Spin Flip TDA/ TDDFT using Multicollinear Approach. (#48)

* pretest

* Add the sftda dic, containing SF-TDA/ TDDFT.

* Add the top comment for files and a spin flip tda calculation example.

* copy __mcfun_fn_eval_xc to numint2c_sftd.py.

* define fucntion get_ab_sf() as TDA_SF class function.

* Modified the mistakes in CI/Lint tests.

* Skip the unittests if mcfun isn't installed.

* lint errors

---------

Co-authored-by: maohaohao <[email protected]>
Co-authored-by: Qiming Sun <[email protected]>

examples/sftda/00-spin_flip_tda.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env/python
+
+'''
+Spin flip TDA/ TDDFT. The kwarg 'extype' is used to control which
+kind of excited energy to be calculated, 0 for spin flip up, 1 for
+spin flip down.
+'''
+
+from pyscf import lib, gto, dft
+from pyscf import sftda
+
+mol = gto.Mole()
+mol.verbose = 6
+mol.output = None
+mol.spin = 2
+mol.atom = 'O 0 0 2.07; O 0 0 0'
+mol.unit = 'B'
+mol.basis = '631g'
+mol.build()
+
+mf = dft.UKS(mol)
+mf.xc = 'svwn' # blyp, b3lyp, tpss
+mf.kernel()
+
+#
+# 1. spin flip up TDA
+#
+mftd1 = sftda.TDA_SF(mf)
+mftd1.nstates = 5 # the number of excited states
+mftd1.extype = 0 # 0 for spin flip up excited energies
+# the spin sample points in multicollinear approach, which
+# can be increased by users.
+mftd1.collinear_samples=200
+mftd1.kernel()
+
+mftd1.e # to get the excited energies
+mftd1.xy # to get the transition vectors
+
+#
+# 2. spin flip down TDA
+#
+mftd2 = sftda.uks_sf.TDA_SF(mf)
+mftd2.nstates = 5 # the number of excited states
+mftd2.extype = 1 # 1 for spin flip down excited energies
+mftd2.collinear_samples=200
+mftd2.kernel()
+
+mftd2.e # to get the excited energies
+mftd2.xy # to get the transition vectors
+
+#
+# 3. get_ab_sf()
+#
+# a, b = sftda.TDA_SF(mf).get_ab_sf()
+a, b = sftda.TDDFT_SF(mf).get_ab_sf()
+# List a has two items: (A_abab, A_baba) with A[i,a,j,b].
+# List b has two items: (B_abba, B_baab) with B[i,a,j,b].
+
+#
+# 4. spin flip up TDDFT, which can not converged.
+#    Just give an input example here.
+#    Users can use get_ab_sf() to construct the whole TDDFT matrix
+#    to get the excited energies, if the system is small.
+# mftd3 = sftda.TDDFT_SF(mf) # equal to mftd3 = mf.TDDFT_SF()
+# mftd3.nstates = 4
+# mftd3.extype = 0
+# mftd3.collinear_samples=200
+# mftd3.kernel()
+
+# mftd3.e
+# mftd3.xy
+
+#
+# 5. spin flip down TDDFT, which can not converged.
+#    Just give an input example here.
+# mftd4 = sftda.uks_sf.CasidaTDDFT(mf)
+# mftd4.nstates = 4
+# mftd4.extype = 1
+# mftd4.collinear_samples=200
+# mftd4.kernel()
+
+# mftd4.e
+# mftd4.xy

pyscf/sftda/__init__.py

@@ -0,0 +1,35 @@
+#/usr/bin/env python
+# Copyright 2014-2024 The PySCF 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 pyscf import scf
+from pyscf.dft import KohnShamDFT
+from pyscf.sftda import uks_sf
+from pyscf.sftda.uks_sf import TDUKS_SF
+
+
+def TDA_SF(mf):
+    mf = mf.remove_soscf()
+    if isinstance(mf, scf.rohf.ROHF):
+        if isinstance(mf, KohnShamDFT):
+            mf = mf.to_uks()
+        else:
+            mf = mf.to_uhf()
+    return mf.TDA_SF()
+
+def TDDFT_SF(mf):
+    print('Warning!!! SF-TDDFT ruining in the slow divergence, ' + \
+          'you can choose get_ab_sf() to construct the full matrix ' + \
+          'to obtain the excited energies.')
+    return mf.TDDFT_SF()

pyscf/sftda/numint2c_sftd.py

@@ -0,0 +1,92 @@
+#/usr/bin/env python
+# Copyright 2014-2024 The PySCF 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.
+
+# This file can be merged into pyscf.dft.numint2c.py
+
+MGGA_DENSITY_LAPL = False # just copy from pyscf.dft.numint2c.py
+
+import functools
+import numpy as np
+from pyscf import lib
+from pyscf.dft import numint, xc_deriv
+
+# This function is copied from pyscf.dft.numint2c.py
+def __mcfun_fn_eval_xc(ni, xc_code, xctype, rho, deriv):
+    evfk = ni.eval_xc_eff(xc_code, rho, deriv=deriv, xctype=xctype)
+    for order in range(1, deriv+1):
+        if evfk[order] is not None:
+            evfk[order] = xc_deriv.ud2ts(evfk[order])
+    return evfk
+
+# This function can be merged with pyscf.dft.numint2c.mcfun_eval_xc_adapter()
+# This function should be a class function in the Numint2c class.
+def mcfun_eval_xc_adapter_sf(ni, xc_code):
+    '''Wrapper to generate the eval_xc function required by mcfun
+
+    Kwargs:
+        dim: int
+            eval_xc_eff_sf is for mc collinear sf tddft/ tda case.add().
+    '''
+
+    try:
+        import mcfun
+    except ImportError:
+        raise ImportError('This feature requires mcfun library.\n'
+                          'Try install mcfun with `pip install mcfun`')
+
+    xctype = ni._xc_type(xc_code)
+    fn_eval_xc = functools.partial(__mcfun_fn_eval_xc, ni, xc_code, xctype)
+    nproc = lib.num_threads()
+
+    def eval_xc_eff(xc_code, rho, deriv=1, omega=None, xctype=None,
+                verbose=None):
+        return mcfun.eval_xc_eff_sf(
+            fn_eval_xc, rho, deriv,
+            collinear_samples=ni.collinear_samples, workers=nproc)
+    return eval_xc_eff
+
+# This function should be a class function in the Numint2c class.
+def cache_xc_kernel_sf(self, mol, grids, xc_code, mo_coeff, mo_occ, spin=1,max_memory=2000):
+    '''Compute the fxc_sf, which can be used in SF-TDDFT/TDA
+    '''
+    xctype = self._xc_type(xc_code)
+    if xctype == 'GGA':
+        ao_deriv = 1
+    elif xctype == 'MGGA':
+        ao_deriv = 2 if MGGA_DENSITY_LAPL else 1
+    else:
+        ao_deriv = 0
+    with_lapl = MGGA_DENSITY_LAPL
+
+    assert mo_coeff[0].ndim == 2
+    assert spin == 1
+
+    nao = mo_coeff[0].shape[0]
+    rhoa = []
+    rhob = []
+
+    ni = numint.NumInt()
+    for ao, mask, weight, coords \
+            in self.block_loop(mol, grids, nao, ao_deriv, max_memory=max_memory):
+        rhoa.append(ni.eval_rho2(mol, ao, mo_coeff[0], mo_occ[0], mask, xctype, with_lapl))
+        rhob.append(ni.eval_rho2(mol, ao, mo_coeff[1], mo_occ[1], mask, xctype, with_lapl))
+    rho_ab = (np.hstack(rhoa), np.hstack(rhob))
+    rho_ab = np.asarray(rho_ab)
+    rho_tmz = np.zeros_like(rho_ab)
+    rho_tmz[0] += rho_ab[0]+rho_ab[1]
+    rho_tmz[1] += rho_ab[0]-rho_ab[1]
+    eval_xc = mcfun_eval_xc_adapter_sf(self,xc_code)
+    fxc_sf = eval_xc(xc_code, rho_tmz, deriv=2, xctype=xctype)
+    return fxc_sf