diff --git a/src/nest/__init__.py b/src/nest/__init__.py index 52cc229..f90fbce 100644 --- a/src/nest/__init__.py +++ b/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"] diff --git a/src/nest/dz0scf/__init__.py b/src/nest/dz0scf/__init__.py new file mode 100644 index 0000000..c72f51c --- /dev/null +++ b/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', +] \ No newline at end of file diff --git a/src/nest/dz0scf/dz0scf.py b/src/nest/dz0scf/dz0scf.py new file mode 100644 index 0000000..943ab77 --- /dev/null +++ b/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 + diff --git a/src/nest/dz0scf/tests/test_dz0scf.py b/src/nest/dz0scf/tests/test_dz0scf.py new file mode 100644 index 0000000..fc88305 --- /dev/null +++ b/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, + ) \ No newline at end of file