diff --git a/.github/workflows/install_and_test.yml b/.github/workflows/install_and_test.yml
index fb695eb..d83ed18 100644
--- a/.github/workflows/install_and_test.yml
+++ b/.github/workflows/install_and_test.yml
@@ -46,6 +46,11 @@ jobs:
pytest -v tests/unit/test_kinetic_density_eq.py
pytest -v tests/unit/test_eigenproblem.py
pytest -v tests/unit/test_loss.py
+ pytest -v tests/unit/test_eri_storage.py
+ pytest -v tests/unit/test_eri_contraction.py
+ pytest -v tests/unit/test_eri_generation.py
+ pytest -v tests/unit/test_eri_serialization.py
+ pytest -v tests/unit/test_eri_consumers.py
- name: Run integration tests
run: |
pytest -v tests/integration/molecules/test_non_xc_energy.py
@@ -53,6 +58,7 @@ jobs:
pytest -v tests/integration/molecules/test_Harris.py
pytest -v tests/integration/molecules/test_predict_B88.py
pytest -v tests/integration/molecules/test_training.py
+ pytest -v tests/integration/molecules/test_eri_storage.py
pytest -v tests/integration/solids/test_training.py
pytest -v tests/integration/solids/test_non_xc_energy.py
pytest -v tests/integration/solids/test_functional_implementations.py
diff --git a/README.md b/README.md
index 5893898..0cdd29c 100644
--- a/README.md
+++ b/README.md
@@ -110,6 +110,60 @@ mf.kernel()
molecule = molecule_from_pyscf(mf)
```
+### Compressed ERI storage (opt-in)
+
+By default `molecule_from_pyscf` stores the electron-repulsion integrals (ERIs)
+as a dense `(n, n, n, n)` tensor, which dominates a `Molecule`'s memory for
+larger basis sets. Pass `eri_storage` to keep PySCF's symmetry-compressed form
+instead:
+
+```python
+molecule = molecule_from_pyscf(mf, eri_storage="s8") # "s4", or omit for "dense"
+```
+
+| `eri_storage` | stored shape | payload | comment |
+| --- | --- | --- | --- |
+| `"dense"` (default) | `(n, n, n, n)` | `n⁴` | current behaviour, unchanged |
+| `"s4"` | `(M, M)` | `≈ n⁴/4` | `M = n (n + 1) / 2` |
+| `"s8"` | `(T,)` | `≈ n⁴/8` | `T = M (M + 1) / 2` |
+
+The format is explicit metadata on the returned `PackedERI` (never inferred
+from the array shape) and the storage dtype is preserved. `save_molecule_data`
+writes the format alongside the ERIs (as dataset attributes), and datasets
+saved before this feature (dense, no metadata) load unchanged.
+
+These operations use the compressed form directly, without ever materializing
+an `n⁴` buffer:
+
+* building the Coulomb (J) potential and the Coulomb energy, i.e. `Molecule.nonXC()`, `Functional.energy`, Harris-functional training and all training losses;
+* `jax.jit` and `jax.grad` through those paths.
+
+An `n⁴` buffer is materialized only explicitly, off the training path:
+
+* `from grad_dft.utils import packed_to_dense; packed_to_dense(molecule.rep_tensor)`;
+* comparing two molecules whose ERIs are stored in different formats.
+
+Unsupported cases raise an informative `ValueError`:
+
+* periodic/PBC mean-field objects (including Gamma point) and complex ERI dtypes — use dense storage for those;
+* `"s2ij"`/`"s2kl"` packings and numeric aliases.
+
+Measured on CPU in float64 with `examples/advanced_scripts/benchmark_eri_storage.py`
+(see that directory's readme to reproduce):
+
+| molecule | basis | dense | s4 | s8 |
+| --- | --- | --- | --- | --- |
+| H2O (n = 24) | cc-pvdz | 2.53 MiB | 0.69 MiB | 0.34 MiB |
+| C6H6 (n = 36) | sto-3g | 12.81 MiB | 3.38 MiB | 1.69 MiB |
+
+Only the storage payload shrinks: the steady-state Coulomb contraction is
+within a small factor of the dense `einsum` (s8 ≈ 2.6x on C6H6,
+at sub-millisecond absolute times) and one training step (loss plus parameter
+gradients) stays comparable, so no speedup is claimed.
+
+Limitations and next steps: periodic compressed formats, `"s2ij"`/`"s2kl"`
+packings and GPU measurements are not implemented yet.
+
### Creating a neural functional
A more complex, neural functional can be created as
diff --git a/examples/advanced_scripts/benchmark_eri_storage.py b/examples/advanced_scripts/benchmark_eri_storage.py
new file mode 100644
index 0000000..b3dcfa4
--- /dev/null
+++ b/examples/advanced_scripts/benchmark_eri_storage.py
@@ -0,0 +1,418 @@
+# Copyright 2023 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+# http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Reproducible benchmark for GradDFT symmetry-compressed ERI storage (#15).
+
+Each measurement runs in its own child process (``--child`` is internal) so
+that peak RSS is reported per ``(format, mode)`` without pollution from the
+other configurations. The parent spawns the children, collects one JSON line
+per child and prints a markdown table.
+
+Notes
+-----
+- Linux/macOS only (``resource.getrusage``): run it on the machine whose
+ numbers you want to quote.
+- ``jax_enable_x64`` is enabled, matching the test suite (float64 storage).
+- "compile" is the first jitted call after tracing (compile + execution);
+ "steady" is the median of ``--repeats`` calls with ``jax.block_until_ready``
+ after warmup. CPU numbers are for guidance, not a hardware-neutral claim.
+- Memory and runtime are reported separately; no speedup is claimed.
+
+Examples
+--------
+ python examples/advanced_scripts/benchmark_eri_storage.py --molecule H2O --basis cc-pvdz
+ python examples/advanced_scripts/benchmark_eri_storage.py --mode generate --formats dense,s8
+"""
+
+import argparse
+import json
+import resource
+import statistics
+import subprocess
+import sys
+import time
+
+from jax import config
+
+import jax
+import jax.numpy as jnp
+from jax.nn import gelu, sigmoid
+from jax.random import PRNGKey
+import numpy as np
+from flax import linen as nn
+from pyscf import dft, gto
+
+from grad_dft import (
+ Molecule,
+ NeuralFunctional,
+ molecule_from_pyscf,
+ non_scf_predictor,
+ simple_energy_loss,
+)
+from grad_dft.utils.eri import PackedERI, packed_to_dense
+
+config.update("jax_enable_x64", True)
+
+RESULT_PREFIX = "ERI_BENCH_RESULT "
+MIB = float(2**20)
+SCF_ITERS = 200
+STEADY_REPEATS = 20
+WARMUP_REPEATS = 3
+FORMATS = ("dense", "s4", "s8")
+
+# Geometries in Angstrom (one atom per line: PySCF's atom parser).
+GEOMETRIES = {
+ "H2": "H 0.0 0.0 -0.3705\nH 0.0 0.0 0.3705",
+ "H2O": "O 0.0 0.0 0.1173\nH 0.0 0.7572 -0.4692\nH 0.0 -0.7572 -0.4692",
+ "CH4": (
+ "C 0.0 0.0 0.0\nH 0.6291 0.6291 0.6291\nH -0.6291 -0.6291 0.6291\n"
+ "H -0.6291 0.6291 -0.6291\nH 0.6291 -0.6291 -0.6291"
+ ),
+ "C6H6": (
+ "C 1.2090 0.0000 0.0000\nC 0.6045 1.0470 0.0000\nC -0.6045 1.0470 0.0000\n"
+ "C -1.2090 0.0000 0.0000\nC -0.6045 -1.0470 0.0000\nC 0.6045 -1.0470 0.0000\n"
+ "H 2.1490 0.0000 0.0000\nH 1.0745 1.8610 0.0000\nH -1.0745 1.8610 0.0000\n"
+ "H -2.1490 0.0000 0.0000\nH -1.0745 -1.8610 0.0000\nH 1.0745 -1.8610 0.0000"
+ ),
+}
+
+
+def converged_mf(molecule: str, basis: str):
+ r"""Converge a spin-restricted Kohn-Sham calculation with no XC energy.
+
+ ``molecule`` is a key of :data:`GEOMETRIES`.
+ """
+ mol = gto.M(atom=GEOMETRIES[molecule], basis=basis, verbose=0)
+ mf = dft.RKS(mol)
+ mf.xc = "0.00*LDA" # no XC energy: isolates the ERI-dependent terms
+ mf.max_cycle = SCF_ITERS
+ mf.kernel()
+ return mf
+
+
+def build_molecule(molecule: str, basis: str, fmt: str):
+ r"""Converge SCF and return ``(mf, molecule, scf_seconds)``."""
+ start = time.perf_counter()
+ mf = converged_mf(molecule, basis)
+ scf_seconds = time.perf_counter() - start
+ mol = molecule_from_pyscf(mf, eri_storage=fmt)
+ jax.block_until_ready(mol.rep_tensor)
+ return mf, mol, scf_seconds
+
+
+def timed_steady(fn, repeats: int):
+ r"""Return ``(compile_seconds, steady_median_seconds)`` for a jitted call."""
+ start = time.perf_counter()
+ jax.block_until_ready(fn())
+ compile_seconds = time.perf_counter() - start
+ for _ in range(WARMUP_REPEATS):
+ jax.block_until_ready(fn())
+ durations = []
+ for _ in range(repeats):
+ start = time.perf_counter()
+ jax.block_until_ready(fn())
+ durations.append(time.perf_counter() - start)
+ return compile_seconds, statistics.median(durations)
+
+
+def storage_stats(rep_tensor):
+ r"""Physical storage shape and bytes of a dense array or PackedERI."""
+ array = rep_tensor.packed if isinstance(rep_tensor, PackedERI) else rep_tensor
+ return {
+ "stored_shape": list(np.shape(array)),
+ "stored_bytes": int(array.size * array.dtype.itemsize),
+ }
+
+
+def peak_rss_mib():
+ r"""Peak resident set size of this process in MiB (Linux: KiB units)."""
+ return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024.0
+
+
+def tracked_peak_mib():
+ r"""Peak device memory in MiB if the backend reports it, else None."""
+ try:
+ stats = jax.devices()[0].memory_stats()
+ except (AttributeError, NotImplementedError, RuntimeError):
+ return None
+ if not stats or "peak_bytes_in_use" not in stats:
+ return None
+ return stats["peak_bytes_in_use"] / MIB
+
+
+def coefficient_inputs(molecule: Molecule, *_, **__):
+ r"""Single-feature density input (clipped), following the test suite."""
+ rho = jnp.clip(molecule.density(), min=1e-30)
+ return jnp.concatenate((rho,), axis=1)
+
+
+def energy_densities(molecule: Molecule, *_, **__):
+ r"""LSDA exchange energy density feature, following the test suite."""
+ rho = jnp.clip(molecule.density(), min=1e-30)
+ return (
+ -3
+ / 2
+ * (3 / (4 * jnp.pi)) ** (1 / 3)
+ * (rho ** (4 / 3)).sum(axis=1, keepdims=True)
+ )
+
+
+def coefficients(_instance, rhoinputs):
+ r"""A minimal parametric energy-density head (one trainable Dense layer)."""
+ x = nn.Dense(features=1)(rhoinputs)
+ x = nn.LayerNorm()(x)
+ return sigmoid(gelu(x))
+
+
+def measure_generate(args) -> dict:
+ r"""ERI generation: SCF time, packing time, storage shape and bytes."""
+ start = time.perf_counter()
+ mf = converged_mf(args.molecule, args.basis)
+ scf_seconds = time.perf_counter() - start
+
+ start = time.perf_counter()
+ mol = molecule_from_pyscf(mf, eri_storage=args.format)
+ jax.block_until_ready(mol.rep_tensor)
+ generate_seconds = time.perf_counter() - start
+
+ rep_tensor = mol.rep_tensor
+ n_orbitals = (
+ rep_tensor.n_orbitals
+ if isinstance(rep_tensor, PackedERI)
+ else rep_tensor.shape[0]
+ )
+ return {
+ "n_orbitals": int(n_orbitals),
+ "scf_seconds": scf_seconds,
+ "generate_seconds": generate_seconds,
+ **storage_stats(rep_tensor),
+ }
+
+
+def measure_contraction(args) -> dict:
+ r"""Coulomb contraction on stored ERIs: to_dense, compile and steady time."""
+ _, mol, scf_seconds = build_molecule(args.molecule, args.basis, args.format)
+ to_dense_seconds = None
+ if isinstance(mol.rep_tensor, PackedERI):
+ start = time.perf_counter()
+ jax.block_until_ready(packed_to_dense(mol.rep_tensor))
+ to_dense_seconds = time.perf_counter() - start
+
+ def contraction(molecule):
+ return molecule.get_coulomb_potential()
+
+ jit_contraction = jax.jit(contraction)
+ compile_seconds, steady_seconds = timed_steady(
+ lambda: jit_contraction(mol), args.repeats
+ )
+ return {
+ "scf_seconds": scf_seconds,
+ "to_dense_seconds": to_dense_seconds,
+ "contraction_compile_seconds": compile_seconds,
+ "contraction_steady_seconds": steady_seconds,
+ }
+
+
+def measure_train(args) -> dict:
+ r"""One training step (loss + parameter gradients) on stored ERIs."""
+ mf, mol, scf_seconds = build_molecule(args.molecule, args.basis, args.format)
+ functional = NeuralFunctional(coefficients, energy_densities, coefficient_inputs)
+ params = functional.init(PRNGKey(args.seed), coefficient_inputs(mol))
+ predictor = non_scf_predictor(functional)
+ truth_energy = mf.e_tot
+
+ def step(params, molecule):
+ (loss, _), grads = simple_energy_loss(params, predictor, molecule, truth_energy)
+ return loss, grads
+
+ jit_step = jax.jit(step)
+ compile_seconds, steady_seconds = timed_steady(
+ lambda: jit_step(params, mol), args.repeats
+ )
+ loss_value = float(jit_step(params, mol)[0])
+ return {
+ "scf_seconds": scf_seconds,
+ "train_compile_seconds": compile_seconds,
+ "train_steady_seconds": steady_seconds,
+ "train_loss": loss_value,
+ }
+
+
+MEASUREMENTS = {
+ "generate": measure_generate,
+ "contraction": measure_contraction,
+ "train": measure_train,
+}
+
+
+def run_child(args) -> int:
+ r"""Run one (format, mode) measurement and emit a single JSON result line."""
+ payload = MEASUREMENTS[args.mode](args)
+ payload.update(
+ {
+ "mode": args.mode,
+ "format": args.format,
+ "molecule": args.molecule,
+ "basis": args.basis,
+ "peak_rss_mib": peak_rss_mib(),
+ "tracked_peak_mib": tracked_peak_mib(),
+ }
+ )
+ print(RESULT_PREFIX + json.dumps(payload))
+ return 0
+
+
+def run_parent(args) -> int:
+ r"""Spawn one child per (format, mode) and print the aggregated table."""
+ modes = ["generate", "contraction", "train"] if args.mode == "all" else [args.mode]
+ results = {}
+ for fmt in args.formats:
+ for mode in modes:
+ command = [
+ sys.executable,
+ __file__,
+ "--child",
+ "--molecule",
+ args.molecule,
+ "--basis",
+ args.basis,
+ "--format",
+ fmt,
+ "--mode",
+ mode,
+ "--repeats",
+ str(args.repeats),
+ "--seed",
+ str(args.seed),
+ ]
+ completed = subprocess.run(
+ command, capture_output=True, text=True, check=False
+ )
+ lines = [
+ line
+ for line in completed.stdout.splitlines()
+ if line.startswith(RESULT_PREFIX)
+ ]
+ if completed.returncode != 0 or not lines:
+ sys.stderr.write(completed.stdout)
+ sys.stderr.write(completed.stderr)
+ raise SystemExit(
+ f"child failed for format={fmt!r} mode={mode!r} "
+ f"(exit {completed.returncode})"
+ )
+ results.setdefault(fmt, {}).update(
+ json.loads(lines[-1][len(RESULT_PREFIX) :])
+ )
+ print(render_table(results))
+ if args.json_out:
+ with open(args.json_out, "w", encoding="utf-8") as handle:
+ json.dump(results, handle, indent=2, sort_keys=True)
+ print(f"\nRaw results written to {args.json_out}")
+ return 0
+
+
+def _cell(value, unit="", precision=4):
+ r"""Format a table cell, using ``n/a`` for unavailable measurements."""
+ if value is None:
+ return "n/a"
+ if isinstance(value, float):
+ return f"{value:.{precision}f}{unit}"
+ return f"{value}{unit}"
+
+
+def render_table(results: dict) -> str:
+ r"""Render the aggregated per-format results as a markdown table."""
+ columns = [
+ ("stored shape", lambda row: _shape(row.get("stored_shape"))),
+ ("stored size", lambda row: _cell(_to_mib(row.get("stored_bytes")), " MiB", 4)),
+ ("generate (s)", lambda row: _cell(row.get("generate_seconds"))),
+ ("to_dense (s)", lambda row: _cell(row.get("to_dense_seconds"))),
+ (
+ "coul. compile (s)",
+ lambda row: _cell(row.get("contraction_compile_seconds")),
+ ),
+ (
+ "coul. steady (s)",
+ lambda row: _cell(row.get("contraction_steady_seconds"), "", 6),
+ ),
+ ("train compile (s)", lambda row: _cell(row.get("train_compile_seconds"))),
+ ("train steady (s)", lambda row: _cell(row.get("train_steady_seconds"), "", 6)),
+ ("peak RSS (MiB)", lambda row: _cell(row.get("peak_rss_mib"), "", 1)),
+ ("tracked peak (MiB)", lambda row: _cell(row.get("tracked_peak_mib"), "", 1)),
+ ]
+ header = "| format | " + " | ".join(name for name, _ in columns) + " |"
+ divider = "| --- | " + " | ".join("---" for _ in columns) + " |"
+ rows = [header, divider]
+ for fmt in FORMATS:
+ row = results.get(fmt)
+ if row is None:
+ continue
+ cells = [fmt] + [render(row) for _, render in columns]
+ rows.append("| " + " | ".join(cells) + " |")
+ return "\n".join(rows)
+
+
+def _shape(shape):
+ r"""Render a stored shape (or ``n/a``), keeping 1-D shapes unambiguous."""
+ if shape is None:
+ return "n/a"
+ if len(shape) == 1:
+ return f"({shape[0]},)"
+ return "(" + ", ".join(str(dim) for dim in shape) + ")"
+
+
+def _to_mib(byte_count):
+ r"""Convert bytes to MiB, or ``None``."""
+ if byte_count is None:
+ return None
+ return byte_count / MIB
+
+
+def parse_args(argv=None):
+ r"""Parse command-line arguments (``--child`` is internal)."""
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--molecule", choices=sorted(GEOMETRIES), default="H2O")
+ parser.add_argument("--basis", default="cc-pvdz")
+ parser.add_argument("--formats", default=",".join(FORMATS))
+ parser.add_argument(
+ "--mode", choices=["generate", "contraction", "train", "all"], default="all"
+ )
+ parser.add_argument("--repeats", type=int, default=STEADY_REPEATS)
+ parser.add_argument("--seed", type=int, default=0)
+ parser.add_argument("--json-out", default=None, help="optional results file")
+ parser.add_argument("--child", action="store_true", help=argparse.SUPPRESS)
+ parser.add_argument("--format", choices=FORMATS, help=argparse.SUPPRESS)
+ args = parser.parse_args(argv)
+ selected = [fmt.strip() for fmt in args.formats.split(",") if fmt.strip()]
+ invalid = [fmt for fmt in selected if fmt not in FORMATS]
+ if invalid:
+ parser.error(
+ f"unknown format(s) {invalid}; expected a comma-separated subset of {FORMATS}"
+ )
+ args.formats = selected
+ if args.child and args.format is None:
+ parser.error("--format is required for child measurements")
+ return args
+
+
+def main(argv=None) -> int:
+ r"""Dispatch to the child measurement or the parent aggregator."""
+ args = parse_args(argv)
+ if args.child:
+ return run_child(args)
+ return run_parent(args)
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/examples/advanced_scripts/readme.md b/examples/advanced_scripts/readme.md
index 261136c..18e090d 100644
--- a/examples/advanced_scripts/readme.md
+++ b/examples/advanced_scripts/readme.md
@@ -14,3 +14,36 @@ molecule = molecule_from_pyscf(mf, omegas = omegas)
```
The `test_constraints.py` file executes an example showing that constraints run. However, note that they have not been tested thoroughly.
+
+## Benchmarking compressed ERI storage
+
+`benchmark_eri_storage.py` measures the memory and runtime footprint of the
+three ERI storage formats (`dense`, `s4`, `s8`) end to end. Each measurement
+runs in its own child process, so peak RSS is reported per `(format, mode)`;
+the parent aggregates the JSON lines into one markdown table.
+
+```bash
+python examples/advanced_scripts/benchmark_eri_storage.py --molecule H2O --basis cc-pvdz
+```
+
+Modes (`--mode`):
+
+- `generate` — SCF time, ERI generation/packing time, physical storage shape and bytes.
+- `contraction` — `to_dense()` time (packed formats) and `Molecule.get_coulomb_potential()` compile vs steady-state time.
+- `train` — one training step (a tiny `NeuralFunctional` with `simple_energy_loss`, including parameter gradients) compile vs steady-state time.
+- `all` (default): the three above.
+
+Measurement notes:
+
+* float64 (`jax_enable_x64`) is enabled, matching the test suite; the tiny
+ functional uses a fixed seed and SCF uses `xc="0.00*LDA"`, so only
+ ERI-dependent terms are timed.
+* "compile" is the first jitted call (tracing + compilation + execution);
+ "steady" is the median of `--repeats` calls after warmup, each forced to
+ completion with `jax.block_until_ready`. `to_dense()` is a plain gather (not
+ jitted) and is timed once.
+* Peak RSS is Linux `ru_maxrss` per child (`device.memory_stats()` is used when
+ the backend reports it). It includes the Python/JAX runtime (roughly 0.8-1.7
+ GiB here), so read it as a build-time footprint; the `stored size` column is
+ the exact ERI payload.
+* Numbers are CPU-only and machine-specific: re-run to get your own.
\ No newline at end of file
diff --git a/grad_dft/__init__.py b/grad_dft/__init__.py
index d95b263..fc87973 100644
--- a/grad_dft/__init__.py
+++ b/grad_dft/__init__.py
@@ -24,6 +24,7 @@
grad_density,
coulomb_energy
)
+from .utils import PackedERI, eri_allclose
from .solid import (
Solid
)
diff --git a/grad_dft/constraints.py b/grad_dft/constraints.py
index b9111bd..949436a 100644
--- a/grad_dft/constraints.py
+++ b/grad_dft/constraints.py
@@ -24,7 +24,8 @@
density,
grad_density,
abs_clip,
- energy_predictor
+ energy_predictor,
+ eri_allclose,
)
from grad_dft.interface.pyscf import (
generate_chi_tensor,
@@ -995,7 +996,7 @@ def xc1(
raise ValueError("The two molecules must have the same omegas")
rep_tensor = molecule1.rep_tensor
- if not jnp.isclose(molecule1.rep_tensor, molecule2.rep_tensor).all():
+ if not eri_allclose(molecule1.rep_tensor, molecule2.rep_tensor):
raise ValueError("The two molecules must have the same rep_tensor")
energy = jnp.empty((0))
basis = molecule1.basis
diff --git a/grad_dft/external/__init__.py b/grad_dft/external/__init__.py
index 963c075..8ba9dc0 100644
--- a/grad_dft/external/__init__.py
+++ b/grad_dft/external/__init__.py
@@ -1,10 +1,10 @@
+"""Adapters around external code that GradDFT builds on.
+
+``Functional``, ``NeuralNumInt`` and ``_SystemState`` come from the installed
+``density_functional_approximation_dm21`` package through the compatibility
+layer in :mod:`grad_dft.external._dm21_compat`; ``_nu_chunk`` is GradDFT's
+own JAX implementation.
+"""
+
+from grad_dft.external._dm21_compat import Functional, NeuralNumInt, _SystemState
from grad_dft.external._hf_density import _nu_chunk
-from grad_dft.external.density_functional_approximation_dm21.density_functional_approximation_dm21.neural_numint import (
- NeuralNumInt,
-)
-from grad_dft.external.density_functional_approximation_dm21.density_functional_approximation_dm21.neural_numint import (
- Functional,
-)
-from grad_dft.external.density_functional_approximation_dm21.density_functional_approximation_dm21.neural_numint import (
- _SystemState,
-)
diff --git a/grad_dft/external/_dm21_compat.py b/grad_dft/external/_dm21_compat.py
new file mode 100644
index 0000000..77239ef
--- /dev/null
+++ b/grad_dft/external/_dm21_compat.py
@@ -0,0 +1,265 @@
+# Copyright 2026 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+# http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Compatibility layer between the installed DM21 package and GradDFT.
+
+``density_functional_approximation_dm21`` is installed from its upstream
+repository (pinned in ``requirements.txt``). Upstream targets TF-Hub's TF1
+Module API and a setuptools that still ships ``pkg_resources``, while this
+environment runs tensorflow-hub 0.16, setuptools 84 (no ``pkg_resources``,
+removed in setuptools 82) and PySCF 2.13. This module carries exactly the
+three deltas that make the installed package run here; everything else is
+used unmodified:
+
+1. A ``pkg_resources`` shim, required because ``tensorflow_hub.__init__``
+ imports ``pkg_resources.parse_version`` at import time while setuptools
+ >= 82 no longer ships ``pkg_resources``.
+2. Shims for the TF1 Module API symbols TF-Hub 0.16 removed
+ (``hub.Module``, ``hub.add_signature``, ``hub.create_module_spec``), so
+ upstream's ``_build_graph`` runs unmodified against the SavedModel
+ checkpoints via ``hub.load``.
+3. The ``eval_xc_eff`` override required on PySCF >= 2.7: nr_rks/nr_uks
+ reach the functional through ``eval_xc_eff``, whose base-class
+ implementation dispatches to the libxc backend (``self.libxc.eval_xc1``)
+ instead of this subclass's ``eval_xc``, and strips the laplacian row that
+ DM21's graph placeholders expect.
+
+The three deltas were developed for the vendored copy of upstream that used
+to live under ``grad_dft/external/`` (tasks #6 and #10) and are carried over
+verbatim here, after task #13 replaced that copy with the installed package
+(git history has the vendored code).
+"""
+
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+
+# Delta 1: tensorflow-hub imports pkg_resources.parse_version at import time,
+# so the shim must run before tensorflow_hub is imported. try/except keeps the
+# real pkg_resources when the environment has one (setuptools < 82).
+try:
+ import pkg_resources # pylint: disable=unused-import
+except ModuleNotFoundError:
+ # setuptools >= 82 removed pkg_resources. Provide the single symbol
+ # tensorflow-hub needs.
+ import sys as _sys
+ from types import SimpleNamespace as _SimpleNamespace
+
+ def _parse_version(version):
+ """Returns a comparable key for a dotted numeric version string."""
+ return tuple(int(part) for part in version.split(".") if part.isdigit())
+
+ _sys.modules.setdefault(
+ "pkg_resources", _SimpleNamespace(parse_version=_parse_version)
+ )
+
+# Delta 2 (and the delta-1 shim above) intentionally come before these
+# imports: the module's whole purpose is to control import order.
+# pylint: disable=wrong-import-position
+import tensorflow_hub as hub
+import density_functional_approximation_dm21.neural_numint as _neural_numint
+
+if not hasattr(hub, "Module"):
+
+ class _SavedModelModule: # pylint: disable=too-few-public-methods
+ """Stand-in for the TF1 ``hub.Module`` that TF-Hub 0.16 removed.
+
+ The DM21 checkpoints are SavedModels, so ``spec`` (a checkpoint
+ directory) is loaded once with ``hub.load`` and called through its
+ ``default`` signature, matching upstream's call convention
+ ``module(tensor_dict, as_dict=True)`` whose keyword names keep the
+ ``tensor_dict$`` prefix.
+ """
+
+ def __init__(self, spec):
+ """Loads the SavedModel's ``default`` signature once."""
+ self._signature = hub.load(spec).signatures["default"]
+
+ def __call__(self, tensor_dict, as_dict=True):
+ """Applies the functional to the feature tensors in ``tensor_dict``."""
+ del as_dict # the signature result already is a dict of tensors
+ return self._signature(**tensor_dict)
+
+ def _add_signature(*args, **kwargs):
+ """No-op replacing the removed ``hub.add_signature``.
+
+ Upstream registers the functional-derivative outputs as a module
+ signature here. The derivatives stay graph attributes consumed by
+ ``eval_xc`` through the session, so only the export path (unsupported,
+ see below) misses the registered signature.
+ """
+ del args, kwargs # unused
+
+ def _create_module_spec(*args, **kwargs):
+ """Raises for upstream's export path, which needs the removed TF1 API.
+
+ ``export_functional_and_derivatives`` builds an exportable TF1 Hub
+ module from these primitives; TF-Hub 0.16 removed them, and a TF2
+ ``tf.saved_model.save`` reimplementation would be needed instead. Only
+ upstream's ``export_saved_model.py`` script and ``neural_numint_test.py``
+ use this path, and neither is part of GradDFT's test suite.
+ """
+ del args, kwargs # unused
+ raise NotImplementedError(
+ "export_functional_and_derivatives requires the TF1 Hub Module API "
+ "(hub.create_module_spec / hub.Module.export) that tensorflow-hub "
+ "0.16 removed; a TF2 tf.saved_model.save reimplementation would be "
+ "needed instead."
+ )
+
+ hub.Module = _SavedModelModule
+ hub.add_signature = _add_signature
+ hub.create_module_spec = _create_module_spec
+
+Functional = _neural_numint.Functional
+_SystemState = _neural_numint._SystemState # pylint: disable=protected-access
+
+
+class NeuralNumInt( # pylint: disable=too-few-public-methods
+ _neural_numint.NeuralNumInt
+):
+ """DM21 numerical integration from the installed package, plus PySCF >= 2.7.
+
+ PySCF's nr_rks/nr_uks grid loops reach the functional through
+ ``eval_xc_eff`` (delta 3 in the module docstring), which upstream does not
+ implement; the override below routes the evaluation back into upstream's
+ ``eval_xc`` and converts its sigma derivatives to the density-parameter
+ derivatives PySCF contracts with the atomic orbitals.
+ """
+
+ def eval_xc_eff( # pylint: disable=too-many-arguments,too-many-positional-arguments
+ self,
+ xc_code: str,
+ rho: np.ndarray,
+ deriv: int = 1,
+ omega: Optional[float] = None,
+ xctype: Optional[str] = None,
+ verbose=None,
+ spin: Optional[int] = None,
+ ) -> Union[List[np.ndarray], Tuple[np.ndarray, np.ndarray, None, None]]:
+ """Evaluates the XC energy and derivatives against density parameters.
+
+ PySCF's nr_rks/nr_uks grid loops call this method rather than eval_xc.
+ It differs from eval_xc only in the derivative convention: eval_xc
+ returns derivatives with respect to sigma = |nabla rho|^2, while
+ eval_xc_eff returns them with respect to the density parameters
+ [rho, nabla rho, tau]. This override applies that chain rule to
+ eval_xc's outputs, so the DM21 evaluation itself (including the local
+ Hartree-Fock side effects on self._vmat_hf) stays in eval_xc.
+
+ See pyscf.dft.numint.NumInt.eval_xc_eff for more details on the
+ interface. The layouts below were verified against that method
+ (max |diff| = 0.0 on mgga_x_tpss for both spin channels).
+
+ Args:
+ xc_code: unused (see eval_xc).
+ rho: density and density derivatives at each grid point, shape
+ (5, N) [rho, nabla_x, nabla_y, nabla_z, tau], or (2, 5, N) for
+ spin-polarized calculations; a laplacian-carrying (6, N) /
+ (2, 6, N) rho is accepted as well. eval_rho runs with
+ MGGA_DENSITY_LAPL off, so rho arrives without the laplacian row
+ the graph's (6, N) placeholders expect; a zero row is inserted
+ before evaluation (the graph unstacks and discards it).
+ deriv: derivative order. 1 (the default) is supported; 0 evaluates
+ the energy and the local Hartree-Fock side effects but returns no
+ potential; greater than 1 raises.
+ omega: RSH parameter. None (the default) is passed through to
+ eval_xc, which raises for an explicit value: DM21 has no
+ range-separation support.
+ xctype: unused. The DM21 functional is always a meta-GGA.
+ verbose: unused.
+ spin: 0 for a spin-unpolarized (restricted) calculation, 1 for
+ spin-polarized. Inferred from rho's leading dimension when None,
+ the same rule as pyscf.dft.numint.NumInt.eval_xc_eff.
+
+ Returns:
+ exc, vxc, fxc, kxc, where:
+ exc is the XC energy density at each grid point, shape (N).
+ vxc has shape (5, N) for spin=0, with rows
+ [dE/d rho, dE/d nabla rho (3 components), dE/d tau], or shape
+ (2, 5, N) for spin=1 with the same rows for the alpha and beta
+ channels (rows [rho, nabla rho, tau] of each).
+ fxc is None. (Second derivatives are not implemented.)
+ kxc is None. (Third derivatives are not implemented.)
+ For deriv=0 the return is [exc, None, None, None].
+
+ Raises:
+ NotImplementedError: if deriv > 1 (DM21 provides no second or higher
+ derivatives; the base class would fail inside transform_xc).
+ """
+ if omega is None:
+ omega = self.omega
+ if deriv > 1:
+ raise NotImplementedError(
+ "DM21 functionals provide no second or higher derivatives, so "
+ "eval_xc_eff cannot be called with deriv > 1."
+ )
+ del xctype # unused: the DM21 functional is always a meta-GGA
+
+ rho = np.asarray(rho, order="C", dtype=np.float64)
+ if spin is None:
+ spin = 1 if (rho.ndim >= 2 and rho.shape[0] == 2) else 0
+ if rho.shape[-2] == 5:
+ rho = np.insert(rho, 4, 0.0, axis=-2)
+
+ exc, (vrho, vsigma, _vlapl, vtau), _, _ = self.eval_xc(
+ xc_code, rho, spin=spin, deriv=deriv, omega=omega, verbose=verbose
+ )
+ if deriv < 1:
+ # Evaluated for exc and the _vmat_hf side effect; no potential is
+ # requested at derivative order 0.
+ return [exc, None, None, None]
+
+ # Chain rule from eval_xc's sigma derivatives (dE/d sigma, with sigma
+ # the PySCF/libxc (sigma_aa, sigma_ab, sigma_bb) components) to the
+ # density-parameter derivatives PySCF contracts with the AOs. To
+ # re-derive it, follow PySCF's MGGA branches: nr_rks does
+ # wv = weight * vxc, wv[0] *= .5, wv[4] *= .5,
+ # _scale_ao_sparse(ao[:4], wv[:4], ...) and _tau_dot_sparse(ao, ao,
+ # wv[4], ...); nr_uks does the same per spin row (wv[:, 0] and
+ # wv[:, 4]) — i.e. row 0 is dE/d rho, rows 1:4 are dE/d nabla rho and
+ # row 4 is dE/d tau. dE/d nabla rho = 2 (dE/d sigma) nabla rho for a
+ # single sigma, and the pairwise expansion below for sigma_ab.
+ if spin == 0:
+ # sigma = nabla rho . nabla rho, so the gradient rows pick up a
+ # factor of 2.
+ vxc = np.concatenate(
+ [
+ vrho[None, :],
+ 2.0 * vsigma[None, :] * rho[1:4],
+ vtau[None, :],
+ ],
+ axis=0,
+ )
+ else:
+ nabla_a, nabla_b = rho[0, 1:4], rho[1, 1:4]
+ vsigma_aa, vsigma_ab, vsigma_bb = vsigma[:, 0], vsigma[:, 1], vsigma[:, 2]
+ vxc_a = np.concatenate(
+ [
+ vrho[:, 0][None, :],
+ 2.0 * vsigma_aa[None, :] * nabla_a + vsigma_ab[None, :] * nabla_b,
+ vtau[:, 0][None, :],
+ ],
+ axis=0,
+ )
+ vxc_b = np.concatenate(
+ [
+ vrho[:, 1][None, :],
+ 2.0 * vsigma_bb[None, :] * nabla_b + vsigma_ab[None, :] * nabla_a,
+ vtau[:, 1][None, :],
+ ],
+ axis=0,
+ )
+ vxc = np.stack([vxc_a, vxc_b], axis=0)
+ return exc, vxc, None, None
diff --git a/grad_dft/external/density_functional_approximation_dm21/.DS_Store b/grad_dft/external/density_functional_approximation_dm21/.DS_Store
deleted file mode 100644
index f5ba6a2..0000000
Binary files a/grad_dft/external/density_functional_approximation_dm21/.DS_Store and /dev/null differ
diff --git a/grad_dft/external/density_functional_approximation_dm21/.bazelrc b/grad_dft/external/density_functional_approximation_dm21/.bazelrc
deleted file mode 100644
index 68d0fcf..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/.bazelrc
+++ /dev/null
@@ -1,609 +0,0 @@
-# The following configuration has been taken from the TensorFlow 2.5.1 release.
-#
-# TensorFlow Bazel configuration file.
-# This file tries to group and simplify build options for TensorFlow
-#
-# ----CONFIG OPTIONS----
-# Android options:
-# android:
-# android_arm:
-# android_arm64:
-# android_x86:
-# android_x86_64:
-#
-# iOS options:
-# ios:
-# ios_armv7:
-# ios_arm64:
-# ios_i386:
-# ios_x86_64:
-# ios_fat:
-#
-# Macosx options
-# darwin_arm64:
-#
-# Compiler options:
-# cuda_clang: Use clang when building CUDA code.
-# c++17: Build with C++17 options (links with libc++)
-# c++1z: Build with C++17 options (links with libc++)
-# c++17_gcc: Build with C++17 options (links with stdlibc++)
-# c++1z_gcc: Build with C++17 options (links with stdlibc++)
-# avx_linux: Build with avx instruction set on linux.
-# avx2_linux: Build with avx2 instruction set on linux.
-# native_arch_linux: Build with instruction sets available to the host machine on linux
-# avx_win: Build with avx instruction set on windows
-# avx2_win: Build with avx2 instruction set on windows
-#
-# Other build options:
-# short_logs: Only log errors during build, skip warnings.
-# verbose_logs: Show all compiler warnings during build.
-# monolithic: Build all TF C++ code into a single shared object.
-# dynamic_kernels: Try to link all kernels dynamically (experimental).
-# libc++: Link against libc++ instead of stdlibc++
-# asan: Build with the clang address sanitizer
-# msan: Build with the clang memory sanitizer
-# ubsan: Build with the clang undefined behavior sanitizer
-#
-#
-# TF version options;
-# v1: Build TF V1 (without contrib)
-# v2: Build TF v2
-#
-# Feature and Third party library support options:
-# xla: Build TF with XLA
-# tpu: Build TF with TPU support
-# cuda: Build with full cuda support.
-# rocm: Build with AMD GPU support (rocm).
-# mkl: Enable full mkl support.
-# tensorrt: Enable Tensorrt support.
-# numa: Enable numa using hwloc.
-# noaws: Disable AWS S3 storage support
-# nogcp: Disable GCS support.
-# nohdfs: Disable hadoop hdfs support.
-# nonccl: Disable nccl support.
-#
-#
-# Remote build execution options (only configured to work with TF team projects for now.)
-# rbe: General RBE options shared by all flavors.
-# rbe_linux: General RBE options used on all linux builds.
-# rbe_win: General RBE options used on all windows builds.
-#
-# rbe_cpu_linux: RBE options to build with only CPU support.
-# rbe_linux_cuda_nvcc_py*: RBE options to build with GPU support using nvcc.
-#
-# rbe_linux_py2: Linux Python 2 RBE config.
-# rbe_linux_py3: Linux Python 3 RBE config
-#
-# rbe_win_py37: Windows Python 3.7 RBE config
-# rbe_win_py38: Windows Python 3.8 RBE config
-#
-# tensorflow_testing_rbe_linux: RBE options to use RBE with tensorflow-testing project on linux
-# tensorflow_testing_rbe_win: RBE options to use RBE with tensorflow-testing project on windows
-#
-# Embedded Linux options (experimental and only tested with TFLite build yet)
-# elinux: General Embedded Linux options shared by all flavors.
-# elinux_aarch64: Embedded Linux options for aarch64 (ARM64) CPU support.
-# elinux_armhf: Embedded Linux options for armhf (ARMv7) CPU support.
-#
-# Release build options (for all operating systems)
-# release_base: Common options for all builds on all operating systems.
-# release_gpu_base: Common options for GPU builds on Linux and Windows.
-# release_cpu_linux: Toolchain and CUDA options for Linux CPU builds.
-# release_cpu_macos: Toolchain and CUDA options for MacOS CPU builds.
-# release_gpu_linux: Toolchain and CUDA options for Linux GPU builds.
-# release_cpu_windows: Toolchain and CUDA options for Windows CPU builds.
-# release_gpu_windows: Toolchain and CUDA options for Windows GPU builds.
-
-# Default build options. These are applied first and unconditionally.
-
-# For projects which use TensorFlow as part of a Bazel build process, putting
-# nothing in a bazelrc will default to a monolithic build. The following line
-# opts in to modular op registration support by default.
-build --define framework_shared_object=true
-
-# For workaround https://github.com/bazelbuild/bazel/issues/8772 with Bazel >= 0.29.1
-build --java_toolchain=@tf_toolchains//toolchains/java:tf_java_toolchain
-build --host_java_toolchain=@tf_toolchains//toolchains/java:tf_java_toolchain
-
-build --define=use_fast_cpp_protos=true
-build --define=allow_oversize_protos=true
-
-build --spawn_strategy=standalone
-build -c opt
-
-# Make Bazel print out all options from rc files.
-build --announce_rc
-
-build --define=grpc_no_ares=true
-
-# See https://github.com/bazelbuild/bazel/issues/7362 for information on what
-# --incompatible_remove_legacy_whole_archive flag does.
-# This flag is set to true in Bazel 1.0 and newer versions. We tried to migrate
-# Tensorflow to the default, however test coverage wasn't enough to catch the
-# errors.
-# There is ongoing work on Bazel team's side to provide support for transitive
-# shared libraries. As part of migrating to transitive shared libraries, we
-# hope to provide a better mechanism for control over symbol exporting, and
-# then tackle this issue again.
-#
-# TODO: Remove this line once TF doesn't depend on Bazel wrapping all library
-# archives in -whole_archive -no_whole_archive.
-build --noincompatible_remove_legacy_whole_archive
-
-# These are bazel 2.0's incompatible flags. Tensorflow needs to use bazel 2.0.0
-# to use cc_shared_library, as part of the Tensorflow Build Improvements RFC:
-# https://github.com/tensorflow/community/pull/179
-build --noincompatible_prohibit_aapt1
-
-build --enable_platform_specific_config
-
-# Enable XLA support by default.
-build --define=with_xla_support=true
-
-build --config=short_logs
-
-build --config=v2
-
-# Default options should come above this line.
-
-# Allow builds using libc++ as a linker library
-# This is mostly for OSSFuzz, so we also pass in the flags from environment to clean build file
-build:libc++ --action_env=CC
-build:libc++ --action_env=CXX
-build:libc++ --action_env=CXXFLAGS=-stdlib=libc++
-build:libc++ --action_env=PATH
-build:libc++ --define force_libcpp=enabled
-build:libc++ --linkopt -fuse-ld=lld
-
-# Android configs. Bazel needs to have --cpu and --fat_apk_cpu both set to the
-# target CPU to build transient dependencies correctly. See
-# https://docs.bazel.build/versions/master/user-manual.html#flag--fat_apk_cpu
-build:android --crosstool_top=//external:android/crosstool
-build:android --host_crosstool_top=@bazel_tools//tools/cpp:toolchain
-build:android_arm --config=android
-build:android_arm --cpu=armeabi-v7a
-build:android_arm --fat_apk_cpu=armeabi-v7a
-build:android_arm64 --config=android
-build:android_arm64 --cpu=arm64-v8a
-build:android_arm64 --fat_apk_cpu=arm64-v8a
-build:android_x86 --config=android
-build:android_x86 --cpu=x86
-build:android_x86 --fat_apk_cpu=x86
-build:android_x86_64 --config=android
-build:android_x86_64 --cpu=x86_64
-build:android_x86_64 --fat_apk_cpu=x86_64
-
-# Sets the default Apple platform to macOS.
-build:macos --apple_platform_type=macos
-
-# gRPC on MacOS requires this #define
-build:macos --copt=-DGRPC_BAZEL_BUILD
-
-# Settings for MacOS on ARM CPUs.
-build:macos_arm64 --cpu=darwin_arm64
-
-# iOS configs for each architecture and the fat binary builds.
-build:ios --apple_platform_type=ios
-build:ios --apple_bitcode=embedded --copt=-fembed-bitcode
-build:ios --copt=-Wno-c++11-narrowing
-build:ios_armv7 --config=ios
-build:ios_armv7 --cpu=ios_armv7
-build:ios_arm64 --config=ios
-build:ios_arm64 --cpu=ios_arm64
-build:ios_i386 --config=ios
-build:ios_i386 --cpu=ios_i386
-build:ios_x86_64 --config=ios
-build:ios_x86_64 --cpu=ios_x86_64
-build:ios_fat --config=ios
-build:ios_fat --ios_multi_cpus=armv7,arm64,i386,x86_64
-
-# Config to use a mostly-static build and disable modular op registration
-# support (this will revert to loading TensorFlow with RTLD_GLOBAL in Python).
-# By default, TensorFlow will build with a dependence on
-# //tensorflow:libtensorflow_framework.so.
-build:monolithic --define framework_shared_object=false
-
-# Please note that MKL on MacOS or windows is still not supported.
-# If you would like to use a local MKL instead of downloading, please set the
-# environment variable "TF_MKL_ROOT" every time before build.
-build:mkl --define=build_with_mkl=true --define=enable_mkl=true
-build:mkl --define=tensorflow_mkldnn_contraction_kernel=0
-build:mkl --define=build_with_openmp=true
-build:mkl -c opt
-
-# config to build OneDNN backend with a user specified threadpool.
-build:mkl_threadpool --define=build_with_mkl=true --define=enable_mkl=true
-build:mkl_threadpool --define=tensorflow_mkldnn_contraction_kernel=0
-build:mkl_threadpool --define=build_with_mkl_opensource=true
-build:mkl_threadpool -c opt
-
-# Config setting to build oneDNN with Compute Library for the Arm Architecture (ACL).
-# This build is for the inference regime only.
-build:mkl_aarch64 --define=build_with_mkl_aarch64=true --define=enable_mkl=true
-build:mkl_aarch64 --define=tensorflow_mkldnn_contraction_kernel=0
-build:mkl_aarch64 --define=build_with_mkl_opensource=true
-build:mkl_aarch64 --define=build_with_openmp=true
-build:mkl_aarch64 -c opt
-
-# This config refers to building CUDA op kernels with nvcc.
-build:cuda --repo_env TF_NEED_CUDA=1
-build:cuda --crosstool_top=@local_config_cuda//crosstool:toolchain
-build:cuda --@local_config_cuda//:enable_cuda
-
-# This config refers to building CUDA op kernels with clang.
-build:cuda_clang --config=cuda
-build:cuda_clang --repo_env TF_CUDA_CLANG=1
-build:cuda_clang --@local_config_cuda//:cuda_compiler=clang
-
-# Debug config
-build:dbg -c dbg
-# for now, disable arm_neon. see: https://github.com/tensorflow/tensorflow/issues/33360
-build:dbg --cxxopt -DTF_LITE_DISABLE_X86_NEON
-# AWS SDK must be compiled in release mode. see: https://github.com/tensorflow/tensorflow/issues/37498
-build:dbg --copt -DDEBUG_BUILD
-
-# Config to build TPU backend
-build:tpu --define=with_tpu_support=true
-
-build:tensorrt --repo_env TF_NEED_TENSORRT=1
-
-build:rocm --crosstool_top=@local_config_rocm//crosstool:toolchain
-build:rocm --define=using_rocm=true --define=using_rocm_hipcc=true
-build:rocm --repo_env TF_NEED_ROCM=1
-
-# Options extracted from configure script
-build:numa --define=with_numa_support=true
-
-# Options to disable default on features
-build:noaws --define=no_aws_support=true
-build:nogcp --define=no_gcp_support=true
-build:nohdfs --define=no_hdfs_support=true
-build:nonccl --define=no_nccl_support=true
-
-build:stackdriver_support --define=stackdriver_support=true
-
-# Modular TF build options
-build:dynamic_kernels --define=dynamic_loaded_kernels=true
-build:dynamic_kernels --copt=-DAUTOLOAD_DYNAMIC_KERNELS
-
-# Build TF with C++ 17 features.
-build:c++17 --cxxopt=-std=c++1z
-build:c++17 --cxxopt=-stdlib=libc++
-build:c++1z --config=c++17
-build:c++17_gcc --cxxopt=-std=c++1z
-build:c++1z_gcc --config=c++17_gcc
-
-# Don't trigger --config= when cross-compiling.
-build:android --noenable_platform_specific_config
-build:ios --noenable_platform_specific_config
-
-# Suppress C++ compiler warnings, otherwise build logs become 10s of MBs.
-build:android --copt=-w
-build:ios --copt=-w
-build:linux --copt=-w
-build:linux --host_copt=-w
-build:macos --copt=-w
-build:windows --copt=/W0
-
-# Tensorflow uses M_* math constants that only get defined by MSVC headers if
-# _USE_MATH_DEFINES is defined.
-build:windows --copt=/D_USE_MATH_DEFINES
-build:windows --host_copt=/D_USE_MATH_DEFINES
-
-# Default paths for TF_SYSTEM_LIBS
-build:linux --define=PREFIX=/usr
-build:linux --define=LIBDIR=$(PREFIX)/lib
-build:linux --define=INCLUDEDIR=$(PREFIX)/include
-build:linux --define=PROTOBUF_INCLUDE_PATH=$(PREFIX)/include
-build:macos --define=PREFIX=/usr
-build:macos --define=LIBDIR=$(PREFIX)/lib
-build:macos --define=INCLUDEDIR=$(PREFIX)/include
-build:macos --define=PROTOBUF_INCLUDE_PATH=$(PREFIX)/include
-# TF_SYSTEM_LIBS do not work on windows.
-
-# By default, build TF in C++ 14 mode.
-build:android --cxxopt=-std=c++14
-build:android --host_cxxopt=-std=c++14
-build:ios --cxxopt=-std=c++14
-build:ios --host_cxxopt=-std=c++14
-build:linux --cxxopt=-std=c++14
-build:linux --host_cxxopt=-std=c++14
-build:macos --cxxopt=-std=c++14
-build:macos --host_cxxopt=-std=c++14
-build:windows --cxxopt=/std:c++14
-build:windows --host_cxxopt=/std:c++14
-
-# On windows, we still link everything into a single DLL.
-build:windows --config=monolithic
-
-# On linux, we dynamically link small amount of kernels
-build:linux --config=dynamic_kernels
-
-# Make sure to include as little of windows.h as possible
-build:windows --copt=-DWIN32_LEAN_AND_MEAN
-build:windows --host_copt=-DWIN32_LEAN_AND_MEAN
-build:windows --copt=-DNOGDI
-build:windows --host_copt=-DNOGDI
-
-# MSVC (Windows): Standards-conformant preprocessor mode
-# See https://docs.microsoft.com/en-us/cpp/preprocessor/preprocessor-experimental-overview
-build:windows --copt=/experimental:preprocessor
-build:windows --host_copt=/experimental:preprocessor
-
-# Misc build options we need for windows.
-build:windows --linkopt=/DEBUG
-build:windows --host_linkopt=/DEBUG
-build:windows --linkopt=/OPT:REF
-build:windows --host_linkopt=/OPT:REF
-build:windows --linkopt=/OPT:ICF
-build:windows --host_linkopt=/OPT:ICF
-build:windows --experimental_strict_action_env=true
-
-# Verbose failure logs when something goes wrong
-build:windows --verbose_failures
-
-# On windows, we never cross compile
-build:windows --distinct_host_configuration=false
-
-# Configure short or long logs
-build:short_logs --output_filter=DONT_MATCH_ANYTHING
-build:verbose_logs --output_filter=
-
-# Instruction set optimizations
-# TODO(gunan): Create a feature in toolchains for avx/avx2 to
-# avoid having to define linux/win separately.
-build:avx_linux --copt=-mavx
-build:avx_linux --host_copt=-mavx
-build:avx2_linux --copt=-mavx2
-build:native_arch_linux --copt=-march=native
-build:avx_win --copt=/arch=AVX
-build:avx2_win --copt=/arch=AVX2
-
-# Options to build TensorFlow 1.x or 2.x.
-build:v1 --define=tf_api_version=1 --action_env=TF2_BEHAVIOR=0
-build:v2 --define=tf_api_version=2 --action_env=TF2_BEHAVIOR=1
-
-# Disable XLA on mobile.
-build:xla --define=with_xla_supprt=true # TODO: remove, it's on by default.
-build:android --define=with_xla_support=false
-build:ios --define=with_xla_support=false
-
-# BEGIN TF REMOTE BUILD EXECUTION OPTIONS
-# Options when using remote execution
-# WARNING: THESE OPTIONS WONT WORK IF YOU DO NOT HAVE PROPER AUTHENTICATION AND PERMISSIONS
-
-# Flag to enable remote config
-common --experimental_repo_remote_exec
-
-build:rbe --repo_env=BAZEL_DO_NOT_DETECT_CPP_TOOLCHAIN=1
-build:rbe --google_default_credentials
-build:rbe --bes_backend=buildeventservice.googleapis.com
-build:rbe --bes_results_url="https://source.cloud.google.com/results/invocations"
-build:rbe --bes_timeout=600s
-build:rbe --define=EXECUTOR=remote
-build:rbe --distinct_host_configuration=false
-build:rbe --flaky_test_attempts=3
-build:rbe --jobs=200
-build:rbe --remote_executor=grpcs://remotebuildexecution.googleapis.com
-build:rbe --remote_timeout=3600
-build:rbe --spawn_strategy=remote,worker,standalone,local
-test:rbe --test_env=USER=anon
-# Attempt to minimize the amount of data transfer between bazel and the remote
-# workers:
-build:rbe --remote_download_toplevel
-
-build:rbe_linux --config=rbe
-build:rbe_linux --action_env=PATH="/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/local/go/bin"
-build:rbe_linux --host_javabase=@bazel_toolchains//configs/ubuntu16_04_clang/1.1:jdk8
-build:rbe_linux --javabase=@bazel_toolchains//configs/ubuntu16_04_clang/1.1:jdk8
-build:rbe_linux --host_java_toolchain=@bazel_tools//tools/jdk:toolchain_hostjdk8
-build:rbe_linux --java_toolchain=@bazel_tools//tools/jdk:toolchain_hostjdk8
-
-# Non-rbe settings we should include because we do not run configure
-build:rbe_linux --config=avx_linux
-# TODO(gunan): Check why we need this specified in rbe, but not in other builds.
-build:rbe_linux --linkopt=-lrt
-build:rbe_linux --host_linkopt=-lrt
-build:rbe_linux --linkopt=-lm
-build:rbe_linux --host_linkopt=-lm
-
-build:rbe_cpu_linux --config=rbe_linux
-build:rbe_cpu_linux --host_crosstool_top="@org_tensorflow//third_party/toolchains/preconfig/ubuntu16.04/gcc7_manylinux2010:toolchain"
-build:rbe_cpu_linux --crosstool_top="@org_tensorflow//third_party/toolchains/preconfig/ubuntu16.04/gcc7_manylinux2010:toolchain"
-build:rbe_cpu_linux --extra_toolchains="@org_tensorflow//third_party/toolchains/preconfig/ubuntu16.04/gcc7_manylinux2010:cc-toolchain-k8"
-build:rbe_cpu_linux --extra_execution_platforms="@ubuntu16.04-manylinux2010-py3_config_platform//:platform"
-build:rbe_cpu_linux --extra_execution_platforms="@ubuntu16.04-manylinux2010-py3_config_platform//:platform"
-build:rbe_cpu_linux --host_platform="@ubuntu16.04-manylinux2010-py3_config_platform//:platform"
-build:rbe_cpu_linux --platforms="@ubuntu16.04-manylinux2010-py3_config_platform//:platform"
-
-build:rbe_linux_cuda_base --config=rbe_linux
-build:rbe_linux_cuda_base --config=cuda
-build:rbe_linux_cuda_base --config=tensorrt
-build:rbe_linux_cuda_base --action_env=TF_CUDA_VERSION=11
-build:rbe_linux_cuda_base --action_env=TF_CUDNN_VERSION=8
-build:rbe_linux_cuda_base --repo_env=REMOTE_GPU_TESTING=1
-test:rbe_linux_cuda_base --test_env=LD_LIBRARY_PATH="/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64"
-
-build:rbe_linux_cuda11.2_nvcc_base --config=rbe_linux_cuda_base
-build:rbe_linux_cuda11.2_nvcc_base --host_crosstool_top="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda//crosstool:toolchain"
-build:rbe_linux_cuda11.2_nvcc_base --crosstool_top="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda//crosstool:toolchain"
-build:rbe_linux_cuda11.2_nvcc_base --extra_toolchains="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda//crosstool:toolchain-linux-x86_64"
-build:rbe_linux_cuda11.2_nvcc_base --extra_execution_platforms="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_platform//:platform"
-build:rbe_linux_cuda11.2_nvcc_base --host_platform="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_platform//:platform"
-build:rbe_linux_cuda11.2_nvcc_base --platforms="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_platform//:platform"
-build:rbe_linux_cuda11.2_nvcc_base --repo_env=TF_CUDA_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda"
-build:rbe_linux_cuda11.2_nvcc_base --repo_env=TF_TENSORRT_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_tensorrt"
-build:rbe_linux_cuda11.2_nvcc_base --repo_env=TF_NCCL_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_nccl"
-build:rbe_linux_cuda11.2_nvcc_py3.6 --config=rbe_linux_cuda11.2_nvcc_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python3.6"
-build:rbe_linux_cuda11.2_nvcc_py3.7 --config=rbe_linux_cuda11.2_nvcc_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python3.7"
-build:rbe_linux_cuda11.2_nvcc_py3.8 --config=rbe_linux_cuda11.2_nvcc_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python3.8"
-build:rbe_linux_cuda11.2_nvcc_py3.9 --config=rbe_linux_cuda11.2_nvcc_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python3.9"
-
-# Map default to CUDA 11.2.
-build:rbe_linux_cuda_nvcc_py36 --config=rbe_linux_cuda11.2_nvcc_py3.6
-build:rbe_linux_cuda_nvcc_py37 --config=rbe_linux_cuda11.2_nvcc_py3.7
-build:rbe_linux_cuda_nvcc_py38 --config=rbe_linux_cuda11.2_nvcc_py3.8
-build:rbe_linux_cuda_nvcc_py39 --config=rbe_linux_cuda11.2_nvcc_py3.9
-
-# Deprecated configs that people might still use.
-build:rbe_linux_cuda_nvcc --config=rbe_linux_cuda_nvcc_py36
-build:rbe_gpu_linux --config=rbe_linux_cuda_nvcc
-
-build:rbe_linux_cuda_clang_base --config=rbe_linux_cuda_base
-build:rbe_linux_cuda_clang_base --repo_env TF_CUDA_CLANG=1
-build:rbe_linux_cuda_clang_base --@local_config_cuda//:cuda_compiler=clang
-build:rbe_linux_cuda_clang_base --crosstool_top="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda//crosstool:toolchain"
-build:rbe_linux_cuda_clang_base --extra_toolchains="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda//crosstool:toolchain-linux-x86_64"
-build:rbe_linux_cuda_clang_base --extra_execution_platforms="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_platform//:platform"
-build:rbe_linux_cuda_clang_base --host_platform="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_platform//:platform"
-build:rbe_linux_cuda_clang_base --platforms="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_platform//:platform"
-build:rbe_linux_cuda_clang_base --repo_env=TF_CUDA_CONFIG_REPO="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_cuda"
-build:rbe_linux_cuda_clang_base --repo_env=TF_TENSORRT_CONFIG_REPO="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_tensorrt"
-build:rbe_linux_cuda_clang_base --repo_env=TF_NCCL_CONFIG_REPO="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_nccl"
-build:rbe_linux_cuda_clang_py27 --config=rbe_linux_cuda_clang_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python2.7"
-build:rbe_linux_cuda_clang_py35 --config=rbe_linux_cuda_clang_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python3.5"
-build:rbe_linux_cuda_clang_py36 --config=rbe_linux_cuda_clang_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python3.6"
-build:rbe_linux_cuda_clang_py37 --config=rbe_linux_cuda_clang_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python3.7"
-build:rbe_linux_cuda_clang_py38 --config=rbe_linux_cuda_clang_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-clang_manylinux2010-cuda11.2-cudnn8.1-tensorrt7.2_config_python3.8"
-
-# ROCm
-build:rbe_linux_rocm_base --config=rbe_linux
-build:rbe_linux_rocm_base --repo_env=TF_NEED_ROCM=1
-build:rbe_linux_rocm_base --crosstool_top="@ubuntu18.04-gcc7_manylinux2010-rocm_config_rocm//crosstool:toolchain"
-build:rbe_linux_rocm_base --extra_toolchains="@ubuntu18.04-gcc7_manylinux2010-rocm_config_rocm//crosstool:toolchain-linux-x86_64"
-build:rbe_linux_rocm_base --extra_execution_platforms="@ubuntu18.04-gcc7_manylinux2010-rocm_config_platform//:platform"
-build:rbe_linux_rocm_base --host_platform="@ubuntu18.04-gcc7_manylinux2010-rocm_config_platform//:platform"
-build:rbe_linux_rocm_base --platforms="@ubuntu18.04-gcc7_manylinux2010-rocm_config_platform//:platform"
-build:rbe_linux_rocm_base --action_env=TF_ROCM_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-rocm_config_rocm"
-build:rbe_linux_rocm_base --define=using_rocm_hipcc=true
-build:rbe_linux_rocm_py2.7 --config=rbe_linux_rocm_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-rocm_config_python2.7"
-build:rbe_linux_rocm_py3.5 --config=rbe_linux_rocm_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-rocm_config_python3.5"
-build:rbe_linux_rocm_py3.6 --config=rbe_linux_rocm_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-rocm_config_python3.6"
-build:rbe_linux_rocm_py3.7 --config=rbe_linux_rocm_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-rocm_config_python3.7"
-build:rbe_linux_rocm_py3.8 --config=rbe_linux_rocm_base --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu18.04-gcc7_manylinux2010-rocm_config_python3.8"
-
-# Linux CPU
-build:rbe_linux_py2 --config=rbe_linux
-build:rbe_linux_py2 --repo_env=PYTHON_BIN_PATH="/usr/bin/python2"
-build:rbe_linux_py2 --python_path="/usr/bin/python2"
-build:rbe_linux_py2 --repo_env=TF_PYTHON_CONFIG_REPO="@org_tensorflow//third_party/toolchains/preconfig/ubuntu16.04/py"
-
-build:rbe_linux_py3 --config=rbe_linux
-build:rbe_linux_py3 --python_path="/usr/bin/python3"
-build:rbe_linux_py3 --repo_env=TF_PYTHON_CONFIG_REPO="@ubuntu16.04-manylinux2010-py3_config_python"
-
-build:rbe_win --config=rbe
-build:rbe_win --crosstool_top="@org_tensorflow//third_party/toolchains/preconfig/win/tf_win_08062020:toolchain"
-build:rbe_win --extra_toolchains="@org_tensorflow//third_party/toolchains/preconfig/win/tf_win_08062020:cc-toolchain-x64_windows"
-build:rbe_win --host_javabase="@org_tensorflow//third_party/toolchains/preconfig/win:windows_jdk8"
-build:rbe_win --javabase="@org_tensorflow//third_party/toolchains/preconfig/win:windows_jdk8"
-build:rbe_win --extra_execution_platforms="@org_tensorflow//third_party/toolchains/preconfig/win:rbe_windows_ltsc2019"
-build:rbe_win --host_platform="@org_tensorflow//third_party/toolchains/preconfig/win:rbe_windows_ltsc2019"
-build:rbe_win --platforms="@org_tensorflow//third_party/toolchains/preconfig/win:rbe_windows_ltsc2019"
-build:rbe_win --shell_executable=C:\\tools\\msys64\\usr\\bin\\bash.exe
-
-# TODO(gunan): Remove once we use MSVC 2019 with latest patches.
-build:rbe_win --define=override_eigen_strong_inline=true
-build:rbe_win --jobs=100
-
-build:rbe_win_py37 --config=rbe
-build:rbe_win_py37 --repo_env=TF_PYTHON_CONFIG_REPO="@windows_py37_config_python"
-build:rbe_win_py37 --python_path=C:\\Python37\\python.exe
-
-build:rbe_win_py38 --config=rbe
-build:rbe_win_py38 --repo_env=PYTHON_BIN_PATH=C:\\Python38\\python.exe
-build:rbe_win_py38 --repo_env=PYTHON_LIB_PATH=C:\\Python38\\lib\\site-packages
-build:rbe_win_py38 --repo_env=TF_PYTHON_CONFIG_REPO=@org_tensorflow//third_party/toolchains/preconfig/win_1803/py38
-build:rbe_win_py38 --python_path=C:\\Python38\\python.exe
-
-# These you may need to change for your own GCP project.
-build:tensorflow_testing_rbe --project_id=tensorflow-testing
-common:tensorflow_testing_rbe_linux --remote_instance_name=projects/tensorflow-testing/instances/default_instance
-build:tensorflow_testing_rbe_linux --config=tensorflow_testing_rbe
-
-common:tensorflow_testing_rbe_win --remote_instance_name=projects/tensorflow-testing/instances/windows
-build:tensorflow_testing_rbe_win --config=tensorflow_testing_rbe
-
-# TFLite build configs for generic embedded Linux
-build:elinux --crosstool_top=@local_config_embedded_arm//:toolchain
-build:elinux --host_crosstool_top=@bazel_tools//tools/cpp:toolchain
-build:elinux_aarch64 --config=elinux
-build:elinux_aarch64 --cpu=aarch64
-build:elinux_armhf --config=elinux
-build:elinux_armhf --cpu=armhf
-# END TF REMOTE BUILD EXECUTION OPTIONS
-
-# Config-specific options should come above this line.
-
-# Load rc file written by ./configure.
-try-import %workspace%/.tf_configure.bazelrc
-
-# Load rc file with user-specific options.
-try-import %workspace%/.bazelrc.user
-
-# Here are bazelrc configs for release builds
-build:release_base --config=v2
-build:release_base --distinct_host_configuration=false
-test:release_base --flaky_test_attempts=3
-test:release_base --test_size_filters=small,medium
-
-build:release_cpu_linux --config=release_base
-build:release_cpu_linux --config=avx_linux
-build:release_cpu_linux --crosstool_top=@org_tensorflow//third_party/toolchains/preconfig/ubuntu16.04/gcc7_manylinux2010-nvcc-cuda11.2:toolchain
-test:release_cpu_linux --test_env=LD_LIBRARY_PATH
-
-build:release_cpu_macos --config=release_base
-build:release_cpu_macos --config=avx_linux
-
-build:release_gpu_base --config=cuda
-build:release_gpu_base --action_env=TF_CUDA_VERSION="11"
-build:release_gpu_base --action_env=TF_CUDNN_VERSION="8"
-build:release_gpu_base --repo_env=TF_CUDA_COMPUTE_CAPABILITIES="sm_35,sm_50,sm_60,sm_70,sm_75,compute_80"
-
-build:release_gpu_linux --config=release_cpu_linux
-build:release_gpu_linux --config=release_gpu_base
-build:release_gpu_linux --config=tensorrt
-build:release_gpu_linux --action_env=CUDA_TOOLKIT_PATH="/usr/local/cuda-11.2"
-build:release_gpu_linux --action_env=LD_LIBRARY_PATH="/usr/local/cuda:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64:/usr/local/tensorrt/lib"
-build:release_gpu_linux --action_env=GCC_HOST_COMPILER_PATH="/usr/bin/gcc-5"
-build:release_gpu_linux --crosstool_top=@org_tensorflow//third_party/toolchains/preconfig/ubuntu16.04/gcc7_manylinux2010-nvcc-cuda11.2:toolchain
-
-build:release_cpu_windows --config=release_base
-build:release_cpu_windows --config=avx_win
-build:release_cpu_windows --define=no_tensorflow_py_deps=true
-# First available in VS 16.4. Speeds Windows compile times by a lot. See
-# https://groups.google.com/a/tensorflow.org/d/topic/build/SsW98Eo7l3o/discussion
-build:release_cpu_windows --copt=/d2ReducedOptimizeHugeFunctions --host_copt=/d2ReducedOptimizeHugeFunctions
-
-build:release_gpu_windows --config=release_cpu_windows
-build:release_gpu_windows --config=release_gpu_base
-
-# Address sanitizer
-# CC=clang bazel build --config asan
-build:asan --strip=never
-build:asan --copt -fsanitize=address
-build:asan --copt -DADDRESS_SANITIZER
-build:asan --copt -g
-build:asan --copt -O3
-build:asan --copt -fno-omit-frame-pointer
-build:asan --linkopt -fsanitize=address
-
-# Memory sanitizer
-# CC=clang bazel build --config msan
-build:msan --strip=never
-build:msan --copt -fsanitize=memory
-build:msan --copt -DADDRESS_SANITIZER
-build:msan --copt -g
-build:msan --copt -O3
-build:msan --copt -fno-omit-frame-pointer
-build:msan --linkopt -fsanitize=memory
-
-# Undefined Behavior Sanitizer
-# CC=clang bazel build --config ubsan
-build:ubsan --strip=never
-build:ubsan --copt -fsanitize=undefined
-build:ubsan --copt -g
-build:ubsan --copt -O3
-build:ubsan --copt -fno-omit-frame-pointer
-build:ubsan --linkopt -fsanitize=undefined
-build:ubsan --linkopt -lubsan
diff --git a/grad_dft/external/density_functional_approximation_dm21/.vscode/launch.json b/grad_dft/external/density_functional_approximation_dm21/.vscode/launch.json
deleted file mode 100644
index 5da9c63..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/.vscode/launch.json
+++ /dev/null
@@ -1,16 +0,0 @@
-{
- // Use IntelliSense para saber los atributos posibles.
- // Mantenga el puntero para ver las descripciones de los existentes atributos.
- // Para más información, visite: https://go.microsoft.com/fwlink/?linkid=830387
- "version": "0.2.0",
- "configurations": [
- {
- "name": "Python: archivo actual",
- "type": "python",
- "request": "launch",
- "program": "${file}",
- "console": "integratedTerminal",
- "justMyCode": true
- }
- ]
-}
\ No newline at end of file
diff --git a/grad_dft/external/density_functional_approximation_dm21/BUILD.bazel b/grad_dft/external/density_functional_approximation_dm21/BUILD.bazel
deleted file mode 100644
index 838a602..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/BUILD.bazel
+++ /dev/null
@@ -1,125 +0,0 @@
-"""DM21 functionals."""
-
-load("@org_tensorflow//tensorflow/python/tools:tools.bzl", "saved_model_compile_aot")
-load("@rules_python//python:defs.bzl", "py_binary", "py_library")
-load("@external_py_deps//:requirements.bzl", "requirement")
-
-licenses(["notice"])
-
-py_library(
- name = "compute_hfx_density",
- srcs = ["density_functional_approximation_dm21/compute_hfx_density.py"],
- srcs_version = "PY3",
- deps = [
- requirement("numpy"),
- requirement("pyscf"),
- ],
-)
-
-py_test(
- name = "compute_hfx_density_test",
- srcs = ["density_functional_approximation_dm21/compute_hfx_density_test.py"],
- python_version = "PY3",
- srcs_version = "PY3",
- deps = [
- ":compute_hfx_density",
- requirement("attrs"),
- requirement("numpy"),
- requirement("scipy"),
- requirement("pyscf"),
- "@io_abseil_py//absl/testing:absltest",
- "@io_abseil_py//absl/testing:parameterized",
- ],
-)
-
-filegroup(
- name = "dm21_checkpoints",
- srcs = glob(["density_functional_approximation_dm21/checkpoints/**"]),
-)
-
-py_library(
- name = "neural_numint",
- srcs = ["density_functional_approximation_dm21/neural_numint.py"],
- data = [":dm21_checkpoints"],
- srcs_version = "PY3",
- deps = [
- ":compute_hfx_density",
- requirement("attrs"),
- requirement("keras"),
- requirement("numpy"),
- requirement("pyscf"),
- "@org_tensorflow//tensorflow:tensorflow_py",
- requirement("tensorflow-hub"),
- ],
-)
-
-py_test(
- name = "neural_numint_test",
- srcs = ["density_functional_approximation_dm21/neural_numint_test.py"],
- python_version = "PY3",
- srcs_version = "PY3",
- deps = [
- ":neural_numint",
- requirement("attrs"),
- requirement("pyscf"),
- "@io_abseil_py//absl/testing:parameterized",
- "@org_tensorflow//tensorflow:tensorflow_py",
- ],
-)
-
-py_binary(
- name = "export_saved_model",
- srcs = ["density_functional_approximation_dm21/export_saved_model.py"],
- data = [":dm21_checkpoints"],
- python_version = "PY3",
- srcs_version = "PY3",
- deps = [
- ":neural_numint",
- ],
-)
-
-EXPORTED_SAVED_MODEL_OBJECTS = [
- "DM21/saved_model.pb",
- "DM21/variables/variables.index",
- "DM21/variables/variables.data-00000-of-00001",
-]
-
-genrule(
- name = "create_model_for_aot_compile",
- outs = EXPORTED_SAVED_MODEL_OBJECTS,
- # The functional used in dm21_aot_compiled_example can be changed by
- # setting the --functional flag to the desired functional.
- cmd = "$(location :export_saved_model) --functional DM21 --batch_size 1000 --out_dir $(@D)/DM21",
- tools = [":export_saved_model"],
-)
-
-filegroup(
- name = "dm21_exported_model",
- srcs = EXPORTED_SAVED_MODEL_OBJECTS,
-)
-
-saved_model_compile_aot(
- name = "aot_compiled_dm21",
- cpp_class = "dm21::functional",
- directory = ":DM21",
- filegroups = [":dm21_exported_model"],
- force_without_xla_support_flag = False,
- multithreading = False,
- signature_def = "default",
- tag_set = "''",
- target_triple = "x86_64-pc-linux",
-)
-
-cc_library(
- name = "dm21_aot_compiled_example",
- srcs = ["cc/dm21_aot_compiled_example.cc"],
- hdrs = ["cc/dm21_aot_compiled_example.h"],
- deps = [":aot_compiled_dm21"],
-)
-
-cc_binary(
- name = "run_dm21_aot_compiled_example",
- srcs = ["cc/run_dm21_aot_compiled_example.cc"],
- copts = ["-DXLA_AVAILABLE"],
- deps = [":dm21_aot_compiled_example"],
-)
diff --git a/grad_dft/external/density_functional_approximation_dm21/README.md b/grad_dft/external/density_functional_approximation_dm21/README.md
deleted file mode 100644
index f8d8bf3..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/README.md
+++ /dev/null
@@ -1,308 +0,0 @@
-# Pushing the Frontiers of Density Functionals by Solving the Fractional Electron Problem
-
-This package provides a PySCF interface to the DM21 (DeepMind 21) family of
-exchange-correlation functionals described in the paper ["Pushing the Frontiers
-of Density Functionals by Solving the Fractional Electron
-Problem"](https://doi.org/10.1126/science.abj6511).
-
-## Installation
-
-`pip install .` installs the interfaces to the DM21 functionals and required
-dependencies. This is best done inside a
-[virtual environment](https://docs.python-guide.org/dev/virtualenvs/).
-
-Note: using PySCF 2.0 (or later) enables substantially more efficient
-calculation of the local Hartree-Fock features, resulting in a large speed
-increase.
-
-### Installing directly
-
-To install without cloning or downloading the deepmind_research repository,
-execute:
-
-```shell
-python3 -m venv ~/venv/DM21
-source ~/venv/DM21/bin/activate
-pip install git+git://github.com/deepmind/deepmind-research.git#subdirectory=density_functional_approximation_dm21
-```
-
-### Downloading and installing from a local git repository
-
-Alternatively, clone the deepmind_research repository and install from a local
-repository:
-
-```shell
-git clone https://github.com/deepmind/deepmind-research.git
-cd deepmind-research/density_functional_approximation_dm21
-python3 -m venv ~/venv/DM21
-source ~/venv/DM21/bin/activate
-pip install .
-```
-
-The tests can be run either by running the test files directly or using
-`py.test`, again from the `density_functional_approximation_dm21` subdirectory:
-
-```shell
-pip install '.[testing]'
-py.test
-```
-
-## PySCF interface
-
-A typical DFT calculation with PySCF is set up and run using:
-
-```python
-from pyscf import gto
-from pyscf import dft
-
-# Create the molecule of interest and select the basis set.
-mol = gto.Mole()
-mol.atom = 'Ne 0.0 0.0 0.0'
-mol.basis = 'cc-pVDZ'
-mol.build()
-
-# Create a DFT solver and select the exchange-correlation functional.
-mf = dft.RKS(mol)
-mf.xc = 'b3lyp'
-
-# Run the DFT calculation.
-mf.kernel()
-```
-
-The DM21 functionals can be used in a very similar way, except we need to
-compute local Hartree-Fock features, which does not fit in with the interface
-used by conventional functionals. Instead, this package supplies a lightweight
-wrapper around PySCF's numerical integration class which evaluates the
-exchange-correlation energy and potential over a real-space grid. To use the
-DM21 functional with PySCF, the above code needs to be changed to:
-
-```python
-import density_functional_approximation_dm21 as dm21
-from pyscf import gto
-from pyscf import dft
-
-# Create the molecule of interest and select the basis set.
-mol = gto.Mole()
-mol.atom = 'Ne 0.0 0.0 0.0'
-mol.basis = 'cc-pVDZ'
-mol.build()
-
-# Create a DFT solver and insert the DM21 functional into the solver.
-mf = dft.RKS(mol)
-mf._numint = dm21.NeuralNumInt(dm21.Functional.DM21)
-
-# Run the DFT calculation.
-mf.kernel()
-```
-
-i.e. instead of specifying the functional with `mf.xc = `, the
-functional is specified using `mf._numint = dm21.NeuralNumInt()`, where `` is the corresponding member of the
-`Functional` enum.
-
-Available functionals are:
-
-* `DM21` - trained on molecules dataset, and fractional charge, and fractional
- spin constraints.
-* `DM21m` - trained on molecules dataset.
-* `DM21mc` - trained on molecules dataset, and fractional charge constraints.
-* `DM21mu` - trained on molecules dataset, and electron gas constraints.
-
-Full details of the network architecture, training method and datasets used can
-be found in the paper (reference below). Note that the results in our paper also
-include D3 corrections, which must be
-[included separately](https://pyscf.org/user/dft.html#dispersion-corrections).
-
-
-### Best practices for using the neural functionals.
-
-In this section, we suggest some tips for using the neural functionals in a way
-similar to how they were used for benchmarking in the paper. The tensorflow
-network that we used is running at single precision, and as such it is very
-hard to converge calculations to the high convergence thresholds which are
-default in pyscf. For example, the following script should allow users to
-run an atomization energy calculation for methane.
-
-```python
-import density_functional_approximation_dm21 as dm21
-from pyscf import gto
-from pyscf import dft
-
-# Create the molecule of interest and select the basis set.
-methane = gto.Mole()
-methane.atom = """H 0.000000000000 0.000000000000 0.000000000000
- C 0.000000000000 0.000000000000 1.087900000000
- H 1.025681956337 0.000000000000 1.450533333333
- H -0.512840978169 0.888266630391 1.450533333333
- H -0.512840978169 -0.888266630391 1.450533333333"""
-methane.basis = 'def2-qzvp'
-methane.verbose = 4
-methane.build()
-
-carbon = gto.Mole()
-carbon.atom = 'C 0.0 0.0 0.0'
-carbon.basis = 'def2-qzvp'
-carbon.spin = 2
-carbon.verbose = 4
-carbon.build()
-
-hydrogen = gto.Mole()
-hydrogen.atom = 'H 0.0 0.0 0.0'
-hydrogen.basis = 'def2-qzvp'
-hydrogen.spin = 1
-hydrogen.verbose = 4
-hydrogen.build()
-
-energies = []
-for mol in [methane, carbon, hydrogen]:
- # Create a DFT solver and insert the DM21 functional into the solver.
- if mol.spin == 0:
- mf = dft.RKS(mol)
- else:
- mf = dft.UKS(mol)
- # It will make SCF faster to start close to the solution with a cheaper
- # functional.
- mf.xc = 'B3LYP'
- mf.run()
- dm0 = mf.make_rdm1()
-
- mf._numint = dm21.NeuralNumInt(dm21.Functional.DM21)
- # It's wise to relax convergence tolerances.
- mf.conv_tol = 1E-6
- mf.conv_tol_grad = 1E-3
- # Run the DFT calculation.
- energy = mf.kernel(dm0=dm0)
- energies.append(energy)
-
-print({'CH4': energies[0], 'C': energies[1], 'H': energies[2]})
-```
-
-This script should produce three energies (in Hartrees) for the water molecule
-and the two atoms of
-`{'CH4': -40.51785372584538, 'C': -37.84542045526023, 'H': -0.5011533955627797}`
-, this leads to an atomization energy of 419.06 kcal/mol, which is then
-corrected with the D3(BJ) correction for methane (1.20 kcal/mol) to yield a
-predicted atomization energy of 420.26 kcal/mol. Comparing this to the
-literature value of 420.42, leads us to deduce an error of around 0.2 kcal/mol.
-
-It should also be noted that if a closed shell system is run unrestricted it
-can give a small difference between spin densities and eigenvalues with a
-negligible effect on the energy.
-
-
-## Using DM21 from C++
-
-There are two options for using the DM21 from C++.
-
-1. Load the SavedModel using
- [TensorFlow's C++ API](https://www.tensorflow.org/guide/saved_model#load_a_savedmodel_in_c).
- This requires a run-time dependency on the TensorFlow library.
-2. Compile the model ahead-of-time into a standalone library. This requires all
- array dimensions to be fixed at compile time, which imposes a minor
- limitation on the interface for using the functional. As the DM21
- functionals act on grid points independently, this does not restrict the
- system size.
-
-The first option is more flexible but trickier to setup. Consequently we
-demonstrate the second option: compiling the functional into a standalone
-library. An example of running the DM21 functional using a standalone compiled
-library is provided in `cc/dm21_aot_compiled_example.cc`. This requires a
-link-time dependency on parts of the `xla_compiled_cpu_runtime_standalone`
-library, which are not included in the compiled functional library. The easiest
-way to build this is to use [Bazel](https://bazel.build). The first step is to
-[install Bazel](https://docs.bazel.build/versions/4.2.0/install.html).
-[Bazelisk](https://docs.bazel.build/versions/main/install-bazelisk.html) is
-another way to install Bazel if a native installer is not available. The
-following has been tested with Bazel 4.2.0. It is best to continue working
-inside a virtual environment.
-
-Assuming the above installation steps using `git clone` have been followed, and
-`Bazel` has been installed, the example can be built and run using:
-
-```
-pip install -r requirements_aot_compilation.txt
-bazel run --experimental_repo_remote_exec :run_dm21_aot_compiled_example
-```
-
-where the `pip install` command is only required if a fresh virtual environment
-is used, and installs required prerequisites for TensorFlow. See the
-[TensorFlow documentation](https://www.tensorflow.org/install/source) for more
-details.
-
-A static library, using the `cc_library` rule, can similarly be built and then
-linked against from a separate program with no additional dependencies required,
-shown in `cc/dm21_aot_compiled_example.cc` and the `dm21_aot_compiled_example`
-build rule.
-
-For calling from C, we recommend wrapping a C++ interface in `extern C { ... }`
-to create a C API. For calling from Fortran, the C API can be used via the
-Fortran 2003 ISO_C_BINDING feature.
-
-### Detailed explanation
-
-The supplied functionals in the `checkpoints` subdirectory are not easy to use
-from C++, as they only contain operations for the forward pass through the
-functional. This means they can only be used to evaluate the energy on a fixed
-density. Self-consistent calculations require various gradients of the
-functional, and it is easiest to create these using TensorFlow's Python API.
-`NeuralNumInt` contains a method for exporting the functional and functional
-derivatives. Assuming the above installation steps have been followed, a
-functional and its derivatives can be exported by:
-
-```shell
-export_saved_model.py --functional=DM21 \
- --out_dir=/path/to/export/DM21 \
- --batch_size=1000
-```
-
-where `--out_dir` specifies the directory to save the model containing the
-functional and functional derivatives to, and `--batch_size` the number of grid
-points the functional will be evaluated on at a time. The functional must be the
-name of a functional in the `neural_numint.Functional` enum. Note that the batch
-size needs only be fixed for exporting models to be used with ahead-of-time
-compilation.
-
-The functional can now be compiled using the `saved_model_cli` tool supplied
-with TensorFlow:
-
-```shell
-$ saved_model_cli aot_compile_cpu \
- --dir /path/to/export/DM21 \
- --output_prefix /path/to/compiled/DM21/dm21 \
- --cpp_class dm21::functional \
- --tag_set '' \
- --signature_def_key default
-```
-
-The output prefix and C++ class can be arbitrarily chosen. The `tag_set` and
-`signature_def_key` arguments must be as given above. This creates the following
-files in the output directory:
-
-- dm21.h: header file defining the C++ interface to the DM21 functional.
-- dm21_makefile.inc: a snippet to be included in a Makefile for setting
- include, library and compilation flags.
-- dm21_metadata.o, dm21.o: object files for running the DM21 functional.
-
-## Reference
-
-If this repository is helpful for your research please cite the following
-publication:
-
-Pushing the Frontiers of Density Functionals by Solving the Fractional Electron
-Problem, James Kirkpatrick, Brendan McMorrow, David H. P. Turban, Alexander L.
-Gaunt, James S. Spencer, Alexander G. D. G. Matthews, Annette Obika, Louis
-Thiry, Meire Fortunato, David Pfau, Lara Román Castellanos, Stig Petersen,
-Alexander W. R. Nelson, Pushmeet Kohli, Paula Mori-Sánchez, Demis Hassabis, Aron
-J. Cohen, Science, DOI: https://doi.org/10.1126/science.abj6511
-
-## License
-
-All code is made available under the Apache 2.0 License. Model parameters
-(contained in the `density_functional_approximation_dm21/checkpoints/`
-subdirectory) are made available under the Creative Commons Attribution 4.0
-International (CC BY 4.0) License. See
-https://creativecommons.org/licenses/by/4.0/legalcode for more details.
-
-## Disclaimer
-
-This is not an official Google product.
diff --git a/grad_dft/external/density_functional_approximation_dm21/__init__.py b/grad_dft/external/density_functional_approximation_dm21/__init__.py
deleted file mode 100644
index 85e4032..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/__init__.py
+++ /dev/null
@@ -1,13 +0,0 @@
-# Copyright 2023 Xanadu Quantum Technologies Inc.
-
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-
-# http://www.apache.org/licenses/LICENSE-2.0
-
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
diff --git a/grad_dft/external/density_functional_approximation_dm21/cc/dm21_aot_compiled_example.cc b/grad_dft/external/density_functional_approximation_dm21/cc/dm21_aot_compiled_example.cc
deleted file mode 100644
index 86059be..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/cc/dm21_aot_compiled_example.cc
+++ /dev/null
@@ -1,97 +0,0 @@
-// Copyright 2021 DeepMind Technologies Limited.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-#include
-
-// Users should adjust header path to match generated location. The generated
-// location is determined by the -output_prefix flags passed to saved_model_cli
-// aot_compile_cpu.
-#include "aot_compiled_dm21.h" // generated
-
-void run_dm21_compiled_functional() {
- // The functional class name is set by a flag passed to saved_model_cli
- // aot_compile_cpu.
- dm21::functional dm21_xc;
-
- // This assumes the functional was compiled using instructions given in the
- // README.md, in particular that the model was exported with a batch size of
- // 1000. We will compute the functional just at a single point and pad the
- // rest of the batch using grid_weight = 0.
-
- constexpr int batch_dim = 1000;
- // See docstring for neural_numint.FunctionalInputs for descriptions on each
- // input feature. Note that, due how the model is compiled, we don't need to
- // pass any values for the grid coordinates.
- float rho_a[6][batch_dim] = {
- {1.238148615359934e-11}, {-5.4047671667795604e-11},
- {-5.4047671667795604e-11}, {-7.2613887530595865e-12},
- {4.436179569857956e-10}, {5.972594378309204e-11},
- };
- float hfx_a[batch_dim][2] = {{-4.33218591e-13, -4.32842821e-13}};
- float grid_weights[batch_dim] = {1.7594189642339968};
-
- // Use same values for both alpha and beta electrons (restricted calculation)
- dm21_xc.set_arg_feed_rho_a_data(rho_a);
- dm21_xc.set_arg_feed_rho_b_data(rho_a);
- dm21_xc.set_arg_feed_hfx_a_data(hfx_a);
- dm21_xc.set_arg_feed_hfx_b_data(hfx_a);
- dm21_xc.set_arg_feed_grid_weights_data(grid_weights);
-
- std::puts("Running functional...");
- bool status = dm21_xc.Run();
-
- if (status) {
- std::puts("Successfully ran functional.");
- // Fetch results.
- // Other methods for fetching results exist which may be more convenient.
- // Please see the generated header.
- // See neural_numint.NeuralNumint._build_graph and
- // See neural_numint.NeuralNumint.eval_xc for more details.
- // XC potential at each grid point, shape (batch_dim).
- const float* vxc = dm21_xc.result_fetch_vxc_data();
- // Derivative of the energy wtih respect to the density.
- // In python, this has shape (2, batch_dim), where the zeroth component is
- // with respect to the alpha density and the first component with respect to
- // the beta density. In C++, a flat 1D array is returned.
- const float* vrho = dm21_xc.result_fetch_vrho_data();
- // Derivative of the energy wtih respect to sigma.
- // In python, this has shape (3, batch_dim), where the zeroth component is
- // with respect to the alpha spin channel, the first component with respect
- // to the spin channel and the third component with respect to the total. In
- // C++, a flat 1D array is returned.
- const float* vsigma = dm21_xc.result_fetch_vsigma_data();
- // Derivative of the energy wtih respect to tau, the kinetic energy density.
- // In python, this has shape (2, batch_dim), where the zeroth component is
- // with respect to the alpha spin channel, and the first component with
- // respect to the spin channel. In C++, a flat 1D array is returned.
- const float* vtau = dm21_xc.result_fetch_vtau_data();
- // Intermediates required for evaluating the contribution of local
- // Hartree-Fock features to the derivative of the Fock matrix. See
- // docstrings and comments in compute_hfx_density.py and neural_numint.py.
- // In python, this has shape (2, batch_dim, nomega), where nomega is the
- // number of omega values used for the Hartree-Fock kernels. The zeroth
- // component is with respect to the Hartree-Fock energy density at each grid
- // point for the alpha-spin density and the first component with respect to
- // the beta-spin density. In C++, a flat 1D array is returned.
- const float* vhf = dm21_xc.result_fetch_vhf_data();
- std::printf("vxc[0] = %.6g\n", vxc[0]);
- std::printf("vrho[0] = %.6g, %.6g\n", vrho[0], vrho[batch_dim]);
- std::printf("vsigma[0] = %.6g, %.6g %.6g\n", vsigma[0], vsigma[batch_dim],
- vsigma[2 * batch_dim]);
- std::printf("vtau[0] = %.6g, %.6g\n", vtau[0], vtau[batch_dim]);
- std::printf("vhf[0] = %.6g, %.6g\n", vhf[0], vhf[1]);
- } else {
- std::puts("Failed to run functional.");
- }
-}
diff --git a/grad_dft/external/density_functional_approximation_dm21/cc/dm21_aot_compiled_example.h b/grad_dft/external/density_functional_approximation_dm21/cc/dm21_aot_compiled_example.h
deleted file mode 100644
index 529e0a4..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/cc/dm21_aot_compiled_example.h
+++ /dev/null
@@ -1,21 +0,0 @@
-// Copyright 2021 DeepMind Technologies Limited.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-#ifndef THIRD_PARTY_DEEPMIND_DEEPMIND_RESEARCH_DENSITY_FUNCTIONAL_APPROXIMATION_DM21_CC_DM21_AOT_COMPILED_EXAMPLE_H_
-#define THIRD_PARTY_DEEPMIND_DEEPMIND_RESEARCH_DENSITY_FUNCTIONAL_APPROXIMATION_DM21_CC_DM21_AOT_COMPILED_EXAMPLE_H_
-
-void run_dm21_compiled_functional();
-
-
-#endif // THIRD_PARTY_DEEPMIND_DEEPMIND_RESEARCH_DENSITY_FUNCTIONAL_APPROXIMATION_DM21_CC_DM21_AOT_COMPILED_EXAMPLE_H_
diff --git a/grad_dft/external/density_functional_approximation_dm21/cc/run_dm21_aot_compiled_example.cc b/grad_dft/external/density_functional_approximation_dm21/cc/run_dm21_aot_compiled_example.cc
deleted file mode 100644
index 17f96a0..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/cc/run_dm21_aot_compiled_example.cc
+++ /dev/null
@@ -1,30 +0,0 @@
-// Copyright 2021 DeepMind Technologies Limited.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-#include
-
-#ifdef XLA_AVAILABLE
-#include "cc/dm21_aot_compiled_example.h"
-#endif
-
-int main(int argc, char** argv) {
-#ifdef XLA_AVAILABLE
- run_dm21_compiled_functional();
- return 0;
-#else
- std::puts("Built without XLA support. Cannot run functional!");
- return 1;
-#endif
-}
-
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/__init__.py b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/__init__.py
deleted file mode 100644
index 4bffc24..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/__init__.py
+++ /dev/null
@@ -1,26 +0,0 @@
-# Copyright 2021 DeepMind Technologies Limited.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""An interface to DM21 family of exchange-correlation functionals for PySCF."""
-
-from grad_dft.external.density_functional_approximation_dm21.density_functional_approximation_dm21.neural_numint import (
- Functional,
-)
-from grad_dft.external.density_functional_approximation_dm21.density_functional_approximation_dm21.neural_numint import (
- NeuralNumInt,
-)
-import grad_dft.external.density_functional_approximation_dm21.density_functional_approximation_dm21.compute_hfx_density as compute_hfx_density
-from grad_dft.external.density_functional_approximation_dm21.density_functional_approximation_dm21.neural_numint import (
- _SystemState,
-)
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/saved_model.pb b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/saved_model.pb
deleted file mode 100644
index 112a158..0000000
Binary files a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/saved_model.pb and /dev/null differ
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/smart_module.pb b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/smart_module.pb
deleted file mode 100644
index b33834c..0000000
Binary files a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/smart_module.pb and /dev/null differ
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/tfhub_module.pb b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/tfhub_module.pb
deleted file mode 100644
index d65dd8f..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/tfhub_module.pb
+++ /dev/null
@@ -1 +0,0 @@
-
\ No newline at end of file
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/variables/variables.data-00000-of-00001 b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/variables/variables.data-00000-of-00001
deleted file mode 100644
index 8bee898..0000000
Binary files a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/variables/variables.data-00000-of-00001 and /dev/null differ
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/variables/variables.index b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/variables/variables.index
deleted file mode 100644
index 245f492..0000000
Binary files a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21/variables/variables.index and /dev/null differ
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/saved_model.pb b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/saved_model.pb
deleted file mode 100644
index 7b1ef6e..0000000
Binary files a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/saved_model.pb and /dev/null differ
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/smart_module.pb b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/smart_module.pb
deleted file mode 100644
index acd48e8..0000000
Binary files a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/smart_module.pb and /dev/null differ
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/tfhub_module.pb b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/tfhub_module.pb
deleted file mode 100644
index d65dd8f..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/tfhub_module.pb
+++ /dev/null
@@ -1 +0,0 @@
-
\ No newline at end of file
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/variables/variables.data-00000-of-00001 b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/variables/variables.data-00000-of-00001
deleted file mode 100644
index 2a9952b..0000000
Binary files a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/variables/variables.data-00000-of-00001 and /dev/null differ
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/variables/variables.index b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/variables/variables.index
deleted file mode 100644
index dfb7433..0000000
Binary files a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21m/variables/variables.index and /dev/null differ
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/saved_model.pb b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/saved_model.pb
deleted file mode 100644
index 47f2664..0000000
Binary files a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/saved_model.pb and /dev/null differ
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/smart_module.pb b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/smart_module.pb
deleted file mode 100644
index adbaff6..0000000
Binary files a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/smart_module.pb and /dev/null differ
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/tfhub_module.pb b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/tfhub_module.pb
deleted file mode 100644
index d65dd8f..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/tfhub_module.pb
+++ /dev/null
@@ -1 +0,0 @@
-
\ No newline at end of file
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/variables/variables.data-00000-of-00001 b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/variables/variables.data-00000-of-00001
deleted file mode 100644
index 684a385..0000000
Binary files a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/variables/variables.data-00000-of-00001 and /dev/null differ
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/variables/variables.index b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/variables/variables.index
deleted file mode 100644
index ac424ff..0000000
Binary files a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mc/variables/variables.index and /dev/null differ
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/saved_model.pb b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/saved_model.pb
deleted file mode 100644
index 5bd3cff..0000000
Binary files a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/saved_model.pb and /dev/null differ
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/smart_module.pb b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/smart_module.pb
deleted file mode 100644
index 0749a4e..0000000
Binary files a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/smart_module.pb and /dev/null differ
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/tfhub_module.pb b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/tfhub_module.pb
deleted file mode 100644
index d65dd8f..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/tfhub_module.pb
+++ /dev/null
@@ -1 +0,0 @@
-
\ No newline at end of file
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/variables/variables.data-00000-of-00001 b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/variables/variables.data-00000-of-00001
deleted file mode 100644
index a8ad26d..0000000
Binary files a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/variables/variables.data-00000-of-00001 and /dev/null differ
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/variables/variables.index b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/variables/variables.index
deleted file mode 100644
index d80edc8..0000000
Binary files a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/checkpoints/DM21mu/variables/variables.index and /dev/null differ
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/compute_hfx_density.py b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/compute_hfx_density.py
deleted file mode 100644
index d1dcbea..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/compute_hfx_density.py
+++ /dev/null
@@ -1,334 +0,0 @@
-# Copyright 2021 DeepMind Technologies Limited.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-r"""Computation of the Hartree-Fock exchange density.
-
-We consider two types of potential:
-
- 1. Coulomb potential v(r,r') = 1/|r-r'|, which results in the full HF exchange
- density and energy.
- 2. Screened (long-range) Coulomb potential v(r,r') = erf(\omega|r-r'|)/|r-r'|,
- which results in the screened HF exchange density energy.
-
-Note that PySCF and libcint treat a value of omega=0 to correspond to the
-Coulomb potential. In the following, HF refers to full HF exchange if the
-Coulomb potential is used and to screened HF exchange if the screened Coulomb
-potential is used.
-
-The Hartree-Fock (HF) exchange energy can be written as:
-
--2 HF_x = \sum_{a,b,c,d} D_{ab} D_{cd} \int dr \int dr'
- [ \chi_a(r) \chi_c(r) v(r, r') \chi_b(r') \chi_d(r') ]
-
-where D is the density matrix, \chi_a the atomic basis functions and r, r' are
-coordinates. For clarity we have dropped the spin-channel label of the density
-matrix.
-
-Defining the following intermediates:
-
-\nu_{bd}(r) = \int dr' (\chi_b(r') v(r, r') \chi_d(r'))
-E_b(r) = \sum_a D_{ab} \chi_a(r)
-
-we get the following expression for HF:
-
--2 HF_x = \int dr \sum_{bd} E_b(r) E_d(r) \nu_{bd}(r)
-
-Therefore the quantity
-
-exx(r) = -0.5 sum_{bd} E_b(r) E_d(r) \nu_{bd}(r)
-
-represents an energy density at location r which integrates to the HF exchange
-energy.
-
-The Fock matrix, F, is the derivative of the energy with respect to the density
-matrix. If the energy depends upon the set of features {x}, then the Fock matrix
-can be evaluated as \sum_x dE/dx dx/dD_{ab}. The derivatives with respect to the
-features can be easily evaluated using automatic differentiation. We hence
-require the derivative of exx with respect to the density matrix:
-
-dexx(r)/dD_{ab} = -D_{cd} \chi_a(r) \chi_c(r) \nu_{bd}(r)
-
-This is too large to store, so we instead compute the following intermediate,
-and evaluate the derivative as required on the fly:
-
-fxx_a(r) = D_{bd} \chi_a(r) \nu_{bd}(r)
-
-Note: we compute exx and fxx for each spin channel for both restricted and
-unrestricted calculations.
-"""
-
-from typing import Generator, Optional, Tuple, Union
-
-import attr
-import numpy as np
-from pyscf.dft import numint
-from pyscf.gto import mole
-from pyscf.lib import logger
-from pyscf.lib import numpy_helper
-
-
-def _evaluate_nu_slow(
- mol: mole.Mole, coords: np.ndarray, omega: float, hermi: int = 1
-) -> np.ndarray:
- """Computes nu integrals for given coordinates using a slow loop."""
- nu = []
- # Use the Gaussian nuclear model in int1e_rinv_sph to evaluate the screened
- # integrals.
- with mol.with_rinv_zeta(zeta=omega * omega):
- # This is going to be slow...
- for coord in coords:
- with mol.with_rinv_origin(coord):
- nu.append(mol.intor("int1e_rinv_sph", hermi=hermi))
- return np.asarray(nu)
-
-
-def _evaluate_nu(mol: mole.Mole, coords: np.ndarray, omega: float, hermi: int = 1) -> np.ndarray:
- """Computes nu integrals for given coordinates."""
- try:
- with mol.with_range_coulomb(omega=omega):
- # grids keyword argument supported in pyscf 2.0.0-alpha.
- nu = mol.intor(
- "int1e_grids_sph", hermi=hermi, grids=coords
- ) # pytype: disable=wrong-keyword-args
- except TypeError:
- logger.info(
- mol,
- "Support for int1e_grids not found (requires libcint 4.4.1 and "
- "pyscf 2.0.0a or later. Falling back to slow loop over individual grid "
- "points.",
- )
- nu = _evaluate_nu_slow(mol, coords, omega)
- return nu
-
-
-def _nu_chunk(
- mol: mole.Mole, coords: np.ndarray, omega: float, chunk_size: int = 1000
-) -> Generator[Tuple[int, int, np.ndarray], None, None]:
- r"""Yields chunks of nu integrals over the grid.
-
- Args:
- mol: pyscf Mole object.
- coords: coordinates, r', at which to evaluate the nu integrals, shape (N,3).
- omega: range separation parameter. A value of 0 disables range-separation
- (i.e. uses the kernel v(r,r') = 1/|r-r'| instead of
- v(r,r') = erf(\omega |r-r'|) / |r-r'|)
- chunk_size: number of coordinates to evaluate the integrals at a time.
-
- Yields:
- start_index, end_index, nu_{ab}(r) where
- start_index, end_index are indices into coords,
- nu is an array of shape (end_index-start_index, nao, nao), where nao is
- the number of atomic orbitals and contains
- nu_{ab}(r) = , where a,b are atomic
- orbitals and r' are the grid coordinates in coords[start_index:end_index].
-
- Raises:
- ValueError: if omega is negative.
- """
- if omega < 0:
- raise ValueError("Range-separated parameter omega must be non-negative!")
- ncoords = len(coords)
- for chunk_index in range(0, ncoords, chunk_size):
- end_index = min(chunk_index + chunk_size, ncoords)
- coords_chunk = coords[chunk_index:end_index]
- nu_chunk = _evaluate_nu(mol, coords_chunk, omega=omega)
- yield chunk_index, end_index, nu_chunk
-
-
-def _compute_exx_block(nu: np.ndarray, e: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
- r"""Computes exx and fxx.
-
- Args:
- nu: batch of integrals, in format (k,i,j) where r_k is the
- position of the k-th grid point, i and j label atomic orbitals.
- e: density matrix in the AO basis at each grid point.
-
- Returns:
- exx and fxx, where
- fxx_{gb} =\sum_c nu_{gbc} e_{gc} and
- exx_{g} = -0.5 \sum_b e_{gb} fxx_{gb}.
- """
- fxx = np.einsum("gbc,gc->gb", nu, e)
- exx = -0.5 * np.einsum("gb,gb->g", e, fxx)
- return exx, fxx
-
-
-def _compute_jk_block(
- nu: np.ndarray, fxx: np.ndarray, dm: np.ndarray, ao_value: np.ndarray, weights: np.ndarray
-) -> Tuple[np.ndarray, np.ndarray]:
- """Computes J and K contributions from the given block of nu integrals."""
- batch_size = nu.shape[0]
- vj = numpy_helper.dot(nu.reshape(batch_size, -1), dm.reshape(-1, 1))
- vj = np.squeeze(vj)
- vj_ao = np.einsum("g,gb->gb", vj * weights, ao_value)
- j = numpy_helper.dot(ao_value.T, vj_ao)
- w_ao = np.einsum("g,gb->gb", weights, ao_value)
- k = numpy_helper.dot(fxx.T, w_ao)
- return j, k
-
-
-@attr.s(auto_attribs=True)
-class HFDensityResult:
- r"""Container for results returned by get_hf_density.
-
- Note that the kernel used in all integrals is defined by the omega input
- argument.
-
- Attributes:
- exx: exchange energy density at position r on the grid for the alpha, beta
- spin channels. Each array is shape (N), where N is the number of grid
- points.
- fxx: intermediate for evaluating dexx/dD^{\sigma}_{ab}, where D is the
- density matrix and \sigma is the spin coordinate. See top-level docstring
- for details. Each array is shape (N, nao), where nao is the number of
- atomic orbitals.
- coulomb: coulomb matrix (restricted calculations) or matrices (unrestricted
- calculations). Each array is shape (nao, nao).
- Restricted calculations: \sum_{} D_{cd} (ab|cd)
- Unrestricted calculations: \sum_{} D^{\sigma}_{cd} (ab|cd)
- exchange: exchange matrix (restricted calculations) or matrices
- (unrestricted calculations). Each array is shape (nao, nao).
- Restricted calculations: \sum_{} D_{cd} (ab|cd)
- Unrestricted calculations: \sum_{} D^{\sigma}_{cd} (ac|bd).
- """
- exx: Tuple[np.ndarray, np.ndarray]
- fxx: Optional[Tuple[np.ndarray, np.ndarray]] = None
- coulomb: Optional[Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]] = None
- exchange: Optional[Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]] = None
-
-
-def get_hf_density(
- mol: mole.Mole,
- dm: Union[Tuple[np.ndarray, np.ndarray], np.ndarray],
- coords: np.ndarray,
- omega: float = 0.0,
- deriv: int = 0,
- ao: Optional[np.ndarray] = None,
- chunk_size: int = 1000,
- weights: Optional[np.ndarray] = None,
-) -> HFDensityResult:
- r"""Computes the (range-separated) HF energy density.
-
- Args:
- mol: PySCF molecule.
- dm: The density matrix. For restricted calculations, an array of shape
- (M, M), where M is the number of atomic orbitals. For unrestricted
- calculations, either an array of shape (2, M, M) or a tuple of arrays,
- each of shape (M, M), where dm[0] is the density matrix for the alpha
- electrons and dm[1] the density matrix for the beta electrons.
- coords: The coordinates to compute the HF density at, shape (N, 3), where N
- is the number of grid points.
- omega: The inverse width of the error function. An omega of 0. means range
- separation and a 1/|r-R| kernel is used in the nu integrals. Otherwise,
- the kernel erf(\omega|r-R|)/|r-R|) is used. Must be non-negative.
- deriv: The derivative order. Only first derivatives (deriv=1) are currently
- implemented. deriv=0 indicates no derivatives are required.
- ao: The atomic orbitals evaluated on the grid, shape (N, M). These are
- computed if not supplied.
- chunk_size: The number of coordinates to compute the HF density for at once.
- Reducing this saves memory since we don't have to keep as many Nus (nbasis
- x nbasis) in memory at once.
- weights: weight of each grid point, shape (N). If present, the Coulomb and
- exchange matrices are also computed semi-numerically, otherwise only the
- HF density and (if deriv=1) its first derivative are computed.
-
- Returns:
- HFDensityResult object with the HF density (exx), the derivative of the HF
- density with respect to the density (fxx) if deriv is 1, and the Coulomb and
- exchange matrices if the weights argument is provided.
-
- Raises:
- NotImplementedError: if a Cartesian basis set is used or if deriv is greater
- than 1.
- ValueError: if omega or deriv are negative.
- """
- if mol.cart:
- raise NotImplementedError(
- "Local HF exchange is not implmented for basis " "sets with Cartesian functions!"
- )
- if deriv < 0:
- raise ValueError(f"`deriv` must be non-negative, got {deriv}")
- if omega < 0:
- raise ValueError(f"`omega` must be non-negative, got {omega}")
- if deriv > 1:
- raise NotImplementedError("Higher order derivatives are not implemented.")
-
- if isinstance(dm, tuple) or dm.ndim == 3:
- dma, dmb = dm
- restricted = False
- else:
- dma = dm / 2
- dmb = dm / 2
- restricted = True
-
- logger.info(mol, "Computing contracted density matrix ...")
- if ao is None:
- ao = numint.eval_ao(mol, coords, deriv=0)
- e_a = np.dot(ao, dma)
- e_b = np.dot(ao, dmb)
-
- exxa = []
- exxb = []
- fxxa = []
- fxxb = []
- ja = np.zeros_like(dma)
- jb = np.zeros_like(dmb)
- ka = np.zeros_like(dma)
- kb = np.zeros_like(dmb)
-
- for start, end, nu in _nu_chunk(mol, coords, omega, chunk_size=chunk_size):
- logger.info(mol, "Computing exx %s / %s ...", end, len(e_a))
- exxa_block, fxxa_block = _compute_exx_block(nu, e_a[start:end])
- exxa.extend(exxa_block)
- if not restricted:
- exxb_block, fxxb_block = _compute_exx_block(nu, e_b[start:end])
- exxb.extend(exxb_block)
- if deriv == 1:
- fxxa.extend(fxxa_block)
- if not restricted:
- fxxb.extend(fxxb_block)
-
- if weights is not None:
- ja_block, ka_block = _compute_jk_block(
- nu, fxxa_block, dma, ao[start:end], weights[start:end]
- )
- ja += ja_block
- ka += ka_block
- if not restricted:
- jb_block, kb_block = _compute_jk_block(
- nu, fxxb_block, dmb, ao[start:end], weights[start:end]
- )
- jb += jb_block
- kb += kb_block
-
- exxa = np.asarray(exxa)
- fxxa = np.asarray(fxxa)
- if restricted:
- exxb = exxa
- fxxb = fxxa
- else:
- exxb = np.asarray(exxb)
- fxxb = np.asarray(fxxb)
-
- result = HFDensityResult(exx=(exxa, exxb))
- if deriv == 1:
- result.fxx = (fxxa, fxxb)
- if weights is not None:
- if restricted:
- result.coulomb = 2 * ja
- result.exchange = 2 * ka
- else:
- result.coulomb = (ja, jb)
- result.exchange = (ka, kb)
- return result
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/compute_hfx_density_test.py b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/compute_hfx_density_test.py
deleted file mode 100644
index e8edf93..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/compute_hfx_density_test.py
+++ /dev/null
@@ -1,173 +0,0 @@
-# Copyright 2021 DeepMind Technologies Limited.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Tests for compute_hfx_density."""
-
-from absl.testing import absltest
-from absl.testing import parameterized
-import numpy as np
-from pyscf import dft
-from pyscf import gto
-from pyscf import lib
-from pyscf import scf
-import scipy
-from grad_dft.external.density_functional_approximation_dm21 import compute_hfx_density
-
-
-class ComputeHfxDensityTest(parameterized.TestCase):
- def setUp(self):
- super().setUp()
- lib.param.TMPDIR = None
- lib.num_threads(1)
-
- @parameterized.named_parameters(
- {"testcase_name": "local_hf", "omega": 0.0},
- {"testcase_name": "range_separated_local_hf_0.5", "omega": 0.5},
- {"testcase_name": "range_separated_local_hf_1.0", "omega": 1.0},
- {"testcase_name": "range_separated_local_hf_2.0", "omega": 2.0},
- )
- def test_closed_shell(self, omega):
- mol = gto.M(atom="He 0. 0. 0.", basis="3-21g")
- solver = dft.RKS(mol)
- solver.grids.level = 2
- solver.grids.build()
- solver.kernel()
- dm = solver.make_rdm1()
- with mol.with_range_coulomb(omega=omega):
- target_j, target_k = scf.hf.get_jk(mol, dm)
- target_hf = -0.25 * np.einsum("ij,ji", dm, target_k)
- target_coulomb = np.einsum("ij,ji", dm, target_j)
-
- coords = solver.grids.coords
- weights = solver.grids.weights
-
- results = compute_hfx_density.get_hf_density(mol, dm, coords, omega=omega, weights=weights)
- coulomb = np.einsum("ij,ji", dm, results.coulomb)
- hf = -0.25 * np.einsum("ij,ji", dm, results.exchange)
- predicted_hf = np.sum((results.exx[0] + results.exx[1]) * weights)
-
- with self.subTest("test_hf_density"):
- self.assertAlmostEqual(target_hf, predicted_hf)
-
- with self.subTest("test_get_jk"):
- np.testing.assert_allclose(results.coulomb, target_j)
- np.testing.assert_allclose(results.exchange, target_k)
- self.assertAlmostEqual(coulomb, target_coulomb)
- self.assertAlmostEqual(hf, target_hf)
-
- @parameterized.named_parameters(
- {"testcase_name": "local_hf", "omega": 0.0},
- {"testcase_name": "range_separated_local_hf_0.5", "omega": 0.5},
- {"testcase_name": "range_separated_local_hf_1.0", "omega": 1.0},
- {"testcase_name": "range_separated_local_hf_2.0", "omega": 2.0},
- )
- def test_hf_density_on_open_shell(self, omega):
- mol = gto.M(atom="He 0. 0. 0.", basis="3-21g", charge=1, spin=1)
- solver = dft.UKS(mol)
- solver.grids.level = 2
- solver.grids.build()
- solver.kernel()
- dm = solver.make_rdm1()
- with mol.with_range_coulomb(omega=omega):
- target_j, target_k = scf.hf.get_jk(mol, dm)
- target_hf = -0.5 * (
- np.einsum("ij,ji", dm[0], target_k[0]) + np.einsum("ij,ji", dm[1], target_k[1])
- )
- target_coulomb = np.einsum("ij,ji", dm[0], target_j[0]) + np.einsum(
- "ij,ji", dm[1], target_j[1]
- )
-
- coords = solver.grids.coords
- weights = solver.grids.weights
-
- results = compute_hfx_density.get_hf_density(mol, dm, coords, omega=omega, weights=weights)
-
- predicted_hf = np.sum((results.exx[0] + results.exx[1]) * weights)
- coulomb = np.einsum("ij,ji", dm[0], results.coulomb[0]) + np.einsum(
- "ij,ji", dm[1], results.coulomb[1]
- )
- hf = -0.5 * (
- np.einsum("ij,ji", dm[0], results.exchange[0])
- + np.einsum("ij,ji", dm[1], results.exchange[1])
- )
-
- with self.subTest("test_hf_density"):
- self.assertAlmostEqual(target_hf, predicted_hf, places=3)
-
- with self.subTest("test_get_jk"):
- np.testing.assert_allclose(results.coulomb[0], target_j[0])
- np.testing.assert_allclose(results.coulomb[1], target_j[1])
- np.testing.assert_allclose(results.exchange[0], target_k[0])
- np.testing.assert_allclose(results.exchange[1], target_k[1])
- self.assertAlmostEqual(coulomb, target_coulomb)
- self.assertAlmostEqual(hf, target_hf)
-
-
-def _nu_test_systems():
- systems = [
- {"atom": "N 0 0 0; N 0 0 2.4", "charge": 0, "spin": 0, "basis": "cc-pVDZ", "num_grids": -1},
- {"atom": "N 0 0 0; N 0 0 2.4", "charge": 0, "spin": 0, "basis": "cc-pVDZ", "num_grids": 1},
- {"atom": "N 0 0 0; N 0 0 2.4", "charge": 0, "spin": 0, "basis": "cc-pVDZ", "num_grids": 2},
- {"atom": "N 0 0 0; N 0 0 2.4", "charge": 0, "spin": 0, "basis": "cc-pVDZ", "num_grids": 10},
- {"atom": "N 0 0 0; N 0 0 2.4", "charge": 0, "spin": 0, "basis": "cc-pVDZ", "num_grids": 32},
- {"atom": "N 0 0 0; N 0 0 2.4", "charge": 0, "spin": 0, "basis": "cc-pVDZ", "num_grids": 33},
- {"atom": "Li 0 0 0", "charge": 0, "spin": 1, "basis": "cc-pVTZ", "num_grids": -1},
- {"atom": "H 0 0 0", "charge": 0, "spin": 1, "basis": "cc-pVQZ", "num_grids": -1},
- ]
- system_names = ["N2", "N2_1", "N2_2", "N2_10", "N2_32", "N2_33", "Li", "H"]
- for name, system in zip(system_names, systems):
- yield {"testcase_name": f"{name}_hermitian", "hermi": 0, **system}
- yield {"testcase_name": f"{name}_non_hermitian", "hermi": 1, **system}
-
-
-class NuTest(parameterized.TestCase):
- def setUp(self):
- super(NuTest, self).setUp()
- lib.param.TMPDIR = None
- lib.num_threads(1)
-
- @parameterized.named_parameters(_nu_test_systems())
- def test_nu_integrals(self, atom, charge, spin, basis, num_grids, hermi):
- mol = gto.M(atom=atom, charge=charge, spin=spin, basis=basis)
- mf = dft.UKS(mol)
- mf.grids.build()
- if num_grids == -1:
- test_coords = mf.grids.coords
- else:
- test_coords = mf.grids.coords[0:num_grids]
- nu_slow = compute_hfx_density._evaluate_nu_slow(mol, test_coords, omega=0.0, hermi=hermi)
- nu_fast = compute_hfx_density._evaluate_nu(mol, test_coords, omega=0.0, hermi=hermi)
- np.testing.assert_allclose(nu_slow, nu_fast, atol=1e-13)
-
- def test_range_separated_nu(self):
- mol = gto.M(atom="He 0 0 0", basis="cc-pVDZ")
- r0 = np.array([[0.1, 0.2, 1.0]])
- omega = 1.0
- result = np.squeeze(compute_hfx_density._evaluate_nu(mol, r0, omega=omega))
-
- solver = dft.RKS(mol)
- solver.grids.level = 2
- solver.grids.build()
- coords = solver.grids.coords
- weights = solver.grids.weights
- ao_value = dft.numint.eval_ao(mol, coords, deriv=0)
- dist = np.linalg.norm(coords - r0, axis=1)
- erf = scipy.special.erf(omega * dist) / dist
- expected_result = np.squeeze(np.einsum("g,ga,gb->ab", weights * erf, ao_value, ao_value))
-
- np.testing.assert_allclose(result, expected_result)
-
-
-if __name__ == "__main__":
- absltest.main()
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/export_saved_model.py b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/export_saved_model.py
deleted file mode 100644
index 0cbf577..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/export_saved_model.py
+++ /dev/null
@@ -1,60 +0,0 @@
-#!/usr/bin/env python3
-# Copyright 2021 DeepMind Technologies Limited.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Helper for exporting a functional and its derivatives to a saved_model."""
-
-from typing import Sequence
-
-from absl import app
-from absl import flags
-
-from density_functional_approximation_dm21 import neural_numint
-
-_OUT_DIR = flags.DEFINE_string("out_dir", None, "Output directory.", required=True)
-_BATCH_SIZE = flags.DEFINE_integer(
- "batch_size",
- 1000,
- "Number of grid points exported functional will process in a single call.",
- lower_bound=0,
-)
-_FUNCTIONAL = flags.DEFINE_enum_class(
- "functional", neural_numint.Functional.DM21, neural_numint.Functional, "Functional to export."
-)
-
-
-def export(
- functional: neural_numint.Functional,
- export_path: str,
- batch_dim: int,
-) -> None:
- """Export a functional and its derivatives to a single saved_model.
-
- Args:
- functional: functional to export.
- export_path: path to saved the model to.
- batch_dim: number of grid points to process in a single call.
- """
- ni = neural_numint.NeuralNumInt(functional)
- ni.export_functional_and_derivatives(export_path=export_path, batch_dim=batch_dim)
-
-
-def main(argv: Sequence[str]) -> None:
- if len(argv) > 1:
- raise app.UsageError("Too many command-line arguments.")
- export(_FUNCTIONAL.value, _OUT_DIR.value, _BATCH_SIZE.value)
-
-
-if __name__ == "__main__":
- app.run(main)
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py
deleted file mode 100644
index 5a8c62e..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py
+++ /dev/null
@@ -1,954 +0,0 @@
-# Copyright 2021 DeepMind Technologies Limited.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""An interface to DM21 family of exchange-correlation functionals for PySCF."""
-
-import enum
-import os
-from typing import Generator, List, Optional, Sequence, Tuple, Union
-
-import attr
-import numpy as np
-from pyscf import dft
-from pyscf import gto
-from pyscf.dft import numint
-import tensorflow.compat.v1 as tf
-
-try:
- import pkg_resources # pylint: disable=unused-import
-except ModuleNotFoundError: # setuptools >= 81 removed pkg_resources
- # tensorflow_hub.__init__ still imports pkg_resources.parse_version to
- # check the installed tensorflow version. Provide the one symbol it needs
- # when pkg_resources is no longer part of the environment.
- import sys as _sys
- from types import SimpleNamespace as _SimpleNamespace
-
- def _parse_version(version):
- """Returns a comparable key for a dotted numeric version string."""
- return tuple(int(part) for part in version.split(".") if part.isdigit())
-
- _sys.modules.setdefault(
- "pkg_resources", _SimpleNamespace(parse_version=_parse_version)
- )
-
-import tensorflow_hub as hub
-
-import grad_dft.external.density_functional_approximation_dm21.density_functional_approximation_dm21.compute_hfx_density as compute_hfx_density
-
-tf.disable_v2_behavior()
-# tf.enable_eager_execution()
-
-# TODO(b/196260242): avoid depending upon private function
-_dot_ao_ao = numint._dot_ao_ao # pylint: disable=protected-access
-
-
-@enum.unique
-class Functional(enum.Enum):
- """Enum for exchange-correlation functionals in the DM21 family.
-
- Attributes:
- DM21: trained on molecules dataset, and fractional charge, and fractional
- spin constraints.
- DM21m: trained on molecules dataset.
- DM21mc: trained on molecules dataset, and fractional charge constraints.
- DM21mu: trained on molecules dataset, and electron gas constraints.
- """
-
- # Break pylint's preferred naming pattern to match the functional names used
- # in the paper.
- # pylint: disable=invalid-name
- DM21 = enum.auto()
- DM21m = enum.auto()
- DM21mc = enum.auto()
- DM21mu = enum.auto()
- # pylint: enable=invalid-name
-
-
-# We use attr.s instead of here instead of dataclasses.dataclass as
-# dataclasses.asdict returns a deepcopy of the attributes. This is wasteful in
-# memory if they are large and breaks (as in the case of tf.Tensors) if they are
-# not serializable. attr.asdict does not perform this copy and so works with
-# both np.ndarrays and tf.Tensors.
-@attr.s(auto_attribs=True)
-class FunctionalInputs:
- r""" "Inputs required for DM21 functionals.
-
- Depending upon the context, this is either a set of numpy arrays (feature
- construction) or TF tensors (constructing placeholders/running functionals).
-
- Attributes:
- rho_a: Density information for the alpha electrons.
- PySCF for meta-GGAs supplies a single array for the total density
- (restricted calculations) and a pair of arrays, one for each spin channel
- (unrestricted calculations).
- Each array/tensor is of shape (6, N) and contains the density and density
- derivatives, where:
- rho(0, :) - density at each grid point
- rho(1, :) - norm of the derivative of the density at each grid point
- along x
- rho(2, :) - norm of the derivative of the density at each grid point
- along y
- rho(3, :) - norm of the derivative of the density at each grid point
- along z
- rho(4, :) - \nabla^2 \rho [not used]
- rho(5, :) - tau (1/2 (\nabla \rho)^2) at each grid point.
- See pyscf.dft.numint.eval_rho for more details.
- We require separate inputs for both alpha- and beta-spin densities, even
- in restricted calculations (where rho_a = rho_b = rho/2, where rho is the
- total density).
- rho_b: as for rho_a for the beta electrons.
- hfx_a: local Hartree-Fock energy density at each grid point for the alpha-
- spin density for each value of omega. Shape [N, len(omega_values)].
- See compute_hfx_density for more details.
- hfx_b: as for hfx_a for the beta-spin density.
- grid_coords: grid coordinates at which to evaluate the density. Shape
- (N, 3), where N is the number of grid points. Note that this is currently
- unused by the functional, but is still a required input.
- grid_weights: weight of each grid point. Shape (N).
- """
- rho_a: Union[tf.Tensor, np.ndarray]
- rho_b: Union[tf.Tensor, np.ndarray]
- hfx_a: Union[tf.Tensor, np.ndarray]
- hfx_b: Union[tf.Tensor, np.ndarray]
- grid_coords: Union[tf.Tensor, np.ndarray]
- grid_weights: Union[tf.Tensor, np.ndarray]
-
-
-@attr.s(auto_attribs=True)
-class _GridState:
- """Internal state required for the numerical grid.
-
- Attributes:
- coords: coordinates of the grid. Shape (N, 3), where N is the number of grid
- points.
- weight: weight associated with each grid point. Shape (N).
- mask: mask indicating whether a shell is zero at a grid point. Shape
- (N, nbas) where nbas is the number of shells in the basis set. See
- pyscf.dft.gen_grids.make_mask.
- ao: atomic orbitals evaluated on the grid. Shape (N, nao), where nao is the
- number of atomic orbitals, or shape (:, N, nao), where the 0-th element
- contains the ao values, the next three elements contain the first
- derivatives, and so on.
- """
-
- coords: np.ndarray
- weight: np.ndarray
- mask: np.ndarray
- ao: np.ndarray
-
-
-@attr.s(auto_attribs=True)
-class _SystemState:
- """Internal state required for system of interest.
-
- Attributes:
- mol: PySCF molecule
- dms: density matrix or matrices (unrestricted calculations only).
- Restricted calculations: shape (nao, nao), where nao is the number of
- atomic orbitals.
- Unrestricted calculations: shape (2, nao, nao) or a sequence (length 2) of
- arrays of shape (nao, nao), and dms[0] and dms[1] are the density matrices
- of the alpha and beta electrons respectively.
- """
-
- mol: gto.Mole
- dms: Union[np.ndarray, Sequence[np.ndarray]]
-
-
-def _get_number_of_density_matrices(dms):
- """Returns the number of density matrices in dms."""
- # See pyscf.numint.NumInt._gen_rho_evaluator
- if isinstance(dms, np.ndarray) and dms.ndim == 2:
- return 1
- return len(dms)
-
-
-class NeuralNumInt(numint.NumInt):
- """A wrapper around pyscf.dft.numint.NumInt for the DM21 functionals.
-
- In order to supply the local Hartree-Fock features required for the DM21
- functionals, we lightly wrap the NumInt class. The actual evaluation of the
- exchange-correlation functional is performed in NeuralNumInt.eval_xc, which
- PySCF's nr_rks/nr_uks grid loops reach through the eval_xc_eff override.
-
- Usage:
- mf = dft.RKS(...) # dft.ROKS and dft.UKS are also supported.
- # Specify the functional by monkey-patching mf._numint rather than using
- # mf._xc or mf._define_xc_.
- mf._numint = NeuralNumInt(Functional.DM21)
- mf.kernel()
- """
-
- def __init__(self, functional: Functional, *, checkpoint_path: Optional[str] = None):
- """Constructs a NeuralNumInt object.
-
- Args:
- functional: member of Functional enum giving the name of the
- functional.
- checkpoint_path: Optional path to specify the directory containing the
- checkpoints of the DM21 family of functionals. If not specified, attempt
- to find the checkpoints using a path relative to the source code.
- """
-
- self._functional_name = functional.name
- if checkpoint_path:
- self._model_path = os.path.join(checkpoint_path, self._functional_name)
- else:
- self._model_path = os.path.join(
- os.path.dirname(__file__), "checkpoints", self._functional_name
- )
-
- # All DM21 functionals use local Hartree-Fock features with a non-range
- # separated 1/r kernel and a range-seperated kernel with \omega = 0.4.
- # Note an omega of 0.0 is interpreted by PySCF and libcint to indicate no
- # range-separation.
- self._omega_values = [0.0, 0.4]
- self._graph = tf.Graph()
- with self._graph.as_default():
- self._build_graph()
- self._session = tf.Session()
- self._session.run(tf.global_variables_initializer())
-
- self._grid_state = None
- self._system_state = None
- self._vmat_hf = None
- super().__init__()
-
- def _build_graph(self, batch_dim: Optional[int] = None):
- """Builds the TensorFlow graph for evaluating the functional.
-
- Args:
- batch_dim: the batch dimension of the grid to use in the model. Default:
- None (determine at runtime). This should only be set if building a model
- in order to export and ahead-of-time compile it into a standalone
- library.
- """
-
- # TF-Hub 0.16 removed the TF1 Module API (hub.Module). The checkpoints
- # are SavedModels, so load them with hub.load; the "default" signature's
- # keyword names keep the "tensor_dict$" prefix used below (only the
- # TensorSpec name= fields are sanitised to "_").
- self._functional = hub.load(self._model_path).signatures["default"]
-
- grid_coords = tf.placeholder(tf.float32, shape=[batch_dim, 3], name="grid_coords")
- grid_weights = tf.placeholder(tf.float32, shape=[batch_dim], name="grid_weights")
-
- # Density information.
- rho_a = tf.placeholder(tf.float32, shape=[6, batch_dim], name="rho_a")
- rho_b = tf.placeholder(tf.float32, shape=[6, batch_dim], name="rho_b")
-
- # Split into corresponding terms.
- rho_only_a, grad_a_x, grad_a_y, grad_a_z, _, tau_a = tf.unstack(rho_a, axis=0)
- rho_only_b, grad_b_x, grad_b_y, grad_b_z, _, tau_b = tf.unstack(rho_b, axis=0)
-
- # Evaluate |\del \rho|^2 for each spin density and for the total density.
- norm_grad_a = grad_a_x**2 + grad_a_y**2 + grad_a_z**2
- norm_grad_b = grad_b_x**2 + grad_b_y**2 + grad_b_z**2
- grad_x = grad_a_x + grad_b_x
- grad_y = grad_a_y + grad_b_y
- grad_z = grad_a_z + grad_b_z
- norm_grad = grad_x**2 + grad_y**2 + grad_z**2
-
- # The local Hartree-Fock energy densities at each grid point for the alpha-
- # and beta-spin densities for each value of omega.
- # Note an omega of 0 indicates no screening of the Coulomb potential.
- hfxa = tf.placeholder(tf.float32, shape=[batch_dim, len(self._omega_values)], name="hfxa")
- hfxb = tf.placeholder(tf.float32, shape=[batch_dim, len(self._omega_values)], name="hfxb")
-
- # Make all features 2D arrays on input for ease of handling inside the
- # functional.
- features = {
- "grid_coords": grid_coords,
- "grid_weights": tf.expand_dims(grid_weights, 1),
- "rho_a": tf.expand_dims(rho_only_a, 1),
- "rho_b": tf.expand_dims(rho_only_b, 1),
- "tau_a": tf.expand_dims(tau_a, 1),
- "tau_b": tf.expand_dims(tau_b, 1),
- "norm_grad_rho_a": tf.expand_dims(norm_grad_a, 1),
- "norm_grad_rho_b": tf.expand_dims(norm_grad_b, 1),
- "norm_grad_rho": tf.expand_dims(norm_grad, 1),
- "hfxa": hfxa,
- "hfxb": hfxb,
- }
- tensor_dict = {f"tensor_dict${k}": v for k, v in features.items()}
-
- predictions = self._functional(**tensor_dict)
- local_xc = predictions["grid_contribution"]
- weighted_local_xc = local_xc * grid_weights
- unweighted_xc = tf.reduce_sum(local_xc, axis=0)
- xc = tf.reduce_sum(weighted_local_xc, axis=0)
-
- # The potential is the local exchange correlation divided by the
- # total density. Add a small constant to deal with zero density.
- self._vxc = local_xc / (rho_only_a + rho_only_b + 1e-12)
-
- # The derivatives of the exchange-correlation (XC) energy with respect to
- # input features. PySCF weights the (standard) derivatives by the grid
- # weights, so we need to compute this with respect to the unweighted sum
- # over grid points.
- self._vrho = tf.gradients(
- unweighted_xc,
- [features["rho_a"], features["rho_b"]],
- name="GRAD_RHO",
- unconnected_gradients=tf.UnconnectedGradients.ZERO,
- )
- self._vsigma = tf.gradients(
- unweighted_xc,
- [features["norm_grad_rho_a"], features["norm_grad_rho_b"], features["norm_grad_rho"]],
- name="GRAD_SIGMA",
- unconnected_gradients=tf.UnconnectedGradients.ZERO,
- )
- self._vtau = tf.gradients(
- unweighted_xc,
- [features["tau_a"], features["tau_b"]],
- name="GRAD_TAU",
- unconnected_gradients=tf.UnconnectedGradients.ZERO,
- )
- # Standard meta-GGAs do not have a dependency on local HF, so we need to
- # compute the contribution to the Fock matrix ourselves. Just use the
- # weighted XC energy to avoid having to weight this later.
- self._vhf = tf.gradients(
- xc,
- [features["hfxa"], features["hfxb"]],
- name="GRAD_HFX",
- unconnected_gradients=tf.UnconnectedGradients.ZERO,
- )
-
- self._placeholders = FunctionalInputs(
- rho_a=rho_a,
- rho_b=rho_b,
- hfx_a=hfxa,
- hfx_b=hfxb,
- grid_coords=grid_coords,
- grid_weights=grid_weights,
- )
-
- # Note: TF-Hub 0.16 removed hub.add_signature/create_module_spec, so no
- # signature is registered here. The derivatives (self._vrho, ...) stay
- # graph attributes consumed by eval_xc via self._session.run; only the
- # export path below needed them in a signature, and it now raises.
-
- def export_functional_and_derivatives(
- self,
- export_path: str,
- batch_dim: Optional[int] = None,
- ):
- """Exports the TensorFlow graph containing the functional and derivatives.
-
- The hub modules supplied contain the TensorFlow operations for the
- evaluation of the exchange-correlation energy. Evaluation of the functional
- derivatives, required for a self-consistent calculation, are added in
- _build_graph. The module created by export_functional_and_derivatives
- contains the evaluation of the functional and the functional derivatives.
- This is much simpler to use from languages other than Python, e.g. using the
- C or C++ TensorFlow API, or using tfcompile to create a standalone C++
- library.
-
- Args:
- export_path: path to write the Hub model to. The exported model can be
- loaded using either TF-Hub or SavedModel APIs.
- batch_dim: the batch dimension of the grid to use in the model. Default:
- None (determine at runtime). This should only be set if the exported
- model is to be ahead-of-time compiled into a standalone library.
-
- Raises:
- NotImplementedError: always. TF-Hub 0.16 removed the TF1 Module API
- (hub.create_module_spec, hub.Module.export) this method was built
- on; exporting would require a TF2 tf.saved_model.save
- reimplementation. Only the vendored export_saved_model.py and
- neural_numint_test.py use this method, and neither is part of the
- test suite.
- """
- del export_path, batch_dim # unused: always raises
- raise NotImplementedError(
- "export_functional_and_derivatives requires the TF1 Hub Module API "
- "(hub.create_module_spec / hub.Module.export) that tensorflow-hub "
- "0.16 removed, and would need a TF2 tf.saved_model.save "
- "reimplementation instead. Only the vendored export_saved_model.py "
- "and neural_numint_test.py use this method; neither is part of the "
- "test suite."
- )
-
- # DM21* functionals include the hybrid term directly, so set the
- # range-separated and hybrid parameters expected by PySCF to 0 so PySCF
- # doesn't also add these contributions in separately.
- def rsh_coeff(self, *args):
- """Returns the range separated parameters, omega, alpha, beta."""
- return [0.0, 0.0, 0.0]
-
- def hybrid_coeff(self, *args, **kwargs):
- """Returns the fraction of Hartree-Fock exchange to include."""
- return 0.0
-
- def _xc_type(self, *args, **kwargs):
- return "MGGA"
-
- def nr_rks(
- self,
- mol: gto.Mole,
- grids: dft.Grids,
- xc_code: str,
- dms: Union[np.ndarray, Sequence[np.ndarray]],
- relativity: int = 0,
- hermi: int = 0,
- max_memory: float = 20000,
- verbose=None,
- ) -> Tuple[float, float, np.ndarray]:
- """Calculates RKS XC functional and potential matrix on a given grid.
-
- Args:
- mol: PySCF molecule.
- grids: grid on which to evaluate the functional.
- xc_code: XC code. Unused. NeuralNumInt hard codes the XC functional
- based upon the functional argument given to the constructor.
- dms: the density matrix or sequence of density matrices. Multiple density
- matrices are not currently supported. Shape (nao, nao), where nao is the
- number of atomic orbitals.
- relativity: Unused. (pyscf.numint.NumInt.nr_rks does not currently use
- this argument.)
- hermi: 0 if the density matrix is Hermitian, 1 if the density matrix is
- non-Hermitian.
- max_memory: the maximum cache to use, in MB.
- verbose: verbosity level. Unused. (PySCF currently does not handle the
- verbosity level passed in here.)
-
- Returns:
- nelec, excsum, vmat, where
- nelec is the number of electrons obtained by numerical integration of
- the density matrix.
- excsum is the functional's XC energy.
- vmat is the functional's XC potential matrix, shape (nao, nao).
-
- Raises:
- NotImplementedError: if multiple density matrices are supplied.
- """
- # Wrap nr_rks so we can store internal variables required to evaluate the
- # contribution to the XC potential from local Hartree-Fock features.
- # See pyscf.dft.numint.nr_rks for more details.
- ndms = _get_number_of_density_matrices(dms)
- if ndms > 1:
- raise NotImplementedError(
- "NeuralNumInt does not support multiple density matrices. "
- "Only ground state DFT calculations are currently implemented."
- )
- nao = mol.nao_nr()
- self._vmat_hf = np.zeros((nao, nao))
- self._system_state = _SystemState(mol=mol, dms=dms)
- nelec, excsum, vmat = super().nr_rks(
- mol=mol,
- grids=grids,
- xc_code=xc_code,
- dms=dms,
- relativity=relativity,
- hermi=hermi,
- max_memory=max_memory,
- verbose=verbose,
- )
- vmat += self._vmat_hf + self._vmat_hf.T
-
- # Clear internal state to prevent accidental re-use.
- self._system_state = None
- self._grid_state = None
- return nelec, excsum, vmat
-
- def nr_uks(
- self,
- mol: gto.Mole,
- grids: dft.Grids,
- xc_code: str,
- dms: Union[Sequence[np.ndarray], Sequence[Sequence[np.ndarray]]],
- relativity: int = 0,
- hermi: int = 0,
- max_memory: float = 20000,
- verbose=None,
- ) -> Tuple[np.ndarray, float, np.ndarray]:
- """Calculates UKS XC functional and potential matrix on a given grid.
-
- Args:
- mol: PySCF molecule.
- grids: grid on which to evaluate the functional.
- xc_code: XC code. Unused. NeuralNumInt hard codes the XC functional
- based upon the functional argument given to the constructor.
- dms: the density matrix or sequence of density matrices for each spin
- channel. Multiple density matrices for each spin channel are not
- currently supported. Each density matrix is shape (nao, nao), where nao
- is the number of atomic orbitals.
- relativity: Unused. (pyscf.dft.numint.NumInt.nr_rks does not currently use
- this argument.)
- hermi: 0 if the density matrix is Hermitian, 1 if the density matrix is
- non-Hermitian.
- max_memory: the maximum cache to use, in MB.
- verbose: verbosity level. Unused. (PySCF currently does not handle the
- verbosity level passed in here.)
-
- Returns:
- nelec, excsum, vmat, where
- nelec is the number of alpha, beta electrons obtained by numerical
- integration of the density matrix as an array of size 2.
- excsum is the functional's XC energy.
- vmat is the functional's XC potential matrix, shape (2, nao, nao), where
- vmat[0] and vmat[1] are the potential matrices for the alpha and beta
- spin channels respectively.
-
- Raises:
- NotImplementedError: if multiple density matrices for each spin channel
- are supplied.
- """
- # Wrap nr_uks so we can store internal variables required to evaluate the
- # contribution to the XC potential from local Hartree-Fock features.
- # See pyscf.dft.numint.nr_uks for more details.
- if isinstance(dms, np.ndarray) and dms.ndim == 2: # RHF DM
- ndms = _get_number_of_density_matrices(dms)
- else:
- ndms = _get_number_of_density_matrices(dms[0])
- if ndms > 1:
- raise NotImplementedError(
- "NeuralNumInt does not support multiple density matrices. "
- "Only ground state DFT calculations are currently implemented."
- )
-
- nao = mol.nao_nr()
- self._vmat_hf = np.zeros((2, nao, nao))
- self._system_state = _SystemState(mol=mol, dms=dms)
- nelec, excsum, vmat = super().nr_uks(
- mol=mol,
- grids=grids,
- xc_code=xc_code,
- dms=dms,
- relativity=relativity,
- hermi=hermi,
- max_memory=max_memory,
- verbose=verbose,
- )
- vmat[0] += self._vmat_hf[0] + self._vmat_hf[0].T
- vmat[1] += self._vmat_hf[1] + self._vmat_hf[1].T
-
- # Clear internal state to prevent accidental re-use.
- self._system_state = None
- self._grid_state = None
- self._vmat_hf = None
- return nelec, excsum, vmat
-
- def block_loop(
- self,
- mol: gto.Mole,
- grids: dft.Grids,
- nao: Optional[int] = None,
- deriv: int = 0,
- max_memory: float = 2000,
- non0tab: Optional[np.ndarray] = None,
- blksize: Optional[int] = None,
- buf: Optional[np.ndarray] = None,
- ) -> Generator[Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray], None, None]:
- """Loops over the grid by blocks. See pyscf.dft.numint.NumInt.block_loop.
-
- Args:
- mol: PySCF molecule.
- grids: grid on which to evaluate the functional.
- nao: number of basis functions. If None, obtained from mol.
- deriv: unused. The first functional derivatives are always computed.
- max_memory: the maximum cache to use for the information on the grid, in
- MB. Determines the size of each block if blksize is None.
- non0tab: mask determining if a shell in the basis set is zero at a grid
- point. Shape (N, nbas), where N is the number of grid points and nbas
- the number of shells in the basis set. Obtained from grids if not
- supplied.
- blksize: size of each block. Calculated from max_memory if None.
- buf: buffer to use for storing ao. If None, a new array for ao is created
- for each block.
-
- Yields:
- ao, mask, weight, coords: information on a block of the grid containing N'
- points, where
- ao: atomic orbitals evaluated on the grid. Shape (N', nao), where nao is
- the number of atomic orbitals.
- mask: mask indicating whether a shell in the basis set is zero at a grid
- point. Shape (N', nbas).
- weight: weight associated with each grid point. Shape (N').
- coords: coordinates of the grid. Shape (N', 3).
- """
- # Wrap block_loop so we can store internal variables required to evaluate
- # the contribution to the XC potential from local Hartree-Fock features.
- for ao, mask, weight, coords in super().block_loop(
- mol=mol,
- grids=grids,
- nao=nao,
- deriv=deriv,
- max_memory=max_memory,
- non0tab=non0tab,
- blksize=blksize,
- buf=buf,
- ):
- # Cache the curent block so we can access it in eval_xc.
- self._grid_state = _GridState(ao=ao, mask=mask, weight=weight, coords=coords)
- yield ao, mask, weight, coords
-
- def construct_functional_inputs(
- self,
- mol: gto.Mole,
- dms: Union[np.ndarray, Sequence[np.ndarray]],
- spin: int,
- coords: np.ndarray,
- weights: np.ndarray,
- rho: Union[np.ndarray, Tuple[np.ndarray, np.ndarray]],
- ao: Optional[np.ndarray] = None,
- ) -> Tuple[FunctionalInputs, Tuple[np.ndarray, np.ndarray]]:
- """Constructs the input features required for the functional.
-
- Args:
- mol: PySCF molecule.
- dms: density matrix of shape (nao, nao) (restricted calculations) or of
- shape (2, nao, nao) (unrestricted calculations) or tuple of density
- matrices for each spin channel, each of shape (nao, nao) (unrestricted
- calculations).
- spin: 0 for a spin-unpolarized (restricted Kohn-Sham) calculation, and
- spin-polarized (unrestricted) otherwise.
- coords: coordinates of the grid. Shape (N, 3), where N is the number of
- grid points.
- weights: weight associated with each grid point. Shape (N).
- rho: density and density derivatives at each grid point. Single array
- containing the total density for restricted calculations, tuple of
- arrays for each spin channel for unrestricted calculations. Each array
- has shape (6, N). See pyscf.dft.numint.eval_rho and comments in
- FunctionalInputs for more details.
- ao: The atomic orbitals evaluated on the grid, shape (N, nao). Computed if
- not supplied.
-
- Returns:
- inputs, fxx, where
- inputs: FunctionalInputs object containing the inputs (as np.ndarrays)
- for the functional.
- fxx: intermediates, shape (N, nao) for the alpha- and beta-spin
- channels, required for computing the first derivative of the local
- Hartree-Fock density with respect to the density matrices. See
- compute_hfx_density for more details.
- """
- if spin == 0:
- # RKS
- rhoa = rho / 2
- rhob = rho / 2
- else:
- # UKS
- rhoa, rhob = rho
-
- # Local HF features.
- exxa, exxb = [], []
- fxxa, fxxb = [], []
- for omega in sorted(self._omega_values):
- hfx_results = compute_hfx_density.get_hf_density(
- mol, dms, coords=coords, omega=omega, deriv=1, ao=ao
- )
- exxa.append(hfx_results.exx[0])
- exxb.append(hfx_results.exx[1])
- fxxa.append(hfx_results.fxx[0])
- fxxb.append(hfx_results.fxx[1])
- exxa = np.stack(exxa, axis=-1)
- fxxa = np.stack(fxxa, axis=-1)
- if spin == 0:
- exx = (exxa, exxa)
- fxx = (fxxa, fxxa)
- else:
- exxb = np.stack(exxb, axis=-1)
- fxxb = np.stack(fxxb, axis=-1)
- exx = (exxa, exxb)
- fxx = (fxxa, fxxb)
-
- return (
- FunctionalInputs(
- rho_a=rhoa,
- rho_b=rhob,
- hfx_a=exx[0],
- hfx_b=exx[1],
- grid_coords=coords,
- grid_weights=weights,
- ),
- fxx,
- )
-
- def eval_xc(
- self,
- xc_code: str,
- rho: Union[np.ndarray, Tuple[np.ndarray, np.ndarray]],
- spin: int = 0,
- relativity: int = 0,
- deriv: int = 1,
- omega: Optional[float] = None,
- verbose=None,
- ) -> Tuple[np.ndarray, Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray], None, None]:
- """Evaluates the XC energy and functional derivatives.
-
- See pyscf.dft.libxc.eval_xc for more details on the interface.
-
- Note: this also sets self._vmat_extra, which contains the contribution the
- the potential matrix from the local Hartree-Fock terms in the functional.
-
- Args:
- xc_code: unused.
- rho: density and density derivatives at each grid point. Single array
- containing the total density for restricted calculations, tuple of
- arrays for each spin channel for unrestricted calculations. Each array
- has shape (6, N), where N is the number of grid points. See
- pyscf.dft.numint.eval_rho and comments in FunctionalInputs for more
- details.
- spin: 0 for a spin-unpolarized (restricted Kohn-Sham) calculation, and
- spin-polarized (unrestricted) otherwise.
- relativity: Not supported.
- deriv: unused. The first functional derivatives are always computed.
- omega: RSH parameter. Not supported.
- verbose: unused.
-
- Returns:
- exc, vxc, fxc, kxc, where:
- exc is the exchange-correlation potential matrix evaluated at each grid
- point, shape (N).
- vxc is (vrho, vgamma, vlapl, vtau), the first-order functional
- derivatives evaluated at each grid point, each shape (N).
- fxc is set to None. (The second-order functional derivatives are not
- computed.)
- kxc is set to None. (The third-order functional derivatives are not
- computed.)
- """
- del xc_code, verbose, deriv # unused
-
- if relativity != 0:
- raise NotImplementedError(
- "Relatistic calculations are not implemented " "for DM21 functionals."
- )
- if omega is not None:
- raise NotImplementedError(
- "User-specifed range seperation parameters are "
- "not implemented for DM21 functionals."
- )
-
- # Retrieve cached state.
- ao = self._grid_state.ao
- if ao.ndim == 3:
- # Just need the AO values, not the gradients.
- ao = ao[0]
- if self._grid_state.weight is None:
- weights = np.array([1.0])
- else:
- weights = self._grid_state.weight
- mask = self._grid_state.mask
-
- inputs, (fxxa, fxxb) = self.construct_functional_inputs(
- mol=self._system_state.mol,
- dms=self._system_state.dms,
- spin=spin,
- rho=rho,
- weights=weights,
- coords=self._grid_state.coords,
- ao=ao,
- )
-
- with self._graph.as_default():
- feed_dict = dict(
- zip(
- attr.asdict(self._placeholders).values(),
- attr.asdict(inputs).values(),
- )
- )
- tensor_list = [
- self._vxc,
- self._vrho,
- self._vsigma,
- self._vtau,
- self._vhf,
- ]
- exc, vrho, vsigma, vtau, vhf = self._session.run(tensor_list, feed_dict=feed_dict)
-
- mol = self._system_state.mol
- shls_slice = (0, mol.nbas)
- ao_loc_nr = mol.ao_loc_nr()
- # Note: tf.gradients returns a list of gradients.
- # vrho, vsigma, vtau are derivatives of objects that had
- # tf.expand_dims(..., 1) applied. The [:, 0] indexing undoes this by
- # selecting the 0-th (and only) element from the second dimension.
- if spin == 0:
- vxc_0 = (vrho[0][:, 0] + vrho[1][:, 0]) / 2.0
- # pyscf expects derivatives with respect to:
- # grad_rho . grad_rho.
- # The functional uses the first and last as inputs, but then has
- # grad_(rho_a + rho_b) . grad_(rho_a + rho_b)
- # as input. The following computes the correct total derivatives.
- vxc_1 = vsigma[0][:, 0] / 4.0 + vsigma[1][:, 0] / 4.0 + vsigma[2][:, 0]
- vxc_3 = (vtau[0][:, 0] + vtau[1][:, 0]) / 2.0
- vxc_2 = np.zeros_like(vxc_3)
- vhfs = (vhf[0] + vhf[1]) / 2.0
- # Local Hartree-Fock terms
- for i in range(len(self._omega_values)):
- # Factor of 1/2 is to account for adding vmat_hf + vmat_hf.T to vmat,
- # which we do to match existing PySCF style. Unlike other terms, vmat_hf
- # is already symmetric though.
- aow = np.einsum("pi,p->pi", fxxa[:, :, i], -0.5 * vhfs[:, i])
- self._vmat_hf += _dot_ao_ao(mol, ao, aow, mask, shls_slice, ao_loc_nr)
- else:
- vxc_0 = np.stack([vrho[0][:, 0], vrho[1][:, 0]], axis=1)
- # pyscf expects derivatives with respect to:
- # grad_rho_a . grad_rho_a
- # grad_rho_a . grad_rho_b
- # grad_rho_b . grad_rho_b
- # The functional uses the first and last as inputs, but then has
- # grad_(rho_a + rho_b) . grad_(rho_a + rho_b)
- # as input. The following computes the correct total derivatives.
- vxc_1 = np.stack(
- [
- vsigma[0][:, 0] + vsigma[2][:, 0],
- 2.0 * vsigma[2][:, 0],
- vsigma[1][:, 0] + vsigma[2][:, 0],
- ],
- axis=1,
- )
- vxc_3 = np.stack([vtau[0][:, 0], vtau[1][:, 0]], axis=1)
- vxc_2 = np.zeros_like(vxc_3)
- vhfs = np.stack([vhf[0], vhf[1]], axis=2)
- for i in range(len(self._omega_values)):
- # Factors of 1/2 are due to the same reason as in the spin=0 case.
- aow = np.einsum("pi,p->pi", fxxa[:, :, i], -0.5 * vhfs[:, i, 0])
- self._vmat_hf[0] += _dot_ao_ao(mol, ao, aow, mask, shls_slice, ao_loc_nr)
- aow = np.einsum("pi,p->pi", fxxb[:, :, i], -0.5 * vhfs[:, i, 1])
- self._vmat_hf[1] += _dot_ao_ao(mol, ao, aow, mask, shls_slice, ao_loc_nr)
-
- fxc = None # Second derivative not implemented
- kxc = None # Second derivative not implemented
- # PySCF C routines expect float64.
- exc = exc.astype(np.float64)
- vxc = tuple(v.astype(np.float64) for v in (vxc_0, vxc_1, vxc_2, vxc_3))
- return exc, vxc, fxc, kxc
-
- def eval_xc_eff( # pylint: disable=too-many-arguments,too-many-positional-arguments
- self,
- xc_code: str,
- rho: np.ndarray,
- deriv: int = 1,
- omega: Optional[float] = None,
- xctype: Optional[str] = None,
- verbose=None,
- spin: Optional[int] = None,
- ) -> Union[List[np.ndarray], Tuple[np.ndarray, np.ndarray, None, None]]:
- """Evaluates the XC energy and derivatives against density parameters.
-
- PySCF's nr_rks/nr_uks grid loops call this method rather than eval_xc.
- It differs from eval_xc only in the derivative convention: eval_xc
- returns derivatives with respect to sigma = |nabla rho|^2, while
- eval_xc_eff returns them with respect to the density parameters
- [rho, nabla rho, tau]. This override applies that chain rule to
- eval_xc's outputs, so the DM21 evaluation itself (including the local
- Hartree-Fock side effects on self._vmat_hf) stays in eval_xc.
-
- See pyscf.dft.numint.NumInt.eval_xc_eff for more details on the
- interface. The layouts below were verified against that method
- (max |diff| = 0.0 on mgga_x_tpss for both spin channels).
-
- Args:
- xc_code: unused (see eval_xc).
- rho: density and density derivatives at each grid point, shape
- (5, N) [rho, nabla_x, nabla_y, nabla_z, tau], or (2, 5, N) for
- spin-polarized calculations; a laplacian-carrying (6, N) /
- (2, 6, N) rho is accepted as well. eval_rho runs with
- MGGA_DENSITY_LAPL off, so rho arrives without the laplacian row
- the graph's (6, N) placeholders expect; a zero row is inserted
- before evaluation (the graph unstacks and discards it).
- deriv: derivative order. 1 (the default) is supported; 0 evaluates
- the energy and the local Hartree-Fock side effects but returns no
- potential; greater than 1 raises.
- omega: RSH parameter. None (the default) is passed through to
- eval_xc, which raises for an explicit value: DM21 has no
- range-separation support.
- xctype: unused. The DM21 functional is always a meta-GGA.
- verbose: unused.
- spin: 0 for a spin-unpolarized (restricted) calculation, 1 for
- spin-polarized. Inferred from rho's leading dimension when None,
- the same rule as pyscf.dft.numint.NumInt.eval_xc_eff.
-
- Returns:
- exc, vxc, fxc, kxc, where:
- exc is the XC energy density at each grid point, shape (N).
- vxc has shape (5, N) for spin=0, with rows
- [dE/d rho, dE/d nabla rho (3 components), dE/d tau], or shape
- (2, 5, N) for spin=1 with the same rows for the alpha and beta
- channels (rows [rho, nabla rho, tau] of each).
- fxc is None. (Second derivatives are not implemented.)
- kxc is None. (Third derivatives are not implemented.)
- For deriv=0 the return is [exc, None, None, None].
-
- Raises:
- NotImplementedError: if deriv > 1 (DM21 provides no second or higher
- derivatives; the base class would fail inside transform_xc).
- """
- if omega is None:
- omega = self.omega
- if deriv > 1:
- raise NotImplementedError(
- "DM21 functionals provide no second or higher derivatives, so "
- "eval_xc_eff cannot be called with deriv > 1."
- )
- del xctype # unused: the DM21 functional is always a meta-GGA
-
- rho = np.asarray(rho, order="C", dtype=np.float64)
- if spin is None:
- spin = 1 if (rho.ndim >= 2 and rho.shape[0] == 2) else 0
- if rho.shape[-2] == 5:
- rho = np.insert(rho, 4, 0.0, axis=-2)
-
- exc, (vrho, vsigma, _vlapl, vtau), _, _ = self.eval_xc(
- xc_code, rho, spin=spin, deriv=deriv, omega=omega, verbose=verbose
- )
- if deriv < 1:
- # Evaluated for exc and the _vmat_hf side effect; no potential is
- # requested at derivative order 0.
- return [exc, None, None, None]
-
- # Chain rule from eval_xc's sigma derivatives (dE/d sigma, with sigma
- # the PySCF/libxc (sigma_aa, sigma_ab, sigma_bb) components) to the
- # density-parameter derivatives PySCF contracts with the AOs. To
- # re-derive it, follow PySCF's MGGA branches: nr_rks does
- # wv = weight * vxc, wv[0] *= .5, wv[4] *= .5,
- # _scale_ao_sparse(ao[:4], wv[:4], ...) and _tau_dot_sparse(ao, ao,
- # wv[4], ...); nr_uks does the same per spin row (wv[:, 0] and
- # wv[:, 4]) — i.e. row 0 is dE/d rho, rows 1:4 are dE/d nabla rho and
- # row 4 is dE/d tau. dE/d nabla rho = 2 (dE/d sigma) nabla rho for a
- # single sigma, and the pairwise expansion below for sigma_ab.
- if spin == 0:
- # sigma = nabla rho . nabla rho, so the gradient rows pick up a
- # factor of 2.
- vxc = np.concatenate(
- [
- vrho[None, :],
- 2.0 * vsigma[None, :] * rho[1:4],
- vtau[None, :],
- ],
- axis=0,
- )
- else:
- nabla_a, nabla_b = rho[0, 1:4], rho[1, 1:4]
- vsigma_aa, vsigma_ab, vsigma_bb = vsigma[:, 0], vsigma[:, 1], vsigma[:, 2]
- vxc_a = np.concatenate(
- [
- vrho[:, 0][None, :],
- 2.0 * vsigma_aa[None, :] * nabla_a + vsigma_ab[None, :] * nabla_b,
- vtau[:, 0][None, :],
- ],
- axis=0,
- )
- vxc_b = np.concatenate(
- [
- vrho[:, 1][None, :],
- 2.0 * vsigma_bb[None, :] * nabla_b + vsigma_ab[None, :] * nabla_a,
- vtau[:, 1][None, :],
- ],
- axis=0,
- )
- vxc = np.stack([vxc_a, vxc_b], axis=0)
- return exc, vxc, None, None
diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint_test.py b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint_test.py
deleted file mode 100644
index 9bcea65..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint_test.py
+++ /dev/null
@@ -1,137 +0,0 @@
-# Copyright 2021 DeepMind Technologies Limited.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Tests for neural_numint."""
-
-import os
-
-
-from absl.testing import parameterized
-import attr
-from pyscf import dft
-from pyscf import gto
-from pyscf import lib
-import tensorflow.compat.v1 as tf
-
-from density_functional_approximation_dm21 import neural_numint
-
-
-class NeuralNumintTest(tf.test.TestCase, parameterized.TestCase):
- def setUp(self):
- super().setUp()
- lib.param.TMPDIR = None
- lib.num_threads(1)
-
- # Golden values were obtained using the version of PySCF (including integral
- # generation) reported in the DM21 paper.
- @parameterized.parameters(
- {"functional": neural_numint.Functional.DM21, "expected_energy": -126.898521},
- {"functional": neural_numint.Functional.DM21m, "expected_energy": -126.907332},
- {"functional": neural_numint.Functional.DM21mc, "expected_energy": -126.922127},
- {"functional": neural_numint.Functional.DM21mu, "expected_energy": -126.898178},
- )
- def test_rks(self, functional, expected_energy):
- ni = neural_numint.NeuralNumInt(functional)
-
- mol = gto.Mole()
- mol.atom = [["Ne", 0.0, 0.0, 0.0]]
- mol.basis = "sto-3g"
- mol.build()
-
- mf = dft.RKS(mol)
- mf.small_rho_cutoff = 1.0e-20
- mf._numint = ni
- mf.run()
- self.assertAlmostEqual(mf.e_tot, expected_energy, delta=2.0e-4)
-
- @parameterized.parameters(
- {"functional": neural_numint.Functional.DM21, "expected_energy": -37.34184876},
- {"functional": neural_numint.Functional.DM21m, "expected_energy": -37.3377766},
- {"functional": neural_numint.Functional.DM21mc, "expected_energy": -37.33489173},
- {"functional": neural_numint.Functional.DM21mu, "expected_energy": -37.34015315},
- )
- def test_uks(self, functional, expected_energy):
- ni = neural_numint.NeuralNumInt(functional)
-
- mol = gto.Mole()
- mol.atom = [["C", 0.0, 0.0, 0.0]]
- mol.spin = 2
- mol.basis = "sto-3g"
- mol.build()
-
- mf = dft.UKS(mol)
- mf.small_rho_cutoff = 1.0e-20
- mf._numint = ni
- mf.run()
- self.assertAlmostEqual(mf.e_tot, expected_energy, delta=2.0e-4)
-
- def test_exported_model(self):
- mol = gto.Mole()
- mol.atom = [["C", 0.0, 0.0, 0.0]]
- mol.spin = 2
- mol.basis = "sto-3g"
- mol.build()
-
- ni = neural_numint.NeuralNumInt(neural_numint.Functional.DM21)
- mf = dft.UKS(mol)
- mf.small_rho_cutoff = 1.0e-20
- mf._numint = ni
- mf.run()
-
- dms = mf.make_rdm1()
- ao = ni.eval_ao(mol, mf.grids.coords, deriv=2)
- rho_a = ni.eval_rho(mol, ao, dms[0], xctype="MGGA")
- rho_b = ni.eval_rho(mol, ao, dms[1], xctype="MGGA")
- inputs, _ = ni.construct_functional_inputs(
- mol=mol,
- dms=dms,
- spin=1,
- coords=mf.grids.coords,
- weights=mf.grids.weights,
- rho=(rho_a, rho_b),
- ao=ao[0],
- )
-
- feed_dict = dict(
- zip(
- attr.asdict(ni._placeholders).values(),
- attr.asdict(inputs).values(),
- )
- )
- with ni._graph.as_default():
- outputs = ni._session.run(
- {
- "vxc": ni._vxc,
- "vrho": ni._vrho,
- "vsigma": ni._vsigma,
- "vtau": ni._vtau,
- "vhf": ni._vhf,
- },
- feed_dict=feed_dict,
- )
-
- export_path = os.path.join(self.get_temp_dir(), "export")
- ni.export_functional_and_derivatives(export_path)
- model = tf.saved_model.load_v2(export_path)
- tensor_inputs = {
- k: tf.constant(v, dtype=tf.float32) for k, v in attr.asdict(inputs).items()
- }
- exported_output_tensors = model.signatures["default"](**tensor_inputs)
- with tf.Session() as session:
- session.run(tf.global_variables_initializer())
- exported_outputs = session.run(exported_output_tensors)
- self.assertAllClose(outputs, exported_outputs, atol=5.0e-5, rtol=1.0e-5)
-
-
-if __name__ == "__main__":
- tf.test.main()
diff --git a/grad_dft/external/density_functional_approximation_dm21/requirements.txt b/grad_dft/external/density_functional_approximation_dm21/requirements.txt
deleted file mode 100644
index ff94956..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/requirements.txt
+++ /dev/null
@@ -1,9 +0,0 @@
-absl-py==0.13.0
-attrs==21.2.0
-h5py==3.1.0
-numpy==1.19.5
-pyscf==1.7.6.post1
-pytest==6.2.4
-scipy==1.7.1
-tensorflow==2.6.0
-tensorflow-hub==0.12.0
diff --git a/grad_dft/external/density_functional_approximation_dm21/requirements_aot_compilation.txt b/grad_dft/external/density_functional_approximation_dm21/requirements_aot_compilation.txt
deleted file mode 100644
index eb9409b..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/requirements_aot_compilation.txt
+++ /dev/null
@@ -1,3 +0,0 @@
-Keras-Preprocessing==1.1.2
-numpy==1.19.5
-tensorflow-estimator==2.7.0
diff --git a/grad_dft/external/density_functional_approximation_dm21/run.sh b/grad_dft/external/density_functional_approximation_dm21/run.sh
deleted file mode 100755
index 319f2a3..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/run.sh
+++ /dev/null
@@ -1,27 +0,0 @@
-#!/bin/bash
-# Copyright 2021 DeepMind Technologies Limited.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# Install density_functional_approximation_dm21 in a clean virtualenv and run
-# the tests. This assumes the working directory is the top-level directory of
-# the deepmind-research repository, i.e.:
-# git clone git@github.com:deepmind/deepmind-research.git
-# cd deepmind_research
-# density_functional_approximation_dm21/run.sh
-
-python3 -m venv /tmp/DM21
-source /tmp/DM21/bin/activate
-pip3 install -r density_functional_approximation_dm21/requirements.txt
-pip install 'density_functional_approximation_dm21/[testing]'
-py.test density_functional_approximation_dm21/
diff --git a/grad_dft/external/density_functional_approximation_dm21/setup.py b/grad_dft/external/density_functional_approximation_dm21/setup.py
deleted file mode 100644
index 7f01b2c..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/setup.py
+++ /dev/null
@@ -1,51 +0,0 @@
-# Copyright 2021 DeepMind Technologies Limited
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-# https://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Setup for DM21 functionals interface to PySCF."""
-
-from setuptools import setup
-
-REQUIRED_PACKAGES = [
- "absl-py",
- "attrs",
- # Note PySCF 1.7.6 and older do not support h5py 3.3.0:
- # https://github.com/pyscf/pyscf/issues/1016
- "h5py",
- "numpy",
- # Note DM21 functionals are compatible with PySCF 1.7.6 if an older version
- # of h5py is used.
- "pyscf>=2.0",
- "tensorflow",
- "tensorflow_hub",
-]
-CHECKPOINT_DATA = ["checkpoints/DM21*/*.pb", "checkpoints/DM21*/variables/*"]
-
-setup(
- name="density_functional_approximation_dm21",
- version="0.1",
- description="An interface to PySCF for the DM21 functionals.",
- url="https://github.com/deepmind/deepmind-research/density_functional_approximation_dm21",
- author="DeepMind",
- author_email="no-reply@google.com",
- # Contained modules and scripts.
- packages=["density_functional_approximation_dm21"],
- package_data={
- "density_functional_approximation_dm21": CHECKPOINT_DATA,
- },
- scripts=["density_functional_approximation_dm21/export_saved_model.py"],
- install_requires=REQUIRED_PACKAGES,
- platforms=["any"],
- license="Apache 2.0",
- extras_require={"testing": ["pytest", "scipy"]},
-)
diff --git a/grad_dft/external/density_functional_approximation_dm21/test.py b/grad_dft/external/density_functional_approximation_dm21/test.py
deleted file mode 100644
index 81b1088..0000000
--- a/grad_dft/external/density_functional_approximation_dm21/test.py
+++ /dev/null
@@ -1,15 +0,0 @@
-import density_functional_approximation_dm21 as dm21
-from openfermion.chem import geometry_from_pubchem
-import tensorflow as tf
-
-from pyscf import gto, dft
-
-mol = gto.Mole()
-mol.atom = geometry_from_pubchem("water")
-mol.basis = "cc-pvdz"
-mol.build()
-
-mf = dft.RKS(mol)
-mf._numint = dm21.NeuralNumInt(dm21.Functional.DM21)
-
-mf.kernel()
diff --git a/grad_dft/interface/pyscf.py b/grad_dft/interface/pyscf.py
index d109f0a..52fe533 100644
--- a/grad_dft/interface/pyscf.py
+++ b/grad_dft/interface/pyscf.py
@@ -37,15 +37,20 @@
# from qdft.reaction import Reaction, make_reaction, get_grad
from grad_dft.molecule import Grid, Molecule, Reaction, make_reaction
from grad_dft.solid import Solid, KPointInfo
-from grad_dft.utils import DType, default_dtype, DensityFunctional, HartreeFock
+from grad_dft.utils import (
+ DType,
+ default_dtype,
+ DensityFunctional,
+ HartreeFock,
+ PackedERI,
+ validate_eri_storage,
+)
from jax.tree_util import tree_map
from grad_dft.external import NeuralNumInt
from grad_dft.external import Functional as ExternalFunctional
-from pyscf.pbc import df
-
import h5py
-from pyscf import cc, dft, scf
+from pyscf import cc, dft, df, scf
from jaxtyping import Array, Scalar, Int, Bool
from grad_dft.external import _nu_chunk
@@ -111,12 +116,54 @@ def molecule_from_pyscf(
scf_iteration: Scalar = jnp.int32(50),
chunk_size: Optional[Scalar] = jnp.int32(1024),
grad_order: Optional[Scalar] = jnp.int32(2),
-) -> Molecule:
+ eri_storage: str = "dense",
+) -> Molecule:
+ r"""Convert a converged PySCF mean field object into a Molecule.
+
+ Parameters
+ ----------
+ mf : DensityFunctional
+ PySCF mean field object (kernel already run).
+ dtype : DType, optional
+ Target floating point dtype.
+ omegas : Array, optional
+ Range-separation parameters; triggers chi-tensor generation.
+ energy : Scalar, optional
+ Energy stored on the molecule.
+ name : Array, optional
+ Molecule name (encoded as integers).
+ scf_iteration : Scalar, optional
+ Number of SCF iterations performed.
+ chunk_size : Scalar, optional
+ Chunk size for chi-tensor generation.
+ grad_order : Scalar, optional
+ Highest derivative order of the atomic orbitals.
+ eri_storage : str, optional
+ ERI storage format for ``rep_tensor``: "dense" (default, current
+ behavior), "s4" or "s8" (symmetry-compressed PySCF storage wrapped
+ in a ``PackedERI``). Isolated molecules only; periodic systems
+ raise ValueError (validated here before any integrals run, and
+ again in :func:`calc_eri_with_pyscf` as defense in depth).
+
+ Returns
+ -------
+ Molecule
+ """
+ # Fail fast on unsupported ERI storage before any SCF or integral work.
+ eri_storage = validate_eri_storage(eri_storage)
+ if hasattr(mf, "cell") and eri_storage != "dense":
+ raise ValueError(
+ f"eri_storage={eri_storage!r} is only supported for isolated "
+ "molecules; periodic systems keep eri_storage='dense' "
+ "(use molecule_from_pyscf without compression)."
+ )
if hasattr(mf, "kpts"):
if not np.array_equal(mf.kpts, np.array([[0.0, 0.0, 0.0]])):
raise RuntimeError("Input was periodic with BZ sampling beyond gamma-point only. Use solid_from_pyscf instead.")
grid = grid_from_pyscf(mf.grids, dtype=dtype)
- outputs = _package_outputs(mf, mf.grids, scf_iteration, grad_order)
+ outputs = _package_outputs(
+ mf, mf.grids, scf_iteration, grad_order, eri_storage=eri_storage
+ )
(
ao,
grad_ao,
@@ -372,6 +419,11 @@ def saver(
|- All the attributes in Molecule class
|- chi (attributes: omegas)
+ ERIs in compressed storage (a PackedERI, eri_storage="s4"/"s8") are saved
+ as the packed array with ``storage_format``/``n_orbitals`` as *dataset*
+ attributes; dense molecules are saved exactly as before (no attributes),
+ so existing HDF5 files keep loading unchanged.
+
Raises:
-------
@@ -454,6 +506,13 @@ def loader(
molecule/reaction : Molecule or Reaction
The molecule or reaction object.
+ Notes
+ -----
+ Compressed ERIs written by :func:`saver` (dataset attributes
+ ``storage_format``/``n_orbitals``) are reconstructed as ``PackedERI``;
+ dense datasets, including all files written before compressed storage
+ existed, load unchanged.
+
todo: randomize input
"""
@@ -493,6 +552,8 @@ def loader(
args[key] = jnp.stack(
[jnp.asarray(value, dtype = jnp.float64)[:, i] for i in indices], axis=1
)
+ elif key == "rep_tensor":
+ args[key] = _load_rep_tensor(value, dtype=jnp.float64)
else:
args[key] = jnp.asarray(value, dtype=jnp.float64)
@@ -553,6 +614,8 @@ def loader(
args[key] = jnp.stack(
[jnp.asarray(value)[:, i] for i in indices], axis=1
)
+ elif key == "rep_tensor":
+ args[key] = _load_rep_tensor(value)
else:
args[key] = jnp.asarray(value)
@@ -580,7 +643,14 @@ def loader(
def save_molecule_data(mol_group: h5py.Group, molecule: Molecule):
- r"""Auxiliary function to save all data except for chi"""
+ r"""Auxiliary function to save all data except for chi
+
+ ``rep_tensor`` in compressed storage (a ``PackedERI``, from
+ ``eri_storage="s4"/"s8"``) is saved as the packed array with its
+ ``storage_format``/``n_orbitals`` as *dataset* attributes; dense
+ molecules are saved exactly as before (no attributes), so existing
+ HDF5 files keep loading unchanged.
+ """
to_numpy = (
lambda arr: arr if (isinstance(arr, str) or isinstance(arr, float)) else np.asarray(arr)
@@ -596,10 +666,33 @@ def save_molecule_data(mol_group: h5py.Group, molecule: Molecule):
d = mol_group.create_group(name)
for k, v in data.items():
d.create_dataset(f"{k}", data=v)
+ elif isinstance(data, PackedERI):
+ dataset = mol_group.create_dataset(name, data=np.asarray(data.packed))
+ dataset.attrs["storage_format"] = data.storage_format
+ dataset.attrs["n_orbitals"] = data.n_orbitals
else:
mol_group.create_dataset(name, data=data)
+def _load_rep_tensor(dataset: h5py.Dataset, dtype: Optional[DType] = None):
+ r"""Load a saved ERI tensor, packed or dense.
+
+ Compressed storage is identified by the ``storage_format`` dataset
+ attribute written by :func:`save_molecule_data` — never by the array's
+ shape; any other value indicates a corrupt file. Old dense files carry
+ no attributes and load exactly as before.
+ """
+ storage_format = str(dataset.attrs.get("storage_format", "dense"))
+ validate_eri_storage(storage_format)
+ if storage_format == "dense":
+ return jnp.asarray(dataset, dtype=dtype)
+ return PackedERI(
+ jnp.asarray(dataset, dtype=dtype),
+ storage_format,
+ int(dataset.attrs["n_orbitals"]),
+ )
+
+
def save_molecule_chi(
molecule: Molecule,
omegas: Union[Sequence[Scalar], Scalar],
@@ -640,7 +733,12 @@ def to_device_arrays(*arrays, dtype: Optional[DType] = None):
out = []
for array in arrays:
- if isinstance(array, dict):
+ if isinstance(array, PackedERI):
+ # Dense arrays below ignore `dtype`; the packed leaf honors it
+ # (a single cheap conversion, no N^4 materialization). With the
+ # default `dtype=None` this matches the dense behavior.
+ out.append(array.replace(packed=jnp.asarray(array.packed, dtype=dtype)))
+ elif isinstance(array, dict):
for k, v in array.items():
array[k] = jnp.asarray(v)
out.append(array)
@@ -741,7 +839,25 @@ def pbc_ao_grads(cell: Cell, coords: Array, order=2, kpts=None) -> Dict:
i += 1
return result
-def calc_eri_with_pyscf(mf, kpts=np.zeros(3)) -> np.ndarray:
+def _packed_exact_eri_from_mol(mol, eri_storage: str) -> PackedERI:
+ r"""Exact molecular integrals in compressed storage (no density fitting)."""
+ return PackedERI(
+ np.asarray(mol.intor("int2e", aosym=eri_storage)), eri_storage, mol.nao_nr()
+ )
+
+
+def _packed_df_eri(eri_compressed, eri_storage: str, n_ao: int) -> PackedERI:
+ r"""Wrap density-fitted molecular ERIs in compressed storage.
+
+ PySCF's molecular ``DF.get_eri()`` has no ``aosym`` argument and always
+ returns s8 storage; the s4 format re-packs it with ``restore(4, ·)``.
+ """
+ if eri_storage == "s4":
+ eri_compressed = restore(4, eri_compressed, n_ao)
+ return PackedERI(np.asarray(eri_compressed), eri_storage, n_ao)
+
+
+def calc_eri_with_pyscf(mf, kpts=np.zeros(3), eri_storage: str = "dense"):
r"""Calculate the ERIs using the method detected from the PySCF mean field object.
Inputs
@@ -751,16 +867,36 @@ def calc_eri_with_pyscf(mf, kpts=np.zeros(3)) -> np.ndarray:
PySCF mean field object
kpts:
Array of k-points (absolute, not fractional).
+ eri_storage:
+ ERI storage format: "dense" (default, current behavior), "s4" or
+ "s8" (PySCF symmetry-compressed storage wrapped in a
+ :class:`~grad_dft.utils.eri.PackedERI`). Only supported for isolated
+ molecules; periodic systems (including gamma-point-only) keep
+ "dense" and raise ValueError otherwise, since GradDFT assumes no
+ ERI symmetry for periodic integrals. With density fitting, the
+ packed formats are taken from the density fitter's ``get_eri()``
+ (always s8 storage), never from exact integrals.
Outputs
----------
- np.ndarray
+ np.ndarray or PackedERI
The ERIs. Output shape is (nao, nao, nao, nao) for isolated molecules and gamma-point only
periodic calculations. For full BZ calculations, the output shape is (nkpt, nkpt, nao, nao, nao, nao).
+ With eri_storage="s4"/"s8" (isolated molecules only) a PackedERI wrapping
+ the compressed storage of the same integrals.
"""
+ # Validate before any integral work so bad configuration fails fast.
+ eri_storage = validate_eri_storage(eri_storage)
+
# Solid or Isolated molecule?
if hasattr(mf, "cell"): # Periodic system
+ if eri_storage != "dense":
+ raise ValueError(
+ "Periodic systems must keep eri_storage='dense': GradDFT "
+ "assumes no ERI symmetry for periodic integrals (complex "
+ f"k-point phases, density fitting), got {eri_storage!r}."
+ )
# Check for the three density fitting methods. DF is always used for periodic calculations
if isinstance(mf.with_df, FFTDF):
@@ -792,11 +928,17 @@ def calc_eri_with_pyscf(mf, kpts=np.zeros(3)) -> np.ndarray:
try:
_ = mf.with_df
except(AttributeError):
+ if eri_storage != "dense":
+ return _packed_exact_eri_from_mol(mf.mol, eri_storage)
eri = mf.mol.intor("int2e")
return eri
# Use default DF method when DF is used on molecules
+ # (`df` is pyscf.df, the molecular module — NOT pyscf.pbc.df, whose
+ # DF class requires a periodic Cell).
density_fitter = df.DF(mf.mol)
eri_compressed = density_fitter.get_eri()
+ if eri_storage != "dense":
+ return _packed_df_eri(eri_compressed, eri_storage, mf.mol.nao_nr())
eri = restore(1, eri_compressed, mf.mol.nao_nr())
return eri
@@ -808,7 +950,31 @@ def _package_outputs(
grids: Optional[Grids] = None,
scf_iteration: Scalar = jnp.int32(50),
grad_order: Scalar = jnp.int32(2),
-):
+ eri_storage: str = "dense",
+):
+ r"""Package a converged PySCF mean field object into GradDFT arrays.
+
+ Parameters
+ ----------
+ mf : DensityFunctional
+ PySCF mean field object (kernel already run).
+ grids : Grids, optional
+ PySCF grids; taken from ``mf.grids`` downstream.
+ scf_iteration : Scalar, optional
+ Number of SCF iterations performed.
+ grad_order : Scalar, optional
+ Highest derivative order of the atomic orbitals.
+ eri_storage : str, optional
+ ERI storage format forwarded to
+ :func:`calc_eri_with_pyscf` ("dense", "s4" or "s8"); the returned
+ ``rep_tensor`` is then a ``PackedERI`` instead of a dense array.
+
+ Returns
+ -------
+ Tuple
+ (ao, grad_ao, grad_n_ao, rdm1, energy_nuc, h1e, vj, mo_coeff,
+ mo_energy, mo_occ, mf_e_tot, s1e, fock, rep_tensor, kpt_info).
+ """
if scf_iteration != 0:
rdm1 = mf.make_rdm1(mf.mo_coeff, mf.mo_occ)
else:
@@ -839,7 +1005,7 @@ def _package_outputs(
) # The 2 is to compensate for the /2 in the definition of the density matrix
dm = mf.make_rdm1(mf.mo_coeff, mf.mo_occ)
fock = np.stack([h1e, h1e], axis=0) + mf.get_veff(mf.mol, dm)
- rep_tensor = calc_eri_with_pyscf(mf)
+ rep_tensor = calc_eri_with_pyscf(mf, eri_storage=eri_storage)
kpt_info = None
# Unrestricted (spin polarized), open boundary conditions
@@ -858,7 +1024,7 @@ def _package_outputs(
) # The 2 is to compensate for the /2 in the definition of the density matrix
dm = mf.make_rdm1(mf.mo_coeff, mf.mo_occ)
fock = np.stack([h1e, h1e], axis=0) + mf.get_veff(mf.mol, dm)
- rep_tensor = calc_eri_with_pyscf(mf)
+ rep_tensor = calc_eri_with_pyscf(mf, eri_storage=eri_storage)
kpt_info = None
# Restricted (non-spin polarized), periodic boundary conditions, full BZ sampling
@@ -890,7 +1056,7 @@ def _package_outputs(
kpt_info = kpt_info_from_pyscf(mf)
# Compute ERIs for all pairs of k-points. Needed for Coulomb energy calculation
- rep_tensor = calc_eri_with_pyscf(mf, kpts=mf.kpts)
+ rep_tensor = calc_eri_with_pyscf(mf, kpts=mf.kpts, eri_storage=eri_storage)
# Unrestricted (spin polarized), periodic boundary conditions, full BZ sampling
elif rdm1.ndim == 4 and hasattr(mf, "cell") and rdm1.shape[1] != 1:
@@ -916,7 +1082,7 @@ def _package_outputs(
kpt_info = kpt_info_from_pyscf(mf)
# Compute ERIs for all pairs of k-points. Needed for Coulomb energy calculation
- rep_tensor = calc_eri_with_pyscf(mf, kpts=mf.kpts)
+ rep_tensor = calc_eri_with_pyscf(mf, kpts=mf.kpts, eri_storage=eri_storage)
# Restricted (non-spin polarized), periodic boundary conditions, gamma point only
elif rdm1.ndim == 3 and hasattr(mf, "cell") and rdm1.shape[0] == 1:
@@ -960,7 +1126,7 @@ def _package_outputs(
fock = np.squeeze(fock, axis=1)
vj = np.squeeze(vj, axis=1)
h1e = np.squeeze(h1e, axis=0)
- rep_tensor = calc_eri_with_pyscf(mf)
+ rep_tensor = calc_eri_with_pyscf(mf, eri_storage=eri_storage)
kpt_info = None
# Unrestricted (spin polarized), periodic boundary conditions, gamma point only
@@ -999,7 +1165,7 @@ def _package_outputs(
fock = np.squeeze(fock, axis=1)
vj = np.squeeze(vj, axis=1)
h1e = np.squeeze(h1e, axis=0)
- rep_tensor = calc_eri_with_pyscf(mf)
+ rep_tensor = calc_eri_with_pyscf(mf, eri_storage=eri_storage)
kpt_info = None
else:
diff --git a/grad_dft/molecule.py b/grad_dft/molecule.py
index ea1f060..d93949a 100644
--- a/grad_dft/molecule.py
+++ b/grad_dft/molecule.py
@@ -16,7 +16,12 @@
from dataclasses import fields
from typeguard import typechecked
-from grad_dft.utils import vmap_chunked
+from grad_dft.utils import (
+ ERITensor,
+ PackedERI,
+ packed_coulomb_potential,
+ vmap_chunked,
+)
from functools import partial
from jax import numpy as jnp
@@ -71,7 +76,12 @@ def integrate(self, vals: Array, axis: int = 0) -> Array:
@struct.dataclass
class Molecule:
- r""" Base class for storing molecule properties and methods."""
+ r""" Base class for storing molecule properties and methods.
+
+ ``rep_tensor`` accepts either a dense ``(n, n, n, n)`` array or a
+ :class:`~grad_dft.utils.eri.PackedERI` with symmetry-compressed ERIs
+ (``eri_storage="s4"``/``"s8"``); contractions dispatch on the type.
+ """
grid: Grid
atom_index: Int[Array, "atom"]
@@ -90,7 +100,7 @@ class Molecule:
s1e: Optional[Float[Array, "orbitals orbitals"]] = None # Not used during training
omegas: Optional[Float[Array, "omega"]] = None
chi: Optional[Float[Array, "grid omega spin orbitals"]] = None
- rep_tensor: Optional[Float[Array, "orbitals orbitals orbitals orbitals"]] = None
+ rep_tensor: Optional[ERITensor] = None
energy: Optional[Scalar] = None
basis: Optional[Int[Array, '...']] = None # The name is saved as a list of integers, JAX does not accept str
name: Optional[Int[Array, '...']] = None # The name is saved as a list of integers, JAX does not accept str
@@ -108,6 +118,9 @@ def grid_size(self):
def get_coulomb_potential(self, *args, **kwargs) -> Float[Array, "orbitals orbitals"]:
r"""Compute the Coulomb potential matrix.
+ ``rep_tensor`` may be a dense ``(n, n, n, n)`` array or a
+ :class:`~grad_dft.utils.eri.PackedERI` (compressed storage).
+
Returns
-------
Float[Array, "spin orbitals orbitals"]
@@ -296,7 +309,11 @@ def HF_coefficient_input_grad_2_Fock(
)
def nonXC(self, *args, **kwargs) -> Scalar:
- r""" Computes the non-XC energy of a DFT functional."""
+ r""" Computes the non-XC energy of a DFT functional.
+
+ ``rep_tensor`` may be a dense ``(n, n, n, n)`` array or a
+ :class:`~grad_dft.utils.eri.PackedERI` (compressed storage).
+ """
return nonXC(self.rdm1.sum(axis = 0), self.h1e, self.rep_tensor, self.nuclear_repulsion, *args, **kwargs)
def make_rdm1(self) -> Array:
@@ -697,7 +714,7 @@ def abs_clip(arr, threshold):
def nonXC(
rdm1: Float[Array, "orbitals orbitals"],
h1e: Float[Array, "orbitals orbitals"],
- rep_tensor: Float[Array, "orbitals orbitals orbitals orbitals"],
+ rep_tensor: ERITensor,
nuclear_repulsion: Scalar,
precision=Precision.HIGHEST,
) -> Scalar:
@@ -711,9 +728,11 @@ def nonXC(
h1e : Float[Array, "orbitals orbitals"]
The 1-electron Hamiltonian.
Equivalent to mf.get_hcore(mf.mol) in pyscf.
- rep_tensor : Float[Array, "orbitals orbitals orbitals orbitals"]
- The repulsion tensor.
- Equivalent to mf.mol.intor('int2e') in pyscf.
+ rep_tensor : ERITensor
+ The repulsion tensor: dense ``(orbitals, orbitals, orbitals,
+ orbitals)`` array (equivalent to mf.mol.intor('int2e') in pyscf) or
+ a :class:`~grad_dft.utils.eri.PackedERI` (``s4``/``s8``), contracted
+ without materializing the dense :math:`N^4` tensor when compressed.
nuclear_repulsion : Scalar
The nuclear repulsion energy.
Equivalent to mf.mol.energy_nuc() in pyscf.
@@ -762,7 +781,7 @@ def one_body_energy(
@partial(jax.jit, static_argnames=["precision"])
def coulomb_energy(
rdm1: Float[Array, "orbitals orbitals"],
- rep_tensor: Float[Array, "orbitals orbitals orbitals orbitals"],
+ rep_tensor: ERITensor,
precision=Precision.HIGHEST,
) -> Scalar:
r"""A function that computes the Coulomb two-body energy of a DFT functional.
@@ -771,8 +790,11 @@ def coulomb_energy(
----------
rdm1 : Float[Array, "spin orbitals orbitals"]
The 1-Reduced Density Matrix.
- rep_tensor : Float[Array, "orbitals orbitals orbitals orbitals"]
- The repulsion tensor.
+ rep_tensor : ERITensor
+ The repulsion tensor: dense ``(orbitals, orbitals, orbitals,
+ orbitals)`` array or a :class:`~grad_dft.utils.eri.PackedERI`
+ (``s4``/``s8``), contracted without materializing the dense
+ :math:`N^4` tensor when compressed.
Returns
-------
@@ -787,7 +809,7 @@ def coulomb_energy(
@partial(jax.jit, static_argnames=["precision"])
def coulomb_potential(
rdm1: Float[Array, "orbitals orbitals"],
- rep_tensor: Float[Array, "orbitals orbitals orbitals orbitals"],
+ rep_tensor: ERITensor,
precision=Precision.HIGHEST,
) -> Float[Array, "orbitals orbitals"]:
r"""
@@ -798,9 +820,16 @@ def coulomb_potential(
rdm1 : Float[Array, "spin orbitals orbitals"]
The 1-Reduced Density Matrix.
Equivalent to mf.make_rdm1() in pyscf.
- rep_tensor : Float[Array, "orbitals orbitals orbitals orbitals"]
- The repulsion tensor.
- Equivalent to mf.mol.intor('int2e') in pyscf.
+ rep_tensor : ERITensor
+ The repulsion tensor: either the dense
+ ``Float[Array, "orbitals orbitals orbitals orbitals"]``
+ (equivalent to mf.mol.intor('int2e') in pyscf) or a
+ :class:`~grad_dft.utils.eri.PackedERI` holding PySCF
+ symmetry-compressed ``s4``/``s8`` storage. The compressed path
+ contracts the packed storage directly and never materializes the
+ dense :math:`N^4` tensor; it is exact for an arbitrary ``rdm1``
+ (no density symmetry assumed), matching the dense einsum up to
+ floating-point summation order.
precision : Precision, optional
The precision to use for the computation, by default Precision.HIGHEST
@@ -808,6 +837,8 @@ def coulomb_potential(
-------
Float[Array, "spin orbitals orbitals"]
"""
+ if isinstance(rep_tensor, PackedERI):
+ return packed_coulomb_potential(rdm1, rep_tensor, precision)
return jnp.einsum("pqrt,rt->pq", rep_tensor, rdm1, precision=precision)
@jaxtyped
diff --git a/grad_dft/utils/__init__.py b/grad_dft/utils/__init__.py
index 32c02b3..2c2a328 100644
--- a/grad_dft/utils/__init__.py
+++ b/grad_dft/utils/__init__.py
@@ -23,8 +23,21 @@
DType,
HartreeFock,
DensityFunctional,
+ ERITensor,
default_dtype,
)
+from .eri import (
+ ERI_STORAGE_FORMATS,
+ PACKED_STORAGE_FORMATS,
+ PackedERI,
+ eri_allclose,
+ fold_density_matrix,
+ pair_indices,
+ packed8_to_packed4,
+ packed_coulomb_potential,
+ packed_to_dense,
+ validate_eri_storage,
+)
from .tree import tree_size, tree_isfinite, tree_randn_like, tree_func, tree_shape
from .utils import to_device_arrays, Utils
from .chunk import vmap_chunked
diff --git a/grad_dft/utils/eri.py b/grad_dft/utils/eri.py
new file mode 100644
index 0000000..db76388
--- /dev/null
+++ b/grad_dft/utils/eri.py
@@ -0,0 +1,458 @@
+# Copyright 2023 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+# http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Symmetry-compressed electron-repulsion integral (ERI) storage.
+
+This module defines the explicit storage configuration for ERIs and the
+dedicated container used when GradDFT stores PySCF symmetry-compressed
+integrals instead of the dense four-index tensor.
+
+Storage formats
+---------------
+``"dense"``
+ Plain real array of shape ``(n, n, n, n)`` (PySCF ``aosym="s1"``).
+ The default. A dense ERI is always a plain array; it is never wrapped
+ in :class:`PackedERI`, so existing behavior is unchanged.
+``"s4"``
+ Both within-pair symmetries applied, physical shape ``(M, M)`` with
+ ``M = n (n + 1) / 2`` unique index pairs. Both index pairs are packed
+ in PySCF's lower-triangular pair order, ``pair(a, b) = a (a + 1) / 2 + b``
+ for ``a >= b`` (verified against ``pyscf.ao2mo.restore``). Roughly 4x
+ smaller than dense.
+``"s8"``
+ Eightfold symmetry ``(ij|kl) = (kl|ij)`` applied on top of ``s4``,
+ physical shape ``(T,)`` with ``T = M (M + 1) / 2``: the ``(M, M)``
+ ``s4`` matrix packed into its lower triangle, row ``i`` occupying
+ ``[i (i + 1) / 2, i (i + 1) / 2 + i + 1)``. Roughly 8x smaller than
+ dense.
+
+The eightfold relations ``(ij|kl) = (ji|kl) = (ij|lk) = (kl|ij)`` hold for
+real ERIs of isolated molecules. They are *not* assumed for complex or
+periodic integrals: those calculation modes are rejected by the generation
+code (``eri_storage != "dense"`` raises) and keep their dense behavior.
+
+Only PySCF's canonical names ``"s4"``/``"s8"`` are accepted as storage
+configuration; PySCF also accepts the aliases ``"4"``/``"8"`` and the
+partial compressions ``"s2ij"``/``"s2kl"``, which GradDFT deliberately
+rejects with an informative error (see :func:`validate_eri_storage`).
+"""
+
+import numpy as np
+
+from jax import lax
+from jax import numpy as jnp
+from flax import struct
+from jaxtyping import Array, Float
+
+#: Maximum number of gathered elements per ``s8`` row-block during packed
+#: Coulomb contraction (see :func:`packed_coulomb_potential`): bounds the
+#: transient index/value block at :math:`2^{20}` elements (~8 MB in float64)
+#: instead of the full virtual ``(M, M)`` matrix. Static — pass a Python int
+#: if overriding under ``jit``.
+PACKED8_BLOCK_ELEMENTS = 2**20
+
+#: Storage formats accepted by the ``eri_storage`` configuration.
+ERI_STORAGE_FORMATS = ("dense", "s4", "s8")
+
+#: Formats physically held by :class:`PackedERI`. ``"dense"`` is a plain array.
+PACKED_STORAGE_FORMATS = ("s4", "s8")
+
+
+def validate_eri_storage(value: str) -> str:
+ r"""Validate an ``eri_storage`` configuration value.
+
+ Parameters
+ ----------
+ value : str
+ Requested ERI storage format.
+
+ Returns
+ -------
+ str
+ The validated format, unchanged.
+
+ Raises
+ ------
+ ValueError
+ If ``value`` is not one of :data:`ERI_STORAGE_FORMATS`. The message
+ lists the supported formats and names the PySCF options that are
+ intentionally not accepted (``"s2ij"``, ``"s2kl"``, ``"4"``, ``"8"``).
+ """
+ if value not in ERI_STORAGE_FORMATS:
+ raise ValueError(
+ f"Unsupported eri_storage format {value!r}. Supported formats are "
+ f"{list(ERI_STORAGE_FORMATS)}: 'dense' is the full (n, n, n, n) "
+ "tensor, 's4' and 's8' are PySCF symmetry-compressed packings. "
+ "PySCF's partial compressions 's2ij'/'s2kl' and aliases '4'/'8' "
+ "are not accepted; periodic and complex ERIs only support 'dense'."
+ )
+ return value
+
+
+def pair_indices(n_orbitals: int) -> Array:
+ r"""Map both members of an index pair to their unique-pair index.
+
+ Uses PySCF's lower-triangular pair order: for :math:`a \ge b` the pair
+ :math:`(a, b)` gets index :math:`a(a+1)/2 + b`, and the map is
+ symmetric, ``pair(a, b) == pair(b, a)``. This is the packing convention
+ behind ``aosym="s4"``/``"s8"`` (verified against ``pyscf.ao2mo.restore``).
+
+ Parameters
+ ----------
+ n_orbitals : int
+ Number of orbitals :math:`n`.
+
+ Returns
+ -------
+ Int[Array, "orbitals orbitals"]
+ Array of shape ``(n, n)`` with values in ``[0, M)``,
+ ``M = n (n + 1) / 2``. The dtype follows the default JAX integer
+ dtype (``int32`` without x64, ``int64`` with x64).
+ """
+ rows, cols = jnp.tril_indices(n_orbitals)
+ pair = rows * (rows + 1) // 2 + cols
+ index = jnp.zeros((n_orbitals, n_orbitals), dtype=pair.dtype)
+ index = index.at[rows, cols].set(pair).at[cols, rows].set(pair)
+ return index
+
+
+def packed8_to_packed4(packed8: Array, n_pairs: int) -> Array:
+ r"""Expand an ``s8``-packed ERI to the ``s4`` ``(M, M)`` representation.
+
+ Row ``i`` of the ``(M, M)`` matrix occupies
+ ``[i (i + 1) / 2, i (i + 1) / 2 + i + 1)`` in the ``s8`` vector; entries
+ above the diagonal are recovered with the pair-exchange symmetry
+ ``(ij|kl) = (kl|ij)``.
+
+ Parameters
+ ----------
+ packed8 : Array
+ Rank-1 ``s8`` ERI of shape ``(T,)``, ``T = M (M + 1) / 2``.
+ n_pairs : int
+ Number of unique index pairs ``M`` (i.e. ``n (n + 1) / 2`` for
+ ``n`` orbitals).
+
+ Returns
+ -------
+ Float[Array, "pair pair"]
+ The symmetric ``(M, M)`` matrix in ``s4`` pair order.
+ """
+ pair = jnp.arange(n_pairs)
+ start = pair * (pair + 1) // 2 # row start offsets in the s8 vector
+ rows, cols = pair[:, None], pair[None, :]
+ source = jnp.where(cols <= rows, start[:, None] + cols, start[None, :] + rows)
+ return packed8[source]
+
+
+def packed_to_dense(eri: "PackedERI") -> Array:
+ r"""Materialize the dense ``(n, n, n, n)`` ERI from a :class:`PackedERI`.
+
+ This is the explicit compatibility boundary of the compressed
+ representation: it allocates a full :math:`N^4` array (and, for
+ ``s8``, an additional ``(M, M)`` intermediate). It is never called on
+ the training path — contractions consume packed storage directly.
+
+ Parameters
+ ----------
+ eri : PackedERI
+ Compressed ERI (``storage_format`` ``"s4"`` or ``"s8"``).
+
+ Returns
+ -------
+ Float[Array, "orbitals orbitals orbitals orbitals"]
+ Dense ERI with the same dtype as ``eri.packed``.
+
+ Raises
+ ------
+ ValueError
+ If ``eri.storage_format`` is not a known packed format.
+ """
+ if eri.storage_format == "s4":
+ packed4 = eri.packed
+ elif eri.storage_format == "s8":
+ n_pairs = eri.n_orbitals * (eri.n_orbitals + 1) // 2
+ packed4 = packed8_to_packed4(eri.packed, n_pairs)
+ else: # unreachable through the constructor; defensive against crafted treedefs
+ raise ValueError(
+ f"Unknown PackedERI storage_format {eri.storage_format!r}; "
+ f"expected one of {list(PACKED_STORAGE_FORMATS)}."
+ )
+ index = pair_indices(eri.n_orbitals)
+ return packed4[index[:, :, None, None], index[None, None, :, :]]
+
+
+@struct.dataclass
+class PackedERI:
+ r"""Symmetry-compressed electron-repulsion integrals.
+
+ A JAX PyTree holding PySCF symmetry-compressed ERIs in :class:`Molecule`
+ in place of the dense ``(n, n, n, n)`` tensor. The packed numerical
+ array is the single dynamic leaf; ``storage_format`` and ``n_orbitals``
+ are static (auxiliary) metadata, so the format is explicit metadata and
+ never inferred from array shape.
+
+ The *logical* shape is always ``(n, n, n, n)``; the *physical* storage
+ shape is ``packed.shape`` (see module docstring for both).
+
+ Parameters
+ ----------
+ packed : Array or ndarray
+ Physical storage: ``(M, M)`` for ``"s4"`` (``M = n (n + 1) / 2``)
+ or ``(T,)`` for ``"s8"`` (``T = M (M + 1) / 2``).
+ storage_format : str
+ One of :data:`PACKED_STORAGE_FORMATS`.
+ n_orbitals : int
+ Number of orbitals :math:`n` of the logical tensor.
+
+ Raises
+ ------
+ ValueError
+ If the format is unknown, ``n_orbitals`` is not a positive integer,
+ the dtype is not a real floating dtype (e.g. complex ERIs are
+ unsupported), or ``packed``'s shape does not match the format and
+ ``n_orbitals``.
+ """
+
+ packed: Array
+ storage_format: str = struct.field(pytree_node=False)
+ n_orbitals: int = struct.field(pytree_node=False)
+
+ def __post_init__(self):
+ if self.storage_format not in PACKED_STORAGE_FORMATS:
+ raise ValueError(
+ f"Unknown PackedERI storage_format {self.storage_format!r}; "
+ f"expected one of {list(PACKED_STORAGE_FORMATS)}. "
+ "Dense ERIs are stored as plain arrays, not as PackedERI."
+ )
+ if not isinstance(self.n_orbitals, (int, np.integer)) or self.n_orbitals < 1:
+ raise ValueError(
+ f"PackedERI n_orbitals must be a positive integer, got {self.n_orbitals!r}."
+ )
+ dtype = np.dtype(self.packed.dtype)
+ if np.issubdtype(dtype, np.complexfloating):
+ raise ValueError(
+ f"PackedERI requires a real dtype, got {dtype}; complex ERIs "
+ "(e.g. multi-k-point periodic calculations) are unsupported and "
+ "must keep eri_storage='dense'."
+ )
+ if not np.issubdtype(dtype, np.floating):
+ raise ValueError(f"PackedERI requires a floating dtype, got {dtype}.")
+ n = int(self.n_orbitals)
+ n_pairs = n * (n + 1) // 2
+ expected = (
+ (n_pairs, n_pairs)
+ if self.storage_format == "s4"
+ else (n_pairs * (n_pairs + 1) // 2,)
+ )
+ if tuple(self.packed.shape) != expected:
+ raise ValueError(
+ f"storage_format={self.storage_format!r} with n_orbitals={n} "
+ f"requires packed shape {expected}, got {tuple(self.packed.shape)}."
+ )
+
+ @property
+ def shape(self) -> tuple:
+ r"""Logical ERI shape ``(n_orbitals, n_orbitals, n_orbitals, n_orbitals)``.
+
+ The physical storage shape is ``self.packed.shape``.
+ """
+ return (self.n_orbitals,) * 4
+
+ @property
+ def dtype(self):
+ r"""Dtype of the physical packed storage (preserved by conversion)."""
+ return self.packed.dtype
+
+ def to_dense(self) -> Array:
+ r"""Materialize the dense ``(n, n, n, n)`` ERI.
+
+ Explicit :math:`N^4` allocation — never used on the training path.
+ See :func:`packed_to_dense` for details.
+ """
+ return packed_to_dense(self)
+
+
+def fold_density_matrix(
+ rdm1: Float[Array, "orbitals orbitals"], index: Array
+) -> Float[Array, "pair"]:
+ r"""Fold a density matrix onto unique index pairs.
+
+ Computes :math:`\tilde D_{\mathrm{pair}(r,t)} = \sum_{(r,t)} D_{rt}`,
+ i.e. ``D̃[pair(r, t)] = D[r, t] + D[t, r]`` for ``r != t`` and
+ ``D̃[pair(r, r)] = D[r, r]``. This segment-fold makes the subsequent
+ packed contraction exact for an **arbitrary** density matrix — no
+ symmetry of ``rdm1`` is assumed anywhere.
+
+ Parameters
+ ----------
+ rdm1 : Float[Array, "orbitals orbitals"]
+ Density matrix (symmetric or not).
+ index : Int[Array, "orbitals orbitals"]
+ Pair-index map from :func:`pair_indices`.
+
+ Returns
+ -------
+ Float[Array, "pair"]
+ Folded vector of length ``M = n (n + 1) / 2``.
+ """
+ n_pairs = index.shape[0] * (index.shape[0] + 1) // 2
+ folded = jnp.zeros((n_pairs,), dtype=rdm1.dtype)
+ return folded.at[index.reshape(-1)].add(rdm1.reshape(-1))
+
+
+def _packed8_matvec(
+ packed8: Array, vector: Array, n_pairs: int, precision, max_block_elements: int
+) -> Array:
+ r"""Contract ``s8`` storage with a folded vector: ``J̃ = e4 @ D̃``.
+
+ Reads the virtual ``(M, M)`` ``s4`` matrix row-block by row-block without
+ ever allocating it: element ``(i, j)`` lives at ``i (i + 1) / 2 + j`` if
+ ``j <= i`` else ``j (j + 1) / 2 + i``. Each scan step gathers at most
+ ``max_block_elements // M`` rows (at least 1, at most ``M``, so small
+ systems degenerate to a single full block, which is optimal there),
+ bounding transient memory instead of the :math:`M^2` matrix.
+
+ Parameters
+ ----------
+ packed8 : Float[Array, "s8"]
+ ``s8``-packed ERI of length ``M (M + 1) / 2``.
+ vector : Float[Array, "pair"]
+ Folded density vector of length ``M``.
+ n_pairs : int
+ Number of unique pairs ``M`` (static).
+ precision : Precision
+ Einsum precision.
+ max_block_elements : int
+ Static cap on gathered elements per scan step (see
+ :data:`PACKED8_BLOCK_ELEMENTS`).
+
+ Returns
+ -------
+ Float[Array, "pair"]
+ Coupled vector ``J̃`` of length ``M``.
+ """
+ pair = jnp.arange(n_pairs)
+ start = pair * (pair + 1) // 2 # row i occupies [start[i], start[i] + i + 1)
+ block_rows = max(1, min(n_pairs, int(max_block_elements) // max(1, n_pairs)))
+ starts = jnp.arange(0, n_pairs, block_rows)
+
+ def block_coupled(i0):
+ rows = jnp.clip(i0 + jnp.arange(block_rows), 0, n_pairs - 1)
+ cols = pair
+ source = jnp.where(
+ cols[None, :] <= rows[:, None],
+ start[rows][:, None] + cols[None, :],
+ start[cols][None, :] + rows[:, None],
+ )
+ return jnp.einsum("ij,j->i", packed8[source], vector, precision=precision)
+
+ # Sequential scan: only one (block_rows, M) block is alive at a time.
+ _, blocks = lax.scan(lambda _, i0: (None, block_coupled(i0)), None, starts)
+ return blocks.reshape(-1)[:n_pairs]
+
+
+def packed_coulomb_potential(
+ rdm1: Float[Array, "orbitals orbitals"],
+ eri: PackedERI,
+ precision=lax.Precision.HIGHEST,
+ max_block_elements: int = PACKED8_BLOCK_ELEMENTS,
+) -> Float[Array, "orbitals orbitals"]:
+ r"""Coulomb potential matrix ``J_pq = Σ_rt (pq|rt) D_rt`` from packed ERIs.
+
+ Exploits the pair packing directly: with :math:`\tilde D` from
+ :func:`fold_density_matrix` the contraction reduces to a matrix-vector
+ product over unique pairs (:math:`\tilde J = E_4 \tilde D`, ``s4``) and a
+ gather ``J[p, q] = J̃[pair(p, q)]``; the ``s8`` variant reads ``E_4``
+ row-block-wise (see :func:`_packed8_matvec`). The dense
+ ``(n, n, n, n)`` tensor is never materialized (and neither is the
+ ``(M, M)`` ``s4`` matrix for large ``s8`` inputs), and the result equals
+ ``jnp.einsum("pqrt,rt->pq", to_dense(eri), rdm1)`` up to floating-point
+ summation order (observed ≤ ~1e-14) for an arbitrary ``rdm1``.
+
+ Parameters
+ ----------
+ rdm1 : Float[Array, "orbitals orbitals"]
+ The 1-Reduced Density Matrix; no symmetry assumed.
+ eri : PackedERI
+ The symmetry-compressed repulsion tensor (``s4`` or ``s8``).
+ precision : Precision, optional
+ Precision for the contractions, by default Precision.HIGHEST.
+ max_block_elements : int, optional
+ Static cap on gathered elements per ``s8`` scan step, by default
+ :data:`PACKED8_BLOCK_ELEMENTS` (mainly a testing/memory-tuning knob;
+ must be a Python int when passed under ``jit``).
+
+ Returns
+ -------
+ Float[Array, "orbitals orbitals"]
+ The Coulomb potential matrix.
+ """
+ index = pair_indices(eri.n_orbitals)
+ folded = fold_density_matrix(rdm1, index)
+ if eri.storage_format == "s4":
+ coupled = jnp.einsum("ij,j->i", eri.packed, folded, precision=precision)
+ elif eri.storage_format == "s8":
+ n_pairs = eri.n_orbitals * (eri.n_orbitals + 1) // 2
+ coupled = _packed8_matvec(
+ eri.packed, folded, n_pairs, precision, max_block_elements
+ )
+ else: # unreachable through the constructor; defensive against crafted treedefs
+ raise ValueError(
+ f"Unknown PackedERI storage_format {eri.storage_format!r}; "
+ f"expected one of {list(PACKED_STORAGE_FORMATS)}."
+ )
+ return coupled[index]
+
+
+def eri_allclose(a, b, rtol: float = 1e-14, atol: float = 1e-14) -> bool:
+ r"""Compare ERI representations for numerical equivalence.
+
+ Both arguments are either ``None``, a dense ``(n, n, n, n)`` array or a
+ :class:`PackedERI`. Semantics: ``None`` equals only ``None``; mixed
+ ``None``/non-``None`` returns ``False``. Same-format packed comparisons
+ (same ``storage_format`` and ``n_orbitals``) run directly on the packed
+ storage; every other combination (packed vs dense, different packed
+ formats) materializes both sides with :func:`packed_to_dense` first — an
+ explicit :math:`N^4` boundary that is never on the training path.
+
+ Parameters
+ ----------
+ a, b : None or Array or PackedERI
+ ERI representations to compare.
+ rtol, atol : float, optional
+ Relative/absolute tolerances passed to ``jnp.allclose``.
+
+ Returns
+ -------
+ bool
+ ``True`` if numerically equal within tolerance.
+ """
+ if a is None and b is None:
+ return True
+ if (a is None) != (b is None):
+ return False
+ if isinstance(a, PackedERI) and isinstance(b, PackedERI):
+ if a.storage_format != b.storage_format or a.n_orbitals != b.n_orbitals:
+ # Mixed packed formats — compare via dense to avoid shape errors.
+ return bool(
+ jnp.allclose(
+ packed_to_dense(a), packed_to_dense(b), rtol=rtol, atol=atol
+ )
+ )
+ return bool(jnp.allclose(a.packed, b.packed, rtol=rtol, atol=atol))
+ if isinstance(a, PackedERI) or isinstance(b, PackedERI):
+ # Mixed with dense or different packed formats: materialize both via dense
+ a_dense = packed_to_dense(a) if isinstance(a, PackedERI) else jnp.asarray(a)
+ b_dense = packed_to_dense(b) if isinstance(b, PackedERI) else jnp.asarray(b)
+ return bool(jnp.allclose(a_dense, b_dense, rtol=rtol, atol=atol))
+ return bool(jnp.allclose(jnp.asarray(a), jnp.asarray(b), rtol=rtol, atol=atol))
diff --git a/grad_dft/utils/types.py b/grad_dft/utils/types.py
index 2e90b1f..9e9d16f 100644
--- a/grad_dft/utils/types.py
+++ b/grad_dft/utils/types.py
@@ -12,14 +12,17 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-from typing import Union
+"""Common type aliases used across GradDFT."""
+from typing import TYPE_CHECKING, Union
import jax
from jax import numpy as jnp
from flax import linen as nn
-import chex, optax
-import pyscf.dft as dft
-import pyscf.scf as scf
+import chex
+import optax
+from pyscf import dft
+from pyscf import scf
+from jaxtyping import Array as JaxtypingArray, Float
Ansatz = nn.Module
Key = chex.PRNGKey
@@ -32,10 +35,23 @@
DensityFunctional = Union[dft.uks.UKS, dft.rks.RKS, dft.roks.ROKS]
Optimizer = optax.GradientTransformation
-Hartree2kcalmol = 627.50947 # http://www.u.arizona.edu/~stefanb/linkpages/conversions.html
+if TYPE_CHECKING:
+ from .eri import PackedERI # noqa: E402
+
+#: Type accepted by functions consuming electron-repulsion integrals: a dense
+#: ``(orbitals, orbitals, orbitals, orbitals)`` tensor (the default) or
+#: symmetry-compressed storage (see :mod:`grad_dft.utils.eri`).
+ERITensor = Union[
+ Float[JaxtypingArray, "orbitals orbitals orbitals orbitals"], "PackedERI"
+]
+
+Hartree2kcalmol = (
+ 627.50947 # http://www.u.arizona.edu/~stefanb/linkpages/conversions.html
+)
Picometers2Angstroms = 0.01
Bohr2Angstroms = 0.52917721092
def default_dtype() -> DType:
+ r"""Return the default floating-point dtype for GradDFT computations."""
return jnp.float64 if jax.config.x64_enabled else jnp.float32
diff --git a/requirements.txt b/requirements.txt
index 212febf..98f84b2 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,6 +5,7 @@ attrs>=23.1.0
flax>=0.7.2
tensorflow>=2.13.0
tensorflow-hub>=0.16.0
+density_functional_approximation_dm21 @ git+https://github.com/google-deepmind/deepmind-research.git@f5de0ede8430809180254ee957abf36ed62579ef#subdirectory=density_functional_approximation_dm21
typeguard==2.13.3
typing_extensions>=4.8.0
jaxtyping
diff --git a/tests/integration/molecules/test_eri_storage.py b/tests/integration/molecules/test_eri_storage.py
new file mode 100644
index 0000000..d1f51fd
--- /dev/null
+++ b/tests/integration/molecules/test_eri_storage.py
@@ -0,0 +1,259 @@
+# Copyright 2023 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+# http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""End-to-end equivalence of compressed ERI storage (S6 of #15).
+
+On real converged SCF molecules (H2/sto-3g and, larger, H2O/sto-3g) this
+verifies that symmetry-compressed ERIs (s4/s8) are numerically equivalent to
+the dense representation through the full GradDFT pipeline: molecule
+construction, dense expansion, Coulomb contractions, nonXC energies, and a
+training loss with gradients wrt the functional parameters — while the
+default (no ``eri_storage``) behavior stays dense and unchanged.
+"""
+
+# The test functions below take pytest fixtures as arguments.
+# pylint: disable=redefined-outer-name
+
+from jax import config
+
+import jax
+import jax.numpy as jnp
+from jax.nn import gelu, sigmoid
+from jax.random import PRNGKey
+import numpy as np
+import pytest
+from flax import linen as nn
+from pyscf import dft, gto
+
+from grad_dft import (
+ Molecule,
+ NeuralFunctional,
+ molecule_from_pyscf,
+ non_scf_predictor,
+ simple_energy_loss,
+)
+from grad_dft.utils.eri import PackedERI, packed_to_dense
+
+config.update("jax_enable_x64", True)
+
+SCF_ITERS = 200
+PACKED_FORMATS = ("s4", "s8")
+ENERGY_TOL = 1e-9 # contraction summation order differs from the dense einsum
+GRAD_TOL = 1e-8 # gradients accumulate the same 1e-14-level differences
+
+H2_MOL = gto.M(
+ atom="H 0.0 0.0 -0.3705\nH 0.0 0.0 0.3705",
+ basis="sto-3g",
+ verbose=0,
+)
+H2O_MOL = gto.M(
+ atom="O 0.0000 0.0000 0.1173\nH 0.0000 0.7572 -0.4692\nH 0.0000 -0.7572 -0.4692",
+ basis="sto-3g",
+ verbose=0,
+)
+
+
+def converged_mf(mol, density_fit: bool = False):
+ r"""Converge a spin-restricted Kohn-Sham calculation with no XC energy."""
+ mf = dft.RKS(mol)
+ mf.xc = "0.00*LDA" # no XC energy, as in test_non_xc_energy.py
+ mf.max_cycle = SCF_ITERS
+ if density_fit:
+ mf = mf.density_fit()
+ mf.kernel()
+ return mf
+
+
+@pytest.fixture(scope="module")
+def h2_data():
+ r"""H2/sto-3g SCF plus its dense and packed Molecules."""
+ mf = converged_mf(H2_MOL)
+ molecules = {"dense": molecule_from_pyscf(mf)}
+ for fmt in PACKED_FORMATS:
+ molecules[fmt] = molecule_from_pyscf(mf, eri_storage=fmt)
+ return mf, molecules
+
+
+@pytest.fixture(scope="module")
+def h2o_data():
+ r"""H2O/sto-3g SCF plus its dense and packed Molecules."""
+ mf = converged_mf(H2O_MOL)
+ molecules = {"dense": molecule_from_pyscf(mf)}
+ for fmt in PACKED_FORMATS:
+ molecules[fmt] = molecule_from_pyscf(mf, eri_storage=fmt)
+ return mf, molecules
+
+
+@pytest.fixture(scope="module")
+def h2_df_data():
+ r"""Density-fitted H2/sto-3g SCF plus dense and packed Molecules."""
+ mf = converged_mf(H2_MOL, density_fit=True)
+ molecules = {"dense": molecule_from_pyscf(mf)}
+ for fmt in PACKED_FORMATS:
+ molecules[fmt] = molecule_from_pyscf(mf, eri_storage=fmt)
+ return mf, molecules
+
+
+def coefficient_inputs(molecule: Molecule, *_, **__):
+ r"""Single-feature density input (clipped), following test_training.py."""
+ rho = jnp.clip(molecule.density(), min=1e-30)
+ return jnp.concatenate((rho,), axis=1)
+
+
+def energy_densities(molecule: Molecule, *_, **__):
+ r"""LSDA exchange energy density feature, following test_training.py."""
+ rho = jnp.clip(molecule.density(), min=1e-30)
+ return (
+ -3
+ / 2
+ * (3 / (4 * jnp.pi)) ** (1 / 3)
+ * (rho ** (4 / 3)).sum(axis=1, keepdims=True)
+ )
+
+
+def coefficients(_instance, rhoinputs):
+ r"""A minimal parametric energy-density head (one trainable Dense layer)."""
+ x = nn.Dense(features=1)(rhoinputs)
+ x = nn.LayerNorm()(x)
+ return sigmoid(gelu(x))
+
+
+def make_functional() -> NeuralFunctional:
+ r"""Tiny neural functional with trainable parameters."""
+ return NeuralFunctional(coefficients, energy_densities, coefficient_inputs)
+
+
+def test_default_eri_storage_stays_dense(h2_data, h2o_data):
+ r"""Without eri_storage, rep_tensor is the unchanged dense (n, n, n, n) array."""
+ for mf, molecules in (h2_data, h2o_data):
+ dense = molecules["dense"].rep_tensor
+ n_orb = mf.mol.nao_nr()
+ assert isinstance(dense, jnp.ndarray)
+ assert not isinstance(dense, PackedERI)
+ assert dense.shape == (n_orb, n_orb, n_orb, n_orb)
+ np.testing.assert_allclose(
+ dense, mf.mol.intor("int2e", aosym="s1"), atol=1e-14, rtol=0
+ )
+
+
+@pytest.mark.parametrize("fmt", PACKED_FORMATS)
+def test_packed_storage_metadata_and_values(h2_data, h2o_data, fmt):
+ r"""Packed molecules carry correct metadata and expand to the PySCF ERI."""
+ for mf, molecules in (h2_data, h2o_data):
+ packed = molecules[fmt].rep_tensor
+ assert isinstance(packed, PackedERI)
+ assert packed.storage_format == fmt
+ assert packed.n_orbitals == mf.mol.nao_nr()
+ reference = jnp.asarray(mf.mol.intor("int2e", aosym="s1"))
+ np.testing.assert_allclose(
+ packed_to_dense(packed), reference, atol=1e-14, rtol=0
+ )
+
+
+@pytest.mark.parametrize("fmt", PACKED_FORMATS)
+def test_nonxc_and_coulomb_match_dense(h2_data, h2o_data, fmt):
+ r"""nonXC energies and Coulomb potentials agree with dense storage."""
+ for _, molecules in (h2_data, h2o_data):
+ dense, packed = molecules["dense"], molecules[fmt]
+ np.testing.assert_allclose(
+ dense.nonXC(), packed.nonXC(), atol=ENERGY_TOL, rtol=1e-10
+ )
+ np.testing.assert_allclose(
+ dense.get_coulomb_potential(),
+ packed.get_coulomb_potential(),
+ atol=ENERGY_TOL,
+ rtol=1e-10,
+ )
+
+
+@pytest.mark.parametrize("fmt", PACKED_FORMATS)
+def test_energy_loss_and_grad_match_dense(h2_data, fmt):
+ r"""A training loss and its parameter gradients are storage-independent.
+
+ `simple_energy_loss` is itself value_and_grad-wrapped (has_aux=True) and
+ returns `((loss, E_predict), grads_wrt_params)`. The loss includes the
+ nonXC (ERI) energy of each storage, while the parameter gradients cover
+ the neural-functional path of the same computation.
+ """
+ mf, molecules = h2_data
+ functional = make_functional()
+ params = functional.init(PRNGKey(42), coefficient_inputs(molecules["dense"]))
+ predictor = non_scf_predictor(functional)
+ truth_energy = mf.e_tot
+
+ def loss_and_grads(params, molecule):
+ (loss, _), grads = simple_energy_loss(params, predictor, molecule, truth_energy)
+ return loss, grads
+
+ dense_loss, dense_grads = loss_and_grads(params, molecules["dense"])
+ packed_loss, packed_grads = loss_and_grads(params, molecules[fmt])
+
+ assert not jnp.isnan(dense_loss) and not jnp.isnan(packed_loss)
+ np.testing.assert_allclose(packed_loss, dense_loss, atol=GRAD_TOL, rtol=1e-10)
+
+ flat_dense, _ = jax.tree_util.tree_flatten(dense_grads)
+ flat_packed, _ = jax.tree_util.tree_flatten(packed_grads)
+ for grad_dense, grad_packed in zip(flat_dense, flat_packed):
+ np.testing.assert_allclose(grad_packed, grad_dense, atol=GRAD_TOL, rtol=1e-8)
+ # guard against a vacuous all-zero comparison
+ assert any(float(jnp.abs(g).max()) > 1e-10 for g in flat_dense)
+
+
+@pytest.mark.parametrize("fmt", PACKED_FORMATS)
+def test_nonxc_rdm1_gradient_matches_dense(h2_data, h2o_data, fmt):
+ r"""d(nonXC)/d(rdm1) is storage-independent: autodiff through the packed ops.
+
+ Unlike the parameter gradients above, this gradient flows *through* the
+ compressed Coulomb contractions (the nonXC energy depends on rdm1), so it
+ exercises backpropagation of the packed8/packed4 matmuls directly.
+ """
+ for _, molecules in (h2_data, h2o_data):
+ dense, packed = molecules["dense"], molecules[fmt]
+
+ def grad_wrt_rdm1(molecule):
+ return jax.grad(lambda r: molecule.replace(rdm1=r).nonXC())(molecule.rdm1)
+
+ grad_dense = grad_wrt_rdm1(dense)
+ grad_packed = grad_wrt_rdm1(packed)
+ assert not jnp.isnan(grad_packed).any()
+ assert float(jnp.abs(grad_dense).max()) > 1e-10 # non-trivial
+ np.testing.assert_allclose(grad_packed, grad_dense, atol=GRAD_TOL, rtol=1e-9)
+
+
+def test_density_fitted_consistent_across_formats(h2_df_data):
+ r"""Density-fitted ERIs are consistent across dense/s4/s8 storage."""
+ _, molecules = h2_df_data
+ reference = np.asarray(molecules["dense"].rep_tensor)
+ for fmt in PACKED_FORMATS:
+ packed = molecules[fmt].rep_tensor
+ assert isinstance(packed, PackedERI)
+ np.testing.assert_allclose(
+ packed_to_dense(packed), reference, atol=1e-14, rtol=0
+ )
+ np.testing.assert_allclose(
+ molecules["dense"].nonXC(),
+ molecules[fmt].nonXC(),
+ atol=ENERGY_TOL,
+ rtol=1e-10,
+ )
+
+
+def test_rep_tensor_none_molecule_path(h2_data):
+ r"""A molecule without ERIs still supports its non-ERI methods."""
+ _, molecules = h2_data
+ no_eri = molecules["dense"].replace(rep_tensor=None)
+ density = no_eri.density()
+ grad_density = no_eri.grad_density()
+ assert not jnp.isnan(density).any()
+ assert not jnp.isnan(grad_density).any()
diff --git a/tests/integration/molecules/test_predict_DM21.py b/tests/integration/molecules/test_predict_DM21.py
index c237bf8..68909d7 100644
--- a/tests/integration/molecules/test_predict_DM21.py
+++ b/tests/integration/molecules/test_predict_DM21.py
@@ -177,7 +177,7 @@ def test_uks():
mol.basis = "sto-3g"
mol.build()
- mf = dft.UKS(mol, grid)
+ mf = dft.UKS(mol)
mf.grids.level = 3
mf.small_rho_cutoff = 1.0e-20
mf._numint = ni
diff --git a/tests/unit/test_eri_consumers.py b/tests/unit/test_eri_consumers.py
new file mode 100644
index 0000000..61ed1df
--- /dev/null
+++ b/tests/unit/test_eri_consumers.py
@@ -0,0 +1,179 @@
+# Copyright 2023 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+# http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for remaining consumers and metadata (S5 of #15).
+
+Covers:
+
+- `eri_allclose`: None/None, None/mixed, same-format packed equal/unequal,
+ mixed representations (packed-vs-packed of different formats, packed-vs-
+ dense) compared via dense materialization, dense pair.
+- `constraints.xc1`: rep_tensor comparison via `eri_allclose`; the full
+ constraint loss is identical for packed (mixed s8/s4) and dense storage.
+- `Molecule`: `rep_tensor=None` default construction and non-ERI methods
+ (`density`, `grad_density`) work.
+"""
+
+from jax import config
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from pyscf import gto
+
+from grad_dft import LSDA, Grid, Molecule
+from grad_dft.constraints import xc1
+from grad_dft.utils.eri import PackedERI, eri_allclose, packed_to_dense as to_dense
+
+config.update("jax_enable_x64", True)
+
+MOL = gto.M(
+ atom=[["H", (0, 0, 0)], ["H", (0.74, 0, 0)]],
+ basis="cc-pvdz",
+ verbose=0,
+)
+N = MOL.nao_nr()
+REP_DENSE = jnp.asarray(MOL.intor("int2e", aosym="s1"))
+
+
+def packed_eri(fmt: str) -> PackedERI:
+ """Build PackedERI storage of the test molecule's integrals."""
+ return PackedERI(np.asarray(MOL.intor("int2e", aosym=fmt)), fmt, N)
+
+
+def test_eri_allclose_none_pairs():
+ """None equals only None; mixed None/non-None is False."""
+ assert eri_allclose(None, None) is True
+ assert eri_allclose(REP_DENSE, None) is False
+ assert eri_allclose(None, REP_DENSE) is False
+
+
+def test_eri_allclose_same_format_equal_unequal():
+ """Same-format packed comparisons run on the packed storage directly."""
+ p4 = packed_eri("s4")
+ p8 = packed_eri("s8")
+ # equal
+ assert eri_allclose(p4, packed_eri("s4")) is True
+ assert eri_allclose(p8, packed_eri("s8")) is True
+ # unequal
+ p4_bad = PackedERI(
+ jnp.array(p4.packed).at[0, 0].set(jnp.array(p4.packed)[0, 0] + 1.0), "s4", N
+ )
+ assert eri_allclose(p4, p4_bad) is False
+
+
+def test_eri_allclose_mixed_formats_via_todense():
+ """Mixed representations (packed-vs-packed, packed-vs-dense) compare equal."""
+ np.testing.assert_allclose(
+ to_dense(packed_eri("s4")), REP_DENSE, atol=1e-14, rtol=0
+ )
+ assert eri_allclose(packed_eri("s4"), packed_eri("s8")) is True
+ assert eri_allclose(packed_eri("s4"), REP_DENSE) is True
+ assert eri_allclose(REP_DENSE, packed_eri("s8")) is True
+
+
+def test_eri_allclose_dense_pair():
+ """Dense arrays compare elementwise within tolerance."""
+ rep2 = jnp.array(REP_DENSE)
+ assert eri_allclose(rep2, REP_DENSE) is True
+ rep_bad = rep2.at[0, 0, 0, 0].set(rep2[0, 0, 0, 0] + 2.0)
+ assert eri_allclose(rep_bad, REP_DENSE) is False
+
+
+def test_eri_allclose_mismatched_packed_formats_compare_via_dense():
+ """Different packed formats holding equal values compare equal (via dense)."""
+ s4 = packed_eri("s4")
+ s8 = packed_eri("s8")
+ assert eri_allclose(s4, s8) is True
+
+
+def make_minimal_molecule(rep_tensor, rdm_seed=17):
+ """Synthetic Molecule with fields as produced by molecule_from_pyscf.
+
+ The grid/ao rng is fixed so that molecules built with different
+ ``rdm_seed`` differ only in their density matrix (xc1 compares ao).
+ """
+ grid_size = 6
+ rng = np.random.default_rng(17) # fixed: xc1 compares ao/grad_ao across molecules
+ grid = Grid(
+ coords=jnp.zeros((grid_size, 3)), weights=jnp.ones((grid_size,)) / grid_size
+ )
+ rdm_rng = np.random.default_rng(rdm_seed)
+ rdm1 = np.eye(N) + 0.2 * rdm_rng.standard_normal((N, N))
+ rdm1 = jnp.asarray(0.5 * (rdm1 + rdm1.T))
+ spin_rdm1 = jnp.stack([0.5 * rdm1, 0.5 * rdm1])
+ return Molecule(
+ grid=grid,
+ atom_index=jnp.array([0, 0]),
+ nuclear_pos=jnp.zeros((2, 3)),
+ ao=jnp.asarray(rng.standard_normal((grid_size, N))),
+ grad_ao=jnp.asarray(rng.standard_normal((grid_size, N, 3))),
+ grad_n_ao={
+ 0: jnp.zeros((grid_size, N)),
+ 1: jnp.zeros((grid_size, N)),
+ 2: jnp.zeros((grid_size, N)),
+ },
+ rdm1=spin_rdm1,
+ nuclear_repulsion=jnp.asarray(1.0),
+ h1e=jnp.eye(N),
+ vj=jnp.zeros((2, N, N)),
+ mo_coeff=jnp.zeros((2, N, N)),
+ mo_occ=jnp.zeros((2, N)),
+ mo_energy=jnp.zeros((2, N)),
+ # fields xc1 compares unguarded (as produced by molecule_from_pyscf):
+ s1e=jnp.eye(N),
+ omegas=jnp.zeros((1,)),
+ basis=jnp.array([ord("H"), ord("2")]),
+ rep_tensor=rep_tensor,
+ name=jnp.array([ord("H"), ord("2")]),
+ )
+
+
+def test_rep_tensor_none_methods_work():
+ """Non-ERI methods (density, grad_density) work with rep_tensor=None."""
+ mol = make_minimal_molecule(rep_tensor=None)
+ # non-ERI methods should work
+ _ = mol.density()
+ _ = mol.grad_density()
+ # jit also works
+ _ = jax.jit(lambda m: (m.density(), m.grad_density()))(mol)
+
+
+def test_xc1_with_packed_molecules():
+ """xc1's constraint loss is identical for packed and dense ERI storage.
+
+ xc1 evaluates the full energy predictor (functional + coulomb path, i.e.
+ the S2 hot dispatch) on both molecules, so this exercises the consumer
+ end-to-end: mixed s8/s4 storage must give the same loss as all-dense.
+ """
+ gamma = 0.3
+ packed_loss = xc1(
+ LSDA,
+ {},
+ make_minimal_molecule(packed_eri("s8")),
+ make_minimal_molecule(packed_eri("s4"), rdm_seed=19),
+ gamma,
+ mol=MOL,
+ )
+ dense_loss = xc1(
+ LSDA,
+ {},
+ make_minimal_molecule(REP_DENSE),
+ make_minimal_molecule(REP_DENSE, rdm_seed=19),
+ gamma,
+ mol=MOL,
+ )
+ assert jnp.isfinite(packed_loss)
+ assert float(packed_loss) > 0 # non-trivial: the two molecules differ
+ np.testing.assert_allclose(packed_loss, dense_loss, atol=1e-10, rtol=1e-10)
diff --git a/tests/unit/test_eri_contraction.py b/tests/unit/test_eri_contraction.py
new file mode 100644
index 0000000..c8f3b2f
--- /dev/null
+++ b/tests/unit/test_eri_contraction.py
@@ -0,0 +1,226 @@
+# Copyright 2023 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+# http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the compressed Coulomb contractions (grad_dft.utils.eri).
+
+Covers, at the contraction level:
+
+(C) numerical equivalence of dense vs ``s4`` vs ``s8`` contractions for
+ symmetric *and asymmetric* density matrices (an asymmetric D proves no
+ density symmetry is assumed);
+(D) jit/eager agreement, gradients wrt ``rdm1`` matching the dense path,
+ and jaxpr inspection showing the packed path never materializes a
+ ``(n, n, n, n)`` tensor;
+(F) the ``s8`` row-chunked scan (forced small blocks) equals the default
+ single-block contraction.
+
+H2O/cc-pvdz is used rather than H2: its ERI has few exact-zero blocks, so
+pair-indexing mistakes cannot be masked by zeros.
+"""
+
+from jax import config
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+from pyscf import gto
+
+from grad_dft.molecule import Grid, Molecule, coulomb_energy, coulomb_potential
+from grad_dft.molecule import nonXC
+from grad_dft.utils.eri import PackedERI, packed_coulomb_potential
+
+config.update("jax_enable_x64", True)
+
+MOL = gto.M(
+ atom=[["O", (0, 0, 0)], ["H", (0, 0, 0.97)], ["H", (0.76, 0, -0.5)]],
+ basis="cc-pvdz",
+ verbose=0,
+)
+N = MOL.nao_nr()
+REP_DENSE = jnp.asarray(MOL.intor("int2e", aosym="s1"))
+PACKED_FORMATS = ("s4", "s8")
+RDM_KINDS = ("symmetric", "asymmetric")
+
+
+def packed_eri(fmt: str) -> PackedERI:
+ """Build PackedERI storage of the test molecule's integrals."""
+ return PackedERI(jnp.asarray(MOL.intor("int2e", aosym=fmt)), fmt, N)
+
+
+def make_rdm1(kind: str) -> jnp.ndarray:
+ """Random (n, n) density matrix; 'symmetric' enforces D = D^T."""
+ rng = np.random.default_rng(7 if kind == "symmetric" else 8)
+ rdm1 = rng.standard_normal((N, N))
+ if kind == "symmetric":
+ rdm1 = 0.5 * (rdm1 + rdm1.T)
+ return jnp.asarray(rdm1)
+
+
+def make_molecule(rep_tensor, rdm1: jnp.ndarray) -> Molecule:
+ """Minimal synthetic Molecule whose nonXC/Coulomb paths are exercised."""
+ grid_size = 8
+ rng = np.random.default_rng(11)
+ grid = Grid(
+ coords=jnp.zeros((grid_size, 3)),
+ weights=jnp.full((grid_size,), 1.0 / grid_size),
+ )
+ spin_rdm1 = jnp.stack([0.6 * rdm1, 0.4 * rdm1])
+ return Molecule(
+ grid=grid,
+ atom_index=jnp.arange(3),
+ nuclear_pos=jnp.zeros((3, 3)),
+ ao=jnp.asarray(rng.standard_normal((grid_size, N))),
+ grad_ao=jnp.asarray(rng.standard_normal((grid_size, N, 3))),
+ grad_n_ao=jnp.zeros((3, grid_size, N)),
+ rdm1=spin_rdm1,
+ nuclear_repulsion=jnp.asarray(0.89),
+ h1e=jnp.asarray(rng.standard_normal((N, N))),
+ vj=jnp.zeros((2, N, N)),
+ mo_coeff=jnp.zeros((2, N, N)),
+ mo_occ=jnp.zeros((2, N)),
+ mo_energy=jnp.zeros((2, N)),
+ rep_tensor=rep_tensor,
+ )
+
+
+def collect_jaxpr_shapes(jaxpr, shapes) -> None:
+ """Recursively collect variable shapes, including inside nested jaxprs."""
+ if hasattr(jaxpr, "eqns") and hasattr(jaxpr, "invars"):
+ for var in list(jaxpr.invars) + list(jaxpr.outvars) + list(jaxpr.constvars):
+ shapes.append(tuple(var.aval.shape))
+ for eqn in jaxpr.eqns:
+ for var in list(eqn.invars) + list(eqn.outvars):
+ shapes.append(tuple(var.aval.shape))
+ for param in eqn.params.values():
+ if isinstance(param, (tuple, list)):
+ for item in param:
+ collect_jaxpr_shapes(item, shapes)
+ else:
+ collect_jaxpr_shapes(param, shapes)
+ elif isinstance(jaxpr, (tuple, list)):
+ for item in jaxpr:
+ collect_jaxpr_shapes(item, shapes)
+
+
+@pytest.mark.parametrize("fmt", PACKED_FORMATS)
+@pytest.mark.parametrize("kind", RDM_KINDS)
+def test_coulomb_potential_matches_dense(fmt, kind):
+ """Packed J equals the dense einsum up to summation order (atol 1e-11)."""
+ rdm1 = make_rdm1(kind)
+ expected = coulomb_potential(rdm1, REP_DENSE)
+ actual = coulomb_potential(rdm1, packed_eri(fmt))
+ assert actual.dtype == expected.dtype == np.float64
+ np.testing.assert_allclose(actual, expected, rtol=1e-10, atol=1e-11)
+
+
+@pytest.mark.parametrize("fmt", PACKED_FORMATS)
+@pytest.mark.parametrize("kind", RDM_KINDS)
+def test_coulomb_energy_matches_dense(fmt, kind):
+ """Packed Coulomb energy equals the dense one for both density kinds."""
+ rdm1 = make_rdm1(kind)
+ expected = coulomb_energy(rdm1, REP_DENSE)
+ actual = coulomb_energy(rdm1, packed_eri(fmt))
+ np.testing.assert_allclose(actual, expected, rtol=1e-10, atol=1e-11)
+
+
+@pytest.mark.parametrize("fmt", PACKED_FORMATS)
+def test_non_xc_matches_dense(fmt):
+ """nonXC is identical between dense and packed rep_tensor."""
+ rdm1 = make_rdm1("asymmetric")
+ h1e = jnp.asarray(np.random.default_rng(5).standard_normal((N, N)))
+ expected = nonXC(rdm1, h1e, REP_DENSE, jnp.asarray(0.89))
+ actual = nonXC(rdm1, h1e, packed_eri(fmt), jnp.asarray(0.89))
+ np.testing.assert_allclose(actual, expected, rtol=1e-10, atol=1e-11)
+
+
+@pytest.mark.parametrize("fmt", PACKED_FORMATS)
+def test_molecule_methods_match_dense(fmt):
+ """Molecule.get_coulomb_potential and Molecule.nonXC dispatch correctly."""
+ rdm1 = make_rdm1("symmetric")
+ dense_mol = make_molecule(REP_DENSE, rdm1)
+ packed_mol = make_molecule(packed_eri(fmt), rdm1)
+ np.testing.assert_allclose(
+ packed_mol.get_coulomb_potential(),
+ dense_mol.get_coulomb_potential(),
+ rtol=1e-10,
+ atol=1e-11,
+ )
+ np.testing.assert_allclose(
+ packed_mol.nonXC(), dense_mol.nonXC(), rtol=1e-10, atol=1e-11
+ )
+
+
+@pytest.mark.parametrize("fmt", PACKED_FORMATS)
+def test_packed_contraction_jit_matches_eager(fmt):
+ """The packed contraction is pure: jit output equals the eager output."""
+ rdm1 = make_rdm1("asymmetric")
+ eri = packed_eri(fmt)
+ eager = packed_coulomb_potential(rdm1, eri)
+ jitted = jax.jit(
+ packed_coulomb_potential, static_argnames=("precision", "max_block_elements")
+ )(rdm1, eri)
+ np.testing.assert_allclose(jitted, eager, rtol=1e-12, atol=1e-12)
+
+
+@pytest.mark.parametrize("fmt", PACKED_FORMATS)
+def test_coulomb_energy_grad_matches_dense(fmt):
+ """d Coulomb energy / d rdm1 matches the dense path through the packed ops."""
+ rdm1 = make_rdm1("asymmetric")
+ grad_fn = jax.grad(coulomb_energy)
+ expected = grad_fn(rdm1, REP_DENSE)
+ actual = grad_fn(rdm1, packed_eri(fmt))
+ np.testing.assert_allclose(actual, expected, rtol=1e-9, atol=1e-10)
+
+
+@pytest.mark.parametrize("fmt", PACKED_FORMATS)
+def test_jaxpr_never_materializes_dense(fmt):
+ """Packed trace contains no (n, n, n, n) aval and no N^4-sized buffer.
+
+ The dense trace is checked in the opposite direction as a sanity check.
+ (Benign rank-3 (n, n, 1) reshapes from JAX's gather lowering are fine —
+ what must never appear is an N^4 allocation.)
+ """
+ rdm1 = make_rdm1("asymmetric")
+
+ def trace(rep):
+ return jax.make_jaxpr(coulomb_potential)(rdm1, rep)
+
+ packed_shapes = []
+ collect_jaxpr_shapes(trace(packed_eri(fmt)).jaxpr, packed_shapes)
+ assert (N, N, N, N) not in packed_shapes
+ assert max(np.prod(shape) for shape in packed_shapes) < N**4
+
+ dense_shapes = []
+ collect_jaxpr_shapes(trace(REP_DENSE).jaxpr, dense_shapes)
+ assert (N, N, N, N) in dense_shapes
+
+
+@pytest.mark.parametrize("block_rows", [1, 16])
+def test_s8_chunked_scan_matches_single_block(block_rows):
+ """Forcing tiny scan blocks (multi-chunk path) keeps the same result.
+
+ For M = 300 pair rows the default cap yields a single full block; small
+ `max_block_elements` exercise the multi-block `lax.scan` branch that
+ large systems use, proving chunk boundaries are value-preserving.
+ """
+ rdm1 = make_rdm1("asymmetric")
+ eri = packed_eri("s8")
+ n_pairs = N * (N + 1) // 2
+ single_block = packed_coulomb_potential(rdm1, eri)
+ chunked = packed_coulomb_potential(
+ rdm1, eri, max_block_elements=block_rows * n_pairs
+ )
+ assert block_rows < n_pairs # the multi-block branch is actually forced
+ np.testing.assert_allclose(chunked, single_block, rtol=1e-12, atol=1e-12)
diff --git a/tests/unit/test_eri_generation.py b/tests/unit/test_eri_generation.py
new file mode 100644
index 0000000..fe23ee6
--- /dev/null
+++ b/tests/unit/test_eri_generation.py
@@ -0,0 +1,175 @@
+# Copyright 2023 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+# http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the PySCF-side generation of compressed ERIs (S3 of #15).
+
+Covers:
+
+(A) `calc_eri_with_pyscf` isolated molecules: packed formats match the PySCF
+ reference packings bitwise (`intor(aosym=...)`) / to machine precision
+ for the density-fitted branch (fresh DF references);
+(A/B) the dense path is behaviorally unchanged (isolated and periodic);
+(F) error handling: invalid formats fail before any integral work, periodic
+ systems reject compression in both `calc_eri_with_pyscf` and
+ `molecule_from_pyscf` (fail-fast, no SCF required);
+(D) `to_device_arrays` preserves PackedERI metadata and converts its leaf;
+(D) end-to-end `molecule_from_pyscf(eri_storage=...)` pass-through.
+"""
+
+from jax import config
+
+import jax.numpy as jnp
+import numpy as np
+import pytest
+from pyscf import df, gto, scf
+from pyscf.ao2mo import restore
+from pyscf.pbc import gto as pbc_gto
+from pyscf.pbc import scf as pbc_scf
+
+from grad_dft.interface.pyscf import (
+ calc_eri_with_pyscf,
+ molecule_from_pyscf,
+ process_mol,
+ to_device_arrays,
+)
+from grad_dft.utils.eri import PackedERI
+
+config.update("jax_enable_x64", True)
+
+PACKED_FORMATS = ("s4", "s8")
+
+MOL = gto.M(
+ atom=[["H", (0, 0, 0)], ["H", (0.74, 0, 0)]],
+ basis="cc-pvdz",
+ verbose=0,
+)
+N = MOL.nao_nr()
+
+CELL = pbc_gto.Cell()
+CELL.atom = "H 0 0 0; H 0.74 0 0"
+CELL.a = np.eye(3) * 4.0
+CELL.basis = "gth-szv"
+CELL.verbose = 0
+CELL.build()
+
+
+@pytest.mark.parametrize("fmt", PACKED_FORMATS)
+def test_calc_eri_isolated_matches_intor(fmt):
+ """Isolated no-DF storage equals PySCF's own intor(aosym=fmt) bitwise."""
+ mf = scf.RHF(MOL)
+ eri = calc_eri_with_pyscf(mf, eri_storage=fmt)
+ assert isinstance(eri, PackedERI)
+ assert eri.storage_format == fmt
+ assert eri.n_orbitals == N
+ assert eri.packed.dtype == np.float64
+ np.testing.assert_array_equal(np.asarray(eri.packed), MOL.intor("int2e", aosym=fmt))
+
+
+def test_calc_eri_isolated_dense_unchanged():
+ """The default dense branch still returns the uncompressed intor output."""
+ mf = scf.RHF(MOL)
+ eri = calc_eri_with_pyscf(mf)
+ assert isinstance(eri, np.ndarray)
+ np.testing.assert_array_equal(eri, MOL.intor("int2e"))
+
+
+@pytest.mark.parametrize("fmt", ("dense",) + PACKED_FORMATS)
+def test_calc_eri_density_fitted(fmt):
+ """DF branch: values come from DF.get_eri() (always s8 storage).
+
+ Packed outputs are never compared against exact integrals — density
+ fitting approximates those by construction.
+ """
+ mf = scf.RHF(MOL).density_fit()
+ n = MOL.nao_nr()
+ eri_ref = df.DF(MOL).get_eri() # fresh reference, same construction as the function
+ eri = calc_eri_with_pyscf(mf, eri_storage=fmt)
+ if fmt == "dense":
+ assert isinstance(eri, np.ndarray)
+ np.testing.assert_allclose(eri, restore(1, eri_ref, n), atol=1e-14, rtol=0)
+ else:
+ assert isinstance(eri, PackedERI)
+ assert eri.storage_format == fmt
+ expected = restore(4, eri_ref, n) if fmt == "s4" else eri_ref
+ np.testing.assert_allclose(np.asarray(eri.packed), expected, atol=1e-14, rtol=0)
+
+
+def test_calc_eri_density_fitted_formats_consistent():
+ """The two packed DF formats reconstruct the same dense tensor."""
+ s4 = calc_eri_with_pyscf(scf.RHF(MOL).density_fit(), eri_storage="s4")
+ s8 = calc_eri_with_pyscf(scf.RHF(MOL).density_fit(), eri_storage="s8")
+ np.testing.assert_allclose(s4.to_dense(), s8.to_dense(), atol=1e-14, rtol=0)
+ np.testing.assert_allclose(
+ s4.to_dense(), restore(1, df.DF(MOL).get_eri(), N), atol=1e-14, rtol=0
+ )
+
+
+@pytest.mark.parametrize("bad", ["s2", "s2ij", "4", "8", "S8", ""])
+def test_invalid_format_fails_before_integral_work(bad):
+ """Invalid formats raise on mf=None: validation precedes any mf access."""
+ with pytest.raises(ValueError, match="Supported formats are"):
+ calc_eri_with_pyscf(None, eri_storage=bad)
+ with pytest.raises(ValueError, match="Supported formats are"):
+ molecule_from_pyscf(None, eri_storage=bad)
+
+
+@pytest.mark.parametrize("fmt", PACKED_FORMATS)
+def test_periodic_rejects_compression(fmt):
+ """Periodic systems reject compressed storage in both entry points."""
+ kmf = pbc_scf.RHF(CELL)
+ with pytest.raises(ValueError, match="Periodic"):
+ calc_eri_with_pyscf(kmf, eri_storage=fmt)
+ with pytest.raises(ValueError, match="eri_storage"):
+ molecule_from_pyscf(kmf, eri_storage=fmt)
+
+
+def test_periodic_dense_unchanged():
+ """The dense periodic (gamma-point) path still returns a dense array."""
+ kmf = pbc_scf.RHF(CELL)
+ eri = calc_eri_with_pyscf(kmf)
+ assert isinstance(eri, np.ndarray)
+ assert eri.shape == (CELL.nao_nr(),) * 4
+
+
+def test_to_device_arrays_preserves_packed():
+ """Packed metadata survives device conversion; the leaf becomes jnp."""
+ eri_np = MOL.intor("int2e", aosym="s8")
+ out = to_device_arrays(PackedERI(eri_np, "s8", N))[0]
+ assert isinstance(out, PackedERI)
+ assert out.storage_format == "s8"
+ assert out.n_orbitals == N
+ assert isinstance(out.packed, jnp.ndarray)
+ assert out.packed.dtype == jnp.float64
+ np.testing.assert_array_equal(np.asarray(out.packed), eri_np)
+
+
+def test_molecule_from_pyscf_pass_through():
+ """eri_storage flows through _package_outputs into Molecule.rep_tensor."""
+ _, mf = process_mol(MOL, compute_energy=False, grid_level=1)
+ mol_dense = molecule_from_pyscf(mf)
+ mol_packed = molecule_from_pyscf(mf, eri_storage="s8")
+
+ rep_dense = np.asarray(mol_dense.rep_tensor)
+ rep_packed = mol_packed.rep_tensor
+ assert isinstance(rep_packed, PackedERI)
+ assert rep_packed.storage_format == "s8"
+ assert rep_packed.packed.dtype == jnp.float64
+ np.testing.assert_array_equal(rep_dense, MOL.intor("int2e"))
+ np.testing.assert_allclose(
+ np.asarray(rep_packed.to_dense()), MOL.intor("int2e"), atol=1e-14, rtol=0
+ )
+ # The end-to-end non-XC energy agrees between the two molecules.
+ np.testing.assert_allclose(
+ mol_packed.nonXC(), mol_dense.nonXC(), atol=1e-11, rtol=1e-10
+ )
diff --git a/tests/unit/test_eri_serialization.py b/tests/unit/test_eri_serialization.py
new file mode 100644
index 0000000..2fceb03
--- /dev/null
+++ b/tests/unit/test_eri_serialization.py
@@ -0,0 +1,239 @@
+# Copyright 2023 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+# http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for serialization of compressed ERIs (S4 of #15).
+
+Covers:
+
+(B/D/E) `saver`/`loader` round trips with a synthetic Molecule (no SCF):
+ dense layout unchanged (no attributes, old files load), packed s4/s8
+ preserves storage_format/n_orbitals/values/dtype, `nonXC` agrees with
+ the original, and `loader(training=True)` still reconstructs packed
+ molecules (metadata rides on dataset attrs, which survive the
+ training-mode group-attr skip);
+(B) the reaction dataset loop of `loader` as well as the molecule loop;
+(F) a corrupt `storage_format` attribute raises an informative ValueError;
+(E) an old-style dense file (plain dataset, no attributes) loads as dense.
+"""
+
+from jax import config
+
+import h5py
+import jax.numpy as jnp
+import numpy as np
+import pytest
+from pyscf import gto
+
+from grad_dft.interface.pyscf import loader, saver
+from grad_dft.molecule import Grid, Molecule, make_reaction
+from grad_dft.utils.eri import PackedERI
+
+config.update("jax_enable_x64", True)
+
+PACKED_FORMATS = ("s4", "s8")
+
+MOL = gto.M(
+ atom=[["H", (0, 0, 0)], ["H", (0.74, 0, 0)]],
+ basis="cc-pvdz",
+ verbose=0,
+)
+N = MOL.nao_nr()
+REP_DENSE = jnp.asarray(MOL.intor("int2e", aosym="s1"))
+
+
+def packed_eri(fmt: str) -> PackedERI:
+ """Build PackedERI storage of the test molecule's integrals (numpy leaf)."""
+ return PackedERI(np.asarray(MOL.intor("int2e", aosym=fmt)), fmt, N)
+
+
+def make_molecule(rep_tensor) -> Molecule:
+ """Minimal synthetic Molecule with real ERIs, ready for save/load."""
+ grid_size = 8
+ rng = np.random.default_rng(13)
+ grid = Grid(
+ coords=jnp.zeros((grid_size, 3)),
+ weights=jnp.full((grid_size,), 1.0 / grid_size),
+ )
+ rdm1 = jnp.asarray(rng.standard_normal((N, N)))
+ rdm1 = 0.5 * (rdm1 + rdm1.T)
+ spin_rdm1 = jnp.stack([0.6 * rdm1, 0.4 * rdm1])
+ return Molecule(
+ grid=grid,
+ atom_index=jnp.array([0, 0]),
+ nuclear_pos=jnp.zeros((2, 3)),
+ ao=jnp.asarray(rng.standard_normal((grid_size, N))),
+ grad_ao=jnp.asarray(rng.standard_normal((grid_size, N, 3))),
+ grad_n_ao={k: jnp.zeros((grid_size, N)) for k in range(3)},
+ rdm1=spin_rdm1,
+ nuclear_repulsion=jnp.asarray(0.35),
+ h1e=jnp.asarray(rng.standard_normal((N, N))),
+ vj=jnp.zeros((2, N, N)),
+ mo_coeff=jnp.zeros((2, N, N)),
+ mo_occ=jnp.zeros((2, N)),
+ mo_energy=jnp.zeros((2, N)),
+ rep_tensor=rep_tensor,
+ name=jnp.array([ord(char) for char in "H2"]),
+ basis=jnp.array([ord(char) for char in "ccpvdz"]),
+ energy=jnp.asarray(-1.16),
+ )
+
+
+def write_old_style_dense_file(path: str, molecule: Molecule) -> None:
+ """Write an HDF5 file with the pre-compression layout (no dataset attrs)."""
+ with h5py.File(f"{path}.hdf5", "w") as file:
+ group = file.create_group("molecule_0")
+ group.create_dataset("coords", data=np.asarray(molecule.grid.coords))
+ group.create_dataset("weights", data=np.asarray(molecule.grid.weights))
+ for key in (
+ "atom_index",
+ "nuclear_pos",
+ "ao",
+ "grad_ao",
+ "rdm1",
+ "nuclear_repulsion",
+ "h1e",
+ "vj",
+ "mo_coeff",
+ "mo_occ",
+ "mo_energy",
+ "energy",
+ ):
+ group.create_dataset(key, data=np.asarray(getattr(molecule, key)))
+ grad_group = group.create_group("grad_n_ao")
+ for k, v in molecule.grad_n_ao.items():
+ grad_group.create_dataset(f"{k}", data=np.asarray(v))
+ group.create_dataset("rep_tensor", data=np.asarray(molecule.rep_tensor))
+ group.create_dataset("name", data="H2")
+ group.create_dataset("basis", data="ccpvdz")
+
+
+def load_one(path: str, training: bool):
+ """Load the single molecule of a file as (kind, molecule)."""
+ return next(iter(loader(path, randomize=False, training=training)))
+
+
+def test_dense_round_trip_unchanged(tmp_path):
+ """Dense files keep the old layout (no attrs) and load as dense arrays."""
+ path = str(tmp_path / "dense")
+ original = make_molecule(REP_DENSE)
+ saver(path, molecules=[original])
+
+ with h5py.File(f"{path}.hdf5", "r") as file:
+ (group_name,) = [key for key in file.keys() if "molecule" in key]
+ dataset = file[group_name]["rep_tensor"]
+ assert dataset.ndim == 4
+ assert not dict(dataset.attrs)
+
+ kind, molecule = load_one(path, training=False)
+ assert kind == "molecule"
+ assert isinstance(molecule.rep_tensor, jnp.ndarray)
+ np.testing.assert_array_equal(
+ np.asarray(molecule.rep_tensor), np.asarray(original.rep_tensor)
+ )
+
+
+@pytest.mark.parametrize("fmt", PACKED_FORMATS)
+def test_packed_round_trip(tmp_path, fmt):
+ """Packed s4/s8 storage survives save/load with metadata, values, dtype."""
+ path = str(tmp_path / f"packed_{fmt}")
+ original = make_molecule(packed_eri(fmt))
+ saver(path, molecules=[original])
+
+ with h5py.File(f"{path}.hdf5", "r") as file:
+ (group_name,) = [key for key in file.keys() if "molecule" in key]
+ dataset = file[group_name]["rep_tensor"]
+ # s4 is stored as the M x M packed matrix, s8 as the flat T-vector.
+ assert dataset.ndim == (2 if fmt == "s4" else 1)
+ assert str(dataset.attrs["storage_format"]) == fmt
+ assert int(dataset.attrs["n_orbitals"]) == N
+
+ kind, molecule = load_one(path, training=False)
+ assert kind == "molecule"
+ loaded = molecule.rep_tensor
+ assert isinstance(loaded, PackedERI)
+ assert loaded.storage_format == fmt
+ assert loaded.n_orbitals == N
+ assert loaded.packed.dtype == jnp.float64
+ np.testing.assert_array_equal(
+ np.asarray(loaded.packed), np.asarray(original.rep_tensor.packed)
+ )
+ np.testing.assert_allclose(
+ molecule.nonXC(), original.nonXC(), atol=1e-12, rtol=1e-12
+ )
+
+
+@pytest.mark.parametrize("fmt", PACKED_FORMATS)
+def test_training_mode_reconstructs_packed(tmp_path, fmt):
+ """Dataset attrs survive the training-mode group-attr skip."""
+ path = str(tmp_path / f"train_{fmt}")
+ original = make_molecule(packed_eri(fmt))
+ saver(path, molecules=[original])
+
+ _, molecule = load_one(path, training=True)
+ loaded = molecule.rep_tensor
+ assert isinstance(loaded, PackedERI)
+ assert loaded.storage_format == fmt
+ assert loaded.n_orbitals == N
+ np.testing.assert_array_equal(
+ np.asarray(loaded.packed), np.asarray(original.rep_tensor.packed)
+ )
+
+
+def test_reaction_round_trip_packed(tmp_path):
+ """The reaction dataset loop reconstructs packed reactants and products."""
+ path = str(tmp_path / "reaction")
+ reactant = make_molecule(packed_eri("s8"))
+ product = make_molecule(packed_eri("s4"))
+ saver(path, reactions=[make_reaction([reactant], [product], energy=-1.25)])
+
+ _, reaction = next(iter(loader(path, randomize=False, training=True)))
+ reactant_loaded = reaction.reactants[0].rep_tensor
+ product_loaded = reaction.products[0].rep_tensor
+ assert isinstance(reactant_loaded, PackedERI)
+ assert reactant_loaded.storage_format == "s8"
+ assert isinstance(product_loaded, PackedERI)
+ assert product_loaded.storage_format == "s4"
+ np.testing.assert_array_equal(
+ np.asarray(reactant_loaded.packed), np.asarray(reactant.rep_tensor.packed)
+ )
+ np.testing.assert_array_equal(
+ np.asarray(product_loaded.packed), np.asarray(product.rep_tensor.packed)
+ )
+
+
+def test_old_style_dense_file_loads_as_dense(tmp_path):
+ """A pre-compression file (plain dense dataset, no attrs) loads as dense."""
+ path = str(tmp_path / "old_style")
+ original = make_molecule(REP_DENSE)
+ write_old_style_dense_file(path, original)
+
+ kind, molecule = load_one(path, training=False)
+ assert kind == "molecule"
+ assert isinstance(molecule.rep_tensor, jnp.ndarray)
+ assert molecule.rep_tensor.shape == (N, N, N, N)
+ np.testing.assert_array_equal(
+ np.asarray(molecule.rep_tensor), np.asarray(original.rep_tensor)
+ )
+
+
+def test_unknown_storage_format_raises(tmp_path):
+ """A corrupt storage_format attribute raises an informative ValueError."""
+ path = str(tmp_path / "corrupt")
+ saver(path, molecules=[make_molecule(packed_eri("s8"))])
+ with h5py.File(f"{path}.hdf5", "a") as file:
+ (group_name,) = [key for key in file.keys() if "molecule" in key]
+ file[group_name]["rep_tensor"].attrs["storage_format"] = "bogus"
+
+ with pytest.raises(ValueError, match="Supported formats are"):
+ load_one(path, training=False)
diff --git a/tests/unit/test_eri_storage.py b/tests/unit/test_eri_storage.py
new file mode 100644
index 0000000..046b4db
--- /dev/null
+++ b/tests/unit/test_eri_storage.py
@@ -0,0 +1,267 @@
+# Copyright 2023 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+# http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the symmetry-compressed ERI storage container (grad_dft.utils.eri).
+
+Covers, at the container level:
+
+(A) the `aosym` options actually accepted by the installed PySCF for
+ `int2e` (shapes, dtypes, aliases, rejected options);
+(B) ERI symmetry relations of real molecular integrals;
+(C) dense expansion against `pyscf.ao2mo.restore` as the trusted reference;
+(F) error handling for invalid formats, dtypes and shapes;
+(D, partial) PyTree registration and jit round-trip of the container.
+"""
+
+from functools import lru_cache
+
+import numpy as np
+import pytest
+
+from jax import config
+import jax
+import jax.numpy as jnp
+from pyscf import gto
+from pyscf.ao2mo import restore
+
+from grad_dft.utils.eri import (
+ ERI_STORAGE_FORMATS,
+ PackedERI,
+ pair_indices,
+ packed8_to_packed4,
+ validate_eri_storage,
+)
+
+# Enable x64 after imports (grad_dft creates no arrays at import time) but
+# before any array is created, so PySCF's float64 ERIs are not downcast.
+config.update("jax_enable_x64", True)
+
+# A symmetric diatomic (H2) and an asymmetric polyatomic (H2O). The latter
+# avoids accidental all-zero ERI blocks that make pair-matching ambiguous.
+MOLECULES = {
+ "H2": [["H", (0, 0, 0)], ["H", (0.74, 0, 0)]],
+ "H2O": [["O", (0, 0, 0)], ["H", (0, 0, 0.97)], ["H", (0.76, 0, -0.5)]],
+}
+
+
+@lru_cache(maxsize=None)
+def build_mol(name: str) -> gto.Mole:
+ """Small cc-pvdz test molecule (cached: integrals are reused across tests)."""
+ return gto.M(atom=MOLECULES[name], basis="cc-pvdz", verbose=0)
+
+
+def n_pairs(n_orbitals: int) -> int:
+ """Number of unique index pairs M = n (n + 1) / 2."""
+ return n_orbitals * (n_orbitals + 1) // 2
+
+
+@pytest.mark.parametrize("name", list(MOLECULES))
+def test_pyscf_aosym_shapes_and_dtypes(name):
+ """Document the layout PySCF returns for each aosym option (task 8A)."""
+ mol = build_mol(name)
+ n = mol.nao_nr()
+ m = n_pairs(n)
+ t = n_pairs(m)
+ dense = mol.intor("int2e", aosym="s1")
+
+ assert dense.shape == (n, n, n, n)
+ assert mol.intor("int2e", aosym="s2ij").shape == (m, n, n)
+ assert mol.intor("int2e", aosym="s2kl").shape == (n, n, m)
+ assert mol.intor("int2e", aosym="s4").shape == (m, m)
+ assert mol.intor("int2e", aosym="s8").shape == (t,)
+
+ for fmt, shape in [("s1", (n, n, n, n)), ("s4", (m, m)), ("s8", (t,))]:
+ eri = mol.intor("int2e", aosym=fmt)
+ assert eri.shape == shape
+ assert eri.dtype == np.float64
+ assert eri.nbytes == eri.size * 8
+
+
+@pytest.mark.parametrize("alias,canonical", [("4", "s4"), ("8", "s8")])
+def test_pyscf_numeric_aosym_aliases(alias, canonical):
+ """PySCF accepts '4'/'8' as aliases of 's4'/'s8'; GradDFT's config takes
+ only the canonical names (see test_validate_eri_storage)."""
+ mol = build_mol("H2")
+ np.testing.assert_array_equal(
+ mol.intor("int2e", aosym=alias), mol.intor("int2e", aosym=canonical)
+ )
+
+
+def test_pyscf_s2_unsupported_for_int2e():
+ """aosym='s2' is not supported for int2e in the installed PySCF 2.13.1
+ (AttributeError from libcgto.so: undefined symbol GTOnr2e_fill_s2).
+ If a future PySCF gains support, this test should fail so the storage
+ audit is re-run."""
+ mol = build_mol("H2")
+ with pytest.raises(Exception): # pylint: disable=broad-except
+ mol.intor("int2e", aosym="s2")
+
+
+@pytest.mark.parametrize("fmt", ["s4", "s8"])
+@pytest.mark.parametrize("name", list(MOLECULES))
+def test_to_dense_matches_pyscf_restore(fmt, name):
+ """Our expansion must reproduce pyscf.ao2mo.restore exactly.
+
+ Bitwise equality is expected: both perform pure index copies of the same
+ stored values, so any difference would indicate a packing-convention bug.
+ """
+ mol = build_mol(name)
+ n = mol.nao_nr()
+ packed = mol.intor("int2e", aosym=fmt)
+ eri = PackedERI(packed, fmt, n)
+
+ trusted = restore(1, packed, n)
+ np.testing.assert_array_equal(np.asarray(eri.to_dense()), trusted)
+ # And the reconstructed tensor equals the direct dense integral up to the
+ # machine-precision difference between independently computed packings.
+ np.testing.assert_allclose(
+ np.asarray(eri.to_dense()), mol.intor("int2e", aosym="s1"), atol=1e-13, rtol=0
+ )
+
+
+@pytest.mark.parametrize("name", list(MOLECULES))
+def test_packed8_to_packed4_matches_intor_s4(name):
+ """s8 storage expands to the same (M, M) matrix as the direct s4 packing.
+
+ Tolerance 1e-14 rather than bitwise: PySCF evaluates the s4 and s8
+ packings independently, and they differ at machine precision (observed
+ max 2.2e-16 on PySCF 2.13.1). The bitwise reference test is
+ `test_to_dense_matches_pyscf_restore`, which compares against a pure
+ index-copy of the same stored values.
+ """
+ mol = build_mol(name)
+ n = mol.nao_nr()
+ m = n_pairs(n)
+ e8 = mol.intor("int2e", aosym="s8")
+ e4 = mol.intor("int2e", aosym="s4")
+ np.testing.assert_allclose(
+ np.asarray(packed8_to_packed4(jnp.asarray(e8), m)), e4, atol=1e-14, rtol=0
+ )
+
+
+def test_pair_indices_convention():
+ """pair(a, b) = a (a + 1) / 2 + b for a >= b, symmetric, bijective onto [0, M)."""
+ n = 6
+ m = n_pairs(n)
+ index = np.asarray(pair_indices(n))
+ assert index.shape == (n, n)
+ for a in range(n):
+ for b in range(n):
+ expected = max(a, b) * (max(a, b) + 1) // 2 + min(a, b)
+ assert index[a, b] == expected
+ np.testing.assert_array_equal(index, index.T)
+ assert set(index.ravel().tolist()) == set(range(m))
+
+
+def test_symmetry_relations_real_molecular_eri():
+ """Real molecular ERIs satisfy (ij|kl) = (ji|kl) = (ij|lk) = (kl|ij).
+
+ Validity scope: real, isolated molecules only — complex/periodic ERIs
+ are not covered by design (they keep dense storage).
+ """
+ dense = build_mol("H2O").intor("int2e", aosym="s1")
+ np.testing.assert_allclose(dense, dense.transpose(1, 0, 2, 3), atol=1e-14, rtol=0)
+ np.testing.assert_allclose(dense, dense.transpose(0, 1, 3, 2), atol=1e-14, rtol=0)
+ np.testing.assert_allclose(dense, dense.transpose(2, 3, 0, 1), atol=1e-14, rtol=0)
+
+
+@pytest.mark.parametrize("fmt", ERI_STORAGE_FORMATS)
+def test_validate_eri_storage_accepts_supported(fmt):
+ """The three supported configuration values are accepted unchanged."""
+ assert validate_eri_storage(fmt) == fmt
+
+
+@pytest.mark.parametrize("value", ["s2", "s2ij", "s2kl", "4", "8", "S8", ""])
+def test_validate_eri_storage_rejects_unsupported(value):
+ """Unsupported values fail with an error listing the valid formats."""
+ with pytest.raises(ValueError, match="Supported formats are"):
+ validate_eri_storage(value)
+
+
+def test_packed_eri_rejects_unknown_format():
+ """PackedERI only holds the packed formats; 'dense' stays a plain array."""
+ with pytest.raises(ValueError, match="storage_format"):
+ PackedERI(np.zeros((6, 6)), "s2", 3)
+
+
+def test_packed_eri_rejects_complex_dtype():
+ """Complex ERIs (e.g. multi-k-point) must keep eri_storage='dense'."""
+ with pytest.raises(ValueError, match="complex"):
+ PackedERI(np.zeros((6, 6), dtype=np.complex128), "s4", 3)
+
+
+def test_packed_eri_rejects_non_floating_dtype():
+ """Only real floating-point dtypes are valid ERI storage."""
+ with pytest.raises(ValueError, match="floating"):
+ PackedERI(np.zeros((6, 6), dtype=np.int32), "s4", 3)
+
+
+def test_packed_eri_rejects_inconsistent_shape():
+ """Physical storage must match format and n_orbitals.
+
+ n = 3 -> M = 6: s4 must be (6, 6), s8 must be (21,).
+ """
+ with pytest.raises(ValueError, match="requires packed shape"):
+ PackedERI(np.zeros((21, 21)), "s4", 3)
+ with pytest.raises(ValueError, match="requires packed shape"):
+ PackedERI(np.zeros((21,)), "s8", 4)
+
+
+@pytest.mark.parametrize("n", [0, -1, 2.5])
+def test_packed_eri_rejects_invalid_n_orbitals(n):
+ """n_orbitals must be a positive integer."""
+ with pytest.raises(ValueError, match="n_orbitals"):
+ PackedERI(np.zeros((6, 6)), "s4", n)
+
+
+def test_packed_eri_logical_vs_physical_shapes():
+ """`.shape` is the logical (n, n, n, n); storage shape lives on `.packed`."""
+ n = 4
+ m = n_pairs(n)
+ t = n_pairs(m)
+ e4 = PackedERI(np.zeros((m, m)), "s4", n)
+ e8 = PackedERI(np.zeros((t,)), "s8", n)
+ assert e4.shape == (n, n, n, n) # logical
+ assert e4.packed.shape == (m, m) # physical
+ assert e8.shape == (n, n, n, n)
+ assert e8.packed.shape == (t,)
+ assert e4.dtype == np.float64
+
+
+def test_packed_eri_is_pytree_with_static_metadata():
+ """The packed array is the single leaf; format and n_orbitals are aux data."""
+ n = 3
+ m = n_pairs(n)
+ eri = PackedERI(jnp.zeros((m, m)), "s4", n)
+ leaves, treedef = jax.tree_util.tree_flatten(eri)
+ assert len(leaves) == 1
+ rebuilt = jax.tree_util.tree_unflatten(treedef, leaves)
+ assert isinstance(rebuilt, PackedERI)
+ assert rebuilt.storage_format == "s4"
+ assert rebuilt.n_orbitals == n
+
+
+def test_packed_eri_jit_roundtrip():
+ """jit round-trips values exactly and accepts mixed formats (one compile
+ per distinct static metadata)."""
+ mol = build_mol("H2")
+ n = mol.nao_nr()
+ identity = jax.jit(lambda x: x)
+ for fmt in ("s4", "s8"):
+ eri = PackedERI(jnp.asarray(mol.intor("int2e", aosym=fmt)), fmt, n)
+ out = identity(eri)
+ assert isinstance(out, PackedERI)
+ assert out.storage_format == fmt
+ assert out.n_orbitals == n
+ np.testing.assert_array_equal(np.asarray(out.packed), np.asarray(eri.packed))