diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 674166e..b9a8d72 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -31,5 +31,7 @@ jobs: python-version: "3.11" cache: pip - run: python -m pip install -e ".[dev]" + - name: Check dispersion and optimization dependencies + run: python -c "from pyscf.dispersion import dftd3, dftd4; import geometric" - run: python -m ruff check . - run: python -m pytest src -q diff --git a/examples/grad/02_sftda_dispersion.py b/examples/grad/02_sftda_dispersion.py new file mode 100644 index 0000000..2fc6cea --- /dev/null +++ b/examples/grad/02_sftda_dispersion.py @@ -0,0 +1,56 @@ +#!/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. + +""" +SF-TDA total energies and gradients with D3/D4 dispersion. +Required dependencies: pyscf-dispersion +""" + +from pyscf import gto +from nest import sftda # Register SFTDA/SFTDDFT on PySCF reference objects. + +mol = gto.M( + atom="O 0 0 0; H 0 -0.757 0.587; H 0 0.757 0.587", + basis="sto-3g", + spin=2, + verbose=3, +) +mf = mol.UKS(xc="PBE") +# Set dispersion BEFORE SCF. Parameters are selected using mf.xc. +# 'd3bj' or 'd3bj2b': D3(BJ), two-body only +# 'd3bjatm': D3(BJ), including the ATM three-body term +# 'd4': D4, including the ATM three-body term +mf.disp = "d4" +mf.kernel() + +# The same gradient interface is available for mf.SFTDDFT(). +td = mf.SFTDA().set(nstates=2, collinear="col").run() +print("TD roots relative to the reference (Hartree):", td.e) +print("Total state energies including dispersion (Hartree):", td.e_tot) + +# state=1 means the first TD root, which can lie below the SCF reference. +# state=0 requests the SCF reference gradient instead. +grad = td.Gradients().kernel(state=1) +print("First-root total gradient including dispersion (Hartree/Bohr):") +print(grad) + +# For geometry optimization, the scanner returns consistent total energy and +# gradient and recomputes dispersion when the geometry changes. +# It reuses the previous SCF density and TD amplitudes as initial guesses. +# The root index does not track state character. +scanner = td.Gradients().as_scanner(state=1) +energy, gradient = scanner(mol) +print("Scanner total energy (Hartree):", energy) +# mf.e_tot, td.e_tot, and the total gradients already include dispersion. diff --git a/examples/grad/03_sftda_geomopt.py b/examples/grad/03_sftda_geomopt.py new file mode 100644 index 0000000..b10654a --- /dev/null +++ b/examples/grad/03_sftda_geomopt.py @@ -0,0 +1,44 @@ +#!/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 the first SF-TDA root with geomeTRIC and D4 dispersion. +Required dependencies: geometric, pyscf-dispersion +""" + +from pyscf import gto +from pyscf.geomopt import geometric_solver +from nest import sftda # Register SF-TDA on PySCF reference objects. + +mol = gto.M( + atom="O 0 0 0; H 0 -0.757 0.587; H 0 0.757 0.587", + basis="6-31g", + spin=2, +) +mf = mol.UKS(xc="B3LYP").set(disp="d4") + +# The TD constructor runs SCF if needed; the optimizer triggers the TD calculation. +td = mf.SFTDA().set(extype=1, collinear='mcol', nstates=3) + +# It reuses the SCF density and projects old TD amplitudes onto the new MO basis. +# A fixed root index does not track electronic character through root crossings. +scanner = td.Gradients().as_scanner(state=1) +optimizer = geometric_solver.GeometryOptimizer(scanner) +optimizer.max_cycle = 50 +mol_eq = optimizer.kernel() + +print("Geometry optimization converged:", optimizer.converged) +print("Optimized geometry (Angstrom):") +print(mol_eq.tostring(format="xyz")) diff --git a/examples/sftda/02_sftddft_roks.py b/examples/sftda/02_sftddft_roks.py index 8583d15..1c8e27d 100644 --- a/examples/sftda/02_sftddft_roks.py +++ b/examples/sftda/02_sftddft_roks.py @@ -16,8 +16,9 @@ ''' Spin-flip TDDFT/TDA examples with ROKS reference. -nest.sftda.SFTDDFT/SFTDA can also be applied to ROKS objects, -or equivalently to mf.to_uks() objects with methods. +ROKS and UKS share the mf.SFTDDFT()/mf.SFTDA() interface. +The nest.sftda.SFTDDFT(mf)/SFTDA(mf) function interface is also available. +Spin-flip TDDFT gradients for RO references are NOT implemented. ''' import numpy as np @@ -43,7 +44,7 @@ def print_header(title): # Spin-flip down TDDFT print_header("CALCULATION 1: Spin-Flip-Down TDDFT") -sfd_tddft = sftda.SFTDDFT(mf) # same as mf.to_uks().SFTDDFT() +sfd_tddft = mf.SFTDDFT() sfd_tddft.nstates = 5 sfd_tddft.extype = 1 # 1 for spin-flip-down excitations sfd_tddft.collinear_samples = 20 @@ -89,7 +90,7 @@ def norm_xy(z): # Spin-flip down TDA print_header("CALCULATION 2: Spin-Flip-Down TDA") -sfd_tda = sftda.SFTDA(mf) # same as mf.to_uks().SFTDA() +sfd_tda = mf.SFTDA() sfd_tda.extype = 1 sfd_tda.nstates = 5 sfd_tda.collinear = 'col' # Use collinear functional diff --git a/pyproject.toml b/pyproject.toml index 8c8cc9f..2c8ba7c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,6 +13,8 @@ license-files = ["LICENSE"] dependencies = [ "mcfun>=0.2.6", "pyscf>=2.13.0", + "pyscf-dispersion", + "geometric", "numpy", "scipy", ] diff --git a/src/nest/grad/tduks_sf.py b/src/nest/grad/tduks_sf.py index 5477da1..2e4027d 100644 --- a/src/nest/grad/tduks_sf.py +++ b/src/nest/grad/tduks_sf.py @@ -357,11 +357,78 @@ def _contract_xc_kernel(td_grad, xc_code, dmvo, dmoo=None, with_vxc=True, with_k return f1vo, f1oo, v1ao, k1ao +class TDSCF_GradScanner(tdrhf_grad.TDSCF_GradScanner): + @property + def converged(self): + return self.base._scf.converged and self.base.converged[self.state - 1] + + class Gradients(tdrhf_grad.Gradients): cphf_max_cycle = tdrhf_grad.Gradients.cphf_max_cycle + 20 + def as_scanner(self, state=1): + if isinstance(self, lib.GradScanner): + return self + if state == 0: + return self.base._scf.nuc_grad_method().as_scanner() + name = self.__class__.__name__ + TDSCF_GradScanner.__name_mixin__ + return lib.set_class(TDSCF_GradScanner(self, state), + (TDSCF_GradScanner, self.__class__), name) + @lib.with_doc(grad_elec.__doc__) - def grad_elec(self, xy, singlet=None, atmlst=None): + def grad_elec(self, xy, atmlst=None): return grad_elec(self, xy, atmlst, self.max_memory, self.verbose) + def dump_flags(self, verbose=None): + super().dump_flags(verbose) + log = logger.new_logger(self, verbose) + td = self.base + log.info('extype = %s', td.extype) + log.info('collinear = %s', td.collinear) + if td.collinear == 'mcol': + log.info('collinear_samples = %s', td.collinear_samples) + mf = td._scf + log.info('dispersion = %s', (mf.disp or mf.xc) if mf.do_disp() else False) + return self + + def kernel(self, xy=None, state=None, atmlst=None): + """Rewrite the `kernel` method to include the reference's D3/D4 correction.""" + if xy is None: + if state is None: + state = self.state + else: + self.state = state + if state == 0: + logger.warn(self, 'state=0: computing the SCF reference gradient.') + # The SCF gradient kernel already includes dispersion. + return self.base._scf.nuc_grad_method().kernel(atmlst=atmlst) + + if xy is None: + if self.base.xy is None: + self.base.run() + xy = self.base.xy[state - 1] + + if atmlst is None: + atmlst = self.atmlst + else: + self.atmlst = atmlst + + if self.verbose >= logger.WARN: + self.check_sanity() + if self.verbose >= logger.INFO: + self.dump_flags() + + de = self.grad_elec(xy, atmlst) + self.de = de + self.grad_nuc(atmlst=atmlst) + if self.mol.symmetry: + self.de = self.symmetrize(self.de, atmlst) + mf = self.base._scf + if mf.do_disp(): + dispersion = mf.nuc_grad_method().get_dispersion() + if atmlst is not None: + dispersion = dispersion[np.asarray(atmlst, dtype=int)] + self.de += dispersion + self._finalize() + return self.de + Grad = Gradients diff --git a/src/nest/grad/tests/test_tduks_sf_grad.py b/src/nest/grad/tests/test_tduks_sf_grad.py index 9764fe3..444b405 100644 --- a/src/nest/grad/tests/test_tduks_sf_grad.py +++ b/src/nest/grad/tests/test_tduks_sf_grad.py @@ -18,6 +18,11 @@ from pyscf import gto from nest import sftda +try: + from pyscf.dispersion import dftd3, dftd4 +except ImportError: + dftd3 = dftd4 = None + def cal_exact_sf_tda_gradient(mf, extype=1, collinear='mcol', collinear_samples=20, state=1): a, b = sftda.uhf_sf.get_ab_sf(mf, collinear=collinear, collinear_samples=collinear_samples) @@ -234,6 +239,67 @@ def test_mcol_tpss(self): self.assertAlmostEqual(abs(grad_iter - ref).max(), 0, delta=1e-5) + @unittest.skipIf(dftd3 is None, 'pyscf-dispersion is not installed') + def test_col_d3bj(self): + mf = self.mol.UKS(xc='PBE').run() + td = mf.SFTDA().set(nstates=2, collinear='col').run() + roots = td.e.copy() + energies = td.e_tot.copy() + grad = td.Gradients().kernel(state=1) + + mf.disp = 'd3bj' + mf.e_tot = mf.energy_tot() + td.run() + disp_energy = mf.get_dispersion() + disp_grad = mf.nuc_grad_method().get_dispersion() + self.assertAlmostEqual(abs(td.e - roots).max(), 0, delta=1e-6) + self.assertAlmostEqual(abs(td.e_tot - energies - disp_energy).max(), 0, delta=1e-6) + corrected = td.Gradients().kernel(state=1) + self.assertAlmostEqual(abs(corrected - grad - disp_grad).max(), 0, delta=1e-6) + + @unittest.skipIf(dftd4 is None, 'pyscf-dispersion is not installed') + def test_mcol_d4(self): + mf = self.mol.UKS(xc='PBE').run() + td = mf.SFTDA().set(nstates=2, collinear='mcol').run() + roots = td.e.copy() + energies = td.e_tot.copy() + grad = td.Gradients().kernel(state=1) + + mf.disp = 'd4' + mf.e_tot = mf.energy_tot() + td.run() + disp_energy = mf.get_dispersion() + disp_grad = mf.nuc_grad_method().get_dispersion() + self.assertAlmostEqual(abs(td.e - roots).max(), 0, delta=1e-6) + self.assertAlmostEqual(abs(td.e_tot - energies - disp_energy).max(), 0, delta=1e-6) + corrected = td.Gradients().kernel(state=1) + self.assertAlmostEqual(abs(corrected - grad - disp_grad).max(), 0, delta=1e-6) + + @unittest.skipIf(dftd4 is None, 'pyscf-dispersion is not installed') + def test_d4_gradient_scanner(self): + # Finite differences require tightly converged energies at both geometries. + mf = self.mol.UKS(xc='PBE').set(disp='d4', conv_tol=1e-14, conv_tol_grad=1e-9) + mf.grids.level = 4 + mf.run() + td = mf.SFTDA().set(nstates=1, collinear='col', conv_tol=1e-9).run() + grad = td.Gradients().kernel(state=1) + scanner = td.Gradients().as_scanner(state=1) + energy, scanner_grad = scanner(self.mol) + self.assertTrue(scanner.converged) + self.assertAlmostEqual(energy, td.e_tot[0], delta=1e-6) + self.assertAlmostEqual(abs(scanner_grad - grad).max(), 0, delta=1e-6) + + # Differentiate the corrected total energy along one nuclear coordinate. + step = 1e-3 # Bohr + coords = self.mol.atom_coords() + coords[1, 1] += step + e_plus = scanner(self.mol.set_geom_(coords, unit='Bohr', inplace=False))[0] + coords[1, 1] -= 2 * step + e_minus = scanner(self.mol.set_geom_(coords, unit='Bohr', inplace=False))[0] + derivative = (e_plus - e_minus) / (2 * step) + # Allow finite-difference and numerical integration errors, in Hartree/Bohr. + self.assertAlmostEqual(derivative, grad[1, 1], delta=1e-5) + if __name__ == '__main__': print('Full tests for SF-TDA and SF-TDDFT analytic gradients') unittest.main() diff --git a/src/nest/sftda/__init__.py b/src/nest/sftda/__init__.py index 5e23338..8a0c08f 100644 --- a/src/nest/sftda/__init__.py +++ b/src/nest/sftda/__init__.py @@ -14,33 +14,50 @@ # See the License for the specific language governing permissions and # limitations under the License. -from pyscf import scf, dft +from pyscf import scf from nest.sftda import uhf_sf -def TDA_SF(mf): +def TDA_SF(mf, extype=1, collinear="mcol", collinear_samples=20): + """Construct SF-TDA from a UHF/UKS or ROHF/ROKS reference.""" mf = mf.remove_soscf() - if isinstance(mf, scf.rohf.ROHF) or isinstance(mf, scf.hf_symm.SymAdaptedROHF): - if isinstance(mf, dft.roks.ROKS) or isinstance(mf, dft.rks_symm.SymAdaptedROKS): + restricted = isinstance(mf, scf.rohf.ROHF) + if restricted: + if mf.mo_coeff is None: + mf.run() + if isinstance(mf, scf.hf.KohnShamDFT): mf = mf.to_uks() else: mf = mf.to_uhf() - return mf.TDA_SF() + td = uhf_sf.TDA_SF(mf, extype, collinear, collinear_samples) + td._ro_reference = restricted + return td -def TDDFT_SF(mf): +def TDDFT_SF(mf, extype=1, collinear="mcol", collinear_samples=20): + """Construct SF-TDDFT from a UHF/UKS or ROHF/ROKS reference.""" mf = mf.remove_soscf() - if isinstance(mf, scf.rohf.ROHF) or isinstance(mf, scf.hf_symm.SymAdaptedROHF): - if isinstance(mf, dft.roks.ROKS) or isinstance(mf, dft.rks_symm.SymAdaptedROKS): + restricted = isinstance(mf, scf.rohf.ROHF) + if restricted: + if mf.mo_coeff is None: + mf.run() + if isinstance(mf, scf.hf.KohnShamDFT): mf = mf.to_uks() else: mf = mf.to_uhf() - return mf.TDDFT_SF() + td = uhf_sf.TDDFT_SF(mf, extype, collinear, collinear_samples) + td._ro_reference = restricted + return td SFTDA = TDA_SF SFTDDFT = TDDFT_SF +scf.uhf.UHF.TDA_SF = scf.uhf.UHF.SFTDA = TDA_SF +scf.uhf.UHF.TDDFT_SF = scf.uhf.UHF.SFTDDFT = TDDFT_SF +scf.rohf.ROHF.TDA_SF = scf.rohf.ROHF.SFTDA = TDA_SF +scf.rohf.ROHF.TDDFT_SF = scf.rohf.ROHF.SFTDDFT = TDDFT_SF + __all__ = [ "SFTDA", "SFTDDFT", diff --git a/src/nest/sftda/tests/test_sftda.py b/src/nest/sftda/tests/test_sftda.py index 148cd58..001c3b0 100644 --- a/src/nest/sftda/tests/test_sftda.py +++ b/src/nest/sftda/tests/test_sftda.py @@ -128,9 +128,9 @@ def test_col_cam_tda(self): self.assertAlmostEqual(abs(e - td.e).max(), 0, delta=1e-6) def test_hf_tda_roks(self): - mf = self.mol.ROKS(xc='HF').run() + mf = self.mol.ROKS(xc='HF') ref = np.array([-0.2204522712, -0.0023966488]) - td = sftda.TDA_SF(mf).set(extype=1, collinear_samples=50, nstates=2).run() + td = mf.SFTDA().set(extype=1, collinear_samples=50, nstates=2).run() self.assertTrue(np.all(td.converged)) self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6) e = diagonalize_tda(mf, extype=1, collinear_samples=50, nstates=2) @@ -190,6 +190,16 @@ def test_col_cam_tda_roks(self): e = diagonalize_tda(mf, extype=1, collinear="col", collinear_samples=50, nstates=2) self.assertAlmostEqual(abs(e - td.e).max(), 0, delta=1e-6) + def test_tda_scanner(self): + mf = self.mol.UKS(xc='B3LYP').run() + for extype in (0, 1): + td = mf.SFTDA().set(extype=extype, nstates=3) + ref = td.kernel()[0].copy() + td_scan = td.as_scanner() + td_scan.max_cycle = 1 + td_scan(self.mol) + self.assertAlmostEqual(abs(td_scan.e - ref).max(), 0, delta=1e-6) + if __name__ == "__main__": print("Full Tests for spin-flip-TDA with UKS and ROKS references") unittest.main() diff --git a/src/nest/sftda/tests/test_sftddft.py b/src/nest/sftda/tests/test_sftddft.py index 9537770..97c2acc 100644 --- a/src/nest/sftda/tests/test_sftddft.py +++ b/src/nest/sftda/tests/test_sftddft.py @@ -128,9 +128,9 @@ def test_col_cam_tddft(self): self.assertAlmostEqual(abs(e - td.e).max(), 0, delta=1e-6) def test_hf_tddft_roks(self): - mf = self.mol.ROKS(xc='HF').run() + mf = self.mol.ROKS(xc='HF') ref = np.array([0.4629613282, 0.5364066167]) - td = sftda.TDDFT_SF(mf).set(extype=0, collinear_samples=50, nstates=2, conv_tol=1e-6).run() + td = mf.SFTDDFT().set(extype=0, collinear_samples=50, nstates=2, conv_tol=1e-6).run() self.assertTrue(np.all(td.converged)) self.assertAlmostEqual(abs(td.e - ref).max(), 0, delta=1e-6) e = diagonalize_tddft(mf, extype=0, collinear_samples=50, nstates=2) @@ -190,6 +190,16 @@ def test_col_cam_tddft_roks(self): e = diagonalize_tddft(mf, extype=1, collinear="col", collinear_samples=50, nstates=2) self.assertAlmostEqual(abs(e - td.e).max(), 0, delta=1e-6) + def test_tddft_scanner(self): + mf = self.mol.UKS(xc='HF').run() + for extype in (0, 1): + td = mf.SFTDDFT().set(extype=extype, nstates=3) + ref = td.kernel()[0].copy() + td_scan = td.as_scanner() + td_scan.max_cycle = 1 + td_scan(self.mol) + self.assertAlmostEqual(abs(td_scan.e - ref).max(), 0, delta=1e-6) + if __name__ == "__main__": print("Full Tests for spin-flip-TDDFT with UKS and ROKS references") unittest.main() diff --git a/src/nest/sftda/uhf_sf.py b/src/nest/sftda/uhf_sf.py index 0c8e6a4..031afb4 100644 --- a/src/nest/sftda/uhf_sf.py +++ b/src/nest/sftda/uhf_sf.py @@ -17,7 +17,7 @@ import numpy as np from pyscf import lib -from pyscf import scf, dft +from pyscf import scf, dft, gto from pyscf import ao2mo from pyscf.lib import logger from pyscf.tdscf import rhf @@ -53,8 +53,9 @@ def get_ab_sf( List A has two items: (A_baba, A_abab). List B has two items: (B_baab, B_abba). ''' - if isinstance(mf, scf.rohf.ROHF) or isinstance(mf, scf.hf_symm.SymAdaptedROHF): - if isinstance(mf, dft.roks.ROKS) or isinstance(mf, dft.rks_symm.SymAdaptedROKS): + ro_reference = isinstance(mf, scf.rohf.ROHF) + if ro_reference: + if isinstance(mf, scf.hf.KohnShamDFT): mf = mf.to_uks() else: mf = mf.to_uhf() @@ -80,8 +81,7 @@ def get_ab_sf( nocc_b = orbo_b.shape[1] nvir_b = orbv_b.shape[1] - if np.allclose(mf.mo_coeff[0], mf.mo_coeff[1]): - logger.info(mf, 'Restricted open-shell detected.') + if ro_reference or np.allclose(mf.mo_coeff[0], mf.mo_coeff[1]): fock_ao_a, fock_ao_b = mf.get_fock() fock_oo_a = orbo_a.T @ fock_ao_a @ orbo_a fock_vv_a = orbv_a.T @ fock_ao_a @ orbv_a @@ -275,14 +275,68 @@ def add_hf_(a, b, hyb=1): return a, b +class TD_Scanner(rhf.TD_Scanner): + def __call__(self, mol_or_geom, **kwargs): + if isinstance(mol_or_geom, gto.MoleBase): + mol = mol_or_geom + else: + mol = self.mol.set_geom_(mol_or_geom, inplace=False) + + basis_prev = np.hstack(self.mol.bas_exps()) + mo_prev = self._scf.mo_coeff + occ_prev = self._scf.mo_occ + self.reset(mol) + mf_e = self._scf(mol) + + x0 = kwargs.pop('x0', None) + if self.xy is not None: + if not np.array_equal(basis_prev, np.hstack(mol.bas_exps())): + self.xy = None + elif x0 is None: + x0 = self._transfer_initial_guess(self.xy, mo_prev, occ_prev) + self.kernel(x0=x0, **kwargs) + return mf_e + self.e + + @lib.with_doc(rhf.TDA.__doc__) class TDA_SF(TDBase): extype = 1 collinear = 'mcol' collinear_samples = 20 + _ro_reference = False _keys = {'extype', 'collinear', 'collinear_samples'} + def as_scanner(self): + if isinstance(self, lib.SinglePointScanner): + return self + name = self.__class__.__name__ + TD_Scanner.__name_mixin__ + return lib.set_class(TD_Scanner(self), (TD_Scanner, self.__class__), name) + + def _transfer_initial_guess(self, xy, mo_coeff, mo_occ): + # Project old amplitudes onto the new MO basis, as in GPU4PySCF. + mf = self._scf + overlap = mf.get_ovlp() + occupied = [] + virtual = [] + for spin in (0, 1): + old_occ = mo_coeff[spin][:, mo_occ[spin] > 0] + old_vir = mo_coeff[spin][:, mo_occ[spin] == 0] + new_occ = mf.mo_coeff[spin][:, mf.mo_occ[spin] > 0] + new_vir = mf.mo_coeff[spin][:, mf.mo_occ[spin] == 0] + occupied.append(new_occ.T @ overlap @ old_occ) + virtual.append(new_vir.T @ overlap @ old_vir) + + source, target = (1, 0) if self.extype == 0 else (0, 1) + x = np.stack([x for x, y in xy]) + x = np.einsum('ui,nij,vj->nuv', occupied[source], x, virtual[target]) + x = x.reshape(len(xy), -1) + if np.isscalar(xy[0][1]): + return x + y = np.stack([y for x, y in xy]) + y = np.einsum('ui,nij,vj->nuv', occupied[target], y, virtual[source]) + return np.hstack((x, y.reshape(len(xy), -1))) + def __init__(self, mf, extype=1, collinear="mcol", collinear_samples=20): TDBase.__init__(self,mf) # extype is used to determine which spin flip excitation will be calculated. @@ -295,6 +349,7 @@ def __init__(self, mf, extype=1, collinear="mcol", collinear_samples=20): def dump_flags(self, verbose=None): TDBase.dump_flags(self, verbose) log = logger.new_logger(self, verbose) + log.info('reference = %s', 'ROHF/ROKS' if self._ro_reference else 'UHF/UKS') log.info("extype = %s", self.extype) log.info("collinear = %s", self.collinear) if self.collinear == "mcol": @@ -306,8 +361,10 @@ def check_sanity(self): raise ValueError("extype must be 0 or 1") if self.collinear not in ("col", "mcol"): raise ValueError("collinear must be 'col' or 'mcol'") - if self.collinear=='mcol' and self.collinear_samples <= 0: - raise ValueError("collinear_samples must be positive") + if self.collinear == 'mcol' and ( + not isinstance(self.collinear_samples, (int, np.integer)) or self.collinear_samples <= 0 + ): + raise ValueError("collinear_samples must be a positive integer") TDBase.check_sanity(self) return self @@ -321,6 +378,8 @@ def gen_vind(self): assert (mo_coeff[0].dtype == np.double) mo_occ = mf.mo_occ + # Keep compatibility with references manually converted by to_uhf/to_uks. + ro_reference = self._ro_reference or np.allclose(mo_coeff[0], mo_coeff[1]) extype = self.extype if extype==0: occidxb = mo_occ[1] > 0 @@ -328,7 +387,7 @@ def gen_vind(self): orbo = mo_coeff[1][:, occidxb] orbv = mo_coeff[0][:, viridxa] ndim = (int(occidxb.sum()), int(viridxa.sum())) - if np.allclose(mo_coeff[0], mo_coeff[1]): + if ro_reference: fock_a, fock_b = mf.get_fock() focko = orbo.conj().T @ fock_b @ orbo fockv = orbv.conj().T @ fock_a @ orbv @@ -342,7 +401,7 @@ def gen_vind(self): orbo = mo_coeff[0][:, occidxa] orbv = mo_coeff[1][:, viridxb] ndim = (int(occidxa.sum()), int(viridxb.sum())) - if np.allclose(mo_coeff[0], mo_coeff[1]): + if ro_reference: fock_a, fock_b = mf.get_fock() focko = orbo.conj().T @ fock_a @ orbo fockv = orbv.conj().T @ fock_b @ orbv @@ -364,7 +423,7 @@ def vind(zs): dms = lib.einsum('xov,pv,qo->xpq', zs, orbv, orbo.conj()) v1ao = vresp(dms) v1mo = lib.einsum('xpq,qo,pv->xov', v1ao, orbo, orbv.conj()) - if np.allclose(mo_coeff[0], mo_coeff[1]): + if ro_reference: v1mo += lib.einsum('ab,xib->xia', fockv, zs) v1mo -= lib.einsum('ji,xja->xia', focko, zs) else: @@ -414,9 +473,6 @@ def kernel(self, x0=None, nstates=None, extype=None): ''' cpu0 = (logger.process_clock(), logger.perf_counter()) - self.check_sanity() - self.dump_flags() - if extype is None: extype = self.extype else: @@ -427,6 +483,10 @@ def kernel(self, x0=None, nstates=None, extype=None): else: self.nstates = nstates + if self.verbose >= logger.WARN: + self.check_sanity() + if self.verbose >= logger.INFO: + self.dump_flags() log = logger.Logger(self.stdout, self.verbose) def all_eigs(w, v, nroots, envs): @@ -437,7 +497,10 @@ def all_eigs(w, v, nroots, envs): x0sym = None if x0 is None: - x0 = self.init_guess() + if self.xy is None: + x0 = self.init_guess() + else: # Reuse the previous amplitudes, including in geometry scans. + x0 = np.asarray([x.ravel() for x, y in self.xy]) self.converged, self.e, x1 = eigh( vind, x0, precond, tol_residual=self.conv_tol, lindep=self.lindep, @@ -516,7 +579,9 @@ def gen_vind(self): nvira = int(viridxa.sum()) nvirb = int(viridxb.sum()) - if np.allclose(mo_coeff[0], mo_coeff[1]): + # Keep compatibility with references manually converted by to_uhf/to_uks. + ro_reference = self._ro_reference or np.allclose(mo_coeff[0], mo_coeff[1]) + if ro_reference: fock_a, fock_b = mf.get_fock() fockoa = orboa.conj().T @ fock_a @ orboa fockva = orbva.conj().T @ fock_a @ orbva @@ -565,7 +630,7 @@ def vind(zs): if self.extype==0: v1_top = lib.einsum('xpq,qo,pv->xov', v1ao, orbob, orbva.conj()) v1_bot = lib.einsum('xpq,po,qv->xov', v1ao, orboa.conj(), orbvb) - if np.allclose(mo_coeff[0], mo_coeff[1]): + if ro_reference: v1_top += lib.einsum('ab,xib->xia', fockva, zs_b2a) v1_top -= lib.einsum('ji,xja->xia', fockob, zs_b2a) v1_bot += lib.einsum('ab,xib->xia', fockvb, zs_a2b) @@ -576,7 +641,7 @@ def vind(zs): elif self.extype==1: v1_top = lib.einsum('xpq,qo,pv->xov', v1ao, orboa, orbvb.conj()) v1_bot = lib.einsum('xpq,po,qv->xov', v1ao, orbob.conj(), orbva) - if np.allclose(mo_coeff[0], mo_coeff[1]): + if ro_reference: v1_top += lib.einsum('ab,xib->xia', fockvb, zs_a2b) v1_top -= lib.einsum('ji,xja->xia', fockoa, zs_a2b) v1_bot += lib.einsum('ab,xib->xia', fockva, zs_b2a) @@ -649,9 +714,6 @@ def kernel(self, x0=None, nstates=None, extype=None): ''' cpu0 = (logger.process_clock(), logger.perf_counter()) - self.check_sanity() - self.dump_flags() - if extype is None: extype = self.extype else: @@ -662,13 +724,20 @@ def kernel(self, x0=None, nstates=None, extype=None): else: self.nstates = nstates + if self.verbose >= logger.WARN: + self.check_sanity() + if self.verbose >= logger.INFO: + self.dump_flags() log = logger.Logger(self.stdout, self.verbose) vind, hdiag = self.gen_vind() precond = self.get_precond(hdiag) if x0 is None: - x0 = self.init_guess() + if self.xy is None: + x0 = self.init_guess() + else: + x0 = np.asarray([np.concatenate((x.ravel(), y.ravel())) for x, y in self.xy]) pickeig = self.gen_pickeig(extype=extype) @@ -961,12 +1030,3 @@ def analyze(tdobj, verbose=None): SFTDA = TDA_SF SFTDDFT = TDDFT_SF - -dft.uks.UKS.TDA_SF = lib.class_as_method(TDA_SF) -dft.uks.UKS.TDDFT_SF = lib.class_as_method(TDDFT_SF) -scf.uhf.UHF.TDA_SF = lib.class_as_method(TDA_SF) -scf.uhf.UHF.TDDFT_SF = lib.class_as_method(TDDFT_SF) -dft.uks.UKS.SFTDA = lib.class_as_method(TDA_SF) -dft.uks.UKS.SFTDDFT = lib.class_as_method(TDDFT_SF) -scf.uhf.UHF.SFTDA = lib.class_as_method(TDA_SF) -scf.uhf.UHF.SFTDDFT = lib.class_as_method(TDDFT_SF)