Add PCM support for derivative coupling calculations (#550)

puzhichen · web-flow · commit 2e1399545015 · 2025-11-04T07:18:10.000-08:00
* finish add the g-e tdrhf with LR-PCM

* finish tdrhf

* finish rks

* finish tests

* fix some bugs

* review the codes
diff --git a/gpu4pyscf/nac/finite_diff.py b/gpu4pyscf/nac/finite_diff.py
@@ -214,7 +214,6 @@ def get_nacv_ee(td_nac, x_yI, x_yJ, nJ, delta=0.001, with_ris=False, singlet=Tru
     s = cp.asarray(s)
     for iatm in range(natm):
         for icart in range(3):
-            print(f"iatm, icart, {iatm} {icart}")
             mol_add = get_new_mol(mol, coords, delta, iatm, icart)
             mf_add, xy_diag_add = get_mf_td(mol_add, mf, s, mo_coeff, with_ris)
             mol_minus = get_new_mol(mol, coords, -delta, iatm, icart)
diff --git a/gpu4pyscf/nac/tdrhf.py b/gpu4pyscf/nac/tdrhf.py
@@ -108,8 +108,6 @@ def get_nacv_ge(td_nac, x_yI, EI, singlet=True, atmlst=None, verbose=logger.INFO
     nvir = nmo - nocc
     orbv = mo_coeff[:, nocc:]
     orbo = mo_coeff[:, :nocc]
-    if getattr(mf, 'with_solvent', None) is not None:
-        raise NotImplementedError('With solvent is not supported yet')
 
     xI, yI = x_yI
     xI = cp.asarray(xI).reshape(nocc, nvir).T
@@ -118,7 +116,7 @@ def get_nacv_ge(td_nac, x_yI, EI, singlet=True, atmlst=None, verbose=logger.INFO
     yI = cp.asarray(yI).reshape(nocc, nvir).T
     LI = xI-yI    # eq.(83) in Ref. [1]
 
-    vresp = mf.gen_response(singlet=None, hermi=1)
+    vresp = td_nac.base.gen_response(singlet=None, hermi=1)
 
     def fvind(x):
         dm = reduce(cp.dot, (orbv, x.reshape(nvir, nocc) * 2, orbo.T)) # double occupency
@@ -152,6 +150,7 @@ def fvind(x):
 
     mf_grad = mf.nuc_grad_method()
     dmz1doo = z1aoS
+    td_nac.dmz1doo = dmz1doo
     oo0 = reduce(cp.dot, (orbo, orbo.T)) * 2.0
 
     if atmlst is None:
@@ -239,8 +238,6 @@ def get_nacv_ee(td_nac, x_yI, x_yJ, EI, EJ, singlet=True, atmlst=None, verbose=l
     nvir = nmo - nocc
     orbv = mo_coeff[:, nocc:]
     orbo = mo_coeff[:, :nocc]
-    if getattr(mf, 'with_solvent', None) is not None:
-        raise NotImplementedError('With solvent is not supported yet')
 
     xI, yI = x_yI
     xJ, yJ = x_yJ
@@ -262,6 +259,8 @@ def get_nacv_ee(td_nac, x_yI, x_yJ, EI, EJ, singlet=True, atmlst=None, verbose=l
     xmyJ = (xJ - yJ)
     dmxpyJ = reduce(cp.dot, (orbv, xpyJ, orbo.T)) 
     dmxmyJ = reduce(cp.dot, (orbv, xmyJ, orbo.T)) 
+    td_nac.dmxpyI = dmxpyI
+    td_nac.dmxpyJ = dmxpyJ
 
     rIJoo =-contract('ai,aj->ij', xJ, xI) - contract('ai,aj->ij', yI, yJ)
     rIJvv = contract('ai,bi->ab', xI, xJ) + contract('ai,bi->ab', yJ, yI)
@@ -297,13 +296,16 @@ def get_nacv_ee(td_nac, x_yI, x_yJ, EI, EJ, singlet=True, atmlst=None, verbose=l
         vk2J = cp.asarray(vk2J)
 
     veff0doo = vj0IJ * 2 - vk0IJ
+    veff0doo += td_nac.solvent_response(dmzooIJ)
     wvo = reduce(cp.dot, (orbv.T, veff0doo, orbo)) * 2
     veffI = vj1I * 2 - vk1I
+    veffI += td_nac.solvent_response(dmxpyI + dmxpyI.T)
     veffI *= 0.5
     veff0mopI = reduce(cp.dot, (mo_coeff.T, veffI, mo_coeff))
     wvo -= contract("ki,ai->ak", veff0mopI[:nocc, :nocc], xpyJ) * 2  
     wvo += contract("ac,ai->ci", veff0mopI[nocc:, nocc:], xpyJ) * 2
     veffJ = vj1J * 2 - vk1J
+    veffJ += td_nac.solvent_response(dmxpyJ + dmxpyJ.T)
     veffJ *= 0.5
     veff0mopJ = reduce(cp.dot, (mo_coeff.T, veffJ, mo_coeff))
     wvo -= contract("ki,ai->ak", veff0mopJ[:nocc, :nocc], xpyI) * 2  
@@ -320,7 +322,7 @@ def get_nacv_ee(td_nac, x_yI, x_yJ, EI, EJ, singlet=True, atmlst=None, verbose=l
     wvo += contract("ac,ai->ci", veff0momJ[nocc:, nocc:], xmyI) * 2
     # The up parts are according to eq. (86) and (86) in Ref. [1]
 
-    vresp = mf.gen_response(singlet=None, hermi=1)
+    vresp = td_nac.base.gen_response(singlet=None, hermi=1)
 
     def fvind(x):
         dm = reduce(cp.dot, (orbv, x.reshape(nvir, nocc) * 2, orbo.T)) # double occupency
@@ -390,6 +392,7 @@ def fvind(x):
     s1 = mf_grad.get_ovlp(mol)
     z1aoS = (z1ao + z1ao.T)*0.5* (EJ - EI)
     dmz1doo = z1aoS + dmzooIJ  # P
+    td_nac.dmz1doo = dmz1doo
     oo0 = reduce(cp.dot, (orbo, orbo.T))*2  # D
 
     if atmlst is None:
@@ -492,7 +495,10 @@ class NAC(lib.StreamObject):
         "de",
         "de_scaled",
         "de_etf",
-        "de_etf_scaled"
+        "de_etf_scaled",
+        "dmz1doo",
+        "dmxpyI",
+        "dmxpyJ"
     }
 
     def __init__(self, td):
@@ -506,6 +512,9 @@ def __init__(self, td):
         self.de_scaled = None # CIS derivative coupling without ETF
         self.de_etf = None  # CIS Force Matrix Element with ETF
         self.de_etf_scaled = None # Knwon as CIS derivative coupling with ETF
+        self.dmz1doo = None
+        self.dmxpyI = None
+        self.dmxpyJ = None
 
     _write      = rhf_grad_cpu.GradientsBase._write
 
@@ -683,7 +692,6 @@ def check_phase_modified(mol0, mo_coeff0, mo1_reordered, xy0, xy1, nocc, s):
             * xy0[idxo_l, idxv_l] * xy1[idxo_r, idxv_r] * 2
 
         total_s_state += s_state_contribution
-        print(total_s_state)
 
     return total_s_state
 
diff --git a/gpu4pyscf/nac/tdrks.py b/gpu4pyscf/nac/tdrks.py
@@ -71,8 +71,6 @@ def get_nacv_ge(td_nac, x_yI, EI, singlet=True, atmlst=None, verbose=logger.INFO
     nvir = nmo - nocc
     orbv = mo_coeff[:, nocc:]
     orbo = mo_coeff[:, :nocc]
-    if getattr(mf, 'with_solvent', None) is not None:
-        raise NotImplementedError('With solvent is not supported yet')
 
     xI, yI = x_yI
     xI = cp.asarray(xI).reshape(nocc, nvir).T
@@ -85,8 +83,10 @@ def get_nacv_ge(td_nac, x_yI, EI, singlet=True, atmlst=None, verbose=logger.INFO
     ni.libxc.test_deriv_order(mf.xc, 3, raise_error=True)
     omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin)
     with_k = ni.libxc.is_hybrid_xc(mf.xc)
-
-    vresp = mf.gen_response(singlet=None, hermi=1)
+    if isinstance(td_nac.base, tdscf.ris.TDDFT) or isinstance(td_nac.base, tdscf.ris.TDA):
+        vresp = td_nac.base._scf.gen_response(singlet=None, hermi=1)
+    else:
+        vresp = td_nac.base.gen_response(singlet=None, hermi=1)
 
     def fvind(x):
         dm = reduce(cp.dot, (orbv, x.reshape(nvir, nocc) * 2, orbo.T)) # double occupency
@@ -121,6 +121,7 @@ def fvind(x):
     mf_grad = mf.nuc_grad_method()
     s1 = mf_grad.get_ovlp(mol)
     dmz1doo = z1aoS
+    td_nac.dmz1doo = dmz1doo
     oo0 = reduce(cp.dot, (orbo, orbo.T)) * 2.0
 
     if atmlst is None:
@@ -242,8 +243,6 @@ def get_nacv_ee(td_nac, x_yI, x_yJ, EI, EJ, singlet=True, atmlst=None, verbose=l
     nvir = nmo - nocc
     orbv = mo_coeff[:, nocc:]
     orbo = mo_coeff[:, :nocc]
-    if getattr(mf, 'with_solvent', None) is not None:
-        raise NotImplementedError('With solvent is not supported yet')
 
     xI, yI = x_yI
     xJ, yJ = x_yJ
@@ -265,6 +264,8 @@ def get_nacv_ee(td_nac, x_yI, x_yJ, EI, EJ, singlet=True, atmlst=None, verbose=l
     xmyJ = (xJ - yJ)
     dmxpyJ = reduce(cp.dot, (orbv, xpyJ, orbo.T)) 
     dmxmyJ = reduce(cp.dot, (orbv, xmyJ, orbo.T)) 
+    td_nac.dmxpyI = dmxpyI
+    td_nac.dmxpyJ = dmxpyJ
 
     rIJoo =-contract('ai,aj->ij', xJ, xI) - contract('ai,aj->ij', yI, yJ)
     rIJvv = contract('ai,bi->ab', xI, xJ) + contract('ai,bi->ab', yJ, yI)
@@ -335,13 +336,16 @@ def get_nacv_ee(td_nac, x_yI, x_yJ, EI, EJ, singlet=True, atmlst=None, verbose=l
             vk2J += vk2J_omega * (alpha - hyb)
 
         veff0doo = vj0IJ * 2 - vk0IJ + f1ooIJ[0] + k1aoIJ[0] * 2
+        veff0doo += td_nac.solvent_response(dmzooIJ)
         wvo = reduce(cp.dot, (orbv.T, veff0doo, orbo)) * 2
         veffI = vj1I * 2 - vk1I + f1voI[0] * 2
+        veffI += td_nac.solvent_response(dmxpyI + dmxpyI.T)
         veffI *= 0.5
         veff0mopI = reduce(cp.dot, (mo_coeff.T, veffI, mo_coeff))
         wvo -= contract("ki,ai->ak", veff0mopI[:nocc, :nocc], xpyJ) * 2  
         wvo += contract("ac,ai->ci", veff0mopI[nocc:, nocc:], xpyJ) * 2
         veffJ = vj1J * 2 - vk1J + f1voJ[0] * 2
+        veffJ += td_nac.solvent_response(dmxpyJ + dmxpyJ.T)
         veffJ *= 0.5
         veff0mopJ = reduce(cp.dot, (mo_coeff.T, veffJ, mo_coeff))
         wvo -= contract("ki,ai->ak", veff0mopJ[:nocc, :nocc], xpyI) * 2  
@@ -369,21 +373,24 @@ def get_nacv_ee(td_nac, x_yI, x_yJ, EI, EJ, singlet=True, atmlst=None, verbose=l
             vj1J = cp.asarray(vj1J)
 
         veff0doo = vj0IJ * 2 + f1ooIJ[0] + k1aoIJ[0] * 2
+        veff0doo += td_nac.solvent_response(dmzooIJ)
         wvo = reduce(cp.dot, (orbv.T, veff0doo, orbo)) * 2
         veffI = vj1I * 2 + f1voI[0] * 2
+        veffI += td_nac.solvent_response(dmxpyI + dmxpyI.T)
         veffI *= 0.5
         veff0mopI = reduce(cp.dot, (mo_coeff.T, veffI, mo_coeff))
         wvo -= contract("ki,ai->ak", veff0mopI[:nocc, :nocc], xpyJ) * 2  
         wvo += contract("ac,ai->ci", veff0mopI[nocc:, nocc:], xpyJ) * 2
         veffJ = vj1J * 2 + f1voJ[0] * 2
+        veffJ += td_nac.solvent_response(dmxpyJ + dmxpyJ.T)
         veffJ *= 0.5
         veff0mopJ = reduce(cp.dot, (mo_coeff.T, veffJ, mo_coeff))
         wvo -= contract("ki,ai->ak", veff0mopJ[:nocc, :nocc], xpyI) * 2  
         wvo += contract("ac,ai->ci", veff0mopJ[nocc:, nocc:], xpyI) * 2
         veff0momI = cp.zeros((nmo, nmo))
         veff0momJ = cp.zeros((nmo, nmo))
 
-    vresp = mf.gen_response(singlet=None, hermi=1)
+    vresp = td_nac.base.gen_response(singlet=None, hermi=1)
 
     def fvind(x):
         dm = reduce(cp.dot, (orbv, x.reshape(nvir, nocc) * 2, orbo.T)) # double occupency
@@ -453,6 +460,7 @@ def fvind(x):
     s1 = mf_grad.get_ovlp(mol)
     z1aoS = (z1ao + z1ao.T)*0.5* (EJ - EI)
     dmz1doo = z1aoS + dmzooIJ  # P
+    td_nac.dmz1doo = dmz1doo
     oo0 = reduce(cp.dot, (orbo, orbo.T))*2  # D
 
     if atmlst is None:
diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py
@@ -62,6 +62,26 @@ def make_tdscf_gradient_object(td_base_method):
                          (WithSolventTDSCFGradient, td_grad.__class__), name)
 
 
+def make_tdscf_nac_object(td_base_method):
+    '''For td_method in vacuum, add td of solvent pcmobj'''
+    # The nuclear gradients of stable exited states should correspond to a
+    # fully relaxed solvent. Strictly, the TDDFT exited states should be
+    # solved using state-specific solvent model. Even if running LR-PCM for
+    # the zeroth order TDDFT, the wavefunction should be comptued using the
+    # same dielectric constant as the ground state (the zero-frequency eps).
+    with_solvent = td_base_method.with_solvent
+    if not with_solvent.equilibrium_solvation:
+        raise RuntimeError(
+            'When computing derivative couplings of PCM-TDDFT, equilibrium solvation should '
+            'be employed. The PCM TDDFT should be initialized as\n'
+            '    mf.TDDFT(equilibrium_solvation=True)')
+    td_nac = td_base_method.undo_solvent().nac_method()
+    td_nac.base = td_base_method
+    name = with_solvent.__class__.__name__ + td_nac.__class__.__name__
+    return lib.set_class(WithSolventTDSCFNacMethod(td_nac),
+                         (WithSolventTDSCFNacMethod, td_nac.__class__), name)
+
+
 class WithSolventTDSCF:
     from gpu4pyscf.lib.utils import to_gpu, device
 
@@ -134,6 +154,7 @@ def undo_solvent(self):
 
     nuc_grad_method = make_tdscf_gradient_object
     Gradients = nuc_grad_method
+    nac_method = make_tdscf_nac_object
 
 
 class WithSolventTDSCFGradient:
@@ -180,3 +201,94 @@ def grad_elec(self, xy, singlet=None, atmlst=None, verbose=logger.INFO):
         de += grad_solver(pcmobj, dmxpy, v_grids=v_grids, v_grids_l=v_grids, q=q) * 2.0
         
         return de
+
+
+class WithSolventTDSCFNacMethod:
+    from gpu4pyscf.lib.utils import to_gpu, device
+
+    def __init__(self, tda_grad_method):
+        self.__dict__.update(tda_grad_method.__dict__)
+
+    def undo_solvent(self):
+        cls = self.__class__
+        name_mixin = self.base.with_solvent.__class__.__name__
+        obj = lib.view(self, lib.drop_class(cls, WithSolventTDSCFNacMethod, name_mixin))
+        del obj.with_solvent
+        return obj
+
+    def solvent_response(self, dm):
+        return self.base.with_solvent._B_dot_x(dm)*2.0 
+        
+    def get_nacv_ge(self, xy, EI, singlet=None, atmlst=None, verbose=logger.INFO):
+        if self.base.with_solvent.frozen:
+            raise RuntimeError('Frozen solvent model is not supported')
+
+        de_tuple = super().get_nacv_ge(xy, EI, singlet, atmlst, verbose) 
+        de, de_scaled, de_etf, de_etf_scaled = de_tuple
+
+        dm = self.base._scf.make_rdm1(ao_repr=True)
+        if dm.ndim == 3:
+            dm = dm[0] + dm[1]
+        dmP = self.dmz1doo  #1.0 * (self.dmz1doo + self.dmz1doo.T)
+        pcmobj = self.base.with_solvent
+        assert pcmobj.equilibrium_solvation
+
+        de_modify = 0
+        q_sym_dm = pcmobj._get_qsym(dm, with_nuc = True)[0]
+        qE_sym_dmP = pcmobj._get_qsym(dmP)[0]
+        de_modify += grad_qv(pcmobj, dm, q_sym = qE_sym_dmP)
+        de_modify += grad_nuc(pcmobj, dm, q_sym = qE_sym_dmP.get())
+        de_modify += grad_qv(pcmobj, dmP, q_sym = q_sym_dm)
+        v_grids_l = pcmobj._get_vgrids(dmP, with_nuc = False)[0]
+        de_modify += grad_solver(pcmobj, dm, v_grids_l = v_grids_l) * 2.0
+
+        de += de_modify
+        de_scaled = de/EI
+        de_etf += de_modify
+        de_etf_scaled = de_etf/EI
+        
+        return de, de_scaled, de_etf, de_etf_scaled
+
+    def get_nacv_ee(self, x_yI, x_yJ, EI, EJ, singlet=None, atmlst=None, verbose=logger.INFO):
+        if self.base.with_solvent.frozen:
+            raise RuntimeError('Frozen solvent model is not supported')
+
+        de_tuple = super().get_nacv_ee(x_yI, x_yJ, EI, EJ, singlet, atmlst, verbose) 
+        de, de_scaled, de_etf, de_etf_scaled = de_tuple
+
+        dm = self.base._scf.make_rdm1(ao_repr=True)
+        if dm.ndim == 3:
+            dm = dm[0] + dm[1]
+        dmP = 0.5 * (self.dmz1doo + self.dmz1doo.T)
+        dmxpyI = self.dmxpyI + self.dmxpyI.T
+        dmxpyJ = self.dmxpyJ + self.dmxpyJ.T
+        pcmobj = self.base.with_solvent
+        assert pcmobj.equilibrium_solvation
+
+        de_modify = 0.0
+        q_sym_dm = pcmobj._get_qsym(dm, with_nuc = True)[0]
+        qE_sym_dmP = pcmobj._get_qsym(dmP)[0]
+        qE_sym_dmxpyI = pcmobj._get_qsym(dmxpyI)[0]
+        qE_sym_dmxpyJ = pcmobj._get_qsym(dmxpyJ)[0]
+        de_modify += grad_qv(pcmobj, dm, q_sym = qE_sym_dmP)
+        de_modify += grad_nuc(pcmobj, dm, q_sym = qE_sym_dmP.get())
+        de_modify += grad_qv(pcmobj, dmP, q_sym = q_sym_dm)
+        v_grids_l = pcmobj._get_vgrids(dmP, with_nuc = False)[0]
+        de_modify += grad_solver(pcmobj, dm, v_grids_l = v_grids_l) * 2.0
+
+        de_modify += grad_qv(pcmobj, dmxpyJ, q_sym = qE_sym_dmxpyI)
+        de_modify += grad_qv(pcmobj, dmxpyI, q_sym = qE_sym_dmxpyJ)
+
+        v_gridsJ = pcmobj._get_vgrids(dmxpyJ, with_nuc = False)[0]
+        v_gridsI = pcmobj._get_vgrids(dmxpyI, with_nuc = False)[0]
+        qI = pcmobj._get_qsym(dmxpyI, with_nuc = False)[1]
+        qJ = pcmobj._get_qsym(dmxpyJ, with_nuc = False)[1]
+        de_modify += grad_solver(pcmobj, dmxpyJ, v_grids=v_gridsI, v_grids_l=v_gridsJ, q=qI)
+        de_modify += grad_solver(pcmobj, dmxpyI, v_grids=v_gridsJ, v_grids_l=v_gridsI, q=qJ)
+
+        de = de + de_modify
+        de_scaled = de/(EJ-EI)
+        de_etf = de_etf + de_modify
+        de_etf_scaled = de_etf/(EJ-EI)
+        
+        return de, de_scaled, de_etf, de_etf_scaled
diff --git a/gpu4pyscf/solvent/tests/test_pcm_tdscf_nac.py b/gpu4pyscf/solvent/tests/test_pcm_tdscf_nac.py
diff --git a/gpu4pyscf/tdscf/ris.py b/gpu4pyscf/tdscf/ris.py
