Source code for pyscf.df.grad.sacasscf

#!/usr/bin/env python
# Copyright 2014-2022 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.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

import numpy as np
import time
import gc
from functools import reduce
from scipy import linalg
from pyscf import gto
from pyscf import mcscf, lib, ao2mo
from pyscf.grad import lagrange
from pyscf.grad import rhf as rhf_grad
from pyscf.grad import sacasscf as sacasscf_grad
from pyscf.grad import casscf as casscf_grad
from pyscf.grad.mp2 import _shell_prange
from pyscf.mcscf import mc1step, mc1step_symm, newton_casscf
from pyscf.mcscf.addons import StateAverageMCSCFSolver
from pyscf.df.grad import casscf as dfcasscf_grad
from pyscf.df.grad import rhf as dfrhf_grad
from pyscf.fci.direct_spin1 import _unpack_nelec
from pyscf.fci.spin_op import spin_square0
from pyscf.fci import cistring
from pyscf.df.grad.casdm2_util import solve_df_rdm2, grad_elec_dferi, grad_elec_auxresponse_dferi

[docs] def Lorb_dot_dgorb_dx (Lorb, mc, mo_coeff=None, ci=None, atmlst=None, mf_grad=None, eris=None, verbose=None, auxbasis_response=True): ''' Modification of pyscf.grad.casscf.kernel to compute instead the orbital Lagrange term nuclear gradient (sum_pq Lorb_pq d2_Ecas/d_lambda d_kpq) This involves removing nuclear-nuclear terms and making the substitution (D_[p]q + D_p[q]) -> D_pq (d_[p]qrs + d_pq[r]s + d_p[q]rs + d_pqr[s]) -> d_pqrs Where [] around an index implies contraction with Lorb from the left, so that the external index (regardless of whether the index on the rdm is bra or ket) is always the first index of Lorb. ''' # dmo = smoT.dao.smo # dao = mo.dmo.moT t0 = (lib.logger.process_clock (), lib.logger.perf_counter ()) if mo_coeff is None: mo_coeff = mc.mo_coeff if ci is None: ci = mc.ci if mf_grad is None: mf_grad = dfrhf_grad.Gradients (mc._scf) if mc.frozen is not None: raise NotImplementedError mol = mc.mol ncore = mc.ncore ncas = mc.ncas nocc = ncore + ncas nelecas = mc.nelecas nao, nmo = mo_coeff.shape mo_core = mo_coeff[:,:ncore] mo_cas = mo_coeff[:,ncore:nocc] # MRH: new 'effective' MO coefficients including contraction from the Lagrange multipliers moL_coeff = np.dot (mo_coeff, Lorb) s0_inv = np.dot (mo_coeff, mo_coeff.T) moL_core = moL_coeff[:,:ncore] moL_cas = moL_coeff[:,ncore:nocc] # MRH: these SHOULD be state-averaged! Use the actual sacasscf object! casdm1, casdm2 = mc.fcisolver.make_rdm12(ci, ncas, nelecas) # gfock = Generalized Fock, Adv. Chem. Phys., 69, 63 # MRH: each index exactly once! dm_core = np.dot(mo_core, mo_core.T) * 2 dm_cas = reduce(np.dot, (mo_cas, casdm1, mo_cas.T)) # MRH: new density matrix terms dmL_core = np.dot(moL_core, mo_core.T) * 2 dmL_cas = reduce(np.dot, (moL_cas, casdm1, mo_cas.T)) dmL_core += dmL_core.T dmL_cas += dmL_cas.T dm1 = dm_core + dm_cas dm1L = dmL_core + dmL_cas # MRH: end new density matrix terms # MRH: wrap the integral instead of the density matrix. I THINK the sign is the same! # mo sets 0 and 2 should be transposed, 1 and 3 should be not transposed; this will lead to correct sign # Except I can't do this for the external index, because the external index is contracted to ovlp matrix, # not the 2RDM aapa = np.zeros ((ncas, ncas, nmo, ncas), dtype=dm_cas.dtype) aapaL = np.zeros ((ncas, ncas, nmo, ncas), dtype=dm_cas.dtype) for i in range (nmo): jbuf = eris.ppaa[i] kbuf = eris.papa[i] aapa[:,:,i,:] = jbuf[ncore:nocc,:,:].transpose (1,2,0) aapaL[:,:,i,:] += np.tensordot (jbuf, Lorb[:,ncore:nocc], axes=((0),(0))) kbuf = np.tensordot (kbuf, Lorb[:,ncore:nocc], axes=((1),(0))).transpose (1,2,0) aapaL[:,:,i,:] += kbuf + kbuf.transpose (1,0,2) # MRH: new vhf terms vj, vk = mc._scf.get_jk(mol, (dm_core, dm_cas)) vjL, vkL = mc._scf.get_jk(mol, (dmL_core, dmL_cas)) h1 = mc.get_hcore() vhf_c = vj[0] - vk[0] * .5 vhf_a = vj[1] - vk[1] * .5 vhfL_c = vjL[0] - vkL[0] * .5 vhfL_a = vjL[1] - vkL[1] * .5 # MRH: I rewrote this Feff calculation completely, double-check it gfock = np.dot (h1, dm1L) # h1e gfock += np.dot ((vhf_c + vhf_a), dmL_core) # core-core and active-core, 2nd 1RDM linked gfock += np.dot ((vhfL_c + vhfL_a), dm_core) # core-core and active-core, 1st 1RDM linked gfock += np.dot (vhfL_c, dm_cas) # core-active, 1st 1RDM linked gfock += np.dot (vhf_c, dmL_cas) # core-active, 2nd 1RDM linked gfock = np.dot (s0_inv, gfock) # Definition of quantity is in MO's; going (AO->MO->AO) incurs an inverse ovlp gfock += reduce (np.dot, (mo_coeff, np.einsum('uviw,uvtw->it', aapaL, casdm2), mo_cas.T)) # active-active # MRH: I have to contract this external 2RDM index explicitly on the 2RDM but fortunately I can do so here gfock += reduce (np.dot, (mo_coeff, np.einsum('uviw,vuwt->it', aapa, casdm2), moL_cas.T)) # MRH: As of 04/18/2019, the two-body part of this is including aapaL is definitely, unambiguously correct dme0 = (gfock+gfock.T)/2 # This transpose is for the overlap matrix later on aapa = vj = vk = vhf_c = vhf_a = None if atmlst is None: atmlst = list (range(mol.natm)) aoslices = mol.aoslice_by_atom() de_hcore = np.zeros((len(atmlst),3)) de_renorm = np.zeros((len(atmlst),3)) de_eri = np.zeros((len(atmlst),3)) de_aux = np.zeros((len(atmlst),3)) de = np.zeros((len(atmlst),3)) #vhf1c, vhf1a, vhf1cL, vhf1aL = mf_grad.get_veff(mol, (dm_core, dm_cas, dmL_core, dmL_cas)) vj, vk = mf_grad.get_jk (mol, (dm_core, dm_cas, dmL_core, dmL_cas)) vhf1c, vhf1a, vhf1cL, vhf1aL = list (vj - vk * 0.5) if auxbasis_response: de_aux = vj.aux - 0.5 * vk.aux # D.T + T.D de_aux = ((de_aux[0,2] + de_aux[2,0]) # core-core + (de_aux[0,3] + de_aux[2,1]) # core-active + (de_aux[1,2] + de_aux[3,0])) # active-core vj = vk = None hcore_deriv = mf_grad.hcore_generator(mol) s1 = mf_grad.get_ovlp(mol) t0 = lib.logger.timer (mc, 'SA-CASSCF Lorb_dot_dgorb 1-electron part', *t0) # I am trying to contract the eris with a notional casdm2 which has four separate terms. casdm2 += casdm2.transpose (1,0,3,2) # Now I should only need 2 separate terms... # The bare 3-center eris and the auxbasis derivatives are always symmetric wrt AOs # grad_elec_dferi is explicitly symmetrized wrt AOs. # If this fails I can always debug it by kludging ncore, ncas -> 0, nmo dfcasdm2 = solve_df_rdm2 (mc, mo_cas=(mo_cas, moL_cas), casdm2=casdm2) de_eri += grad_elec_dferi (mc, mo_cas=mo_cas, dfcasdm2=dfcasdm2, atmlst=atmlst, max_memory=mc.max_memory)[0] if auxbasis_response: de_aux += grad_elec_auxresponse_dferi (mc, mo_cas=mo_cas, dfcasdm2=dfcasdm2, atmlst=atmlst, max_memory=mc.max_memory)[0] dfcasdm2 = solve_df_rdm2 (mc, mo_cas=mo_cas, casdm2=casdm2) de_eri += grad_elec_dferi (mc, mo_cas=(mo_cas, moL_cas), dfcasdm2=dfcasdm2, atmlst=atmlst, max_memory=mc.max_memory)[0] if auxbasis_response: de_aux += grad_elec_auxresponse_dferi (mc, mo_cas=(mo_cas, moL_cas), dfcasdm2=dfcasdm2, atmlst=atmlst, max_memory=mc.max_memory)[0] dfcasdm2 = casdm2 = None for k, ia in enumerate(atmlst): shl0, shl1, p0, p1 = aoslices[ia] h1ao = hcore_deriv(ia) # MRH: h1e and Feff terms de_hcore[k] += np.einsum('xij,ij->x', h1ao, dm1L) de_renorm[k] -= np.einsum('xij,ij->x', s1[:,p0:p1], dme0[p0:p1]) * 2 # MRH: core-core and core-active 2RDM terms de_eri[k] += np.einsum('xij,ij->x', vhf1c[:,p0:p1], dm1L[p0:p1]) * 2 de_eri[k] += np.einsum('xij,ij->x', vhf1cL[:,p0:p1], dm1[p0:p1]) * 2 # MRH: active-core 2RDM terms de_eri[k] += np.einsum('xij,ij->x', vhf1a[:,p0:p1], dmL_core[p0:p1]) * 2 de_eri[k] += np.einsum('xij,ij->x', vhf1aL[:,p0:p1], dm_core[p0:p1]) * 2 # MRH: deleted the nuclear-nuclear part to avoid double-counting # lesson learned from debugging - mol.intor computes -1 * the derivative and only # for one index # on the other hand, mf_grad.hcore_generator computes the actual derivative of # h1 for both indices and with the correct sign lib.logger.debug (mc, "Orb lagrange hcore component:\n{}".format (de_hcore)) lib.logger.debug (mc, "Orb lagrange renorm component:\n{}".format (de_renorm)) lib.logger.debug (mc, "Orb lagrange eri component:\n{}".format (de_eri)) lib.logger.debug (mc, "Orb lagrange aux component:\n{}".format (de_aux)) de = de_hcore + de_renorm + de_eri + de_aux return de
[docs] def Lci_dot_dgci_dx (Lci, weights, mc, mo_coeff=None, ci=None, atmlst=None, mf_grad=None, eris=None, verbose=None, auxbasis_response=True): ''' Modification of pyscf.grad.casscf.kernel to compute instead the CI Lagrange term nuclear gradient (sum_IJ Lci_IJ d2_Ecas/d_lambda d_PIJ) This involves removing all core-core and nuclear-nuclear terms and making the substitution sum_I w_I<L_I|p'q|I> + c.c. -> <0|p'q|0> sum_I w_I<L_I|p'r'sq|I> + c.c. -> <0|p'r'sq|0> The active-core terms (sum_I w_I<L_I|x'iyi|I>, sum_I w_I <L_I|x'iiy|I>, c.c.) must be retained.''' if mo_coeff is None: mo_coeff = mc.mo_coeff if ci is None: ci = mc.ci if mf_grad is None: mf_grad = dfrhf_grad.Gradients (mc._scf) if mc.frozen is not None: raise NotImplementedError t0 = (lib.logger.process_clock (), lib.logger.perf_counter ()) mol = mc.mol ncore = mc.ncore ncas = mc.ncas nocc = ncore + ncas nelecas = mc.nelecas nao, nmo = mo_coeff.shape mo_occ = mo_coeff[:,:nocc] mo_core = mo_coeff[:,:ncore] mo_cas = mo_coeff[:,ncore:nocc] # MRH: TDMs + c.c. instead of RDMs; 06/30/2020: new interface in mcscf.addons makes this much more transparent casdm1, casdm2 = mc.fcisolver.trans_rdm12 (Lci, ci, ncas, nelecas) casdm1 += casdm1.transpose (1,0) casdm2 += casdm2.transpose (1,0,3,2) # gfock = Generalized Fock, Adv. Chem. Phys., 69, 63 dm_core = np.dot(mo_core, mo_core.T) * 2 dm_cas = reduce(np.dot, (mo_cas, casdm1, mo_cas.T)) aapa = np.zeros ((ncas, ncas, nmo, ncas), dtype=dm_cas.dtype) for i in range (nmo): aapa[:,:,i,:] = eris.ppaa[i][ncore:nocc,:,:].transpose (1,2,0) vj, vk = mc._scf.get_jk(mol, (dm_core, dm_cas)) h1 = mc.get_hcore() vhf_c = vj[0] - vk[0] * .5 vhf_a = vj[1] - vk[1] * .5 # MRH: delete h1 + vhf_c from the first line below (core and core-core stuff) # Also extend gfock to span the whole space gfock = np.zeros_like (dm_cas) gfock[:,:nocc] = reduce(np.dot, (mo_coeff.T, vhf_a, mo_occ)) * 2 gfock[:,ncore:nocc] = reduce(np.dot, (mo_coeff.T, h1 + vhf_c, mo_cas, casdm1)) gfock[:,ncore:nocc] += np.einsum('uvpw,vuwt->pt', aapa, casdm2) dme0 = reduce(np.dot, (mo_coeff, (gfock+gfock.T)*.5, mo_coeff.T)) aapa = vj = vk = vhf_c = vhf_a = h1 = gfock = None if atmlst is None: atmlst = range(mol.natm) aoslices = mol.aoslice_by_atom() de_hcore = np.zeros((len(atmlst),3)) de_renorm = np.zeros((len(atmlst),3)) de_eri = np.zeros((len(atmlst),3)) de_aux = np.zeros((len(atmlst),3)) de = np.zeros((len(atmlst),3)) #vhf1c, vhf1a = mf_grad.get_veff(mol, (dm_core, dm_cas)) vj, vk = mf_grad.get_jk (mol, (dm_core, dm_cas)) if auxbasis_response: de_aux = vj.aux - 0.5 * vk.aux de_aux = de_aux[0,1] + de_aux[1,0] # ^ de_aux[0,0] not included b/c this is CAS lagrange multipliers vhf1c, vhf1a = list (vj - vk * 0.5) vj = vk = None hcore_deriv = mf_grad.hcore_generator(mol) s1 = mf_grad.get_ovlp(mol) dfcasdm2 = casdm2 = solve_df_rdm2 (mc, mo_cas=mo_cas, casdm2=casdm2) de_eri = grad_elec_dferi (mc, mo_cas=mo_cas, dfcasdm2=dfcasdm2, atmlst=atmlst, max_memory=mc.max_memory)[0] if auxbasis_response: de_aux += grad_elec_auxresponse_dferi (mc, mo_cas=mo_cas, dfcasdm2=dfcasdm2, atmlst=atmlst, max_memory=mc.max_memory)[0] dfcasdm2 = casdm2 = None t0 = lib.logger.timer (mc, 'SA-CASSCF Lci_dot_dgci 1-electron part', *t0) for k, ia in enumerate(atmlst): shl0, shl1, p0, p1 = aoslices[ia] h1ao = hcore_deriv(ia) # MRH: dm1 -> dm_cas in the line below de_hcore[k] += np.einsum('xij,ij->x', h1ao, dm_cas) de_renorm[k] -= np.einsum('xij,ij->x', s1[:,p0:p1], dme0[p0:p1]) * 2 # MRH: dm1 -> dm_cas in the line below. Also eliminate core-core terms de_eri[k] += np.einsum('xij,ij->x', vhf1c[:,p0:p1], dm_cas[p0:p1]) * 2 de_eri[k] += np.einsum('xij,ij->x', vhf1a[:,p0:p1], dm_core[p0:p1]) * 2 lib.logger.debug (mc, "CI lagrange hcore component:\n{}".format (de_hcore)) lib.logger.debug (mc, "CI lagrange renorm component:\n{}".format (de_renorm)) lib.logger.debug (mc, "CI lagrange eri component:\n{}".format (de_eri)) lib.logger.debug (mc, "CI lagrange aux component:\n{}".format (de_aux)) de = de_hcore + de_renorm + de_eri + de_aux return de
[docs] def as_scanner(mcscf_grad, state=None): '''Generating a nuclear gradients scanner/solver (for geometry optimizer). The returned solver is a function. This function requires one argument "mol" as input and returns energy and first order nuclear derivatives. The solver will automatically use the results of last calculation as the initial guess of the new calculation. All parameters assigned in the nuc-grad object and SCF object (DIIS, conv_tol, max_memory etc) are automatically applied in the solver. Note scanner has side effects. It may change many underlying objects (_scf, with_df, with_x2c, ...) during calculation. Examples: >>> from pyscf import gto, scf, mcscf >>> mol = gto.M(atom='N 0 0 0; N 0 0 1.1', verbose=0) >>> mc_grad_scanner = mcscf.CASSCF(scf.RHF(mol), 4, 4).nuc_grad_method().as_scanner() >>> etot, grad = mc_grad_scanner(gto.M(atom='N 0 0 0; N 0 0 1.1')) >>> etot, grad = mc_grad_scanner(gto.M(atom='N 0 0 0; N 0 0 1.5')) ''' from pyscf import gto if isinstance(mcscf_grad, lib.GradScanner): return mcscf_grad if state is None and (not hasattr (mcscf_grad, 'state') or (mcscf_grad.state is None)): return dfcasscf_grad.as_scanner (mcscf_grad) lib.logger.info(mcscf_grad, 'Create scanner for %s', mcscf_grad.__class__) name = mcscf_grad.__class__.__name__ + CASSCF_GradScanner.__name_mixin__ return lib.set_class(CASSCF_GradScanner(mcscf_grad, state), (CASSCF_GradScanner, mcscf_grad.__class__), name)
[docs] class CASSCF_GradScanner(lib.GradScanner): def __init__(self, g, state): lib.GradScanner.__init__(self, g) if state is not None: self.state = state def __call__(self, mol_or_geom, **kwargs): if isinstance(mol_or_geom, gto.MoleBase): assert mol_or_geom.__class__ == gto.Mole mol = mol_or_geom else: mol = self.mol.set_geom_(mol_or_geom, inplace=False) mc_scanner = self.base e_tot = mc_scanner(mol) if hasattr (mc_scanner, 'e_mcscf'): self.e_mcscf = mc_scanner.e_mcscf #if isinstance (e_tot, (list, tuple, np.ndarray)): e_tot = e_tot[self.state] if hasattr (mc_scanner, 'e_states'): e_tot = mc_scanner.e_states[self.state] self.mol = mol if not ('state' in kwargs): kwargs['state'] = self.state de = self.kernel(**kwargs) return e_tot, de
[docs] class Gradients (sacasscf_grad.Gradients): _keys = {'with_df', 'auxbasis_response'} def __init__(self, mc, state=None): self.auxbasis_response = True sacasscf_grad.Gradients.__init__(self, mc, state=state)
[docs] def kernel (self, **kwargs): mf_grad = kwargs['mf_grad'] if 'mf_grad' in kwargs else None if mf_grad is None: kwargs['mf_grad'] = dfrhf_grad.Gradients (self.base._scf) # The below only works because dfcasscf_grad is NOT a child of casscf_grad # For instance, I can't monkeypatch rhf_grad this way b/c dfrhf_grad refers to rhf_grad # Maybe it should be, in which case I will have to change this # But on the other hand maybe it can be even simpler? with lib.temporary_env (casscf_grad, Gradients=dfcasscf_grad.Gradients): return sacasscf_grad.Gradients.kernel (self, **kwargs)
[docs] def get_LdotJnuc (self, Lvec, **kwargs): with lib.temporary_env (sacasscf_grad, Lci_dot_dgci_dx=Lci_dot_dgci_dx, Lorb_dot_dgorb_dx=Lorb_dot_dgorb_dx): return sacasscf_grad.Gradients.get_LdotJnuc (self, Lvec, **kwargs)
to_gpu = lib.to_gpu