|
| 1 | +# Copyright 2021-2025 The PySCF Developers. All Rights Reserved. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | +""" |
| 15 | +Nonadiabatic derivetive coupling matrix element calculation is now in experiment. |
| 16 | +This module is under development. |
| 17 | +""" |
| 18 | + |
| 19 | +from functools import reduce |
| 20 | +import cupy as cp |
| 21 | +import numpy as np |
| 22 | +from pyscf import lib |
| 23 | +import pyscf |
| 24 | +from gpu4pyscf.lib import logger |
| 25 | +from gpu4pyscf.grad import rhf as rhf_grad |
| 26 | +from gpu4pyscf.df import int3c2e |
| 27 | +from gpu4pyscf.lib.cupy_helper import contract |
| 28 | +from gpu4pyscf.scf import cphf |
| 29 | +from pyscf import __config__ |
| 30 | +from gpu4pyscf.lib import utils |
| 31 | +from gpu4pyscf import tdscf |
| 32 | +from pyscf.scf import _vhf |
| 33 | + |
| 34 | + |
| 35 | +def get_nacv(td_nac, x_yI, EI, singlet=True, atmlst=None, verbose=logger.INFO): |
| 36 | + """ |
| 37 | + Only supports for ground-excited states. |
| 38 | + Ref: |
| 39 | + [1] 10.1063/1.4903986 main reference |
| 40 | + [2] 10.1021/acs.accounts.1c00312 |
| 41 | + [3] 10.1063/1.4885817 |
| 42 | + """ |
| 43 | + if singlet is False: |
| 44 | + raise NotImplementedError('Only supports for singlet states') |
| 45 | + mol = td_nac.mol |
| 46 | + mf = td_nac.base._scf |
| 47 | + mf_grad = mf.nuc_grad_method() |
| 48 | + mo_coeff = cp.asarray(mf.mo_coeff) |
| 49 | + mo_energy = cp.asarray(mf.mo_energy) |
| 50 | + mo_occ = cp.asarray(mf.mo_occ) |
| 51 | + nao, nmo = mo_coeff.shape |
| 52 | + nocc = int((mo_occ > 0).sum()) |
| 53 | + nvir = nmo - nocc |
| 54 | + orbv = mo_coeff[:, nocc:] |
| 55 | + orbo = mo_coeff[:, :nocc] |
| 56 | + |
| 57 | + xI, yI = x_yI |
| 58 | + xI = cp.asarray(xI).reshape(nocc, nvir).T |
| 59 | + if not isinstance(yI, np.ndarray) and not isinstance(yI, cp.ndarray): |
| 60 | + yI = cp.zeros_like(xI) |
| 61 | + yI = cp.asarray(yI).reshape(nocc, nvir).T |
| 62 | + LI = xI-yI # eq.(83) in Ref. [1] |
| 63 | + |
| 64 | + vresp = mf.gen_response(singlet=None, hermi=1) |
| 65 | + |
| 66 | + def fvind(x): |
| 67 | + dm = reduce(cp.dot, (orbv, x.reshape(nvir, nocc) * 2, orbo.T)) # double occupency |
| 68 | + v1ao = vresp(dm + dm.T) |
| 69 | + return reduce(cp.dot, (orbv.T, v1ao, orbo)).ravel() |
| 70 | + |
| 71 | + z1 = cphf.solve( |
| 72 | + fvind, |
| 73 | + mo_energy, |
| 74 | + mo_occ, |
| 75 | + -LI*1.0*EI, # only one spin, negative in cphf |
| 76 | + max_cycle=td_nac.cphf_max_cycle, |
| 77 | + tol=td_nac.cphf_conv_tol)[0] # eq.(83) in Ref. [1] |
| 78 | + |
| 79 | + z1 = z1.reshape(nvir, nocc) |
| 80 | + z1ao = reduce(cp.dot, (orbv, z1, orbo.T)) * 2 # double occupency |
| 81 | + # eq.(50) in Ref. [1] |
| 82 | + z1aoS = (z1ao + z1ao.T)*0.5 # 0.5 is in the definition of z1aoS |
| 83 | + # eq.(73) in Ref. [1] |
| 84 | + GZS = vresp(z1aoS) # generate the double occupency |
| 85 | + GZS_mo = reduce(cp.dot, (mo_coeff.T, GZS, mo_coeff)) |
| 86 | + W = cp.zeros((nmo, nmo)) # eq.(75) in Ref. [1] |
| 87 | + W[:nocc, :nocc] = GZS_mo[:nocc, :nocc] |
| 88 | + zeta0 = mo_energy[nocc:, cp.newaxis] |
| 89 | + zeta0 = z1 * zeta0 |
| 90 | + W[:nocc, nocc:] = GZS_mo[:nocc, nocc:] + 0.5*yI.T*EI + 0.5*zeta0.T #* eq.(43), (56), (28) in Ref. [1] |
| 91 | + zeta1 = mo_energy[cp.newaxis, :nocc] |
| 92 | + zeta1 = z1 * zeta1 |
| 93 | + W[nocc:, :nocc] = 0.5*xI*EI + 0.5*zeta1 |
| 94 | + W = reduce(cp.dot, (mo_coeff, W , mo_coeff.T)) * 2.0 |
| 95 | + |
| 96 | + mf_grad = mf.nuc_grad_method() |
| 97 | + s1 = mf_grad.get_ovlp(mol) |
| 98 | + dmz1doo = z1aoS |
| 99 | + oo0 = reduce(cp.dot, (orbo, orbo.T)) * 2.0 |
| 100 | + |
| 101 | + if atmlst is None: |
| 102 | + atmlst = range(mol.natm) |
| 103 | + |
| 104 | + hcore_deriv = pyscf.grad.rhf.hcore_generator(mf_grad.to_cpu(), mol) |
| 105 | + offsetdic = mol.offset_nr_by_atom() |
| 106 | + de = cp.zeros((len(atmlst),3)) |
| 107 | + de_etf = cp.zeros((len(atmlst),3)) |
| 108 | + xIao = reduce(cp.dot, (orbo, xI.T, orbv.T)) * 2 |
| 109 | + yIao = reduce(cp.dot, (orbv, yI, orbo.T)) * 2 |
| 110 | + eri1 = -mol.intor('int2e_ip1', aosym='s1', comp=3) |
| 111 | + eri1 = eri1.reshape(3,nao,nao,nao,nao) |
| 112 | + for k, ia in enumerate(atmlst): # eq.(58) in Ref. [1] |
| 113 | + shl0, shl1, p0, p1 = offsetdic[ia] |
| 114 | + h1ao = hcore_deriv(ia) |
| 115 | + h1ao = cp.asarray(h1ao) |
| 116 | + s1_tmp = s1.copy() |
| 117 | + s1_tmp[:,:p0] = 0 |
| 118 | + s1_tmp[:,p1:] = 0 |
| 119 | + s1_tmp = s1_tmp + s1_tmp.transpose(0,2,1) |
| 120 | + |
| 121 | + eri1a = eri1.copy() |
| 122 | + eri1a[:,:p0] = 0 |
| 123 | + eri1a[:,p1:] = 0 |
| 124 | + eri1a = eri1a + eri1a.transpose(0,2,1,3,4) |
| 125 | + eri1a = eri1a + eri1a.transpose(0,3,4,1,2) |
| 126 | + erijk = eri1a - eri1a.transpose(0,1,4,3,2)*0.5 |
| 127 | + |
| 128 | + ds1_tmp = s1.copy() |
| 129 | + ds1_tmp = ds1_tmp.transpose(0,2,1) |
| 130 | + ds1_tmp[:,:,:p0] = 0 |
| 131 | + ds1_tmp[:,:,p1:] = 0 |
| 132 | + |
| 133 | + de[k] = cp.einsum('xpq,pq->x', h1ao, dmz1doo) |
| 134 | + de[k] -= cp.einsum('xpq,pq->x', s1_tmp, W) |
| 135 | + de[k] += cp.einsum('xijkl,ij,kl->x', erijk, oo0, dmz1doo) |
| 136 | + de_etf[k] = de[k] |
| 137 | + de[k] += cp.einsum('xij,ij->x', ds1_tmp, xIao*EI) |
| 138 | + de[k] += cp.einsum('xij,ij->x', ds1_tmp, yIao*EI) |
| 139 | + de_etf[k] += cp.einsum('xij,ij->x', s1_tmp, xIao*EI) * 0.5 |
| 140 | + de_etf[k] += cp.einsum('xij,ij->x', s1_tmp, yIao*EI) * 0.5 |
| 141 | + |
| 142 | + de = de.get() |
| 143 | + de_etf = de_etf.get() |
| 144 | + return de, de/EI, de_etf, de_etf/EI |
| 145 | + |
| 146 | + |
| 147 | +class NAC(lib.StreamObject): |
| 148 | + |
| 149 | + cphf_max_cycle = getattr(__config__, "grad_tdrhf_Gradients_cphf_max_cycle", 20) |
| 150 | + cphf_conv_tol = getattr(__config__, "grad_tdrhf_Gradients_cphf_conv_tol", 1e-8) |
| 151 | + |
| 152 | + to_cpu = utils.to_cpu |
| 153 | + to_gpu = utils.to_gpu |
| 154 | + device = utils.device |
| 155 | + |
| 156 | + _keys = { |
| 157 | + "cphf_max_cycle", |
| 158 | + "cphf_conv_tol", |
| 159 | + "mol", |
| 160 | + "base", |
| 161 | + "chkfile", |
| 162 | + "states", |
| 163 | + "atmlst", |
| 164 | + "de", |
| 165 | + "de_scaled", |
| 166 | + "de_etf", |
| 167 | + "de_etf_scaled", |
| 168 | + } |
| 169 | + |
| 170 | + def __init__(self, td): |
| 171 | + self.verbose = td.verbose |
| 172 | + self.stdout = td.stdout |
| 173 | + self.mol = td.mol |
| 174 | + self.base = td |
| 175 | + self.states = (0, 1) # of which the gradients to be computed. |
| 176 | + self.atmlst = None |
| 177 | + self.de = None |
| 178 | + self.de_scaled = None |
| 179 | + self.de_etf = None |
| 180 | + self.de_etf_scaled = None |
| 181 | + |
| 182 | + def dump_flags(self, verbose=None): |
| 183 | + log = logger.new_logger(self, verbose) |
| 184 | + log.info("\n") |
| 185 | + log.info( |
| 186 | + "******** LR %s gradients for %s ********", |
| 187 | + self.base.__class__, |
| 188 | + self.base._scf.__class__, |
| 189 | + ) |
| 190 | + log.info("cphf_conv_tol = %g", self.cphf_conv_tol) |
| 191 | + log.info("cphf_max_cycle = %d", self.cphf_max_cycle) |
| 192 | + log.info("chkfile = %s", self.chkfile) |
| 193 | + log.info(f"States ID = {self.states}") |
| 194 | + log.info("\n") |
| 195 | + return self |
| 196 | + |
| 197 | + @lib.with_doc(get_nacv.__doc__) |
| 198 | + def get_nacv(self, x_yI, EI, singlet, atmlst=None, verbose=logger.INFO): |
| 199 | + return get_nacv(self, x_yI, EI, singlet, atmlst, verbose) |
| 200 | + |
| 201 | + def kernel(self, xy_I=None, xy_J=None, E_I=None, E_J=None, singlet=None, atmlst=None): |
| 202 | + |
| 203 | + logger.warn(self, "This module is under development!!") |
| 204 | + |
| 205 | + if singlet is None: |
| 206 | + singlet = self.base.singlet |
| 207 | + if atmlst is None: |
| 208 | + atmlst = self.atmlst |
| 209 | + else: |
| 210 | + self.atmlst = atmlst |
| 211 | + |
| 212 | + if self.verbose >= logger.WARN: |
| 213 | + self.check_sanity() |
| 214 | + if self.verbose >= logger.INFO: |
| 215 | + self.dump_flags() |
| 216 | + |
| 217 | + if xy_I is None or xy_J is None: |
| 218 | + states = sorted(self.states) |
| 219 | + I, J = states |
| 220 | + if I < 0 or J < 0: |
| 221 | + raise ValueError("Excited states ID should be non-negetive integers.") |
| 222 | + elif I > 0: |
| 223 | + raise NotImplementedError("Only for ground-excited states nonadiabatic coupling.") |
| 224 | + elif I == 0: |
| 225 | + xy_I = self.base.xy[J-1] |
| 226 | + E_I = self.base.e[J-1] |
| 227 | + self.de, self.de_scaled, self.de_etf, self.de_etf_scaled \ |
| 228 | + = self.get_nacv(xy_I, E_I, singlet, atmlst, verbose=self.verbose) |
| 229 | + self._finalize() |
| 230 | + else: |
| 231 | + raise NotImplementedError("Only for ground-excited states nonadiabatic coupling.") |
| 232 | + return self.de |
| 233 | + |
| 234 | + def get_veff(self, mol=None, dm=None, j_factor=1.0, k_factor=1.0, omega=0.0, hermi=0, verbose=None): |
| 235 | + """ |
| 236 | + Computes the first-order derivatives of the energy contributions from |
| 237 | + Veff per atom. |
| 238 | +
|
| 239 | + NOTE: This function is incompatible to the one implemented in PySCF CPU version. |
| 240 | + In the CPU version, get_veff returns the first order derivatives of Veff matrix. |
| 241 | + """ |
| 242 | + if mol is None: |
| 243 | + mol = self.mol |
| 244 | + if dm is None: |
| 245 | + dm = self.base.make_rdm1() |
| 246 | + if omega == 0.0: |
| 247 | + vhfopt = self.base._scf._opt_gpu.get(None, None) |
| 248 | + return rhf_grad._jk_energy_per_atom(mol, dm, vhfopt, j_factor=j_factor, k_factor=k_factor, verbose=verbose) |
| 249 | + else: |
| 250 | + vhfopt = self.base._scf._opt_gpu.get(omega, None) |
| 251 | + with mol.with_range_coulomb(omega): |
| 252 | + return rhf_grad._jk_energy_per_atom( |
| 253 | + mol, dm, vhfopt, j_factor=j_factor, k_factor=k_factor, verbose=verbose) |
| 254 | + |
| 255 | + |
| 256 | + def _finalize(self): |
| 257 | + if self.verbose >= logger.NOTE: |
| 258 | + logger.note( |
| 259 | + self, |
| 260 | + "--------- %s nonadiabatic derivative coupling for states %d and %d----------", |
| 261 | + self.base.__class__.__name__, |
| 262 | + self.states[0], |
| 263 | + self.states[1], |
| 264 | + ) |
| 265 | + self._write(self.mol, self.de, self.atmlst) |
| 266 | + logger.note( |
| 267 | + self, |
| 268 | + "--------- %s nonadiabatic derivative coupling for states %d and %d after E scaled (divided by E)----------", |
| 269 | + self.base.__class__.__name__, |
| 270 | + self.states[0], |
| 271 | + self.states[1], |
| 272 | + ) |
| 273 | + self._write(self.mol, self.de_scaled, self.atmlst) |
| 274 | + logger.note( |
| 275 | + self, |
| 276 | + "--------- %s nonadiabatic derivative coupling for states %d and %d with ETF----------", |
| 277 | + self.base.__class__.__name__, |
| 278 | + self.states[0], |
| 279 | + self.states[1], |
| 280 | + ) |
| 281 | + self._write(self.mol, self.de_etf, self.atmlst) |
| 282 | + logger.note( |
| 283 | + self, |
| 284 | + "--------- %s nonadiabatic derivative coupling for states %d and %d with ETF after E scaled (divided by E)----------", |
| 285 | + self.base.__class__.__name__, |
| 286 | + self.states[0], |
| 287 | + self.states[1], |
| 288 | + ) |
| 289 | + self._write(self.mol, self.de_etf_scaled, self.atmlst) |
| 290 | + logger.note(self, "----------------------------------------------") |
| 291 | + |
| 292 | + def solvent_response(self, dm): |
| 293 | + return 0.0 |
| 294 | + |
| 295 | + as_scanner = NotImplemented |
| 296 | + |
| 297 | + to_gpu = lib.to_gpu |
| 298 | + |
| 299 | + |
| 300 | +tdscf.rhf.TDA.NAC = tdscf.rhf.TDHF.NAC = lib.class_as_method(NAC) |
0 commit comments