Feature/dz0scf gradient clean (reopened for review)

.gitignore

@@ -220,3 +220,6 @@ __marimo__/
 
 # Streamlit
 .streamlit/secrets.toml
+
+# Quarantined local research-code snapshots
+.codex_trash/

examples/grad/02_dz0scf_grad.py

@@ -0,0 +1,58 @@
+#!/usr/bin/env python
+# Copyright 2026 The NEST Developers. All Rights Reserved.
+#
+# 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.
+
+'''
+Analytic nuclear gradient of the Dz0SCF (average-occupation) reference.
+
+Dz0SCF drives one set of orbitals with average occupations 2/1/0 for a
+high-spin open-shell reference, and takes the high-spin ROKS energy evaluated
+on those orbitals as the reference energy.  ``nuc_grad_method()`` returns the
+analytic gradient of that reference energy, which is the zero state used by
+the NTTDA excited-state gradients (see examples/nttda/02_nttda_dz0scf_grad.py).
+
+The reference is non-stationary on the average-occupation orbitals, so the
+driver solves a Z-vector equation for the orbital response; a diffuse enough
+integration grid is required for the force sum to vanish.
+'''
+
+from pyscf import gto
+from nest import dz0scf  # necessary import
+from nest.dz0scf import DZ0SCF
+
+atom = '''
+N   0.000000  -0.040000   0.000000
+H   0.000000   0.780000   0.590000
+H   0.000000  -0.860000   0.520000
+'''
+mol = gto.M(atom=atom, charge=0, spin=1, basis='6-31g', verbose=3)
+fun = 'PBE'  # try also 'SVWN', 'B3LYP', 'M06-2X', etc.
+mf = DZ0SCF(mol, xc=fun)
+mf.conv_tol = 1e-12
+mf.conv_tol_grad = 1e-9
+mf.max_cycle = 120
+mf.grids.level = 5      # dense grid: the force sum is grid-sensitive
+mf.grids.prune = None
+mf.small_rho_cutoff = 0.0
+mf.kernel()
+
+print('Dz0SCF reference energy: %.12f' % mf.high_spin_energy())
+
+grad = mf.nuc_grad_method().kernel()
+print('Analytic reference gradient (Eh/Bohr):\n', grad)
+print('Force sum (should be ~0):\n', grad.sum(axis=0))
+
+# Gradients can also be restricted to selected atoms:
+grad_n = mf.nuc_grad_method().kernel(atmlst=[0])
+print('Gradient on the nitrogen atom only:\n', grad_n)

examples/nttda/02_nttda_dz0scf_grad.py

@@ -0,0 +1,71 @@
+#!/usr/bin/env python
+# Copyright 2026 The NEST Developers. All Rights Reserved.
+#
+# 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.
+
+'''
+NTTDA excited-state gradients on a Dz0SCF (average-occupation) reference.
+
+A Dz0SCF reference gives a common average-occupation orbital set for every
+spin channel.  NTTDA built on that reference can target the same-spin channel
+(``deltaS=0``) and the spin-lowering channel (``deltaS=-1``); the total state
+energy is the high-spin reference energy plus the NTTDA excitation energy, so
+its gradient is the sum of the reference gradient and the excitation-energy
+gradient.
+
+``td.Gradients().kernel(state=n)`` returns the analytic gradient of
+``E_reference + omega_n`` for state ``n`` (1 for the lowest root); ``state=0``
+returns the reference gradient.  Analytic gradients are available for
+``deltaS = -1`` and ``0``; ``deltaS = +1`` is not implemented.
+'''
+
+from pyscf import gto
+from nest import dz0scf, nttda  # necessary imports
+from nest.dz0scf import DZ0SCF
+
+atom = '''
+C    0.020000  -0.030000   0.010000
+H   -0.020000   0.800000   0.620000
+H    0.030000  -0.910000   0.500000
+'''
+mol = gto.M(atom=atom, charge=0, spin=2, basis='sto-3g', verbose=3)
+fun = 'B3LYP'
+mf = DZ0SCF(mol, xc=fun)
+mf.conv_tol = 1e-12
+mf.conv_tol_grad = 1e-9
+mf.max_cycle = 150
+mf.grids.level = 5      # dense grid: the force sum is grid-sensitive
+mf.grids.prune = None
+mf.small_rho_cutoff = 0.0
+mf.kernel()
+
+for delta_s in (-1, 0):
+    td = mf.NTTDA().set(
+        deltaS=delta_s,     # Sf = Si + deltaS
+        nstates=3,
+        conv_tol=1e-9,
+        max_cycle=200,
+        verbose=0,
+    ).run()
+    print('deltaS = %+d' % delta_s)
+    print('  NTTDA excitation energies:', td.e)
+    print('  total energies (E_ref + omega):', td.total_energies())
+
+    grad = td.Gradients().kernel(state=1)
+    print('  state-1 analytic gradient (Eh/Bohr):\n', grad)
+    print('  force sum (should be ~0):\n', grad.sum(axis=0))
+
+    ref_grad = td.Gradients().kernel(state=0)
+    print('  reference (state-0) gradient:\n', ref_grad)
+    print('  excitation-only contribution (state 1 - state 0):\n',
+          grad - ref_grad)

src/nest/dz0scf/dz0scf.py

@@ -55,7 +55,54 @@ def evaluate_high_spin_energy(mf):
     )
 
 class _DZ0VeffMixin:
-    def get_veff(
+    reference_energy_semantics = 'high_spin_roks_energy_on_dz0_orbitals'
+    reference_energy_stationary = False
+    is_average_occupation_reference = True
+
+    def _charge_rks(self):
+        """Return a fresh RKS view used for the spin-unpolarized charge response.
+
+        The view is rebuilt on every call so that a reused mean-field object
+        (``reset(new_mol)``, ``xc`` change, new geometry) never feeds a stale
+        molecule or functional into the response.
+        """
+        charge = dft.rks.RKS(self.mol)
+        for name in (
+                'xc', 'nlc', 'grids', 'nlcgrids', '_numint',
+                'max_memory', 'small_rho_cutoff'):
+            if hasattr(self, name):
+                setattr(charge, name, getattr(self, name))
+        charge.mo_coeff = np.asarray(self.mo_coeff)
+        charge.mo_occ = np.asarray(self.mo_occ)
+        charge.mo_energy = np.asarray(self.mo_energy)
+        charge.verbose = 0
+        return charge
+
+    def make_rdm1s(self, mo_coeff=None, mo_occ=None):
+        """Return equal spin densities ``D/2`` for the spin-unpolarized reference."""
+        if mo_coeff is None:
+            mo_coeff = self.mo_coeff
+        if mo_occ is None:
+            mo_occ = self.mo_occ
+        mo_coeff = np.asarray(mo_coeff)
+        occupation = np.asarray(mo_occ)
+        dm0 = (mo_coeff * occupation) @ mo_coeff.conj().T
+        return 0.5 * dm0, 0.5 * dm0
+
+    def gen_response(self, mo_coeff=None, mo_occ=None, hermi=1, max_memory=None):
+        """Charge-only (spin-unpolarized) linear response of the reference."""
+        if mo_coeff is None:
+            mo_coeff = self.mo_coeff
+        if mo_occ is None:
+            mo_occ = self.mo_occ
+        return self._charge_rks().gen_response(
+            mo_coeff=mo_coeff,
+            mo_occ=mo_occ,
+            hermi=hermi,
+            max_memory=max_memory,
+        )
+
+    def get_veff(
         self,
         mol=None,
         dm=None,
@@ -81,8 +128,11 @@ def get_veff(
             vhf_last,
             hermi,
         )
-    def high_spin_energy(self):
-        return evaluate_high_spin_energy(self)
+    def high_spin_energy(self):
+        return evaluate_high_spin_energy(self)
+
+    def reference_energy(self):
+        return self.high_spin_energy()
 
     def nuc_grad_method(self):
         """Return the Dz0SCF analytic nuclear-gradient driver."""

src/nest/dz0scf/tests/test_dz0scf.py

@@ -79,12 +79,18 @@ def test_svwn_dz0scf(self):
         ])
 
         self.assertTrue(np.all(np.asarray(td_s.converged)))
-        np.testing.assert_allclose(
-            np.asarray(omega_s),
-            omega_s_ref,
-            rtol=0.0,
-            atol=1e-6,
-        )
+        np.testing.assert_allclose(
+            np.asarray(omega_s),
+            omega_s_ref,
+            rtol=0.0,
+            atol=1e-6,
+        )
+        np.testing.assert_allclose(
+            td_s.total_energies(),
+            mf.high_spin_energy() + omega_s,
+            rtol=0.0,
+            atol=1e-12,
+        )
 
         td_t = NTTDA(mf)
         td_t.deltaS = 0
@@ -176,4 +182,4 @@ def test_b3lyp_dz0scf(self):
             omega_t_ref,
             rtol=0.0,
             atol=1e-6,
-        )    
+        )