refactor(aocscf): replace Dz0SCF with average-occupation ROKS

examples/aocscf/01_average_occupation_scf.py

@@ -0,0 +1,37 @@
+#!/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.
+
+"""Optimize orbitals with average occupations and report the high-spin energy."""
+
+from pyscf import gto
+from nest import aocscf  # noqa: F401 - registers ROKS.average_occ()
+
+mol = gto.M(
+    atom="""
+    H  0.000000  0.934473 -0.588078
+    H  0.000000 -0.934473 -0.588078
+    C  0.000000  0.000000  0.000000
+    O  0.000000  0.000000  1.221104
+    """,
+    basis="6-31g",
+    spin=2,
+    symmetry=True
+)
+mf = mol.ROKS(xc="SVWN").average_occ().run()
+
+print(f"SCF converged: {mf.converged}")
+print(f"Average-occupation SCF energy: {mf.e_avg_occ:.12f} Ha")
+print(f"High-spin energy: {mf.e_tot:.12f} Ha")
+mf.analyze(verbose=4)  # same as other SCF methods

examples/aocscf/02_nttda.py

@@ -0,0 +1 @@
+../nttda/04_nttda_aocscf.py

examples/nttda/01_nttda.py

@@ -36,4 +36,8 @@
 td.deltaS = -1  # Final spin: Sf = Si + deltaS. Valid values are -1, 0, and +1.
 td.nobeta = False
 td.nstates = 5
-td.run().analyze(verbose=4)
+td.kernel()
+
+td.analyze(verbose=4)
+print(f"Total energies (Ha): {td.e_tot}", end=",\n")
+print(f"which equals to mf.e_tot + td.e: {mf.e_tot + td.e}")

examples/nttda/04_nttda_aocscf.py

@@ -0,0 +1,41 @@
+#!/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.
+
+"""Run NTTDA from an average-occupation ROKS reference."""
+
+from pyscf import gto
+
+from nest import aocscf, nttda  # noqa: F401 - registers ROKS.average_occ() and NTTDA()
+
+mol = gto.M(
+    atom="""
+    H  0.000000  0.934473 -0.588078
+    H  0.000000 -0.934473 -0.588078
+    C  0.000000  0.000000  0.000000
+    O  0.000000  0.000000  1.221104
+    """,
+    basis="6-31g",
+    spin=2,
+    symmetry=True
+)
+mf = mol.ROKS(xc="CAM-B3LYP").average_occ()
+mf.kernel()
+td = mf.NTTDA().set(nstates=5, deltaS=-1)
+# nobeta does not work for NTTDA with a average-occupation ROKS reference
+td.kernel()
+
+td.analyze(verbose=4)
+print(f"Total energies (Ha): {td.e_tot}", end=",\n")
+print(f"which equals to mf.e_tot + td.e: {mf.e_tot + td.e}")

examples/sftda/01_sftddft.py

@@ -67,6 +67,9 @@ def print_header(title):
 
 sfd_tddft.kernel()
 sfd_tddft.analyze(verbose=4)  # Verbose=4 shows orbital composition
+print(f"Total energies (Ha): {sfd_tddft.e_tot}", end=",\n")
+print(f"which equals to mf.e_tot + sfd_tddft.e: {mf.e_tot + sfd_tddft.e}")
+
 
 # -------------------------------------------------------------------
 # 3. Validation: Full Diagonalization of Casida Matrix

src/nest/__init__.py

@@ -15,8 +15,8 @@
 
 """nest package."""
 
-from nest import dz0scf, nttda, sftda
+from nest import aocscf, nttda, sftda
 
 __version__ = "0.1.0"
 
-__all__ = ["__version__", "dz0scf", "nttda", "sftda"]
+__all__ = ["__version__", "aocscf", "nttda", "sftda"]

src/nest/dz0scf/__init__.py → src/nest/aocscf/__init__.py

@@ -1,28 +1,23 @@
-#!/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.
-
-from .dz0scf import (
-    DZ0SCF,
-    EnsembleROKS,
-    SymAdaptedEnsembleROKS,
-    evaluate_high_spin_energy,
-)
-
-__all__ = [
-    'DZ0SCF',
-    'EnsembleROKS',
-    'SymAdaptedEnsembleROKS',
-    'evaluate_high_spin_energy',
-]
+# 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.
+
+"""Average-occupation SCF with a high-spin reported energy.
+
+Importing this module registers ``average_occ()`` on molecular PySCF ROKS
+and SymAdaptedROKS objects.
+"""
+
+from .aocscf import AverageOccupationROKS, SymAdaptedAverageOccupationROKS, average_occ
+
+__all__ = ['AverageOccupationROKS', 'SymAdaptedAverageOccupationROKS', 'average_occ']

src/nest/aocscf/aocscf.py

@@ -0,0 +1,136 @@
+# 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.
+#
+#
+# Author: Yue Yu & Tai Wang
+#
+"""ROKS orbitals optimized with spin-averaged occupations."""
+
+import numpy as np
+from pyscf import dft, lib
+from pyscf.lib import logger
+
+
+def get_veff(mf, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
+    if mol is None:
+        mol = mf.mol
+    if dm is None:
+        dm = mf.make_rdm1()
+    if np.ndim(dm) == 3:
+        dm = dm[0] + dm[1]
+    if np.ndim(dm_last) == 3:
+        dm_last = dm_last[0] + dm_last[1]
+    veff = dft.rks.get_veff(mf, mol, dm, dm_last, vhf_last, hermi)
+    # ROHF.get_fock expects separate alpha and beta blocks, even when equal.
+    return lib.tag_array(np.asarray((veff, veff)), ecoul=veff.ecoul, exc=veff.exc, vj=veff.vj, vk=veff.vk)
+
+
+def _high_spin_energy_and_summary(mf):
+    """Evaluate high-spin energy with an independent energy summary."""
+    if mf.mo_coeff is None or mf.mo_occ is None:
+        raise RuntimeError('Run SCF before evaluating the high-spin energy.')
+    # A view bypasses the spin-averaged get_veff while retaining the original
+    # object's XC integrator, grids, and other calculation settings.
+    mf_hs = lib.view(mf, dft.roks.ROKS)
+    mf_hs.scf_summary = mf.scf_summary.copy()
+    dm_hs = mf_hs.make_rdm1(mf.mo_coeff, mf.mo_occ)
+    veff_hs = mf_hs.get_veff(mf.mol, dm_hs)
+    energy = mf_hs.energy_tot(dm=dm_hs, h1e=mf_hs.get_hcore(), vhf=veff_hs)
+    return energy, mf_hs.scf_summary
+
+
+def nuc_grad_method(mf):
+    """Return the analytic gradient of the high-spin reported energy."""
+    from nest.grad.aocscf import Gradients
+
+    return Gradients(mf)
+
+
+def dump_flags(mf, verbose=None):
+    if isinstance(mf, dft.rks_symm.SymAdaptedROKS):
+        dft.rks_symm.SymAdaptedROKS.dump_flags(mf, verbose)
+    else:
+        dft.roks.ROKS.dump_flags(mf, verbose)
+    log = logger.new_logger(mf, verbose)
+    log.info('SCF potential: average occupation (equal alpha/beta density)')
+    log.info('Reported total energy and nuclear gradient: high spin')
+    return mf
+
+
+def check_sanity(mf):
+    dft.roks.ROKS.check_sanity(mf)
+    if mf.mo_occ is not None and np.ndim(mf.mo_occ) != 1:
+        raise ValueError('Average-occupation ROKS requires one-dimensional MO occupations')
+    return mf
+
+
+class AverageOccupationROKS(dft.roks.ROKS):
+    """ROKS with average-occupation SCF orbitals and a high-spin ``e_tot``."""
+
+    _keys = {'e_avg_occ'}
+
+    def _finalize(self):
+        # PySCF's SCF driver needs energy_tot() to remain the average-occupation
+        # energy during iterations and convergence checks.
+        # The final reported energy and nuclear gradient are the high-spin values.
+        # e_tot is set to the high-spin energy with `_finalize()` after SCF
+        self.e_avg_occ = self.e_tot
+        if self.mo_coeff is not None and self.mo_occ is not None:
+            self.e_tot, self.scf_summary = _high_spin_energy_and_summary(self)
+        logger.note(self, 'average-occupation SCF energy = %.15g', self.e_avg_occ)
+        return dft.roks.ROKS._finalize(self)
+
+    get_veff = get_veff
+    nuc_grad_method = nuc_grad_method
+    Gradients = nuc_grad_method
+    dump_flags = dump_flags
+    check_sanity = check_sanity
+
+class SymAdaptedAverageOccupationROKS(dft.rks_symm.SymAdaptedROKS):
+    """Symmetry-adapted average-occupation ROKS."""
+
+    _keys = {'e_avg_occ'}
+
+    def _finalize(self):
+        self.e_avg_occ = self.e_tot
+        if self.mo_coeff is not None and self.mo_occ is not None:
+            self.e_tot, self.scf_summary = _high_spin_energy_and_summary(self)
+        logger.note(self, 'average-occupation SCF energy = %.15g', self.e_avg_occ)
+        return dft.rks_symm.SymAdaptedROKS._finalize(self)
+
+    get_veff = get_veff
+    nuc_grad_method = nuc_grad_method
+    Gradients = nuc_grad_method
+    dump_flags = dump_flags
+    check_sanity = check_sanity
+
+
+def average_occ(mf):
+    """Return an average-occupation SCF object from a plain PySCF ROKS."""
+    if type(mf) is dft.roks.ROKS:
+        cls = AverageOccupationROKS
+    elif type(mf) is dft.rks_symm.SymAdaptedROKS:
+        cls = SymAdaptedAverageOccupationROKS
+    else:
+        raise TypeError('average_occ() requires a plain ROKS or SymAdaptedROKS object')
+    averaged = lib.view(mf, cls)
+    averaged.scf_summary = mf.scf_summary.copy()
+    averaged.converged = False
+    averaged.e_tot = None
+    averaged.e_avg_occ = None
+    return averaged
+
+
+dft.roks.ROKS.average_occ = average_occ
+dft.rks_symm.SymAdaptedROKS.average_occ = average_occ

src/nest/aocscf/tests/test_aocscf.py

@@ -0,0 +1,68 @@
+# 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.
+
+import unittest
+
+import numpy as np
+from pyscf import gto
+from nest import aocscf
+
+
+class KnownValues(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        mol = gto.Mole()
+        mol.verbose = 0
+        mol.output = '/dev/null'
+        mol.atom = '''
+        O   0.64372820   0.14077399  -0.04477253
+        O  -0.64862595  -0.12779073  -0.05445498
+        H   1.16027512  -0.65947800   0.36730132
+        H  -1.12109306   0.55561188   0.42651873
+        '''
+        mol.spin = 2
+        mol.basis = '631g'
+        cls.mol = mol.build()
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.mol.stdout.close()
+
+    def test_svwn_aocscf(self):
+        mf = self.mol.ROKS(xc='SVWN').average_occ()
+        mf.conv_tol = 1e-11
+        mf.conv_tol_grad = 1e-8
+        mf.max_cycle = 200
+        mf.grids.level = 3
+        mf.grids.prune = None
+        mf.small_rho_cutoff = 0.0
+        mf.run()
+
+        self.assertTrue(mf.converged)
+        self.assertAlmostEqual(mf.e_avg_occ, -150.15324131943828, delta=1e-7)
+        self.assertAlmostEqual(mf.e_tot, -150.18135492533739, delta=1e-7)
+
+    def test_b3lyp_aocscf(self):
+        mf = self.mol.ROKS(xc='B3LYP').average_occ()
+        mf.conv_tol = 1e-11
+        mf.conv_tol_grad = 1e-8
+        mf.max_cycle = 200
+        mf.grids.level = 3
+        mf.grids.prune = None
+        mf.small_rho_cutoff = 0.0
+        mf.run()
+
+        self.assertTrue(mf.converged)
+        self.assertAlmostEqual(mf.e_avg_occ, -151.18245418239550, delta=1e-7)
+        self.assertAlmostEqual(mf.e_tot, -151.25619865161033, delta=1e-7)