Full migration of moleintor to moleintor_lite (#136) - #145
Draft
claude[bot] wants to merge 13 commits into
Draft
Conversation
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Migrate the jvp branch for int1e_nuc and int2e from the legacy gto/moleintor.py to the jittable/batchable gto/moleintor_lite.py. - int1e_nuc: AO-center (bra/ket) derivative via int1e_nuc_dr10/dr01 plus the operator (nuclear-position) term via the rinv-at-nucleus trick (int1e_rinv_dr10/dr01 scaled by -Z_A), summed over nuclei. - int2e: aosym='s1' only; the gradient is built from the bra-1 derivative int2e_dr1000 and the 8-fold ERI index symmetry (no packed storage). - Higher-order derivatives raise NotImplementedError (out of scope). Verified vs finite difference (int1e_nuc ~1e-9, int2e ~1e-11), reverse vs forward mode to machine precision, and jit==eager. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- scf/diis.py: add a jittable Pulay DIIS class (PytreeNode mirroring
scf/anderson.py) using the commutator error vector and the shared
minimize_residual solver; keep the legacy PySCF CDIIS (SCF_DIIS alias).
- scf/hf_lite.py:
* get_hcore now assembles int1e_kin + int1e_nuc directly (jittable on
MoleLite) instead of delegating to PySCF (which needs _ecpbas).
* add a jittable core-Hamiltonian get_init_guess (the PySCF minao/atom
guesses are not traceable).
* wire mf.diis='diis' into the SCF loop and get_fock.
* add a jittable UHF class (spin-resolved get_occ/make_rdm1/eig/get_veff/
energy_elec/get_grad/init guess); RHF=SCF alias added.
Verified vs PySCF: RHF/UHF energies ~1e-13, nuclear gradients ~1e-7
(implicit-diff GMRES tol), UHF closed-shell == RHF, jit==eager.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Refactor the padding-aware diagonalization/MO-masking into a _SCFPadMixin shared by the restricted (SCF/SCFPad) and new unrestricted (UHF/UHFPad) padded SCF classes. The padding methods act on a single spin block, so they compose with hf_lite.UHF (eig calls _eigh per spin; get_occ masks each spin). Export RHF/UHF/UHFPad from ml.scf. Verified batched vs PySCF (jit(vmap(value_and_grad))): RHF (H2O+H2) and UHF (OH, NH2 doublets) energies match (~1e-13) and forces (~1e-8) at the matching hcore-guess solution, with padding-atom gradients identically 0. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- tests/test_moleintor_lite.py: int1e_ovlp/kin/nuc and int2e coordinate jvps vs finite difference, reverse==forward mode, jit==eager, and higher-order derivatives raise NotImplementedError. - tests/test_hf_lite.py: RHF/UHF energy and nuclear gradient vs PySCF (UHF compared at the matching hcore-guess solution), UHF closed-shell == RHF, DIIS/Anderson/None agree, and jit(value_and_grad)==eager. - tests/test_hf_pad.py: batched RHF (H2O+H2) and UHF (OH+NH2 doublets) via jit(vmap(value_and_grad)) vs PySCF, with padding-atom gradient = 0. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The higher-order guard was too broad and broke the existing gto/test/test_mole_lite.py::test_int1e Hessian test for int1e_ovlp/kin, which the lite jvp supports via the '_dr..' name recursion. Restrict the gate to the newly migrated int1e_nuc (rinv operator term) and int2e branches; plain int1e (ovlp/kin/r/rinv) again recurse to arbitrary order. 2nd-order int1e_nuc still fails loudly via its int1e_nuc_dr10 sub-integral. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ents Land the generalized jvp branch: int1e derivatives recurse through the dr-naming to the order libcintad provides (total order 4, checked loudly by _check_deriv_available), and _int2e_jvp_r0 handles any derivative order while evaluating sub-integrals with the highest packing their remaining pair symmetry allows (s2ij/s2kl), emitting packed tangents for s1/s2ij/s2kl/s4/s8 primals. The s1 primal ERI is built s8-packed on the host and unpacked with ao2mo.restore (~8x cheaper fill). Fixes on top: - _getints_impl now coerces atm/bas/env(/ao_loc) to numpy before calling PySCF's engine; pure_callback hands over jax arrays and make_loc dies on jnp.cumsum(out=...) otherwise. - getints_jvp returns tangents directly instead of accumulating into a zeros_like(primal) buffer, saving one full-size temporary per differentiation level. Replace the stale higher-order NotImplementedError test with real coverage: second derivatives vs finite differences (ovlp/kin/nuc/int2e), packed-aosym primal+jvp vs s1, s4 second derivative (+jit), a dr20-seeded order-4 control (full 4-fold nesting kept as _high_cost), and the order-5 loud-failure boundary. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Replace the unrolled per-atom Python loop in _gen_int1e_nuc_jvp_rc with
a vmap over nuclei: the body traces once, so jaxpr size and compile time
stay constant in the atom count (the unrolled version staged 1-2 host
callbacks per atom per differentiation level), while the callback still
executes sequentially per atom at run time (vmap_method="sequential").
vmap rather than lax.scan: a scan stages the body as an opaque jaxpr in
which the inner getints custom_jvp does not survive forward-over-reverse
differentiation ("Pure callbacks do not support JVP" from
jacfwd(grad(...))), while vmap batches inline and composes in every
order and mode. The transient cost is the (natm, 3*ycomp, nao, nao)
stack of per-nucleus rinv integrals, negligible next to the ERI.
Padding atoms (charge 0) contribute zero terms automatically, and the
dr-name recursion for higher orders is unchanged since env's rinv-origin
slots keep their tangent dependence inside the batched body.
Add jit-vs-eager and jacfwd(grad(...)) second-derivative tests for
int1e_nuc to pin both compositions.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Add pyscfad/scf/_eri_lite.py with a J/K contraction for the s4-packed
(npair, npair) ERI, four times smaller than the dense s1 tensor -- and
since getints emits coordinate tangents in the same packing, the saving
holds for every jacfwd direction of a hessian as well:
- J packs the symmetrized density with off-diagonal weight 2 and reduces
to one (x, npair) @ (npair, npair) matmul.
- K scans AO row blocks with lax.map: each step gathers the s4 rows of
a block, unpacks the column pairs and contracts with the density, so
the transient stays O(block * nao^3) (block sized to VK_BLOCK_BYTES)
and nao^4 is never materialized.
All ops are linear gathers/matmuls, so jvps pass through, reverse mode
transposes them to scatter-adds, and derivatives compose to any order;
scan-inside-vmap keeps hf_pad batching intact. dot_eri_dm dispatches on
the packing and falls back to the dense s1 einsums.
hf_lite.SCF now owns get_jk (previously aliased to the legacy hf.SCF
method, whose packed branch calls the non-traceable _vhf.incore): it
stores mol.intor("int2e", aosym=self.eri_aosym) with a new static
eri_aosym="s4" attribute ("s1" opts back into the dense path) and
raises loudly on omega instead of silently entering a broken
range-separation path. UHF (stacked spin densities) and the padded
classes inherit it unchanged. On CPU the packed contraction also
measures faster than the dense einsums (bandwidth-bound).
Tests: dot_eri_dm_s4 vs the dense einsums (non-symmetric and stacked
densities, with_j/with_k flags, ragged K blocks), s4-vs-s1 energy/grad
(hessian variant as _high_cost), RHF and UHF hessians vs PySCF's
analytic Hessian, and jit-vs-eager hessian.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
test_batched_rhf_hess_high_cost drives jit(vmap(jacfwd(grad))) over the H2O/H2 batch against PySCF's analytic Hessians and checks that padding atoms decouple at second order as well (excluded from the default run by the _high_cost suffix). test_rhf_verbose_jit_smoke pins that a verbose SCF completes under jit, guarding the logger's staging of traced values through jax.debug.callback. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The thirdparty/ tree is a build-output directory: CMake fetches the libcint fork at the pinned commit and installs it there during a source build, and config.h is generated from config.h.in. The earlier WIP commit accidentally captured these local build products, including the 7.9 MB libcintad.so.6.1.1 binary; the build recipe on main already produces the same dr-derivative-enabled library, so nothing depends on the committed copies. Untrack them and extend .gitignore (versioned *.so.*, config.h, thirdparty/) so they cannot come back. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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
1 participant
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.
Fixes #136
Migrates the legacy
gto/moleintor.pyautodiff (jvp branch) to the jittable/batchablegto/moleintor_lite.py, and extends the*_lite/*_padSCF stack to RHF + UHF — with high-order derivatives, packed-ERI memory savings, and end-to-end jit/vmap.Scope (from #136)
moleintor_lite:int1e_nucjvp (rinv-at-nucleus trick) andint2ejvp, forward viajax.pure_callbackscf/diis.py: jittable Pulay DIIS mirroringscf/anderson.py(fixed-size circular buffers,lax.cond-gated extrapolation, ridge-regularized solve)scf/hf_lite.py+ml/scf/hf_pad.py: RHF and UHFHigh-order derivatives
Derivatives recurse through the
_drintegral naming to the orderlibcintadprovides (total order 4 forint1e_ovlp/kin/nuc/rinv/r/rr/rrr/r2,int2c2e,int2e); order 5 fails loudly withNotImplementedError(_check_deriv_available).jacfwd/jacrev/gradcompose in any order and mode through the whole stack, includinghf_lite'slax.custom_root+ gmres tangent solve — no changes to the implicit-diff wiring were needed.Memory
s1ERI primal is builts8-packed on the host and unpacked withao2mo.restore(~8× cheaper fill).hf_litenow stores the ERI s4-packed by default (eri_aosym="s4",(npair, npair)— 4× smaller thans1, and sincegetintsemits tangents in the same packing, the saving holds for everyjacfwddirection of a hessian). J/K come from the new jittablescf/_eri_lite.py: J is one packed matmul; K is a blockedlax.maprow scan with anO(block·nao³)transient that never materializesnao⁴. On CPU the packed contraction also measures faster than the dense einsums;mf.eri_aosym = "s1"opts back into the dense path.int2ejvp evaluates derivative sub-integrals with the highest packing their remaining pair symmetry allows (s2ij/s2kl) and recovers partner centers by index transposition instead of extra integral evaluations; packed primals get packed tangents (s1/s2ij/s2kl/s4/s8at order 0).vmap(constant jaxpr size and compile time innatm; alax.scanversion breaksjacfwd(grad(...))because the scanned body hides thegetintscustom_jvp — pinned by a regression test).getints_jvpreturns tangents directly instead of accumulating into azeros_likebuffer (one fewer full-size temporary per differentiation level).int2e_dr1010) has no residual pair symmetry, so itss1evaluation transiently holds9·nao⁴doubles; avoiding that needs blockedshls_sliceevaluation (future work).Out of scope (unchanged from #136)
ECP, basis-parameter derivatives (
trace_basis),int3cdifferentiation, partialshls_slicein theint2ejvp, and range-separated J/K (omeganow raises loudly).Note: an earlier WIP commit had accidentally captured local pyscfadlib build artifacts (the 7.9 MB
libcintad.so.6.1.1, generated headers,config.h); these are now removed and gitignored —thirdparty/is populated by CMake's libcint ExternalProject at build time, and the pinned recipe onmainalready produces the dr-derivative-enabled library.Verification
All run locally (CPU, float64):
pytest tests/test_moleintor_lite.py tests/test_hf_lite.py tests/test_hf_pad.py— 38 passed (3_high_costdeselected by default; run explicitly, all pass). Covers: jvp vs finite differences, second derivatives vs finite differences, packed-aosym primal/jvp vss1(exact),s4second derivative (+jit), order-4 positive control, order-5 raise,jacfwd(grad)fwd-over-rev composition, RHF/UHF energies+gradients vs PySCF, RHF and UHF hessians vs PySCF's analyticHessian().kernel()(≤1e-6, measured ~6e-7), jit-vs-eager hessian (≤1e-10),dot_eri_dm_s4vs dense einsums (non-symmetric + stacked spin densities, ragged K blocks),s4-vs-s1SCF energy/grad/hessian, batchedjit(vmap(value_and_grad))RHF/UHF with padding-atom gradients exactly 0, batchedjit(vmap(jacfwd(grad)))hessians vs PySCF with padding decoupled at second order, and a verbose-under-jit smoke test.pyscfad/gto/test/test_mole_lite.py tests/xtb— 15 passed.tests/test_scf.py pyscfad/gto— 22 passed.pylint pyscfad— 10.00/10, no warnings.🤖 Generated with Claude Code