Feature/dz0scf gradient clean - #10
Merged
Merged
Conversation
Port the average/roks dual-semantics EnsembleRKS reference (fixed 2/1/0 occupations, Dz = 0) and align Dz0SCF with the shared reference_energy semantics so NTTDA can treat both reference types uniformly.
Port the NTTDA response/gradient stack (deltaS = -1/0 channels, XC ledger kernels, ensemble orbital-response backend) and wire NTTDA.nuc_grad_method to nest.grad.nttda.Gradients. The linear solver uses pyscf.tdscf._lr_eig: the nest._lr_eig fork does not converge on the ensemble orbital Hessian.
Add the ensemble FD/response, scalar-ledger, gradient-layer, and known-value NTTDA energy suites.
Drop the EnsembleRKS reference path and drive the NTTDA gradient from Dz0SCF (average-occupation) and ROKS only. Dz0SCF now exposes the average-occupation marker plus a charge-RKS companion for the spin-unpolarized response, rebuilt on every call so a reused mean-field object (reset/new geometry/xc change) never reuses a stale reference. Port the ensemble suites to Dz0SCF and add a cache-invalidation regression test.
Remove the stale spin-lowering-only transition methods that shadow the upstream all-channel implementation after merging main. Restore nest._lr_eig so oscillator strengths retain upstream convergence behavior. The existing all-channel oscillator-strength regression passes without changing its tolerance.
Keep spin-channel amplitudes and coefficients explicit while sharing orbital partitions, reference responses, J/K derivative batching, Fz exchange and adjoint assembly. Consolidate identical GGA and meta-GGA quadrature flows with explicit density dimensions and kernel dispatch. Remove 1216 net lines from the gradient package; retain existing numerical checks and tolerances.
Retain the branch's PySCF response solver: the shared NEST solver fails the existing small-system Dz0SCF convergence checks. Explicitly converge oscillator-strength amplitudes to 1e-9 with lindep=1e-18 rather than relying on the default 1e-5 residual; reference values and assertion tolerances are unchanged. Seven targeted tests and eight subtests pass.
Build channel projections once per gradient and reuse the shared transition-potential projection. Share adjoint assembly and Krylov mechanics while retaining reference-specific Hessians and preconditioners. Use a singleton density-feature axis for LDA to share the GGA/MGGA quadrature loops, retaining distinct pair kernels. Consolidate HF/DFT Fock derivatives and remove unreferenced channel projection wrappers. Net reduction: 604 implementation lines; existing numerical tests and tolerances are unchanged.
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Critical ROKS/XC gradient correctness issues and missing validation remain unresolved.
Get a fresh assessment by requesting another Copilot review.
Review effort: Lite
Findings: 2
Open (2)
What changed in this PR
Adds analytic NTTDA excited-state gradients for ROKS and Dz0SCF references, including response handling, regression tests, and examples.
Changes:
- Adds gradient, XC response, orbital-response, and Z-vector machinery.
- Extends Dz0SCF reference-energy and response support.
- Expands tests and gradient usage examples.
| File | Description |
|---|---|
src/nest/nttda/tests/test_nttda.py |
Expands NTTDA functional and response tests. |
src/nest/nttda/tests/test_nttda_oscillator_strength.py |
Tightens oscillator-strength test convergence. |
src/nest/nttda/tests/test_nttda_dz0scf.py |
Tests NTTDA Dz0SCF total energies. |
src/nest/nttda/nttda.py |
Adds gradient APIs and Dz0SCF support. |
src/nest/grad/tests/test_nttda_scalar_ledger.py |
Tests scalar response ledgers. |
src/nest/grad/tests/test_nttda_gradient_layers.py |
Tests layered gradient components. |
src/nest/grad/tests/test_nttda_grad.py |
Tests public gradient behavior. |
src/nest/grad/tests/test_nttda_dz0scf_response.py |
Tests Dz0SCF orbital response. |
src/nest/grad/tests/test_nttda_dz0scf_fd.py |
Tests Dz0SCF finite differences. |
src/nest/grad/nttda/xc.py |
Implements XC gradient contractions. |
src/nest/grad/nttda/roks.py |
Implements ROKS adjoint and gradient assembly. |
src/nest/grad/nttda/ensemble.py |
Implements Dz0SCF orbital response. |
src/nest/grad/nttda/delta_s_zero.py |
Implements analytic deltaS=0 gradients. |
src/nest/grad/nttda/delta_s_minus_one.py |
Implements analytic deltaS=-1 gradients. |
src/nest/grad/nttda/common.py |
Provides shared gradient and J/K derivative assembly. |
src/nest/grad/nttda/__init__.py |
Provides analytic and finite-difference gradient drivers. |
src/nest/dz0scf/tests/test_dz0scf.py |
Validates Dz0SCF total energies. |
src/nest/dz0scf/dz0scf.py |
Adds reference-energy and response support. |
examples/nttda/02_nttda_dz0scf_grad.py |
Demonstrates NTTDA gradients. |
examples/grad/02_dz0scf_grad.py |
Demonstrates Dz0SCF gradients. |
.gitignore |
Ignores local research snapshots. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+114
to
+115
| if (isinstance(mf, dft.KohnShamDFT) | ||
| and mf._numint._xc_type(mf.xc) != "HF"): |
Comment on lines
+598
to
+600
| def _reference_fref_kref(mf, rho0, xctype): | ||
| fxc, kxc = mf._numint.eval_xc_eff( | ||
| mf.xc, (rho0, rho0), deriv=3, xctype=xctype, spin=1, |
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.
addd analytic dz0scf-gradient