Repository navigation
Install DM21 package as dependency + compressed rep_tensor - #14
Merged
Merged
Conversation
Declare density_functional_approximation_dm21 as a pinned VCS dependency (upstream HEAD f5de0ede, no PyPI distribution exists) and re-export its Functional/NeuralNumInt/_SystemState through a new grad_dft/external/_dm21_compat.py that carries the three local deltas the installed package needs in this environment: - pkg_resources shim (tensorflow-hub imports parse_version; setuptools 81+ no longer ships it), - only-if-missing shims for the TF1 Hub API symbols tensorflow-hub 0.16 removed (hub.Module via hub.load, hub.add_signature, hub.create_module_spec) so upstream's _build_graph runs unmodified, - the eval_xc_eff override PySCF 2.3+ requires (its base implementation dispatches to the libxc backend instead of the subclass's eval_xc). Verified: Ne/sto-3g RKS DM21 energy matches the vendored baseline bit for bit (-126.898914169039 Ha), CI unit files pass, pre-commit gate clean on touched files. The vendored copy is no longer imported but still on disk; removed in the next subtask.
Remove grad_dft/external/density_functional_approximation_dm21/ (40 tracked files, 3012 lines). The DM21 code and its checkpoints now come from the pinned density_functional_approximation_dm21 package (requirements.txt, installed in a4aba3e), and the tree's only consumer, grad_dft.external, imports through grad_dft.external._dm21_compat, which carries the three local patches the deleted copy had over upstream f5de0ede843. A repo-wide sweep finds no remaining imports or path references to the removed path, and grad_dft.egg-info no longer lists it. The full CI test list runs in the background after this removal and must pass before integration; a4aba3e already verified the unit files and a bit-identical DM21 smoke test. Closes #13
Code review of #13 found three statements in _dm21_compat.py that did not survive verification: - "setuptools >= 81 (no pkg_resources)": the 80.9 and 81.0 wheels still ship pkg_resources; 82.0.0 is the first release without it (checked wheel contents). The threshold is now 82 in all four places. - "PySCF >= 2.3 requires the eval_xc_eff override": nr_rks/nr_uks already called eval_xc_eff in 2.1.1, and the base class still dispatched to self.eval_xc through 2.6.2; the switch to self.eval_xc1 (libxc, bypassing the subclass) landed in 2.7.0 (checked numint.py at every tag in between). The threshold is now 2.7. - The provenance paragraph named the vendored directory as if it still existed; it now refers to grad_dft/external/ plus git history. Docstrings and comments only; no behavior change.
test_uks passed the module-level `grid` (a pyscf Grids object) as the second positional argument of dft.UKS, but that argument is `xc` (true at least since pyscf v1.7). mf.xc therefore held a Grids object and dump_flags crashed in libxc.is_nlc before any DM21 code ran — the file was never green, which is why it never made the CI list. Construct UKS the way the neighboring test_rks already does: default grids configured via mf.grids.level. The whole file now passes (3 tests), which also exercises the spin-polarized DM21 path through the compat layer's eval_xc_eff for the first time. This file is part of the #13 gate: the plan runs it alongside the CI list as the DM21-specific coverage. Style: black passes. The file's pylint findings are pre-existing debt tracked in #7 (none introduced here; the import order is deliberate — the jax x64 config and the sys.path bootstrap must run before the grad_dft imports — and the _numint injection is the point of the test).
New grad_dft/utils/eri.py: storage-format registry with validation, the flax PyTree container (logical (n,n,n,n) shape vs physical packed.shape, static format/n_orbitals metadata, __post_init__ validation including complex-dtype rejection) and exact dense expansion verified bitwise against pyscf.ao2mo.restore. Adds the ERIStorage alias ERITensor in utils/types.py, exports in utils/__init__.py and grad_dft/__init__.py, and tests/unit/test_eri_storage.py (33 tests: PySCF aosym shapes/dtypes/aliases, real-ERI symmetry relations, restore equivalence, validation errors, PyTree/jit round-trip). Part of #15 (plan: tasks/detailed_plan.md).
grad_dft/utils/eri.py gains fold_density_matrix (segment-fold of an arbitrary, possibly asymmetric rdm1 onto unique pairs), a row-chunked lax.scan s8 matvec (no (M,M) allocation beyond the documented PACKED8_BLOCK_ELEMENTS cap) and packed_coulomb_potential. molecule.py: coulomb_potential dispatches on PackedERI at trace time; coulomb_potential/coulomb_energy/nonXC annotations widen to ERITensor, dense einsum untouched. tests/unit/test_eri_contraction.py (20 tests): dense/s4/s8 equivalence for symmetric and asymmetric densities, jit==eager, grad wrt rdm1 matches, synthetic Molecule dispatch, forced multi-block chunking, and jaxpr inspection proving no (n,n,n,n) or N^4-sized buffer on the packed trace. Part of #15.
calc_eri_with_pyscf gains a validated eri_storage argument (dense/s4/s8): isolated molecules return PackedERI from intor(aosym=...) or, with density fitting, from DF.get_eri() (always s8 storage; s4 re-packed via restore(4,·)); periodic systems reject compression with informative ValueErrors and keep the dense path byte-identical. _package_outputs and molecule_from_pyscf forward eri_storage (validated up-front, before any SCF or integral work); to_device_arrays preserves PackedERI metadata and converts its single leaf. Fixes a pre-existing bug: the molecular-DF branch of calc_eri_with_pyscf used the name df which was bound to pyscf.pbc.df (periodic DF requiring a Cell), so df.DF(mf.mol) raised AttributeError for density-fitted molecules on PySCF 2.13; df now binds the molecular pyscf.df. 18 new tests: bitwise equivalence with PySCF reference packings, DF reference consistency (never exact integrals), unchanged dense isolated+periodic paths, fail-fast validation on mf=None, periodic rejection in both entry points, to_device_arrays metadata, and end-to-end molecule_from_pyscf pass-through (nonXC equality). Part of #15.
save_molecule_data writes a PackedERI rep_tensor as the packed array with storage_format/n_orbitals as *dataset* attributes (never group attributes, which training mode drops); the dense path is byte-identical, so existing HDF5 files keep loading unchanged. loader's molecule and reaction loops route rep_tensor through the new _load_rep_tensor helper: no attribute means dense (old files load exactly as before), a known packed format reconstructs a PackedERI with the loop's existing dtype treatment, and a corrupt storage_format raises the informative validate_eri_storage ValueError. 8 round-trip tests (synthetic Molecule, no SCF): dense layout unchanged, packed s4/s8 preserve format/n_orbitals/values/dtype with nonXC equality, loader(training=True) still reconstructs packed molecules, the reaction loop included, an old-style file loads as dense, and a corrupt attribute errors informatively. Part of #15.
eri_allclose(a, b) in eri.py accepts None/dense/PackedERI in any combination: None equals only None, same-format packed compares on the packed storage directly, and every mixed combination (packed-vs-dense or different packed formats) materializes both sides via packed_to_dense — the explicit N^4 boundary, documented and never on the training path. xc1's rep_tensor equality check now uses it (it previously died with a JAX type error on PackedERI inputs); Molecule.rep_tensor is annotated Optional[ERITensor] and the class/method docstrings state that both representations are accepted. The three constraints.py pass-through sites were verified to need no change (Molecule constructor calls, no comparison or conversion). 7 tests: the eri_allclose matrix (None pairs, same-format equal/unequal, mixed via dense, dense pair), rep_tensor=None Molecule construction with density/grad_density under jit, and xc1's full constraint loss identical for mixed s8/s4 packed storage vs all-dense — exercising the S2 hot dispatch inside a complete functional energy evaluation. Part of #15.
New tests/integration/molecules/test_eri_storage.py runs the full pipeline on real converged SCF molecules (H2/sto-3g and H2O/sto-3g): default eri_storage stays dense (n,n,n,n) and equals intor('int2e', aosym='s1'); s4/s8 molecules carry correct metadata and expand to the same integrals (1e-14); nonXC and get_coulomb_potential match dense (1e-9); the simple_energy_loss value and its parameter gradients match dense across storages (1e-8, guarded against a vacuous all-zero comparison); d(nonXC)/d(rdm1) — which backpropagates through the packed8/packed4 contractions themselves — matches dense (1e-8), since the parameter gradients alone do not exercise the ERI term; density-fitted molecules are consistent across all three formats; and rep_tensor=None molecules still support density/grad_density. install_and_test.yml now lists the five unit ERI files and the integration file. Periodic solids remain covered by the unchanged existing solids files. Part of #15.
examples/advanced_scripts/benchmark_eri_storage.py measures the three storage formats end to end. The parent process spawns one child per (format, mode) so peak RSS is per configuration; children emit JSON lines that the parent aggregates into a markdown table (--json-out optional). Modes: generate (SCF time, packing time, physical shape/bytes), contraction (to_dense time and get_coulomb_potential compile vs steady state), train (one jitted training step with a tiny NeuralFunctional + simple_energy_loss, compile vs steady state). Timings use jax.block_until_ready with warmup plus a median over --repeats; float64 is enabled; the SCF uses xc='0.00*LDA' so only ERI-dependent terms are timed; peak RSS comes from resource.getrusage, with device.memory_stats() reported when available. The readme documents the modes, flags and measurement caveats (CPU-only, machine-specific, RSS includes the JAX runtime). Measured on this machine (CPU): H2O/cc-pvdz stored ERI 2.531 MiB dense vs 0.687 s4 vs 0.344 s8; C6H6/sto-3g 12.81 vs 3.38 vs 1.69 MiB; s8 steady-state Coulomb contraction ~2.6x the dense einsum on C6H6 while training-step steady times are comparable. Part of #15.
Adds a 'Compressed ERI storage (opt-in)' subsection to the README under the molecule-creation flow: the eri_storage argument and the exact stored shapes/meaning for dense (n,n,n,n), s4 (M,M) and s8 (T,), what runs directly on the compressed form (Coulomb potential/energy, nonXC, Functional.energy, Harris and the training losses, jit/grad) versus what explicitly materializes dense storage (packed_to_dense, cross-format comparison), and the unsupported cases (periodic/PBC, complex dtype, s2ij/s2kl) that raise an informative ValueError. Quotes the measured CPU/float64 payloads from the S7 benchmark (H2O/cc-pvdz: 2.53/0.69/0.34 MiB; C6H6/sto-3g: 12.81/3.38/1.69 MiB) and states that only memory is reduced while contraction/training times stay comparable, plus the remaining limitations. Documentation-only subtask: no code changes and therefore no new tests; the S6 integration run and the S7 benchmark stand as evidence. Closes #15
Anya497
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
DM21 as dependency
Compressed rep_tensor