From 5a2fca5446ff9d497aa475c0d4456cdf0cd513c2 Mon Sep 17 00:00:00 2001 From: wtpeter Date: Fri, 25 Sep 2026 22:05:41 +0800 Subject: [PATCH 1/3] refactor: change dz0scf to aocscf --- examples/aocscf/01_average_occupation_scf.py | 36 ++++ examples/nttda/04_nttda_aocscf.py | 41 ++++ src/nest/__init__.py | 4 +- src/nest/{dz0scf => aocscf}/__init__.py | 51 +++-- src/nest/aocscf/aocscf.py | 123 ++++++++++++ src/nest/aocscf/tests/test_aocscf.py | 67 +++++++ src/nest/dz0scf/dz0scf.py | 108 ----------- src/nest/dz0scf/tests/test_dz0scf.py | 179 ------------------ src/nest/grad/{dz0scf.py => aocscf.py} | 102 +++++----- ...est_dz0scf_grad.py => test_aocscf_grad.py} | 32 ++-- src/nest/nttda/tests/test_nttda_aocscf.py | 80 ++++++++ 11 files changed, 433 insertions(+), 390 deletions(-) create mode 100644 examples/aocscf/01_average_occupation_scf.py create mode 100644 examples/nttda/04_nttda_aocscf.py rename src/nest/{dz0scf => aocscf}/__init__.py (63%) create mode 100644 src/nest/aocscf/aocscf.py create mode 100644 src/nest/aocscf/tests/test_aocscf.py delete mode 100644 src/nest/dz0scf/dz0scf.py delete mode 100644 src/nest/dz0scf/tests/test_dz0scf.py rename src/nest/grad/{dz0scf.py => aocscf.py} (85%) rename src/nest/grad/tests/{test_dz0scf_grad.py => test_aocscf_grad.py} (85%) create mode 100644 src/nest/nttda/tests/test_nttda_aocscf.py diff --git a/examples/aocscf/01_average_occupation_scf.py b/examples/aocscf/01_average_occupation_scf.py new file mode 100644 index 0000000..f65d4cf --- /dev/null +++ b/examples/aocscf/01_average_occupation_scf.py @@ -0,0 +1,36 @@ +#!/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=""" + 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", + spin=2, +) +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") diff --git a/examples/nttda/04_nttda_aocscf.py b/examples/nttda/04_nttda_aocscf.py new file mode 100644 index 0000000..79debdb --- /dev/null +++ b/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().run() +td = mf.NTTDA().set(nstates=2, deltaS=-1).run() + +print(f"SCF converged: {mf.converged}") +print(f"NTTDA converged: {td.converged}") +print(f"Excitation energies (Ha): {td.e}") +print(f"Total energies (Ha): {td.e_tot}", end=", ") +print(f"which equals to mf.e_tot + td.e: {mf.e_tot + td.e}") +td.analyze(verbose=4) diff --git a/src/nest/__init__.py b/src/nest/__init__.py index f90fbce..ec61da6 100644 --- a/src/nest/__init__.py +++ b/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"] diff --git a/src/nest/dz0scf/__init__.py b/src/nest/aocscf/__init__.py similarity index 63% rename from src/nest/dz0scf/__init__.py rename to src/nest/aocscf/__init__.py index c72f51c..868daa4 100644 --- a/src/nest/dz0scf/__init__.py +++ b/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', -] \ No newline at end of file +# 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'] diff --git a/src/nest/aocscf/aocscf.py b/src/nest/aocscf/aocscf.py new file mode 100644 index 0000000..72299be --- /dev/null +++ b/src/nest/aocscf/aocscf.py @@ -0,0 +1,123 @@ +# 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 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 + + +def _finalize(mf): + # 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 + mf.e_avg_occ = mf.e_tot + if mf.mo_coeff is not None and mf.mo_occ is not None: + mf_hf = lib.view(mf, dft.roks.ROKS) + mf_hf.scf_summary = mf.scf_summary.copy() + dm_hs = mf_hf.make_rdm1(mf.mo_coeff, mf.mo_occ) + veff_hs = mf_hf.get_veff(mf.mol, dm_hs) + energy = mf_hf.energy_tot(dm=dm_hs, h1e=mf_hf.get_hcore(), vhf=veff_hs) + mf.e_tot = energy + mf.scf_summary = mf_hf.scf_summary + logger.note(mf, 'average-occupation SCF energy = %.15g', mf.e_avg_occ) + return dft.roks.ROKS._finalize(mf) + + +class AverageOccupationROKS(dft.roks.ROKS): + """ROKS with average-occupation SCF orbitals and a high-spin ``e_tot``.""" + + _keys = {'e_avg_occ'} + + get_veff = get_veff + nuc_grad_method = nuc_grad_method + dump_flags = dump_flags + check_sanity = check_sanity + _finalize = _finalize + + +class SymAdaptedAverageOccupationROKS(dft.rks_symm.SymAdaptedROKS): + """Symmetry-adapted average-occupation ROKS.""" + + _keys = {'e_avg_occ'} + + get_veff = get_veff + nuc_grad_method = nuc_grad_method + dump_flags = dump_flags + check_sanity = check_sanity + _finalize = _finalize + + +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 diff --git a/src/nest/aocscf/tests/test_aocscf.py b/src/nest/aocscf/tests/test_aocscf.py new file mode 100644 index 0000000..ee47283 --- /dev/null +++ b/src/nest/aocscf/tests/test_aocscf.py @@ -0,0 +1,67 @@ +# 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 + +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) diff --git a/src/nest/dz0scf/dz0scf.py b/src/nest/dz0scf/dz0scf.py deleted file mode 100644 index 7d32fa3..0000000 --- a/src/nest/dz0scf/dz0scf.py +++ /dev/null @@ -1,108 +0,0 @@ -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) - - def nuc_grad_method(self): - """Return the Dz0SCF analytic nuclear-gradient driver.""" - from nest.grad.dz0scf import Gradients - return Gradients(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 deleted file mode 100644 index fc88305..0000000 --- a/src/nest/dz0scf/tests/test_dz0scf.py +++ /dev/null @@ -1,179 +0,0 @@ -# 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 diff --git a/src/nest/grad/dz0scf.py b/src/nest/grad/aocscf.py similarity index 85% rename from src/nest/grad/dz0scf.py rename to src/nest/grad/aocscf.py index 7033b69..9c7772b 100644 --- a/src/nest/grad/dz0scf.py +++ b/src/nest/grad/aocscf.py @@ -1,3 +1,17 @@ +# 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 __future__ import annotations import warnings @@ -66,7 +80,7 @@ def _block_pairs(rows: np.ndarray, cols: np.ndarray) -> tuple[np.ndarray, np.nda def _rotation_space(mo_occ: np.ndarray) -> RotationSpace: f = np.asarray(mo_occ, dtype=float) if f.ndim != 1: - raise ValueError(f"Dz0 occupations must be a one-dimensional array; got {f.shape}.") + raise ValueError(f"Spin-averaged occupations must be a one-dimensional array; got {f.shape}.") is_c = np.isclose(f, 2.0, atol=_OCC_TOL, rtol=0.0) is_o = np.isclose(f, 1.0, atol=_OCC_TOL, rtol=0.0) @@ -74,7 +88,7 @@ def _rotation_space(mo_occ: np.ndarray) -> RotationSpace: if not np.all(is_c | is_o | is_v): bad = np.where(~(is_c | is_o | is_v))[0] raise NotImplementedError( - "This implementation requires Dz0 occupations 0, 1, or 2. " + "This implementation requires occupations 0, 1, or 2. " f"Nonstandard occupations were found at MO indices {bad.tolist()}." ) @@ -90,24 +104,6 @@ def _rotation_space(mo_occ: np.ndarray) -> RotationSpace: return RotationSpace(p=p, q=q, f=f, nalpha=nalpha, nbeta=nbeta) -def _copy_ks_settings(source, target) -> None: - """Copy numerical-integration settings without copying unsupported wrappers.""" - for name in ( - "xc", - "nlc", - "grids", - "nlcgrids", - "_numint", - "max_memory", - "verbose", - "stdout", - "direct_scf_tol", - "small_rho_cutoff", - ): - if hasattr(source, name): - setattr(target, name, getattr(source, name)) - - def _ao_density(mo_coeff: np.ndarray, occupation: np.ndarray) -> np.ndarray: """Return D_AO = C occupation C^dagger.""" return (mo_coeff * occupation) @ mo_coeff.conj().T @@ -140,7 +136,7 @@ def _fractional_rks_fock_skeleton( This is PySCF's RKS ``make_h1`` construction with the density corrected from ``2 C_occ C_occ^T`` to the fractional-occupation density - ``C mo_occ C^T`` required by Dz0SCF. Range-separated exchange is included + ``C mo_occ C^T`` required by spin-averaged SCF. Range-separated exchange is included through PySCF's ``(omega, alpha, hyb)`` decomposition. """ mol = charge_mf.mol @@ -215,8 +211,8 @@ def _fractional_rks_fock_skeleton( return h1ao -class DZ0Gradients(lib.StreamObject): - """Analytic gradient driver for the Dz0SCF high-spin reference energy.""" +class AverageOccupationGradients(lib.StreamObject): + """Analytic gradient of the high-spin energy at spin-averaged SCF orbitals.""" _keys = { "base", @@ -230,7 +226,7 @@ class DZ0Gradients(lib.StreamObject): "de", "z", "g_hs", - "g_dz0", + "g_avg_occ", "b", "e_hs_unrelaxed", } @@ -251,7 +247,7 @@ def __init__(self, mf): self.de = None self.z = None self.g_hs = None - self.g_dz0 = None + self.g_avg_occ = None self.b = None self.e_hs_unrelaxed = None @@ -269,7 +265,7 @@ def __init__(self, mf): def dump_flags(self, verbose=None): log = logger.new_logger(self, verbose) - log.info("******** Dz0SCF high-spin-reference analytic gradient ********") + log.info("******** Spin-averaged SCF high-spin analytic gradient ********") log.info("Z-vector tolerance = %.3g", self.conv_tol) log.info("Z-vector max cycles = %d", self.max_cycle) log.info("GMRES restart = %d", self.restart) @@ -280,26 +276,26 @@ def _validate(self) -> None: mf = self.base mol = self.mol if getattr(mf, "mo_coeff", None) is None or getattr(mf, "mo_occ", None) is None: - raise RuntimeError("Run Dz0SCF before requesting its analytic gradient.") + raise RuntimeError("Run spin-averaged SCF before requesting its analytic gradient.") if hasattr(mf, "converged") and not mf.converged: - warnings.warn("Dz0SCF is not converged; its analytic gradient is not stationary.") + warnings.warn("Spin-averaged SCF is not converged; its analytic gradient is not stationary.") if np.iscomplexobj(mf.mo_coeff) and np.max(np.abs(np.asarray(mf.mo_coeff).imag)) > 1e-12: - raise NotImplementedError("Complex-orbital Dz0SCF gradients are not implemented.") + raise NotImplementedError("Complex-orbital spin-averaged gradients are not implemented.") if self.grid_response: raise NotImplementedError( "Moving-grid response is not implemented consistently in B^(0,A); " "use grid_response=False." ) if getattr(mf, "with_df", None) is not None: - raise NotImplementedError("Density-fitted Dz0SCF gradients are not implemented.") + raise NotImplementedError("Density-fitted spin-averaged gradients are not implemented.") if getattr(mf, "with_x2c", None) is not None: - raise NotImplementedError("X2C Dz0SCF gradients are not implemented.") + raise NotImplementedError("X2C spin-averaged gradients are not implemented.") if getattr(mf, "with_solvent", None) is not None: - raise NotImplementedError("Solvent-response Dz0SCF gradients are not implemented.") + raise NotImplementedError("Solvent-response spin-averaged gradients are not implemented.") if hasattr(mf, "do_nlc") and mf.do_nlc(): - raise NotImplementedError("Nonlocal-correlation (NLC/VV10) Dz0SCF gradients are not implemented.") + raise NotImplementedError("Nonlocal-correlation (NLC/VV10) spin-averaged gradients are not implemented.") if hasattr(mf, "do_disp") and mf.do_disp(): - raise NotImplementedError("Dispersion-corrected Dz0SCF gradients are not implemented.") + raise NotImplementedError("Dispersion-corrected spin-averaged gradients are not implemented.") if getattr(mol, "dimension", 3) != 3: raise NotImplementedError("Only molecular (three-dimensional) calculations are supported.") @@ -309,10 +305,9 @@ def _build_intermediates(self) -> None: c0 = np.asarray(mf.mo_coeff).real space = _rotation_space(np.asarray(mf.mo_occ)) - # Do not call the dft.RKS factory here: for mol.spin != 0 it returns - # ROKS, whereas R^(0) is the spin-unpolarized, charge-only RKS kernel. - charge_mf = dft.rks.RKS(mol) - _copy_ks_settings(mf, charge_mf) + # The factory dft.RKS returns ROKS when mol.spin != 0. A view keeps + # the reference's numerical settings but selects the RKS charge kernel. + charge_mf = lib.view(mf, dft.rks.RKS) charge_mf.mo_coeff = c0 charge_mf.mo_occ = space.f @@ -331,8 +326,7 @@ def _build_intermediates(self) -> None: with_nlc=False, ) - hs_mf = dft.ROKS(mol) - _copy_ks_settings(mf, hs_mf) + hs_mf = lib.view(mf, dft.roks.ROKS) hs_mf.mo_coeff = c0 hs_mf.mo_occ = space.f dm_hs = hs_mf.make_rdm1(c0, space.f) @@ -359,7 +353,7 @@ def _build_intermediates(self) -> None: self._w_hs_mo = w_hs_mo gap = space.occupation_gap - self.g_dz0 = 2.0 * gap * space.pack(f0mo) + self.g_avg_occ = 2.0 * gap * space.pack(f0mo) self.g_hs = 2.0 * ( (space.nalpha[space.q] - space.nalpha[space.p]) * space.pack(f_hs_mo[0]) + (space.nbeta[space.q] - space.nbeta[space.p]) * space.pack(f_hs_mo[1]) @@ -393,7 +387,7 @@ def _solve_z(self) -> np.ndarray: dtype=float, ) - # The exact real-orbital Dz0 Hessian is symmetric, so A^T z = g_HS + # The exact real-orbital spin-averaged Hessian is symmetric, so A^T z = g_HS # is solved with the same matrix-free action. This diagonal contains # the one-electron commutator part and is used only as a preconditioner. diagonal = 2.0 * space.occupation_gap * ( @@ -428,9 +422,9 @@ def _solve_z(self) -> np.ndarray: if info != 0: last = residuals[-1] if residuals else np.nan raise RuntimeError( - "Dz0 Z-vector GMRES did not converge: " + "Spin-averaged Z-vector GMRES did not converge: " f"info={info}, last preconditioned residual={last:.3e}. " - "Increase max_cycle/restart or inspect a near-singular Dz0 orbital Hessian." + "Increase max_cycle/restart or inspect a near-singular orbital Hessian." ) return z @@ -513,18 +507,18 @@ def _build_b(self) -> np.ndarray: return b.real def kernel(self, atmlst=None, verbose=None) -> np.ndarray: - """Compute and return the Dz0SCF high-spin-reference nuclear gradient.""" + """Compute the high-spin nuclear gradient at spin-averaged orbitals.""" log = logger.new_logger(self, verbose) self._validate() self.dump_flags(verbose) self._build_intermediates() - max_g0 = float(np.max(np.abs(self.g_dz0))) if self.g_dz0.size else 0.0 - log.info("max |g_Dz0| = %.6g", max_g0) + max_g0 = float(np.max(np.abs(self.g_avg_occ))) if self.g_avg_occ.size else 0.0 + log.info("max |g_avg_occ| = %.6g", max_g0) scf_grad_tol = getattr(self.base, "conv_tol_grad", 0.0) or 0.0 if max_g0 > max(1e-6, 100.0 * scf_grad_tol): warnings.warn( - f"The packed Dz0 orbital gradient is not small (max={max_g0:.3e}); " + f"The packed spin-averaged orbital gradient is not small (max={max_g0:.3e}); " "the analytic-gradient stationarity equation may be inaccurate." ) @@ -542,7 +536,7 @@ def kernel(self, atmlst=None, verbose=None) -> np.ndarray: self.de = result if log.verbose >= logger.NOTE: - logger.note(self, "--------------- Dz0SCF reference gradients ---------------") + logger.note(self, "------------ Spin-averaged reference gradients ------------") rhf_grad._write(log, self.mol, result, atmlst) logger.note(self, "----------------------------------------------------------") return result @@ -550,13 +544,7 @@ def kernel(self, atmlst=None, verbose=None) -> np.ndarray: grad = kernel -Gradients = DZ0Gradients -Grad = DZ0Gradients - - -def nuc_grad_method(mf) -> DZ0Gradients: - """Functional constructor useful before binding the class as a method.""" - return DZ0Gradients(mf) +Gradients = AverageOccupationGradients -__all__ = ["DZ0Gradients", "Gradients", "Grad", "RotationSpace", "nuc_grad_method"] +__all__ = ["AverageOccupationGradients", "Gradients"] diff --git a/src/nest/grad/tests/test_dz0scf_grad.py b/src/nest/grad/tests/test_aocscf_grad.py similarity index 85% rename from src/nest/grad/tests/test_dz0scf_grad.py rename to src/nest/grad/tests/test_aocscf_grad.py index 6cef964..82b82ae 100644 --- a/src/nest/grad/tests/test_dz0scf_grad.py +++ b/src/nest/grad/tests/test_aocscf_grad.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Regression tests for the Dz0SCF high-spin-reference analytic gradient. +"""Regression tests for the average-occupation SCF high-spin analytic gradient. The reference values were generated with PySCF 2.13.0. The analytic-gradient implementation was independently checked against central finite differences. @@ -27,8 +27,8 @@ import numpy as np from pyscf import gto -from nest.dz0scf import DZ0SCF -from nest.grad.dz0scf import Gradients +from nest import aocscf +from nest.grad.aocscf import Gradients GRADIENT_ATOL = 1.0e-6 @@ -55,8 +55,8 @@ ) -def run_dz0_gradient(atom, spin, xc, basis): - """Run one tightly converged and reproducible Dz0SCF gradient.""" +def run_aocscf_gradient(atom, spin, xc, basis): + """Run one tightly converged and reproducible average-occupation gradient.""" mol = gto.M( atom=atom, unit="Angstrom", @@ -68,7 +68,7 @@ def run_dz0_gradient(atom, spin, xc, basis): output="/dev/null", ) - mf = DZ0SCF(mol, xc=xc) + mf = mol.ROKS(xc=xc).average_occ() mf.conv_tol = 1.0e-12 mf.conv_tol_grad = 1.0e-9 mf.max_cycle = 120 @@ -78,15 +78,15 @@ def run_dz0_gradient(atom, spin, xc, basis): mf.kernel() if not mf.converged: mol.stdout.close() - raise RuntimeError(f"{xc}/{basis} Dz0SCF did not converge") + raise RuntimeError(f"{xc}/{basis} average-occupation SCF did not converge") - # Exercise the public API installed on the Dz0SCF class/mixin. + # Exercise the public API installed on the average-occupation class. grad_obj = mf.nuc_grad_method() if not isinstance(grad_obj, Gradients): mol.stdout.close() raise TypeError( - "DZ0SCF.nuc_grad_method() did not return " - "nest.grad.dz0scf.Gradients" + "AverageOccupationROKS.nuc_grad_method() did not return " + "nest.grad.aocscf.Gradients" ) grad_obj.conv_tol = 1.0e-10 @@ -97,7 +97,7 @@ def run_dz0_gradient(atom, spin, xc, basis): class KnownValues(unittest.TestCase): - def assert_dz0_result( + def assert_aocscf_result( self, mf, grad_obj, @@ -114,7 +114,7 @@ def assert_dz0_result( self.assertEqual(observed_space, expected_space) self.assertLess( - float(np.max(np.abs(grad_obj.g_dz0))), + float(np.max(np.abs(grad_obj.g_avg_occ))), STATIONARITY_TOL, ) @@ -138,7 +138,7 @@ def assert_dz0_result( def test_pbe_nh2_gradient(self): """Check all NH2/PBE/6-31G components in a 44-D response space.""" - mol, mf, grad_obj, gradient = run_dz0_gradient( + mol, mf, grad_obj, gradient = run_aocscf_gradient( atom=""" N 0.000000 -0.040000 0.000000 H 0.000000 0.780000 0.590000 @@ -149,7 +149,7 @@ def test_pbe_nh2_gradient(self): basis="6-31g", ) try: - self.assert_dz0_result( + self.assert_aocscf_result( mf, grad_obj, gradient, @@ -161,7 +161,7 @@ def test_pbe_nh2_gradient(self): def test_b3lyp_ch2_two_open_orbitals_gradient(self): """Check all CH2 components with a two-dimensional open-shell space.""" - mol, mf, grad_obj, gradient = run_dz0_gradient( + mol, mf, grad_obj, gradient = run_aocscf_gradient( atom=""" C 0.020000 -0.030000 0.010000 H -0.020000 0.800000 0.620000 @@ -172,7 +172,7 @@ def test_b3lyp_ch2_two_open_orbitals_gradient(self): basis="sto-3g", ) try: - self.assert_dz0_result( + self.assert_aocscf_result( mf, grad_obj, gradient, diff --git a/src/nest/nttda/tests/test_nttda_aocscf.py b/src/nest/nttda/tests/test_nttda_aocscf.py new file mode 100644 index 0000000..413bdcb --- /dev/null +++ b/src/nest/nttda/tests/test_nttda_aocscf.py @@ -0,0 +1,80 @@ +# 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, nttda + + +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_nttda_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() + + ref = np.array([-0.21222618958794592, 0.022735913574159522]) + td = mf.NTTDA().set(nstates=2, deltaS=-1, nobeta=True, conv_tol=1e-5, max_cycle=200).run() + self.assertTrue(np.all(td.converged)) + self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6) + + ref = np.array([-0.001800257693000168, 0.030755390462627187]) + td = mf.NTTDA().set(nstates=2, deltaS=0, nobeta=True, conv_tol=1e-5, max_cycle=200).run() + self.assertTrue(np.all(td.converged)) + self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6) + + def test_b3lyp_nttda_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() + + ref = np.array([-0.22131467106409972, 0.020196490053532357]) + td = mf.NTTDA().set(nstates=2, deltaS=-1, nobeta=True, conv_tol=1e-5, max_cycle=200).run() + self.assertTrue(np.all(td.converged)) + self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6) + + ref = np.array([-0.006072490213890671, 0.034052405714217956]) + td = mf.NTTDA().set(nstates=2, deltaS=0, nobeta=True, conv_tol=1e-5, max_cycle=200).run() + self.assertTrue(np.all(td.converged)) + self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6) From 0be8cb1ebb1bc00601d20383329cb8156ec41759 Mon Sep 17 00:00:00 2001 From: wtpeter Date: Fri, 25 Sep 2026 22:24:07 +0800 Subject: [PATCH 2/3] modify examples and _finalize of aocscf --- examples/aocscf/01_average_occupation_scf.py | 11 ++-- examples/aocscf/02_nttda.py | 1 + examples/nttda/01_nttda.py | 6 ++- examples/nttda/04_nttda_aocscf.py | 14 +++--- examples/sftda/01_sftddft.py | 3 ++ src/nest/aocscf/aocscf.py | 53 ++++++++++++-------- src/nest/aocscf/tests/test_aocscf.py | 14 ++++++ 7 files changed, 68 insertions(+), 34 deletions(-) create mode 120000 examples/aocscf/02_nttda.py diff --git a/examples/aocscf/01_average_occupation_scf.py b/examples/aocscf/01_average_occupation_scf.py index f65d4cf..04dbfa3 100644 --- a/examples/aocscf/01_average_occupation_scf.py +++ b/examples/aocscf/01_average_occupation_scf.py @@ -16,21 +16,22 @@ """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=""" - 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 + 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 diff --git a/examples/aocscf/02_nttda.py b/examples/aocscf/02_nttda.py new file mode 120000 index 0000000..ba70ac2 --- /dev/null +++ b/examples/aocscf/02_nttda.py @@ -0,0 +1 @@ +../nttda/04_nttda_aocscf.py \ No newline at end of file diff --git a/examples/nttda/01_nttda.py b/examples/nttda/01_nttda.py index 16674dd..a5b216b 100644 --- a/examples/nttda/01_nttda.py +++ b/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}") diff --git a/examples/nttda/04_nttda_aocscf.py b/examples/nttda/04_nttda_aocscf.py index 79debdb..97bffec 100644 --- a/examples/nttda/04_nttda_aocscf.py +++ b/examples/nttda/04_nttda_aocscf.py @@ -30,12 +30,12 @@ spin=2, symmetry=True ) -mf = mol.ROKS(xc="CAM-B3LYP").average_occ().run() -td = mf.NTTDA().set(nstates=2, deltaS=-1).run() +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() -print(f"SCF converged: {mf.converged}") -print(f"NTTDA converged: {td.converged}") -print(f"Excitation energies (Ha): {td.e}") -print(f"Total energies (Ha): {td.e_tot}", end=", ") -print(f"which equals to mf.e_tot + td.e: {mf.e_tot + td.e}") 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}") diff --git a/examples/sftda/01_sftddft.py b/examples/sftda/01_sftddft.py index 217755e..d20b641 100644 --- a/examples/sftda/01_sftddft.py +++ b/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 diff --git a/src/nest/aocscf/aocscf.py b/src/nest/aocscf/aocscf.py index 72299be..f847fdc 100644 --- a/src/nest/aocscf/aocscf.py +++ b/src/nest/aocscf/aocscf.py @@ -36,6 +36,20 @@ def get_veff(mf, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1): 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 @@ -61,46 +75,43 @@ def check_sanity(mf): return mf -def _finalize(mf): - # 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 - mf.e_avg_occ = mf.e_tot - if mf.mo_coeff is not None and mf.mo_occ is not None: - mf_hf = lib.view(mf, dft.roks.ROKS) - mf_hf.scf_summary = mf.scf_summary.copy() - dm_hs = mf_hf.make_rdm1(mf.mo_coeff, mf.mo_occ) - veff_hs = mf_hf.get_veff(mf.mol, dm_hs) - energy = mf_hf.energy_tot(dm=dm_hs, h1e=mf_hf.get_hcore(), vhf=veff_hs) - mf.e_tot = energy - mf.scf_summary = mf_hf.scf_summary - logger.note(mf, 'average-occupation SCF energy = %.15g', mf.e_avg_occ) - return dft.roks.ROKS._finalize(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 dump_flags = dump_flags check_sanity = check_sanity - _finalize = _finalize - 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 dump_flags = dump_flags check_sanity = check_sanity - _finalize = _finalize def average_occ(mf): diff --git a/src/nest/aocscf/tests/test_aocscf.py b/src/nest/aocscf/tests/test_aocscf.py index ee47283..9a2639b 100644 --- a/src/nest/aocscf/tests/test_aocscf.py +++ b/src/nest/aocscf/tests/test_aocscf.py @@ -14,6 +14,7 @@ import unittest +import numpy as np from pyscf import gto from nest import aocscf @@ -65,3 +66,16 @@ def test_b3lyp_aocscf(self): 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) + + def test_symmetry_adapted_orbital_order(self): + mol = gto.M( + atom='H 0 .934473 -.588078; H 0 -.934473 -.588078; C 0 0 0; O 0 0 1.221104', + basis='6-31g', spin=2, symmetry=True, verbose=0, + ) + mf = mol.ROKS(xc='CAM-B3LYP').average_occ().run() + + self.assertTrue(mf.converged) + self.assertAlmostEqual(mf.e_tot, -114.285284450421, delta=1e-6) + np.testing.assert_array_equal(mf.mo_occ[:7], 2) + np.testing.assert_array_equal(mf.mo_occ[7:9], 1) + np.testing.assert_array_equal(mf.mo_occ[9:], 0) From 9a7b25caa99edf2e55e1e558a43f993e4ef01856 Mon Sep 17 00:00:00 2001 From: wtpeter Date: Sat, 26 Sep 2026 08:50:13 +0800 Subject: [PATCH 3/3] register Gradients for aocscf --- src/nest/aocscf/aocscf.py | 2 ++ src/nest/aocscf/tests/test_aocscf.py | 13 ------------- 2 files changed, 2 insertions(+), 13 deletions(-) diff --git a/src/nest/aocscf/aocscf.py b/src/nest/aocscf/aocscf.py index f847fdc..3ad5258 100644 --- a/src/nest/aocscf/aocscf.py +++ b/src/nest/aocscf/aocscf.py @@ -93,6 +93,7 @@ def _finalize(self): get_veff = get_veff nuc_grad_method = nuc_grad_method + Gradients = nuc_grad_method dump_flags = dump_flags check_sanity = check_sanity @@ -110,6 +111,7 @@ def _finalize(self): get_veff = get_veff nuc_grad_method = nuc_grad_method + Gradients = nuc_grad_method dump_flags = dump_flags check_sanity = check_sanity diff --git a/src/nest/aocscf/tests/test_aocscf.py b/src/nest/aocscf/tests/test_aocscf.py index 9a2639b..7d66b23 100644 --- a/src/nest/aocscf/tests/test_aocscf.py +++ b/src/nest/aocscf/tests/test_aocscf.py @@ -66,16 +66,3 @@ def test_b3lyp_aocscf(self): 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) - - def test_symmetry_adapted_orbital_order(self): - mol = gto.M( - atom='H 0 .934473 -.588078; H 0 -.934473 -.588078; C 0 0 0; O 0 0 1.221104', - basis='6-31g', spin=2, symmetry=True, verbose=0, - ) - mf = mol.ROKS(xc='CAM-B3LYP').average_occ().run() - - self.assertTrue(mf.converged) - self.assertAlmostEqual(mf.e_tot, -114.285284450421, delta=1e-6) - np.testing.assert_array_equal(mf.mo_occ[:7], 2) - np.testing.assert_array_equal(mf.mo_occ[7:9], 1) - np.testing.assert_array_equal(mf.mo_occ[9:], 0)