From adceb31f8d199a32101bb1719190bb51c91f4185 Mon Sep 17 00:00:00 2001 From: wtpeter Date: Tue, 22 Sep 2026 20:50:06 +0800 Subject: [PATCH 1/3] fix(nttda): project out the redundant spin-lowering mode --- src/nest/nttda/nttda.py | 36 +++++++++++---- src/nest/nttda/tests/test_nttda.py | 73 ++++++++++++++++++++++++++++++ 2 files changed, 99 insertions(+), 10 deletions(-) diff --git a/src/nest/nttda/nttda.py b/src/nest/nttda/nttda.py index 7724a13..fdc2bf0 100644 --- a/src/nest/nttda/nttda.py +++ b/src/nest/nttda/nttda.py @@ -783,8 +783,6 @@ def kernel(self, x0=None, nstates=None): nstates = self.nstates else: self.nstates = nstates - if self.deltaS == -1: - nstates += 1 log = logger.Logger(self.stdout, self.verbose) def all_eigs(w, v, nroots, envs): @@ -794,11 +792,30 @@ def all_eigs(w, v, nroots, envs): vind, hdiag = self.gen_vind_sc() precond = self.get_precond(hdiag) elif self.deltaS == -1: - vind, hdiag = self.gen_vind_sfd() - precond = self.get_precond(hdiag) + base_vind, hdiag = self.gen_vind_sfd() + base_precond = self.get_precond(hdiag) csidx, osidx, vsidx = _orbital_indices(self) - nocc = len(csidx) + len(osidx) - nvir = len(osidx) + len(vsidx) + ncs, nos = len(csidx), len(osidx) + nocc = ncs + nos + nvir = nos + len(vsidx) + nstates = min(nstates, hdiag.size - 1) + self.nstates = nstates + open_diag = np.arange(nos) + + def project(zs): + # Remove only the OO identity direction; keep the stored layout. + projected = np.array(zs, copy=True) + oo = projected.reshape(-1, nocc, nvir)[:, ncs:, :nos] + trace = np.trace(oo, axis1=1, axis2=2) / nos + oo[:, open_diag, open_diag] -= trace[:, None] + return projected + + def vind(zs): + return project(base_vind(project(zs))) + + def precond(residual, energy): + # Unequal diagonal denominators can reintroduce the OO trace. + return project(base_precond(residual, energy)) elif self.deltaS == 1: vind, hdiag = self.gen_vind_sfu() precond = self.get_precond(hdiag) @@ -811,6 +828,9 @@ def all_eigs(w, v, nroots, envs): x0sym = None if x0 is None: x0 = self.init_guess(hdiag) + if self.deltaS == -1: + # Davidson's initial QR removes dependencies introduced by projection. + x0 = project(x0) self.converged, self.e, x1 = lr_eigh( vind, @@ -830,10 +850,6 @@ def all_eigs(w, v, nroots, envs): self.xy = [(xi, 0) for xi in x1] elif self.deltaS == -1: self.xy = [(xi.reshape(nocc, nvir), 0) for xi in x1] - mask = abs(self.e) > 1e-8 - self.converged = np.asarray(self.converged)[mask] - self.e = self.e[mask] - self.xy = [xy for xy, keep in zip(self.xy, mask) if keep] self.nstates = len(self.e) elif self.deltaS == 1: self.xy = [(xi.reshape(ncs, nvs), 0) for xi in x1] diff --git a/src/nest/nttda/tests/test_nttda.py b/src/nest/nttda/tests/test_nttda.py index f807a32..56b8073 100644 --- a/src/nest/nttda/tests/test_nttda.py +++ b/src/nest/nttda/tests/test_nttda.py @@ -13,6 +13,7 @@ # limitations under the License. import unittest +from unittest.mock import patch import numpy as np from pyscf import gto from nest import nttda @@ -111,6 +112,78 @@ def test_cam_b3lyp_nttda(self): self.assertTrue(np.all(td.converged)) self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6) + def test_hf_nttda_projected_guess(self): + mf = self.mol.ROKS(xc='HF').run() + td = mf.NTTDA().set(nstates=2, deltaS=-1, nobeta=False, conv_tol=1e-6) + vind, hdiag = td.gen_vind_sfd() + nc = np.count_nonzero(mf.mo_occ == 2) + no = np.count_nonzero(mf.mo_occ == 1) + nv = np.count_nonzero(mf.mo_occ == 0) + + # The normalized OO identity is a null vector of the original operator. + q = np.zeros((nc + no, no + nv)) + q[nc:, :no] = np.eye(no) / np.sqrt(no) + self.assertAlmostEqual(np.linalg.norm(vind(q.ravel())), 0, delta=1e-12) + + # Deliberately add this redundant direction to every initial guess. + x0 = td.init_guess(hdiag) + q.ravel() + original = x0.copy() + td.kernel(x0=x0) + + ref = np.array([-0.25588162251385815, 0.03179164805915535]) + self.assertTrue(np.all(td.converged)) + self.assertEqual(td.nstates, 2) + self.assertEqual(len(td.xy), 2) + self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6) + self.assertEqual(abs(x0 - original).max(), 0) + for energy, (x, _) in zip(td.e, td.xy): + self.assertAlmostEqual(np.trace(x[nc:, :no]), 0, delta=1e-12) + residual = vind(x.ravel()).ravel() - energy * x.ravel() + self.assertLess(np.linalg.norm(residual), td.conv_tol) + + def test_nttda_physical_zero_root(self): + # Two open orbitals, no core or virtual orbitals: amplitudes are OO only. + mol = gto.M(atom='H 0 0 0; H 0 0 1', spin=2, basis='sto-3g', verbose=0) + mf = mol.ROKS(xc='HF').run() + td = mf.NTTDA().set(deltaS=-1) + + # Coordinates: (X_00, X_01, X_10, X_11). + # (1, 0, 0, 1) is the redundant zero mode. + # The traceless directions have eigenvalues -2, -1, 0, respectively. + matrix = np.array([ + [-1., 0., 0., 1.], + [ 0., -1., 0., 0.], + [ 0., 0., 0., 0.], + [ 1., 0., 0., -1.], + ]) + + def vind(zs): + return np.asarray(zs) @ matrix.T + + with patch.object(td, 'gen_vind_sfd', return_value=(vind, matrix.diagonal().copy())): + # A negative target root must not cause an extra root to be returned. + td.kernel(nstates=1) + self.assertTrue(np.all(td.converged)) + self.assertEqual(td.nstates, 1) + self.assertEqual(len(td.e), 1) + self.assertAlmostEqual(td.e[0], -2, delta=1e-12) + + # Keep the physical zero root; remove only the OO identity direction. + ref = np.array([-2., -1., 0.]) + td.kernel(nstates=3) + self.assertTrue(np.all(td.converged)) + self.assertEqual(td.nstates, 3) + self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-12) + for x, _ in td.xy: + self.assertAlmostEqual(np.trace(x), 0, delta=1e-12) + + # Four stored coordinates contain only three physical directions. + td.kernel(nstates=4) + self.assertTrue(np.all(td.converged)) + self.assertEqual(td.nstates, 3) + self.assertEqual(len(td.xy), 3) + self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-12) + if __name__ == '__main__': print('Full tests for noncollinear tensor TDA based on ROKS reference') From 5155ab78507cc38edcfc8dbc6155a88051107d6f Mon Sep 17 00:00:00 2001 From: wtpeter Date: Tue, 22 Sep 2026 21:04:59 +0800 Subject: [PATCH 2/3] fix(nttda): improve sanity checks and calculation logging --- src/nest/nttda/nttda.py | 35 ++++++++++++++++---- src/nest/nttda/tests/test_nttda.py | 52 +++++++++++++++++++++++++++++- 2 files changed, 79 insertions(+), 8 deletions(-) diff --git a/src/nest/nttda/nttda.py b/src/nest/nttda/nttda.py index fdc2bf0..135a93a 100644 --- a/src/nest/nttda/nttda.py +++ b/src/nest/nttda/nttda.py @@ -633,7 +633,7 @@ def gen_vind_sfd(td): virt_cols = slice(nos, None) s = nos * 0.5 - assert s >= 0.5, 'NTTDA for Sf=Si-1 only supports case that Si>=1.' + assert s >= 1, 'NTTDA for Sf=Si-1 only supports case that Si>=1.' assert s == (mf.mol.nelec[0] - mf.mol.nelec[1]) * 0.5 log = logger.new_logger(td) @@ -761,9 +761,30 @@ class NTTDA(TDBase): deltaS = -1 nobeta = False + singlet = None _keys = {'deltaS', 'nobeta'} + def dump_flags(self, verbose=None): + TDBase.dump_flags(self, verbose) + log = logger.new_logger(self, verbose) + s = self.mol.spin * 0.5 + log.info('deltaS = %s (Si = %g -> Sf = %g)', self.deltaS, s, s + self.deltaS) + if self.nobeta: + log.info('Numerical stabilization enabled to avoid potential divergence ' + 'at low local beta-electron density') + if self.deltaS == -1: + log.info('Redundant zero-energy state excluded from the excitation calculation') + return self + + def check_sanity(self): + if self.deltaS not in (-1, 0, 1): + raise ValueError('deltaS must be -1, 0, or 1') + if not isinstance(self.nstates, (int, np.integer)) or self.nstates <= 0: + raise ValueError('nstates must be a positive integer') + TDBase.check_sanity(self) + return self + def init_guess(self, hdiag, nstates=None): if nstates is None: nstates = self.nstates @@ -776,14 +797,16 @@ def init_guess(self, hdiag, nstates=None): def kernel(self, x0=None, nstates=None): cpu0 = (logger.process_clock(), logger.perf_counter()) - self.check_sanity() - self.dump_flags() - if nstates is None: nstates = self.nstates else: self.nstates = nstates - log = logger.Logger(self.stdout, self.verbose) + + self.check_sanity() + if self.verbose >= logger.INFO: + self.dump_flags() + + log = logger.new_logger(self) def all_eigs(w, v, nroots, envs): return w, v, np.arange(w.size) @@ -822,8 +845,6 @@ def precond(residual, energy): csidx, _, vsidx = _orbital_indices(self) ncs = len(csidx) nvs = len(vsidx) - else: - raise ValueError('deltaS should be -1, 0, or 1') x0sym = None if x0 is None: diff --git a/src/nest/nttda/tests/test_nttda.py b/src/nest/nttda/tests/test_nttda.py index 56b8073..32241ec 100644 --- a/src/nest/nttda/tests/test_nttda.py +++ b/src/nest/nttda/tests/test_nttda.py @@ -12,10 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import io import unittest from unittest.mock import patch import numpy as np -from pyscf import gto +from pyscf import dft, gto +from pyscf.lib import logger from nest import nttda @@ -141,6 +143,54 @@ def test_hf_nttda_projected_guess(self): residual = vind(x.ravel()).ravel() - energy * x.ravel() self.assertLess(np.linalg.norm(residual), td.conv_tol) + def test_nttda_check_sanity(self): + mf = self.mol.ROKS(xc='HF').run() + td = mf.NTTDA() + self.assertIs(td.check_sanity(), td) + + td.deltaS = 2 + with self.assertRaisesRegex(ValueError, 'deltaS must be'): + td.check_sanity() + with self.assertRaisesRegex(ValueError, 'deltaS must be'): + td.kernel() + td.deltaS = -1 + with self.assertRaisesRegex(ValueError, 'nstates must be a positive integer'): + td.kernel(nstates=0) + with self.assertRaisesRegex(ValueError, 'nstates must be a positive integer'): + td.kernel(nstates=1.5) + + mol = gto.M(atom='H 0 0 0', spin=1, basis='sto-3g', verbose=0) + td = mol.ROKS(xc='HF').run().NTTDA() + with self.assertRaisesRegex(AssertionError, 'Si>=1'): + td.kernel() + + mol = gto.M(atom='He 0 0 0', basis='sto-3g', verbose=0) + td = dft.roks.ROKS(mol, xc='HF').run().NTTDA().set(deltaS=0) + with self.assertRaisesRegex(AssertionError, 'Si>=1/2'): + td.kernel() + + def test_nttda_dump_flags(self): + mf = self.mol.ROKS(xc='HF').run() + td = mf.NTTDA().set(deltaS=-1, nobeta=True) + with io.StringIO() as output: + td.stdout = output + self.assertIs(td.dump_flags(verbose=logger.INFO), td) + text = output.getvalue() + self.assertIn('deltaS = -1 (Si = 1 -> Sf = 0)', text) + self.assertIn('Numerical stabilization enabled', text) + self.assertIn('low local beta-electron density', text) + self.assertIn('Redundant zero-energy state excluded from the excitation calculation', text) + self.assertNotIn('singlet', text) + + output.seek(0) + output.truncate() + td.nobeta = False + td.verbose = logger.INFO + td.kernel(nstates=1) + self.assertTrue(np.all(td.converged)) + self.assertIn('nstates = 1', output.getvalue()) + self.assertNotIn('Numerical stabilization enabled', output.getvalue()) + def test_nttda_physical_zero_root(self): # Two open orbitals, no core or virtual orbitals: amplitudes are OO only. mol = gto.M(atom='H 0 0 0; H 0 0 1', spin=2, basis='sto-3g', verbose=0) From 0958a6381070f406405bdd42e79d7423ac6e820e Mon Sep 17 00:00:00 2001 From: wtpeter Date: Tue, 22 Sep 2026 21:55:07 +0800 Subject: [PATCH 3/3] feat(nttda): add explicit reconstructed TDA matrices --- examples/nttda/03_get_ab.py | 63 ++++++ src/nest/nttda/get_ab.py | 287 ++++++++++++++++++++++++++++ src/nest/nttda/nttda.py | 4 + src/nest/nttda/tests/test_get_ab.py | 94 +++++++++ 4 files changed, 448 insertions(+) create mode 100644 examples/nttda/03_get_ab.py create mode 100644 src/nest/nttda/get_ab.py create mode 100644 src/nest/nttda/tests/test_get_ab.py diff --git a/examples/nttda/03_get_ab.py b/examples/nttda/03_get_ab.py new file mode 100644 index 0000000..61d7e53 --- /dev/null +++ b/examples/nttda/03_get_ab.py @@ -0,0 +1,63 @@ +#!/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. + +""" +Build the dense NTTDA A matrix and compare with Davidson. + +Use a small system: get_ab stores a dense excitation matrix and the full +four-index reference-kernel tensor in the MO basis. +""" + +import numpy as np +from pyscf import gto + +from nest import nttda +from nest.nttda.get_ab import get_ab + +mol = gto.M( + atom="O 0 0 0; H 0 0 1; H 0 1 0", + basis="sto-3g", + spin=2, + verbose=0, +) +mf = mol.ROKS(xc="CAM-B3LYP").run(conv_tol=1e-11) +assert mf.converged + +for delta_s in (-1, 0, 1): + td = mf.NTTDA().set(deltaS=delta_s, nstates=2) + a = get_ab(mf, deltaS=delta_s, nobeta=td.nobeta) # Also available as td.get_ab(). + + if delta_s == -1: + # A uses the flattened [[CO, CV], [OO, OV]] amplitude layout. + # Remove the known OO identity direction by changing basis, rather + # than discarding every eigenvalue near zero. + nc = np.count_nonzero(mf.mo_occ == 2) + no = np.count_nonzero(mf.mo_occ == 1) + nv = np.count_nonzero(mf.mo_occ == 0) + q = np.zeros((nc + no, no + nv)) + q[nc:, :no] = np.eye(no) / np.sqrt(no) + basis = np.linalg.qr(q.reshape(-1, 1), mode="complete")[0][:, 1:] + energies = np.linalg.eigvalsh(basis.T @ a @ basis) + else: + # deltaS=0: CO(1), CV(1), scalar OO(1), OV(1), CV(0). + # deltaS=+1: flattened CV amplitudes only. + energies = np.linalg.eigvalsh(a) + + td.run() + assert np.all(td.converged) + np.testing.assert_allclose(energies[:td.nstates], td.e, atol=1e-8, rtol=0) + print(f"deltaS={delta_s:+d}, A shape={a.shape}") + print("Dense energies (Hartree): ", energies[:td.nstates]) + print("Davidson energies (Hartree):", td.e) diff --git a/src/nest/nttda/get_ab.py b/src/nest/nttda/get_ab.py new file mode 100644 index 0000000..97df488 --- /dev/null +++ b/src/nest/nttda/get_ab.py @@ -0,0 +1,287 @@ +#!/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. + +"""Explicit NT-TDA matrix elements after reference-kernel reconstruction.""" + +import numpy as np +from pyscf import ao2mo, dft, lib +from pyscf.lib import logger + + +def get_ab(mf, deltaS=-1, nobeta=False): + r"""Return the dense A matrix for a ROKS reference (there is no B matrix). + + i,j denote closed-shell orbitals; a,b virtual orbitals; u,v,t,w open-shell + orbitals. Amplitudes store the hole before the particle: X[i,a] for CV, + X[i,u] for CO, X[u,a] for OV, and X[u,t] for the OO amplitude in A_{tu,vw}. + + deltaS=-1 uses [[CO, CV], [OO, OV]].ravel(), including its redundant + OO identity zero mode. deltaS=0 concatenates CO, CV, scalar OO, OV, CV0. + deltaS=+1 uses CV.ravel(). For a batch x of flattened row vectors, + vind(x) = (A @ x.T).T. + + K[p,q,r,s] is K^Ref_{pq,rs}, and Fz[p,q] = 1/2 sum_u K[p,q,u,u]. + XC kernels are contracted directly with MO values and derivatives; + neither response functions nor gen_vind are used in this construction. + This dense implementation stores full four-index MO tensors and is + intended for small systems. With density fitting, both J and K must use + the same fitted integrals (only_dfj=True is not supported). + """ + if not isinstance(mf, (dft.roks.ROKS, dft.rks_symm.SymAdaptedROKS)): + raise TypeError('NTTDA get_ab requires a ROKS reference') + if deltaS not in (-1, 0, 1): + raise ValueError('deltaS must be -1, 0, or 1') + with_df = getattr(mf, 'with_df', None) + if with_df and getattr(mf, 'only_dfj', False): + raise NotImplementedError('NTTDA get_ab does not support only_dfj=True') + + mol = mf.mol + csidx = np.where(mf.mo_occ == 2)[0] + osidx = np.where(mf.mo_occ == 1)[0] + vsidx = np.where(mf.mo_occ == 0)[0] + nc, no, nv = len(csidx), len(osidx), len(vsidx) + s = no * 0.5 + if deltaS == -1: + assert s >= 1, 'NTTDA for Sf=Si-1 only supports case that Si>=1.' + elif deltaS == 0: + assert s >= 0.5, 'NTTDA only supports case that Sf=Si>=1/2.' + assert s == mol.spin * 0.5 + order = np.concatenate((csidx, osidx, vsidx)) + mo = mf.mo_coeff[:, order] + assert mo.dtype == np.double + nmo = nc + no + nv + c, o, v = slice(0, nc), slice(nc, nc + no), slice(nc + no, nmo) + ic, io, iv = np.eye(nc), np.eye(no), np.eye(nv) + + ni = mf._numint + ni.libxc.test_deriv_order(mf.xc, 2, raise_error=True) + xctype = ni._xc_type(mf.xc) + if xctype not in ('HF', 'LDA', 'GGA', 'MGGA'): + raise NotImplementedError('NTTDA get_ab does not support XC type %s' % xctype) + if mf.do_nlc(): + logger.warn(mf, 'NLC contribution in get_ab is NOT included') + omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin) + + # Reference kernel: exact exchange K_{pq,rs} = -(pr|qs). + kref = np.zeros((nmo, nmo, nmo, nmo)) + if hyb != 0: + if with_df: + eri = with_df.ao2mo(mo, compact=False) + else: + eri = ao2mo.kernel(mol, mo, compact=False) + eri = eri.reshape(nmo, nmo, nmo, nmo) + kref -= hyb * lib.einsum('prqs->pqrs', eri) + if omega != 0: + if with_df: + with with_df.range_coulomb(omega) as rsh_df: + eri = rsh_df.ao2mo(mo, compact=False) + else: + with mol.with_range_coulomb(omega): + eri = ao2mo.kernel(mol, mo, compact=False) + eri = eri.reshape(nmo, nmo, nmo, nmo) + kref -= (alpha - hyb) * lib.einsum('prqs->pqrs', eri) + + # The reference XC Hessian is evaluated at rho_alpha = rho_beta = rho/2. + # Pair features are phi_p phi_q, grad(phi_p phi_q), and (for MGGA) + # 1/2 grad(phi_p).grad(phi_q). Contract both pairs with fxc and grid weights. + if xctype != 'HF': + if mf.grids.coords is None: + mf.grids.build(with_non0tab=True) + max_memory = max(2000, mf.max_memory * .8 - lib.current_memory()[0]) + ao_deriv = 0 if xctype == 'LDA' else 1 + for ao, mask, weight, coords in ni.block_loop(mol, mf.grids, mol.nao_nr(), ao_deriv, max_memory): + rho = ni.eval_rho2(mol, ao, mo, mf.mo_occ[order], mask, xctype, with_lapl=False) * .5 + fxc = ni.eval_xc_eff(mf.xc, (rho, rho), deriv=2, xctype=xctype, spin=1)[2] + fxc_ref = .5 * (fxc[0, :, 0] - fxc[0, :, 1] - fxc[1, :, 0] + fxc[1, :, 1]) + if xctype == 'LDA': + phi = lib.einsum('gp,pi->gi', ao, mo) + pair = lib.einsum('gp,gq->gpq', phi, phi) + weighted_pair = pair * (fxc_ref[0, 0] * weight)[:, None, None] + kref += lib.einsum('gpq,grs->pqrs', pair, weighted_pair) + else: + phi = lib.einsum('xgp,pi->xgi', ao, mo) + pair = lib.einsum('xgp,gq->xgpq', phi, phi[0]) + pair[1:4] += lib.einsum('gp,xgq->xgpq', phi[0], phi[1:4]) + if xctype == 'MGGA': + tau_pair = .5 * lib.einsum('xgp,xgq->gpq', phi[1:4], phi[1:4]) + pair = np.concatenate((pair, tau_pair[None]), axis=0) + weighted_pair = lib.einsum('xyg,ygpq->xgpq', fxc_ref * weight, pair) + kref += lib.einsum('xgpq,xgrs->pqrs', pair, weighted_pair) + + # Reconstructed Fz and spin-averaged F0, both in the (C, O, V) MO basis. + fz = .5 * lib.einsum('pquu->pq', kref[:, :, o, o]) + if nobeta: + dma, dmb = mf.make_rdm1() + dm0 = .5 * (dma + dmb) + fock = mf.get_fock(dm=np.array([dm0, dm0])) + else: + fock = mf.get_fock() + f0 = mo.T @ (.5 * (fock.focka + fock.fockb)) @ mo + fp, fm = f0 + fz, f0 - fz + + if deltaS == 1: + # CV-CV: delta_ij (F0+Fz)_ab - delta_ab (F0-Fz)_ji + K_ai,bj. + a = lib.einsum('ij,ab->iajb', ic, fp[v, v]) + a -= lib.einsum('ab,ji->iajb', iv, fm[c, c]) + a += lib.einsum('aibj->iajb', kref[v, c, v, c]) + return a.reshape(nc * nv, nc * nv) + + if deltaS == -1: + # CV-CV + a_cvcv = lib.einsum('ij,ab->iajb', ic, fm[v, v]) + a_cvcv -= lib.einsum('ab,ji->iajb', iv, fp[c, c]) + a_cvcv -= lib.einsum('ij,ab->iajb', ic, fz[v, v]) / s + a_cvcv -= lib.einsum('ab,ji->iajb', iv, fz[c, c]) / s + a_cvcv += lib.einsum('aibj->iajb', kref[v, c, v, c]) + + # CV-CO and CV-OV + a_cvco = lib.einsum('ij,av->iajv', ic, fm[v, o]) + a_cvco += lib.einsum('aivj->iajv', kref[v, c, o, c]) + a_cvco *= np.sqrt((2 * s + 1) / (2 * s)) + a_cvov = -lib.einsum('ab,vi->iavb', iv, fp[o, c]) + a_cvov += lib.einsum('aibv->iavb', kref[v, c, v, o]) + a_cvov *= np.sqrt((2 * s + 1) / (2 * s)) + + # CO-CO + a_coco = lib.einsum('ij,uv->iujv', ic, fm[o, o]) + a_coco -= lib.einsum('uv,ji->iujv', io, fp[c, c]) + a_coco -= 2 / (2 * s - 1) * lib.einsum('uv,ji->iujv', io, fz[c, c]) + a_coco += lib.einsum('uivj->iujv', kref[o, c, o, c]) + a_coco += lib.einsum('uvij->iujv', kref[o, o, c, c]) / (2 * s - 1) + + # CO-OV + a_coov = -lib.einsum('ubiv->iuvb', kref[o, v, c, o]) / (2 * s - 1) + a_coov += 2 * s / (2 * s - 1) * lib.einsum('uibv->iuvb', kref[o, c, v, o]) + + # OV-OV + a_ovov = lib.einsum('uv,ab->uavb', io, fm[v, v]) + a_ovov -= lib.einsum('ab,vu->uavb', iv, fp[o, o]) + a_ovov -= 2 / (2 * s - 1) * lib.einsum('uv,ab->uavb', io, fz[v, v]) + a_ovov += lib.einsum('aubv->uavb', kref[v, o, v, o]) + a_ovov += lib.einsum('abuv->uavb', kref[v, v, o, o]) / (2 * s - 1) + + # OO-CV: row indices are (u,t), corresponding to A_tu,bj. + a_oocv = -lib.einsum('ut,jb->utjb', io, fz[c, v]) / s + a_oocv += lib.einsum('tubj->utjb', kref[o, o, v, c]) + a_oocv *= np.sqrt((2 * s + 1) / (2 * s - 1)) + + # OO-CO + a_ooco = -np.sqrt(2 * s / (2 * s - 1)) * lib.einsum('vt,ju->utjv', io, fp[c, o]) + a_ooco += lib.einsum('ut,jv->utjv', io, fm[c, o]) / np.sqrt(2 * s * (2 * s - 1)) + a_ooco += np.sqrt(2 * s / (2 * s - 1)) * lib.einsum('tuvj->utjv', kref[o, o, o, c]) + + # OO-OV + a_ooov = np.sqrt(2 * s / (2 * s - 1)) * lib.einsum('uv,tb->utvb', io, fm[o, v]) + a_ooov -= lib.einsum('tu,vb->utvb', io, fp[o, v]) / np.sqrt(2 * s * (2 * s - 1)) + a_ooov += np.sqrt(2 * s / (2 * s - 1)) * lib.einsum('tubv->utvb', kref[o, o, v, o]) + + # OO-OO + a_oooo = lib.einsum('wu,tv->utwv', io, fm[o, o]) + a_oooo -= lib.einsum('tv,wu->utwv', io, fp[o, o]) + a_oooo += lib.einsum('tuvw->utwv', kref[o, o, o, o]) + + # Native [[CO, CV], [OO, OV]] layout; transpose the reverse blocks. + nocc, nvir = nc + no, no + nv + a = np.empty((nocc, nvir, nocc, nvir)) + a[:nc, :no, :nc, :no] = a_coco + a[:nc, :no, :nc, no:] = a_cvco.transpose(2, 3, 0, 1) + a[:nc, :no, nc:, :no] = a_ooco.transpose(2, 3, 0, 1) + a[:nc, :no, nc:, no:] = a_coov + a[:nc, no:, :nc, :no] = a_cvco + a[:nc, no:, :nc, no:] = a_cvcv + a[:nc, no:, nc:, :no] = a_oocv.transpose(2, 3, 0, 1) + a[:nc, no:, nc:, no:] = a_cvov + a[nc:, :no, :nc, :no] = a_ooco + a[nc:, :no, :nc, no:] = a_oocv + a[nc:, :no, nc:, :no] = a_oooo + a[nc:, :no, nc:, no:] = a_ooov + a[nc:, no:, :nc, :no] = a_coov.transpose(2, 3, 0, 1) + a[nc:, no:, :nc, no:] = a_cvov.transpose(2, 3, 0, 1) + a[nc:, no:, nc:, :no] = a_ooov.transpose(2, 3, 0, 1) + a[nc:, no:, nc:, no:] = a_ovov + return a.reshape(nocc * nvir, nocc * nvir) + + # deltaS=0: CV0-CV0 + a_cv0cv0 = lib.einsum('ij,ab->iajb', ic, f0[v, v]) + a_cv0cv0 -= lib.einsum('ab,ji->iajb', iv, f0[c, c]) + a_cv0cv0 += lib.einsum('aibj->iajb', kref[v, c, v, c]) + a_cv0cv0 -= 2 * lib.einsum('abij->iajb', kref[v, v, c, c]) + + # CV0-CV + a_cv0cv = -lib.einsum('ij,ab->iajb', ic, fz[v, v]) + a_cv0cv += lib.einsum('ab,ji->iajb', iv, fz[c, c]) + a_cv0cv *= np.sqrt((s + 1) / s) + + # CV0-CO + a_cv0co = lib.einsum('ij,av->iajv', ic, fm[v, o]) / np.sqrt(2) + a_cv0co += lib.einsum('aivj->iajv', kref[v, c, o, c]) / np.sqrt(2) + a_cv0co -= np.sqrt(2) * lib.einsum('avij->iajv', kref[v, o, c, c]) + + # CV0-OV + a_cv0ov = lib.einsum('ab,vi->iavb', iv, fp[o, c]) / np.sqrt(2) + a_cv0ov -= lib.einsum('aibv->iavb', kref[v, c, v, o]) / np.sqrt(2) + a_cv0ov += np.sqrt(2) * lib.einsum('abiv->iavb', kref[v, v, c, o]) + + # CV-CV + a_cvcv = lib.einsum('ij,ab->iajb', ic, f0[v, v] - fz[v, v] / s) + a_cvcv -= lib.einsum('ab,ji->iajb', iv, f0[c, c] + fz[c, c] / s) + a_cvcv += lib.einsum('aibj->iajb', kref[v, c, v, c]) + + # CV-CO and CV-OV + a_cvco = lib.einsum('ij,av->iajv', ic, fm[v, o]) + a_cvco += lib.einsum('aivj->iajv', kref[v, c, o, c]) + a_cvco *= np.sqrt((s + 1) / (2 * s)) + a_cvov = -lib.einsum('ab,vi->iavb', iv, fp[o, c]) + a_cvov += lib.einsum('aibv->iavb', kref[v, c, v, o]) + a_cvov *= np.sqrt((s + 1) / (2 * s)) + + # CO-CO and CO-OV + a_coco = lib.einsum('ij,uv->iujv', ic, fm[o, o]) + a_coco -= lib.einsum('uv,ji->iujv', io, fm[c, c]) + a_coco += lib.einsum('uivj->iujv', kref[o, c, o, c]) + a_coco -= lib.einsum('uvij->iujv', kref[o, o, c, c]) + a_coov = lib.einsum('ubiv->iuvb', kref[o, v, c, o]) + + # OV-OV + a_ovov = lib.einsum('uv,ab->uavb', io, fp[v, v]) + a_ovov -= lib.einsum('ab,vu->uavb', iv, fp[o, o]) + a_ovov += lib.einsum('aubv->uavb', kref[v, o, v, o]) + a_ovov -= lib.einsum('abuv->uavb', kref[v, v, o, o]) + + # The OO block is a scalar in this sector. + a_oocv0 = -np.sqrt(2) * f0[c, v].reshape(1, -1) + a_oocv = 2 * np.sqrt((s + 1) / (2 * s)) * fz[c, v].reshape(1, -1) + a_ooco = -fm[c, o].reshape(1, -1) + a_ooov = fp[o, v].reshape(1, -1) + + # Native order: CO, CV, scalar OO, OV, CV0. + nco, ncv, nov = nc * no, nc * nv, no * nv + a_coco = a_coco.reshape(nco, nco) + a_cvcv = a_cvcv.reshape(ncv, ncv) + a_ovov = a_ovov.reshape(nov, nov) + a_cvco = a_cvco.reshape(ncv, nco) + a_cvov = a_cvov.reshape(ncv, nov) + a_coov = a_coov.reshape(nco, nov) + a_cv0cv0 = a_cv0cv0.reshape(ncv, ncv) + a_cv0cv = a_cv0cv.reshape(ncv, ncv) + a_cv0co = a_cv0co.reshape(ncv, nco) + a_cv0ov = a_cv0ov.reshape(ncv, nov) + return np.block([ + [a_coco, a_cvco.T, a_ooco.T, a_coov, a_cv0co.T], + [a_cvco, a_cvcv, a_oocv.T, a_cvov, a_cv0cv.T], + [a_ooco, a_oocv, np.zeros((1, 1)), a_ooov, a_oocv0], + [a_coov.T, a_cvov.T, a_ooov.T, a_ovov, a_cv0ov.T], + [a_cv0co, a_cv0cv, a_oocv0.T, a_cv0ov, a_cv0cv0], + ]) diff --git a/src/nest/nttda/nttda.py b/src/nest/nttda/nttda.py index 135a93a..e378397 100644 --- a/src/nest/nttda/nttda.py +++ b/src/nest/nttda/nttda.py @@ -883,6 +883,10 @@ def precond(residual, energy): self._finalize() return self.e, self.xy + def get_ab(self): + from nest.nttda.get_ab import get_ab + return get_ab(self._scf, deltaS=self.deltaS, nobeta=self.nobeta) + gen_vind_sfu = gen_vind_sfu gen_vind_sc = gen_vind_sc gen_vind_sfd = gen_vind_sfd diff --git a/src/nest/nttda/tests/test_get_ab.py b/src/nest/nttda/tests/test_get_ab.py new file mode 100644 index 0000000..15453bb --- /dev/null +++ b/src/nest/nttda/tests/test_get_ab.py @@ -0,0 +1,94 @@ +# 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 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.charge = 0 + mol.spin = 2 + mol.basis = '631g' + cls.mol = mol.build() + + @classmethod + def tearDownClass(cls): + cls.mol.stdout.close() + + def test_nttda_get_ab_df(self): + for xc in ('HF', 'PBE', 'CAM-B3LYP'): + mf = self.mol.ROKS(xc=xc).density_fit(auxbasis='weigend').run(conv_tol=1e-11) + self.assertTrue(mf.converged) + td = mf.NTTDA() + for nobeta in (False, True): + td.nobeta = nobeta + for delta_s, generator in ((-1, td.gen_vind_sfd), (0, td.gen_vind_sc), (1, td.gen_vind_sfu)): + td.deltaS = delta_s + vind, hdiag = generator() + a = td.get_ab() + self.assertAlmostEqual(abs(a - a.T).max(), 0, delta=1e-11) + x = np.random.default_rng(12).standard_normal((10, hdiag.size)) + self.assertAlmostEqual(abs(vind(x) - (a @ x.T).T).max(), 0, delta=1e-11) + + # Mixing fitted J with exact K cannot be represented by one K^Ref tensor. + mf.only_dfj = True + with self.assertRaises(NotImplementedError): + td.get_ab() + + def test_nttda_get_ab(self): + for xc in ('HF', 'SVWN', 'PBE', 'M062X', 'CAM-B3LYP'): + mf = self.mol.ROKS(xc=xc).run(conv_tol=1e-11) + self.assertTrue(mf.converged) + for nobeta in (False, True): + for delta_s in (-1, 0, 1): + td = mf.NTTDA().set(deltaS=delta_s, nobeta=nobeta) + if delta_s == -1: + vind, hdiag = td.gen_vind_sfd() + elif delta_s == 0: + vind, hdiag = td.gen_vind_sc() + else: + vind, hdiag = td.gen_vind_sfu() + + a = td.get_ab() + self.assertEqual(a.shape, (hdiag.size, hdiag.size)) + self.assertAlmostEqual(abs(a - a.T).max(), 0, delta=1e-11) + x = np.random.default_rng(12).standard_normal((10, hdiag.size)) + self.assertAlmostEqual(abs(vind(x) - (a @ x.T).T).max(), 0, delta=1e-11) + + if delta_s == -1: + nc = np.count_nonzero(mf.mo_occ == 2) + no = np.count_nonzero(mf.mo_occ == 1) + nv = np.count_nonzero(mf.mo_occ == 0) + q = np.zeros((nc + no, no + nv)) + q[nc:, :no] = np.eye(no) / np.sqrt(no) + self.assertAlmostEqual(np.linalg.norm(a @ q.ravel()), 0, delta=1e-11) + + +if __name__ == '__main__': + unittest.main()