Add Dz0SCF reference implementation

src/nest/__init__.py

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

src/nest/dz0scf/__init__.py

@@ -0,0 +1,28 @@
+#!/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',
+]

src/nest/dz0scf/dz0scf.py

@@ -0,0 +1,103 @@
+import numpy as np
+
+from pyscf import dft,  lib
+from pyscf.dft import uks
+
+def _as_spin_unpolarized_dm(dm):
+    arr = np.asarray(dm)
+    if arr.ndim == 2:
+        dm0 = arr
+    elif arr.ndim == 3 and arr.shape[0] == 2:
+        dm0 = arr[0] + arr[1]
+    else:
+        raise ValueError(
+            f'Expected a 2-D density or two spin densities; got {arr.shape}'
+        )
+
+    dm_ens = np.asarray((0.5 * dm0, 0.5 * dm0))
+
+    mo_coeff = getattr(dm, 'mo_coeff', None)
+    mo_occ = getattr(dm, 'mo_occ', None)
+    if mo_coeff is not None and mo_occ is not None:
+        coeff = mo_coeff
+        if isinstance(coeff, (tuple, list)) or np.asarray(coeff).ndim == 3:
+            coeff = coeff[0]
+
+        occ = np.asarray(mo_occ)
+        if occ.ndim == 2 and occ.shape[0] == 2:
+            occ = occ[0] + occ[1]
+
+        dm_ens = lib.tag_array(
+            dm_ens,
+            mo_coeff=(coeff, coeff),
+            mo_occ=(0.5 * occ, 0.5 * occ),
+        )
+
+    return dm_ens
+
+def evaluate_high_spin_energy(mf):
+    evaluator = dft.ROKS(mf.mol)
+    evaluator.xc = mf.xc
+    evaluator.max_memory = mf.max_memory
+
+    evaluator.grids = mf.grids
+    if hasattr(mf, 'nlcgrids'):
+        evaluator.nlcgrids = mf.nlcgrids
+
+    dm_hs = evaluator.make_rdm1(mf.mo_coeff, mf.mo_occ)
+    hcore = evaluator.get_hcore()
+    veff = evaluator.get_veff(mf.mol, dm_hs)
+
+    return evaluator.energy_tot(
+        dm=dm_hs,
+        h1e=hcore,
+        vhf=veff,
+    )
+
+class _DZ0VeffMixin:
+    def get_veff(
+        self,
+        mol=None,
+        dm=None,
+        dm_last=0,
+        vhf_last=0,
+        hermi=1,
+    ):
+        if mol is None:
+            mol = self.mol
+        if dm is None:
+            dm = self.make_rdm1()
+
+        dm_ens = _as_spin_unpolarized_dm(dm)
+
+        if np.ndim(dm_last) >= 2:
+            dm_last = _as_spin_unpolarized_dm(dm_last)
+
+        return uks.get_veff(
+            self,
+            mol,
+            dm_ens,
+            dm_last,
+            vhf_last,
+            hermi,
+        )
+    def high_spin_energy(self):
+        return evaluate_high_spin_energy(self)
+
+class EnsembleROKS(_DZ0VeffMixin, dft.roks.ROKS):
+    pass
+
+class SymAdaptedEnsembleROKS(_DZ0VeffMixin, dft.rks_symm.SymAdaptedROKS):
+    pass
+
+def DZ0SCF(mol, xc=None):
+    if mol.symmetry:
+        mf = SymAdaptedEnsembleROKS(mol)
+    else:
+        mf = EnsembleROKS(mol)
+
+    if xc is not None:
+        mf.xc = xc
+
+    return mf
+

src/nest/dz0scf/tests/test_dz0scf.py

@@ -0,0 +1,179 @@
+# 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.dz0scf import DZ0SCF
+from nest.nttda import NTTDA
+
+
+class KnownValues(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.mol = gto.M(
+            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
+""",
+            basis="6-31g",
+            unit="Angstrom",
+            charge=0,
+            spin=2,
+            symmetry=False,
+            verbose=0,
+        )
+
+    def test_svwn_dz0scf(self):
+        mf = DZ0SCF(self.mol, xc="SVWN")
+        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.kernel()
+
+        self.assertTrue(mf.converged)
+
+        e_dz0_ref = -150.15324131943828
+        e_high_spin_ref = -150.18135492533739
+
+        self.assertAlmostEqual(
+            mf.e_tot,
+            e_dz0_ref,
+            delta=1e-7,
+        )
+        self.assertAlmostEqual(
+            mf.high_spin_energy(),
+            e_high_spin_ref,
+            delta=1e-7,
+        )
+
+        td_s = NTTDA(mf)
+        td_s.deltaS = -1
+        td_s.nstates = 2
+        td_s.nobeta = True
+        td_s.conv_tol = 1e-5
+        td_s.max_cycle = 200
+
+        omega_s, _ = td_s.kernel()
+
+        omega_s_ref = np.array([
+            -0.21222618958794592,
+             0.022735913574159522,
+        ])
+
+        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,
+        )
+
+        td_t = NTTDA(mf)
+        td_t.deltaS = 0
+        td_t.nstates = 2
+        td_t.nobeta = True
+        td_t.conv_tol = 1e-5
+        td_t.max_cycle = 200
+
+        omega_t, _ = td_t.kernel()
+
+        omega_t_ref = np.array([
+            -0.001800257693000168,
+             0.030755390462627187,
+        ])
+
+        self.assertTrue(np.all(np.asarray(td_t.converged)))
+        np.testing.assert_allclose(
+            np.asarray(omega_t),
+            omega_t_ref,
+            rtol=0.0,
+            atol=1e-6,
+        )
+
+    def test_b3lyp_dz0scf(self):
+        mf = DZ0SCF(self.mol, xc="B3LYP")
+        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.kernel()
+
+        self.assertTrue(mf.converged)
+
+        e_dz0_ref = -151.18245418239550
+        e_high_spin_ref = -151.25619865161033
+
+        self.assertAlmostEqual(
+            mf.e_tot,
+            e_dz0_ref,
+            delta=1e-7,
+        )
+        self.assertAlmostEqual(
+            mf.high_spin_energy(),
+            e_high_spin_ref,
+            delta=1e-7,
+        )
+
+        td_s = NTTDA(mf)
+        td_s.deltaS = -1
+        td_s.nstates = 2
+        td_s.nobeta = True
+        td_s.conv_tol = 1e-5
+        td_s.max_cycle = 200
+
+        omega_s, _ = td_s.kernel()
+
+        omega_s_ref = np.array([
+            -0.22131467106409972,
+             0.020196490053532357,
+        ])
+
+        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,
+        )
+
+        td_t = NTTDA(mf)
+        td_t.deltaS = 0
+        td_t.nstates = 2
+        td_t.nobeta = True
+        td_t.conv_tol = 1e-5
+        td_t.max_cycle = 200
+
+        omega_t, _ = td_t.kernel()
+
+        omega_t_ref = np.array([
+            -0.006072490213890671,
+             0.034052405714217956,
+        ])
+
+        self.assertTrue(np.all(np.asarray(td_t.converged)))
+        np.testing.assert_allclose(
+            np.asarray(omega_t),
+            omega_t_ref,
+            rtol=0.0,
+            atol=1e-6,
+        )