#!/usr/bin/env python
# Copyright 2025 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: Chenghan Li <lch004218@gmail.com>
#
import numpy as np
import pyscf
from pyscf import lib
from pyscf import gto
from pyscf import df
from pyscf import scf
from pyscf import grad
from pyscf.lib import logger
from pyscf.qmmm.pbc import mm_mole
from scipy.special import erfc, erf
[docs]
def add_mm_charges(scf_method, atoms_or_coords, a, charges, radii=None,
rcut_ewald=None, rcut_hcore=None, unit=None):
'''Embedding the one-electron (non-relativistic) potential generated by MM
point (or gaussian-distributed) charges into QM Hamiltonian.
The total energy includes the regular QM energy, the interaction between
the nuclei in QM region and the MM charges, and the static Coulomb
interaction between the electron density and the MM charges. The electrostatic
interactions between reference cell and periodic images are also computed.
It does not include the static Coulomb interactions of the MM charges,
the MM energy, the vdw interaction or other
bonding/non-bonding effects between QM region and MM particles.
Args:
scf_method : a HF or DFT object
atoms_or_coords : 2D array, shape (N,3)
MM particle coordinates
charges : 1D array
MM particle charges
a : 2D array, shape (3,3)
Lattice vectors
Kwargs:
radii : 1D array
The Gaussian charge distribution radii of MM atoms.
rcut_ewald: Float
Ewald real-space cutoff
rcut_hcore: Float
Real-space cutoff for exact QM-MM coupling
unit : str
Bohr, AU, Ang (case insensitive). Default is the same to mol.unit
Returns:
Same method object as the input scf_method with modified 1e Hamiltonian
Examples:
>>> mol = gto.M(atom='H 0 0 0; F 0 0 1', basis='ccpvdz', verbose=0)
>>> mf = add_mm_charges(dft.RKS(mol), [(0.5,0.6,0.8)], np.eye(3)*10, [-0.3])
>>> mf.kernel()
'''
mol = scf_method.mol
if unit is None:
unit = mol.unit
mm_mol = mm_mole.create_mm_mol(atoms_or_coords, a, charges, radii=radii,
rcut_ewald=rcut_ewald, rcut_hcore=rcut_hcore, unit=unit)
return qmmm_for_scf(scf_method, mm_mol)
[docs]
def qmmm_for_scf(method, mm_mol):
'''Add the potential of MM particles to SCF (HF and DFT) method
then generate the corresponding QM/MM method for the QM system.
Args:
mm_mol : MM Mole object
'''
if isinstance(method, scf.hf.SCF):
# Avoid to initialize QMMM twice
if isinstance(method, QMMM):
method.mm_mol = mm_mol
method.s1r = None
method.s1rr = None
method.mm_ewald_pot = None
method.qm_ewald_hess = None
method.e_nuc = None
return method
cls = QMMMSCF
else:
# post-HF methods
raise NotImplementedError()
return lib.set_class(cls(method, mm_mol), (cls, method.__class__))
[docs]
class QMMM:
__name_mixin__ = 'QMMM'
[docs]
class QMMMSCF(QMMM):
_keys = {'mm_mol', 's1r', 's1rr', 'mm_ewald_pot', 'qm_ewald_hess', 'e_nuc'}
to_gpu = NotImplemented
as_scanner = NotImplemented
def __init__(self, method, mm_mol):
self.__dict__.update(method.__dict__)
self.mm_mol = mm_mol
self.s1r = None
self.s1rr = None
self.mm_ewald_pot = None
self.qm_ewald_hess = None
self.e_nuc = None
[docs]
def dump_flags(self, verbose=None):
super().dump_flags(verbose)
logger.info(self, '** Add background charges for %s **',
self.__class__.__name__)
_a = self.mm_mol.lattice_vectors()
logger.info(self, 'lattice vectors a1 [%.9f, %.9f, %.9f]', *_a[0])
logger.info(self, ' a2 [%.9f, %.9f, %.9f]', *_a[1])
logger.info(self, ' a3 [%.9f, %.9f, %.9f]', *_a[2])
if self.verbose >= logger.DEBUG2:
logger.debug2(self, 'Charge Location')
coords = self.mm_mol.atom_coords()
charges = self.mm_mol.atom_charges()
for i, z in enumerate(charges):
logger.debug2(self, '%.9g %s', z, coords[i])
return self
[docs]
def get_mm_ewald_pot(self, mol, mm_mol):
return self.mm_mol.get_ewald_pot(
mol.atom_coords(),
mm_mol.atom_coords(), mm_mol.atom_charges())
[docs]
def get_qm_ewald_pot(self, mol, dm, qm_ewald_hess=None):
# hess = d^2 E / dQ_i dQ_j, d^2 E / dQ_i dD_ja, d^2 E / dDia dDjb, d^2 E/ dQ_i dO_jab
if qm_ewald_hess is None:
qm_ewald_hess = self.mm_mol.get_ewald_pot(mol.atom_coords())
self.qm_ewald_hess = qm_ewald_hess
charges = self.get_qm_charges(dm)
dips = self.get_qm_dipoles(dm)
quads = self.get_qm_quadrupoles(dm)
ewpot0 = lib.einsum('ij,j->i', qm_ewald_hess[0], charges)
ewpot0 += lib.einsum('ijx,jx->i', qm_ewald_hess[1], dips)
ewpot0 += lib.einsum('ijxy,jxy->i', qm_ewald_hess[3], quads)
ewpot1 = lib.einsum('ijx,i->jx', qm_ewald_hess[1], charges)
ewpot1 += lib.einsum('ijxy,jy->ix', qm_ewald_hess[2], dips)
ewpot2 = lib.einsum('ijxy,j->ixy', qm_ewald_hess[3], charges)
return ewpot0, ewpot1, ewpot2
[docs]
def get_hcore(self, mol=None):
cput0 = (logger.process_clock(), logger.perf_counter())
mm_mol = self.mm_mol
if mol is None:
mol = self.mol
rcut_hcore = mm_mol.rcut_hcore
h1e = super().get_hcore(mol)
Ls = mm_mol.get_lattice_Ls()
qm_center = np.mean(mol.atom_coords(), axis=0)
mask = np.linalg.norm(Ls, axis=-1) < 1e-12
Ls[mask] = [np.inf] * 3
r_qm = (mol.atom_coords() - qm_center)[None,:,:] - Ls[:,None,:]
r_qm = np.einsum('Lix,Lix->Li', r_qm, r_qm)
assert rcut_hcore**2 < np.min(r_qm), \
"QM image is within rcut_hcore of QM center. " + \
f"rcut_hcore = {rcut_hcore} >= min(r_qm) = {np.sqrt(np.min(r_qm))}"
Ls[Ls == np.inf] = 0.0
r_qm = mol.atom_coords() - qm_center
r_qm = np.einsum('ix,ix->i', r_qm, r_qm)
assert rcut_hcore**2 > np.max(r_qm), \
"Not all QM atoms are within rcut_hcore of QM center. " + \
f"rcut_hcore = {rcut_hcore} <= max(r_qm) = {np.sqrt(np.max(r_qm))}"
r_qm = None
all_coords = lib.direct_sum('ix+Lx->Lix',
mm_mol.atom_coords(), Ls).reshape(-1,3)
all_charges = np.hstack([mm_mol.atom_charges()] * len(Ls))
dist2 = all_coords - qm_center
dist2 = lib.einsum('ix,ix->i', dist2, dist2)
# charges within rcut_hcore exactly go into hcore
mask = dist2 <= rcut_hcore**2
charges = all_charges[mask]
coords = all_coords[mask]
logger.note(self, '%d MM charges see directly QM density'%charges.shape[0])
nao = mol.nao
max_memory = self.max_memory - lib.current_memory()[0]
blksize = int(min(max_memory*1e6/8/nao**2, 200))
blksize = max(blksize, 1)
if mm_mol.charge_model == 'gaussian' and len(coords) != 0:
expnts = np.hstack([mm_mol.get_zetas()] * len(Ls))[mask]
if mol.cart:
intor = 'int3c2e_cart'
else:
intor = 'int3c2e_sph'
cintopt = gto.moleintor.make_cintopt(mol._atm, mol._bas,
mol._env, intor)
v = 0
for i0, i1 in lib.prange(0, charges.size, blksize):
fakemol = gto.fakemol_for_charges(coords[i0:i1], expnts[i0:i1])
j3c = df.incore.aux_e2(mol, fakemol, intor=intor,
aosym='s2ij', cintopt=cintopt)
v += lib.einsum('xk,k->x', j3c, -charges[i0:i1])
if type(v) is not int:
v = lib.unpack_tril(v)
h1e += v
elif mm_mol.charge_model == 'point' and len(coords) != 0:
for i0, i1 in lib.prange(0, charges.size, blksize):
j3c = mol.intor('int1e_grids', hermi=1, grids=coords[i0:i1])
h1e += lib.einsum('kpq,k->pq', j3c, -charges[i0:i1])
else: # no MM charges
pass
logger.timer(self, 'get_hcore', *cput0)
return h1e
[docs]
def get_qm_charges(self, dm):
aoslices = self.mol.aoslice_by_atom()
chg = self.mol.atom_charges()
dmS = lib.dot(dm, self.get_ovlp())
qm_charges = list()
for iatm in range(self.mol.natm):
p0, p1 = aoslices[iatm, 2:]
qm_charges.append(chg[iatm] - np.trace(dmS[p0:p1, p0:p1]))
return np.asarray(qm_charges)
[docs]
def get_s1r(self):
if self.s1r is None:
cput0 = (logger.process_clock(), logger.perf_counter())
self.s1r = list()
mol = self.mol
bas_atom = mol._bas[:,gto.ATOM_OF]
for i in range(self.mol.natm):
b0, b1 = np.where(bas_atom == i)[0][[0,-1]]
shls_slice = (0, mol.nbas, b0, b1+1)
with mol.with_common_orig(mol.atom_coord(i)):
self.s1r.append(
mol.intor('int1e_r', shls_slice=shls_slice))
logger.timer(self, 'get_s1r', *cput0)
return self.s1r
[docs]
def get_qm_dipoles(self, dm, s1r=None):
if s1r is None:
s1r = self.get_s1r()
aoslices = self.mol.aoslice_by_atom()
qm_dipoles = list()
for iatm in range(self.mol.natm):
p0, p1 = aoslices[iatm, 2:]
qm_dipoles.append(
-lib.einsum('uv,xvu->x', dm[p0:p1], s1r[iatm]))
return np.asarray(qm_dipoles)
[docs]
def get_s1rr(self):
'''
int phi_u phi_v [3(r-Rc)otimes(r-Rc) - |r-Rc|^2] /2 dr
'''
if self.s1rr is None:
cput0 = (logger.process_clock(), logger.perf_counter())
self.s1rr = list()
mol = self.mol
nao = mol.nao
bas_atom = mol._bas[:,gto.ATOM_OF]
for i in range(self.mol.natm):
b0, b1 = np.where(bas_atom == i)[0][[0,-1]]
shls_slice = (0, mol.nbas, b0, b1+1)
with mol.with_common_orig(mol.atom_coord(i)):
s1rr_ = mol.intor('int1e_rr', shls_slice=shls_slice)
s1rr_ = s1rr_.reshape((3,3,nao,-1))
s1rr_trace = lib.einsum('xxuv->uv', s1rr_)
s1rr_ = 3/2 * s1rr_
for k in range(3):
s1rr_[k,k] -= 0.5 * s1rr_trace
self.s1rr.append(s1rr_)
logger.timer(self, 'get_s1rr', *cput0)
return self.s1rr
[docs]
def get_qm_quadrupoles(self, dm, s1rr=None):
if s1rr is None:
s1rr = self.get_s1rr()
aoslices = self.mol.aoslice_by_atom()
qm_quadrupoles = list()
for iatm in range(self.mol.natm):
p0, p1 = aoslices[iatm, 2:]
qm_quadrupoles.append(
-lib.einsum('uv,xyvu->xy', dm[p0:p1], s1rr[iatm]))
return np.asarray(qm_quadrupoles)
[docs]
def get_vdiff(self, mol, ewald_pot):
'''
vdiff_uv = d Q_I / d dm_uv ewald_pot[0]_I
+ d D_Ix / d dm_uv ewald_pot[1]_Ix
+ d O_Ixy / d dm_uv ewald_pot[2]_Ixy
'''
vdiff = np.zeros((mol.nao, mol.nao))
ovlp = self.get_ovlp()
s1r = self.get_s1r()
s1rr = self.get_s1rr()
aoslices = mol.aoslice_by_atom()
for iatm in range(mol.natm):
v0 = ewald_pot[0][iatm]
v1 = ewald_pot[1][iatm]
v2 = ewald_pot[2][iatm]
p0, p1 = aoslices[iatm, 2:]
vdiff[:,p0:p1] -= v0 * ovlp[:,p0:p1]
vdiff[:,p0:p1] -= lib.einsum('x,xuv->uv', v1, s1r[iatm])
vdiff[:,p0:p1] -= lib.einsum('xy,xyuv->uv', v2, s1rr[iatm])
vdiff = (vdiff + vdiff.T) / 2
return vdiff
[docs]
def get_veff(self, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1,
mm_ewald_pot=None, qm_ewald_pot=None):
if mol is None:
mol = self.mol
mm_mol = self.mm_mol
if mm_ewald_pot is None:
if self.mm_ewald_pot is not None:
mm_ewald_pot = self.mm_ewald_pot
else:
cput0 = (logger.process_clock(), logger.perf_counter())
mm_ewald_pot = self.get_mm_ewald_pot(mol, mm_mol)
self.mm_ewald_pot = mm_ewald_pot
logger.timer(self, 'get_mm_ewald_pot', *cput0)
if qm_ewald_pot is None:
if self.qm_ewald_hess is not None:
qm_ewald_pot = self.get_qm_ewald_pot(
mol, dm, self.qm_ewald_hess)
else:
cput0 = (logger.process_clock(), logger.perf_counter())
qm_ewald_pot = self.get_qm_ewald_pot(mol, dm)
logger.timer(self, 'get_qm_ewald_pot', *cput0)
ewald_pot = \
mm_ewald_pot[0] + qm_ewald_pot[0], \
mm_ewald_pot[1] + qm_ewald_pot[1], \
mm_ewald_pot[2] + qm_ewald_pot[2]
vdiff = self.get_vdiff(mol, ewald_pot)
if vhf_last is not None and isinstance(vhf_last, lib.NPArrayWithTag):
vhf_last = vhf_last.veff_rs
veff = super().get_veff(mol, dm, dm_last, vhf_last, hermi)
if isinstance(veff, lib.NPArrayWithTag):
metadata = veff.__dict__
veff = lib.tag_array(veff + vdiff, veff_rs=veff, **metadata)
else:
veff = lib.tag_array(veff + vdiff, veff_rs=veff)
return veff
[docs]
def energy_elec(self, dm=None, h1e=None, vhf=None):
if vhf is None:
# use the original veff to compute energy
vhf = super().get_veff(self.mol, dm)
return super().energy_elec(dm=dm, h1e=h1e, vhf=vhf)
else:
return super().energy_elec(dm=dm, h1e=h1e, vhf=vhf.veff_rs)
[docs]
def energy_ewald(self, dm=None, mm_ewald_pot=None, qm_ewald_pot=None):
# QM-QM and QM-MM pbc correction
if dm is None:
dm = self.make_rdm1()
if mm_ewald_pot is None:
if self.mm_ewald_pot is not None:
mm_ewald_pot = self.mm_ewald_pot
else:
mm_ewald_pot = self.get_mm_ewald_pot(self.mol, self.mm_mol)
if qm_ewald_pot is None:
qm_ewald_pot = self.get_qm_ewald_pot(self.mol, dm, self.qm_ewald_hess)
ewald_pot = mm_ewald_pot[0] + qm_ewald_pot[0] / 2
e = lib.einsum('i,i->', ewald_pot, self.get_qm_charges(dm))
ewald_pot = mm_ewald_pot[1] + qm_ewald_pot[1] / 2
e += lib.einsum('ix,ix->', ewald_pot, self.get_qm_dipoles(dm))
ewald_pot = mm_ewald_pot[2] + qm_ewald_pot[2] / 2
e += lib.einsum('ixy,ixy->', ewald_pot, self.get_qm_quadrupoles(dm))
# TODO add energy correction if sum(charges) !=0 ?
return e
[docs]
def energy_nuc(self):
if self.e_nuc is not None:
return self.e_nuc
else:
cput0 = (logger.process_clock(), logger.perf_counter())
# gas phase nuc energy
nuc = self.mol.energy_nuc() # qm_nuc - qm_nuc
# select mm atoms within rcut_hcore
mol = self.mol
Ls = self.mm_mol.get_lattice_Ls()
qm_center = np.mean(mol.atom_coords(), axis=0)
all_coords = lib.direct_sum('ix+Lx->Lix',
self.mm_mol.atom_coords(), Ls).reshape(-1,3)
all_charges = np.hstack([self.mm_mol.atom_charges()] * len(Ls))
all_expnts = np.hstack([np.sqrt(self.mm_mol.get_zetas())] * len(Ls))
dist2 = all_coords - qm_center
dist2 = lib.einsum('ix,ix->i', dist2, dist2)
mask = dist2 <= self.mm_mol.rcut_hcore**2
charges = all_charges[mask]
coords = all_coords[mask]
expnts = all_expnts[mask]
# qm_nuc - mm_pc
for j in range(self.mol.natm):
q2, r2 = self.mol.atom_charge(j), self.mol.atom_coord(j)
r = lib.norm(r2-coords, axis=1)
nuc += q2*(charges * erf(expnts*r) /r).sum()
logger.timer(self, 'energy_nuc', *cput0)
self.e_nuc = nuc
return nuc
[docs]
def energy_tot(self, dm=None, h1e=None, vhf=None, mm_ewald_pot=None, qm_ewald_pot=None):
nuc = self.energy_nuc()
ewald = self.energy_ewald(dm=dm, mm_ewald_pot=mm_ewald_pot, qm_ewald_pot=qm_ewald_pot)
e_tot = self.energy_elec(dm, h1e, vhf)[0] + nuc + ewald
self.scf_summary['nuc'] = nuc.real
self.scf_summary['ewald'] = ewald
return e_tot
[docs]
def nuc_grad_method(self):
scf_grad = super().nuc_grad_method()
return qmmm_grad_for_scf(scf_grad)
Gradients = nuc_grad_method
[docs]
def add_mm_charges_grad(scf_grad, atoms_or_coords, a, charges, radii=None,
rcut_ewald=None, rcut_hcore=None, unit=None):
'''Apply the MM charges in the QM gradients' method. It affects both the
electronic and nuclear parts of the QM fragment.
Args:
scf_grad : a HF or DFT gradient object (grad.HF or grad.RKS etc)
coords : 2D array, shape (N,3)
MM particle coordinates
charges : 1D array
MM particle charges
a : 2D array, shape (3,3)
Lattice vectors
Kwargs:
radii : 1D array
The Gaussian charge distribution radii of MM atoms.
unit : str
Bohr, AU, Ang (case insensitive). Default is the same to mol.unit
Returns:
Same gradeints method object as the input scf_grad method
'''
assert (isinstance(scf_grad, grad.rhf.Gradients))
mol = scf_grad.mol
if unit is None:
unit = mol.unit
mm_mol = mm_mole.create_mm_mol(atoms_or_coords, a, charges, radii=radii,
rcut_ewald=rcut_ewald, rcut_hcore=rcut_hcore, unit=unit)
mm_grad = qmmm_grad_for_scf(scf_grad)
mm_grad.base.mm_mol = mm_mol
return mm_grad
# Define method mm_charge_grad for backward compatibility
mm_charge_grad = add_mm_charges_grad
[docs]
def qmmm_grad_for_scf(scf_grad):
'''Add the potential of MM particles to SCF (HF and DFT) object and then
generate the corresponding QM/MM gradients method for the QM system.
'''
if getattr(scf_grad.base, 'with_x2c', None):
raise NotImplementedError('X2C with QM/MM charges')
# Avoid to initialize QMMMGrad twice
if isinstance(scf_grad, QMMMGrad):
return scf_grad
assert (isinstance(scf_grad.base, scf.hf.SCF) and
isinstance(scf_grad.base, QMMM))
scf_grad.de_ewald_mm = None
scf_grad.de_nuc_mm = None
return scf_grad.view(lib.make_class((QMMMGrad, scf_grad.__class__)))
[docs]
class QMMMGrad:
__name_mixin__ = 'QMMM'
_keys = {'de_ewald_mm', 'de_nuc_mm'}
to_gpu = NotImplemented
def __init__(self, scf_grad):
self.__dict__.update(scf_grad.__dict__)
[docs]
def dump_flags(self, verbose=None):
super().dump_flags(verbose)
logger.info(self, '** Add background charges for %s **', self.__class__.__name__)
_a = self.base.mm_mol.lattice_vectors()
logger.info(self, 'lattice vectors a1 [%.9f, %.9f, %.9f]', *_a[0])
logger.info(self, ' a2 [%.9f, %.9f, %.9f]', *_a[1])
logger.info(self, ' a3 [%.9f, %.9f, %.9f]', *_a[2])
if self.verbose >= logger.DEBUG2:
logger.debug2(self, 'Charge Location')
coords = self.base.mm_mol.atom_coords()
charges = self.base.mm_mol.atom_charges()
for i, z in enumerate(charges):
logger.debug2(self, '%.9g %s', z, coords[i])
return self
[docs]
def grad_ewald(self, dm=None, with_mm=False, mm_ewald_pot=None, qm_ewald_pot=None):
'''
pbc correction energy grad w.r.t. qm and mm atom positions
'''
cput0 = (logger.process_clock(), logger.perf_counter())
if dm is None: dm = self.base.make_rdm1()
mol = self.base.mol
cell = self.base.mm_mol
assert cell.dimension == 3
qm_charges = self.base.get_qm_charges(dm)
qm_dipoles = self.base.get_qm_dipoles(dm)
qm_quads = self.base.get_qm_quadrupoles(dm)
qm_coords = self.base.mol.atom_coords()
mm_charges = self.base.mm_mol.atom_charges()
mm_coords = self.base.mm_mol.atom_coords()
aoslices = mol.aoslice_by_atom()
# nuc grad due to qm multipole change due to ovlp change
qm_multipole_grad = np.zeros_like(qm_coords)
if mm_ewald_pot is None:
if self.base.mm_ewald_pot is not None:
mm_ewald_pot = self.base.mm_ewald_pot
else:
mm_ewald_pot = self.base.get_mm_ewald_pot(mol, cell)
if qm_ewald_pot is None:
qm_ewald_pot = self.base.get_qm_ewald_pot(mol, dm, self.base.qm_ewald_hess)
ewald_pot = \
mm_ewald_pot[0] + qm_ewald_pot[0], \
mm_ewald_pot[1] + qm_ewald_pot[1], \
mm_ewald_pot[2] + qm_ewald_pot[2]
dEds = np.zeros((mol.nao, mol.nao))
dEdsr = np.zeros((3, mol.nao, mol.nao))
dEdsrr = np.zeros((3, 3, mol.nao, mol.nao))
s1 = self.get_ovlp(mol) # = -mol.intor('int1e_ipovlp')
s1r = list()
s1rr = list()
bas_atom = mol._bas[:,gto.ATOM_OF]
for iatm in range(mol.natm):
v0 = ewald_pot[0][iatm]
v1 = ewald_pot[1][iatm]
v2 = ewald_pot[2][iatm]
p0, p1 = aoslices[iatm, 2:]
dEds[p0:p1] -= v0 * dm[p0:p1]
dEdsr[:,p0:p1] -= lib.einsum('x,uv->xuv', v1, dm[p0:p1])
dEdsrr[:,:,p0:p1] -= lib.einsum('xy,uv->xyuv', v2, dm[p0:p1])
b0, b1 = np.where(bas_atom == iatm)[0][[0,-1]]
shlslc = (b0, b1+1, 0, mol.nbas)
with mol.with_common_orig(qm_coords[iatm]):
# s1r[a,x,u,v] = \int phi_u (r_a-Ri_a) (-\nabla_x phi_v) dr
s1r.append(
np.asarray(-mol.intor('int1e_irp', shls_slice=shlslc).
reshape(3, 3, -1, mol.nao)))
# s1rr[a,b,x,u,v] =
# \int phi_u [3/2*(r_a-Ri_a)(r_b-Ri_b)-1/2*(r-Ri)^2 delta_ab] (-\nable_x phi_v) dr
s1rr_ = np.asarray(-mol.intor('int1e_irrp', shls_slice=shlslc).
reshape(3, 3, 3, -1, mol.nao))
s1rr_trace = lib.einsum('aaxuv->xuv', s1rr_)
s1rr_ *= 3 / 2
for k in range(3):
s1rr_[k,k] -= 0.5 * s1rr_trace
s1rr.append(s1rr_)
for jatm in range(mol.natm):
p0, p1 = aoslices[jatm, 2:]
# d E_qm_pc / d Ri with fixed ewald_pot
qm_multipole_grad[jatm] += \
lib.einsum('uv,xuv->x', dEds[p0:p1], s1[:,p0:p1]) \
- lib.einsum('uv,xuv->x', dEds[:,p0:p1], s1[:,:,p0:p1])
# d E_qm_dip / d Ri
qm_multipole_grad[jatm] -= \
lib.einsum('auv,axuv->x', dEdsr[:,p0:p1], s1r[jatm])
s1r_ = list()
for iatm in range(mol.natm):
s1r_.append(s1r[iatm][...,p0:p1])
s1r_ = np.concatenate(s1r_, axis=-2)
qm_multipole_grad[jatm] += lib.einsum('auv,axuv->x', dEdsr[...,p0:p1], s1r_)
# d E_qm_quad / d Ri
qm_multipole_grad[jatm] -= \
lib.einsum('abuv,abxuv->x', dEdsrr[:,:,p0:p1], s1rr[jatm])
s1rr_ = list()
for iatm in range(mol.natm):
s1rr_.append(s1rr[iatm][...,p0:p1])
s1rr_ = np.concatenate(s1rr_, axis=-2)
qm_multipole_grad[jatm] += lib.einsum('abuv,abxuv->x', dEdsrr[...,p0:p1], s1rr_)
cput1 = logger.timer(self, 'grad_ewald pulay', *cput0)
s1 = s1r = s1rr = dEds = dEdsr = dEdsrr = None
ew_eta, ew_cut = cell.get_ewald_params()
# ---------------------------------------------- #
# -------- Ewald real-space gradient ----------- #
# ---------------------------------------------- #
Lall = cell.get_lattice_Ls()
from pyscf import pbc
rmax_qm = max(np.linalg.norm(qm_coords - np.mean(qm_coords, axis=0), axis=-1))
qm_ewovrl_grad = np.zeros_like(qm_coords)
def grad_Tij(R, r):
Rij = R
Tij = 1 / r
# Tija = \nabla_a Tij
# Tijab = \nabla_a Tijb
# Tijabc = \nabla_a Tijbc
Tija = -lib.einsum('ijx,ij->ijx', Rij, Tij**3)
Tijab = 3 * lib.einsum('ija,ijb->ijab', Rij, Rij)
Tijab = lib.einsum('ijab,ij->ijab', Tijab, Tij**5)
Tijab -= lib.einsum('ij,ab->ijab', Tij**3, np.eye(3))
Tijabc = -15 * lib.einsum('ija,ijb,ijc->ijabc', Rij, Rij, Rij)
Tijabc = lib.einsum('ijabc,ij->ijabc', Tijabc, Tij**7)
Tijabc += 3 * lib.einsum('ija,bc,ij->ijabc', Rij, np.eye(3), Tij**5)
Tijabc += 3 * lib.einsum('ijb,ac,ij->ijabc', Rij, np.eye(3), Tij**5)
Tijabc += 3 * lib.einsum('ijc,ab,ij->ijabc', Rij, np.eye(3), Tij**5)
return Tija, Tijab, Tijabc
def grad_kTij(R, r, eta):
if isinstance(eta, float):
Tij = erfc(eta * r) / r
ekR = np.exp(-eta**2 * r**2)
Tij = erfc(eta * r) / r
invr3 = (Tij + 2*eta/np.sqrt(np.pi) * ekR) / r**2
Tija = -lib.einsum('ijx,ij->ijx', R, invr3)
Tijab = 3 * lib.einsum('ija,ijb,ij->ijab', R, R, 1/r**2)
Tijab -= lib.einsum('ij,ab->ijab', np.ones_like(r), np.eye(3))
invr5 = invr3 + 4/3*eta**3/np.sqrt(np.pi) * ekR # NOTE this is invr5 * r**2
Tijab = lib.einsum('ijab,ij->ijab', Tijab, invr5)
Tijab += 4/3*eta**3/np.sqrt(np.pi)*lib.einsum('ij,ab->ijab', ekR, np.eye(3))
invr7 = invr5 / r**2 + 8/15 / np.sqrt(np.pi) * eta**5 * ekR # NOTE this is invr7 * r**2
Tijabc = -15 * lib.einsum('ija,ijb,ijc->ijabc', R, R, R)
Tijabc = lib.einsum('ijabc,ij->ijabc', Tijabc, 1/r**2)
Tijabc += 3 * lib.einsum('ija,bc->ijabc', R, np.eye(3))
Tijabc += 3 * lib.einsum('ijb,ac->ijabc', R, np.eye(3))
Tijabc += 3 * lib.einsum('ijc,ab->ijabc', R, np.eye(3))
Tijabc = lib.einsum('ijabc,ij->ijabc', Tijabc, invr7)
Tijabc -= 8/5 / np.sqrt(np.pi) * eta**5 * lib.einsum('ij,ija,bc->ijabc', ekR, R, np.eye(3))
Tijabc -= 8/5 / np.sqrt(np.pi) * eta**5 * lib.einsum('ij,ijb,ac->ijabc', ekR, R, np.eye(3))
Tijabc -= 8/5 / np.sqrt(np.pi) * eta**5 * lib.einsum('ij,ijc,ab->ijabc', ekR, R, np.eye(3))
return Tija, Tijab, Tijabc
else:
Tij = erfc(eta * r) / r
ekR = np.exp(-eta**2 * r**2)
Tij = erfc(eta * r) / r
invr3 = (Tij + 2*eta/np.sqrt(np.pi) * ekR) / r**2
Tija = -lib.einsum('ijx,ij->ijx', R, invr3)
Tijab = 3 * lib.einsum('ija,ijb,ij->ijab', R, R, 1/r**2)
Tijab -= lib.einsum('ij,ab->ijab', np.ones_like(r), np.eye(3))
invr5 = invr3 + 4/3*eta**3/np.sqrt(np.pi) * ekR # NOTE this is invr5 * r**2
Tijab = lib.einsum('ijab,ij->ijab', Tijab, invr5)
Tijab += 4/3/np.sqrt(np.pi)*lib.einsum('j,ij,ab->ijab', eta**3, ekR, np.eye(3))
invr7 = invr5 / r**2 + 8/15 / np.sqrt(np.pi) * eta**5 * ekR # NOTE this is invr7 * r**2
Tijabc = -15 * lib.einsum('ija,ijb,ijc->ijabc', R, R, R)
Tijabc = lib.einsum('ijabc,ij->ijabc', Tijabc, 1/r**2)
Tijabc += 3 * lib.einsum('ija,bc->ijabc', R, np.eye(3))
Tijabc += 3 * lib.einsum('ijb,ac->ijabc', R, np.eye(3))
Tijabc += 3 * lib.einsum('ijc,ab->ijabc', R, np.eye(3))
Tijabc = lib.einsum('ijabc,ij->ijabc', Tijabc, invr7)
Tijabc -= 8/5 / np.sqrt(np.pi) * lib.einsum('j,ij,ija,bc->ijabc', eta**5, ekR, R, np.eye(3))
Tijabc -= 8/5 / np.sqrt(np.pi) * lib.einsum('j,ij,ijb,ac->ijabc', eta**5, ekR, R, np.eye(3))
Tijabc -= 8/5 / np.sqrt(np.pi) * lib.einsum('j,ij,ijc,ab->ijabc', eta**5, ekR, R, np.eye(3))
return Tija, Tijab, Tijabc
#------ qm - mm classical ewald energy gradient ------#
all_mm_coords = lib.direct_sum('jx-Lx->Ljx', mm_coords, Lall).reshape(-1,3)
all_mm_charges = np.hstack([mm_charges] * len(Lall))
dist2 = lib.direct_sum('jx-x->jx', all_mm_coords, np.mean(qm_coords, axis=0))
dist2 = lib.einsum('jx,jx->j', dist2, dist2)
if with_mm:
mm_ewovrl_grad = np.zeros_like(all_mm_coords)
mem_avail = self.max_memory - lib.current_memory()[0] # in MB
blksize = int(mem_avail/81/(8e-6*len(all_mm_coords)))
if blksize == 0:
raise RuntimeError(f"Not enough memory, mem_avail = {mem_avail}, blkszie = {blksize}")
for i0, i1 in lib.prange(0, mol.natm, blksize):
R = qm_coords[i0:i1,None,:] - all_mm_coords[None,:,:]
r = np.linalg.norm(R, axis=-1)
r[r<1e-16] = np.inf
# substract the real-space Coulomb within rcut_hcore
mask = dist2 <= cell.rcut_hcore**2
Tija, Tijab, Tijabc = grad_Tij(R[:,mask], r[:,mask])
qm_ewovrl_grad[i0:i1] -= lib.einsum('i,ijx,j->ix', qm_charges[i0:i1], Tija, all_mm_charges[mask])
qm_ewovrl_grad[i0:i1] -= lib.einsum('ia,ijxa,j->ix', qm_dipoles[i0:i1], Tijab, all_mm_charges[mask])
qm_ewovrl_grad[i0:i1] -= lib.einsum('iab,ijxab,j->ix', qm_quads[i0:i1], Tijabc, all_mm_charges[mask]) / 3
if with_mm:
mm_ewovrl_grad[mask] += lib.einsum('i,ijx,j->jx', qm_charges[i0:i1], Tija, all_mm_charges[mask])
mm_ewovrl_grad[mask] += lib.einsum('ia,ijxa,j->jx', qm_dipoles[i0:i1], Tijab, all_mm_charges[mask])
mm_ewovrl_grad[mask] += lib.einsum('iab,ijxab,j->jx', qm_quads[i0:i1], Tijabc, all_mm_charges[mask]) / 3
# difference between MM gaussain charges and MM point charges
mask = dist2 > cell.rcut_hcore**2
min_expnt = np.min(cell.get_zetas())
max_ewrcut = pbc.gto.cell._estimate_rcut(min_expnt, 0, 1., cell.precision)
cut2 = (max_ewrcut + rmax_qm)**2
mask = mask & (dist2 <= cut2)
expnts = np.hstack([np.sqrt(cell.get_zetas())] * len(Lall))[mask]
Tija, Tijab, Tijabc = grad_kTij(R[:,mask], r[:,mask], expnts)
qm_ewovrl_grad[i0:i1] -= lib.einsum('i,ijx,j->ix', qm_charges[i0:i1], Tija, all_mm_charges[mask])
qm_ewovrl_grad[i0:i1] -= lib.einsum('ia,ijxa,j->ix', qm_dipoles[i0:i1], Tijab, all_mm_charges[mask])
qm_ewovrl_grad[i0:i1] -= lib.einsum('iab,ijxab,j->ix', qm_quads[i0:i1], Tijabc, all_mm_charges[mask]) / 3
if with_mm:
mm_ewovrl_grad[mask] += lib.einsum('i,ijx,j->jx', qm_charges[i0:i1], Tija, all_mm_charges[mask])
mm_ewovrl_grad[mask] += lib.einsum('ia,ijxa,j->jx', qm_dipoles[i0:i1], Tijab, all_mm_charges[mask])
mm_ewovrl_grad[mask] += lib.einsum('iab,ijxab,j->jx', qm_quads[i0:i1], Tijabc, all_mm_charges[mask]) / 3
# ewald real-space sum
cut2 = (ew_cut + rmax_qm)**2
mask = dist2 <= cut2
Tija, Tijab, Tijabc = grad_kTij(R[:,mask], r[:,mask], ew_eta)
qm_ewovrl_grad[i0:i1] += lib.einsum('i,ijx,j->ix', qm_charges[i0:i1], Tija, all_mm_charges[mask])
qm_ewovrl_grad[i0:i1] += lib.einsum('ia,ijxa,j->ix', qm_dipoles[i0:i1], Tijab, all_mm_charges[mask])
qm_ewovrl_grad[i0:i1] += lib.einsum('iab,ijxab,j->ix', qm_quads[i0:i1], Tijabc, all_mm_charges[mask]) / 3
if with_mm:
mm_ewovrl_grad[mask] -= lib.einsum('i,ijx,j->jx', qm_charges[i0:i1], Tija, all_mm_charges[mask])
mm_ewovrl_grad[mask] -= lib.einsum('ia,ijxa,j->jx', qm_dipoles[i0:i1], Tijab, all_mm_charges[mask])
mm_ewovrl_grad[mask] -= lib.einsum('iab,ijxab,j->jx', qm_quads[i0:i1], Tijabc, all_mm_charges[mask]) / 3
if with_mm:
mm_ewovrl_grad = mm_ewovrl_grad.reshape(len(Lall), -1, 3)
mm_ewovrl_grad = np.sum(mm_ewovrl_grad, axis=0)
all_mm_coords = all_mm_charges = None
#------ qm - qm classical ewald energy gradient ------#
R = lib.direct_sum('ix-jx->ijx', qm_coords, qm_coords)
r = np.sqrt(lib.einsum('ijx,ijx->ij', R, R))
r[r<1e-16] = 1e100
# substract the real-space Coulomb within rcut_hcore
Tija, Tijab, Tijabc = grad_Tij(R, r)
qm_ewovrl_grad -= lib.einsum('i,ijx,j->ix', qm_charges, Tija, qm_charges)
qm_ewovrl_grad += lib.einsum('i,ijxa,ja->ix', qm_charges, Tijab, qm_dipoles)
qm_ewovrl_grad -= lib.einsum('i,ijxa,ja->jx', qm_charges, Tijab, qm_dipoles) #
qm_ewovrl_grad += lib.einsum('ia,ijxab,jb->ix', qm_dipoles, Tijabc, qm_dipoles)
qm_ewovrl_grad -= lib.einsum('i,ijxab,jab->ix', qm_charges, Tijabc, qm_quads) / 3
qm_ewovrl_grad += lib.einsum('i,ijxab,jab->jx', qm_charges, Tijabc, qm_quads) / 3 #
# ewald real-space sum
# NOTE here I assume ewald real-space sum is over all qm images
# consistent with mm_mole.get_ewald_pot
R = (R[:,:,None,:] - Lall[None,None]).reshape(len(qm_coords), len(Lall)*len(qm_coords), 3)
r = np.sqrt(lib.einsum('ijx,ijx->ij', R, R))
r[r<1e-16] = 1e100
Tija, Tijab, Tijabc = grad_kTij(R, r, ew_eta)
Tija = Tija.reshape(len(qm_coords), len(qm_coords), len(Lall), 3)
Tijab = Tijab.reshape(len(qm_coords), len(qm_coords), len(Lall), 3, 3)
Tijabc = Tijabc.reshape(len(qm_coords), len(qm_coords), len(Lall), 3, 3, 3)
qm_ewovrl_grad += lib.einsum('i,ijLx,j->ix', qm_charges, Tija, qm_charges)
qm_ewovrl_grad -= lib.einsum('i,ijLxa,ja->ix', qm_charges, Tijab, qm_dipoles)
qm_ewovrl_grad += lib.einsum('i,ijLxa,ja->jx', qm_charges, Tijab, qm_dipoles) #
qm_ewovrl_grad -= lib.einsum('ia,ijLxab,jb->ix', qm_dipoles, Tijabc, qm_dipoles)
qm_ewovrl_grad += lib.einsum('i,ijLxab,jab->ix', qm_charges, Tijabc, qm_quads) / 3
qm_ewovrl_grad -= lib.einsum('i,ijLxab,jab->jx', qm_charges, Tijabc, qm_quads) / 3 #
cput2 = logger.timer(self, 'grad_ewald real-space', *cput1)
# ---------------------------------------------- #
# ---------- Ewald k-space gradient ------------ #
# ---------------------------------------------- #
mesh = cell.mesh
Gv, Gvbase, weights = cell.get_Gv_weights(mesh)
absG2 = lib.einsum('gi,gi->g', Gv, Gv)
absG2[absG2==0] = 1e200
coulG = 4*np.pi / absG2
coulG *= weights
Gpref = np.exp(-absG2/(4*ew_eta**2)) * coulG
GvRmm = lib.einsum('gx,ix->ig', Gv, mm_coords)
cosGvRmm = np.cos(GvRmm)
sinGvRmm = np.sin(GvRmm)
zcosGvRmm = lib.einsum("i,ig->g", mm_charges, cosGvRmm)
zsinGvRmm = lib.einsum("i,ig->g", mm_charges, sinGvRmm)
GvRqm = lib.einsum('gx,ix->ig', Gv, qm_coords)
cosGvRqm = np.cos(GvRqm)
sinGvRqm = np.sin(GvRqm)
zcosGvRqm = lib.einsum("i,ig->g", qm_charges, cosGvRqm)
zsinGvRqm = lib.einsum("i,ig->g", qm_charges, sinGvRqm)
DGcosGvRqm = lib.einsum("ia,ga,ig->g", qm_dipoles, Gv, cosGvRqm)
DGsinGvRqm = lib.einsum("ia,ga,ig->g", qm_dipoles, Gv, sinGvRqm)
TGGcosGvRqm = lib.einsum("iab,ga,gb,ig->g", qm_quads, Gv, Gv, cosGvRqm)
TGGsinGvRqm = lib.einsum("iab,ga,gb,ig->g", qm_quads, Gv, Gv, sinGvRqm)
qm_ewg_grad = np.zeros_like(qm_coords)
if with_mm:
mm_ewg_grad = np.zeros_like(mm_coords)
# qm pc - mm pc
p = ['einsum_path', (3, 4), (1, 3), (1, 2), (0, 1)]
qm_ewg_grad -= lib.einsum('i,gx,ig,g,g->ix', qm_charges, Gv, sinGvRqm, zcosGvRmm, Gpref, optimize=p)
qm_ewg_grad += lib.einsum('i,gx,ig,g,g->ix', qm_charges, Gv, cosGvRqm, zsinGvRmm, Gpref, optimize=p)
if with_mm:
p = ['einsum_path', (0, 2), (1, 2), (0, 2), (0, 1)]
mm_ewg_grad -= lib.einsum('i,gx,ig,g,g->ix', mm_charges, Gv, sinGvRmm, zcosGvRqm, Gpref, optimize=p)
mm_ewg_grad += lib.einsum('i,gx,ig,g,g->ix', mm_charges, Gv, cosGvRmm, zsinGvRqm, Gpref, optimize=p)
# qm dip - mm pc
p = ['einsum_path', (4, 5), (1, 4), (0, 1), (0, 2), (0, 1)]
qm_ewg_grad -= lib.einsum('ia,gx,ga,ig,g,g->ix', qm_dipoles, Gv, Gv, sinGvRqm, zsinGvRmm, Gpref, optimize=p)
qm_ewg_grad -= lib.einsum('ia,gx,ga,ig,g,g->ix', qm_dipoles, Gv, Gv, cosGvRqm, zcosGvRmm, Gpref, optimize=p)
if with_mm:
p = ['einsum_path', (1, 3), (0, 2), (0, 2), (0, 1)]
mm_ewg_grad += lib.einsum('g,j,gx,jg,g->jx', DGcosGvRqm, mm_charges, Gv, cosGvRmm, Gpref, optimize=p)
mm_ewg_grad += lib.einsum('g,j,gx,jg,g->jx', DGsinGvRqm, mm_charges, Gv, sinGvRmm, Gpref, optimize=p)
# qm quad - mm pc
p = ['einsum_path', (5, 6), (0, 5), (0, 2), (2, 3), (1, 2), (0, 1)]
qm_ewg_grad += lib.einsum('ga,gb,iab,gx,ig,g,g->ix', Gv, Gv, qm_quads,
Gv, sinGvRqm, zcosGvRmm, Gpref, optimize=p) / 3
qm_ewg_grad -= lib.einsum('ga,gb,iab,gx,ig,g,g->ix', Gv, Gv, qm_quads,
Gv, cosGvRqm, zsinGvRmm, Gpref, optimize=p) / 3
if with_mm:
p = ['einsum_path', (1, 3), (0, 2), (0, 2), (0, 1)]
mm_ewg_grad += lib.einsum('g,j,gx,jg,g->jx', TGGcosGvRqm, mm_charges, Gv, sinGvRmm, Gpref, optimize=p) / 3
mm_ewg_grad -= lib.einsum('g,j,gx,jg,g->jx', TGGsinGvRqm, mm_charges, Gv, cosGvRmm, Gpref, optimize=p) / 3
# qm pc - qm pc
p = ['einsum_path', (3, 4), (1, 3), (1, 2), (0, 1)]
qm_ewg_grad -= lib.einsum('i,gx,ig,g,g->ix', qm_charges, Gv, sinGvRqm, zcosGvRqm, Gpref, optimize=p)
qm_ewg_grad += lib.einsum('i,gx,ig,g,g->ix', qm_charges, Gv, cosGvRqm, zsinGvRqm, Gpref, optimize=p)
# qm pc - qm dip
qm_ewg_grad += lib.einsum('i,gx,ig,g,g->ix', qm_charges, Gv, cosGvRqm, DGcosGvRqm, Gpref, optimize=p)
qm_ewg_grad += lib.einsum('i,gx,ig,g,g->ix', qm_charges, Gv, sinGvRqm, DGsinGvRqm, Gpref, optimize=p)
p = ['einsum_path', (3, 5), (1, 4), (1, 3), (1, 2), (0, 1)]
qm_ewg_grad -= lib.einsum('ja,ga,gx,g,jg,g->jx', qm_dipoles, Gv, Gv, zsinGvRqm, sinGvRqm, Gpref, optimize=p)
qm_ewg_grad -= lib.einsum('ja,ga,gx,g,jg,g->jx', qm_dipoles, Gv, Gv, zcosGvRqm, cosGvRqm, Gpref, optimize=p)
# qm dip - qm dip
p = ['einsum_path', (4, 5), (1, 4), (1, 3), (1, 2), (0, 1)]
qm_ewg_grad -= lib.einsum('ia,ga,gx,ig,g,g->ix', qm_dipoles, Gv, Gv, sinGvRqm, DGcosGvRqm, Gpref, optimize=p)
qm_ewg_grad += lib.einsum('ia,ga,gx,ig,g,g->ix', qm_dipoles, Gv, Gv, cosGvRqm, DGsinGvRqm, Gpref, optimize=p)
# qm pc - qm quad
p = ['einsum_path', (3, 4), (1, 3), (1, 2), (0, 1)]
qm_ewg_grad += lib.einsum('i,gx,ig,g,g->ix', qm_charges, Gv, sinGvRqm, TGGcosGvRqm, Gpref, optimize=p) / 3
qm_ewg_grad -= lib.einsum('i,gx,ig,g,g->ix', qm_charges, Gv, cosGvRqm, TGGsinGvRqm, Gpref, optimize=p) / 3
p = ['einsum_path', (4, 6), (1, 5), (1, 2), (2, 3), (1, 2), (0, 1)]
qm_ewg_grad += lib.einsum('jab,ga,gb,gx,g,jg,g->jx', qm_quads, Gv, Gv,
Gv, zcosGvRqm, sinGvRqm, Gpref, optimize=p) / 3
qm_ewg_grad -= lib.einsum('jab,ga,gb,gx,g,jg,g->jx', qm_quads, Gv, Gv,
Gv, zsinGvRqm, cosGvRqm, Gpref, optimize=p) / 3
logger.timer(self, 'grad_ewald k-space', *cput2)
logger.timer(self, 'grad_ewald', *cput0)
if not with_mm:
return qm_multipole_grad + qm_ewovrl_grad + qm_ewg_grad
else:
return qm_multipole_grad + qm_ewovrl_grad + qm_ewg_grad, \
mm_ewovrl_grad + mm_ewg_grad
[docs]
def get_hcore(self, mol=None):
cput0 = (logger.process_clock(), logger.perf_counter())
mm_mol = self.base.mm_mol
if mol is None:
mol = self.mol
rcut_hcore = mm_mol.rcut_hcore
coords = mm_mol.atom_coords()
charges = mm_mol.atom_charges()
Ls = mm_mol.get_lattice_Ls()
qm_center = np.mean(mol.atom_coords(), axis=0)
all_coords = lib.direct_sum('ix+Lx->Lix',
mm_mol.atom_coords(), Ls).reshape(-1,3)
all_charges = np.hstack([mm_mol.atom_charges()] * len(Ls))
dist2 = all_coords - qm_center
dist2 = lib.einsum('ix,ix->i', dist2, dist2)
# charges within rcut_hcore exactly go into hcore
mask = dist2 <= rcut_hcore**2
charges = all_charges[mask]
coords = all_coords[mask]
nao = mol.nao
max_memory = self.max_memory - lib.current_memory()[0]
blksize = int(min(max_memory*1e6/8/nao**2/3, 200))
blksize = max(blksize, 1)
g_qm = super().get_hcore(mol)
if mm_mol.charge_model == 'gaussian':
expnts = np.hstack([mm_mol.get_zetas()] * len(Ls))[mask]
if mol.cart:
intor = 'int3c2e_ip1_cart'
else:
intor = 'int3c2e_ip1_sph'
cintopt = gto.moleintor.make_cintopt(mol._atm, mol._bas,
mol._env, intor)
v = 0
for i0, i1 in lib.prange(0, charges.size, blksize):
fakemol = gto.fakemol_for_charges(coords[i0:i1], expnts[i0:i1])
j3c = df.incore.aux_e2(mol, fakemol, intor, aosym='s1',
comp=3, cintopt=cintopt)
v += lib.einsum('ipqk,k->ipq', j3c, charges[i0:i1])
g_qm += v
else:
for i0, i1 in lib.prange(0, charges.size, blksize):
j3c = mol.intor('int1e_grids_ip', grids=coords[i0:i1])
g_qm += lib.einsum('ikpq,k->ipq', j3c, charges[i0:i1])
logger.timer(self, 'get_hcore', *cput0)
return g_qm
[docs]
def grad_hcore_mm(self, dm, mol=None):
r'''Nuclear gradients of the electronic energy
'''
cput0 = (logger.process_clock(), logger.perf_counter())
mm_mol = self.base.mm_mol
if mol is None:
mol = self.mol
rcut_hcore = mm_mol.rcut_hcore
coords = mm_mol.atom_coords()
charges = mm_mol.atom_charges()
expnts = mm_mol.get_zetas()
Ls = mm_mol.get_lattice_Ls()
qm_center = np.mean(mol.atom_coords(), axis=0)
all_coords = lib.direct_sum('ix+Lx->Lix',
mm_mol.atom_coords(), Ls).reshape(-1,3)
all_charges = np.hstack([mm_mol.atom_charges()] * len(Ls))
all_expnts = np.hstack([expnts] * len(Ls))
dist2 = all_coords - qm_center
dist2 = lib.einsum('ix,ix->i', dist2, dist2)
# charges within rcut_hcore exactly go into hcore
mask = dist2 <= rcut_hcore**2
charges = all_charges[mask]
coords = all_coords[mask]
expnts = all_expnts[mask]
intor = 'int3c2e_ip2'
nao = mol.nao
max_memory = self.max_memory - lib.current_memory()[0]
blksize = int(min(max_memory*1e6/8/nao**2/3, 200))
blksize = max(blksize, 1)
cintopt = gto.moleintor.make_cintopt(mol._atm, mol._bas,
mol._env, intor)
g = np.zeros_like(all_coords)
g_ = np.zeros_like(coords)
for i0, i1 in lib.prange(0, charges.size, blksize):
fakemol = gto.fakemol_for_charges(coords[i0:i1], expnts[i0:i1])
j3c = df.incore.aux_e2(mol, fakemol, intor, aosym='s1',
comp=3, cintopt=cintopt)
g_[i0:i1] = lib.einsum('ipqk,qp->ik', j3c * charges[i0:i1], dm).T
g[mask] = g_
g = g.reshape(len(Ls), -1, 3)
g = np.sum(g, axis=0)
logger.timer(self, 'grad_hcore_mm', *cput0)
return g
contract_hcore_mm = grad_hcore_mm
[docs]
def grad_nuc(self, mol=None, atmlst=None):
cput0 = (logger.process_clock(), logger.perf_counter())
assert atmlst is None # atmlst needs to be full for computing g_mm
if mol is None: mol = self.mol
mm_mol = self.base.mm_mol
coords = mm_mol.atom_coords()
charges = mm_mol.atom_charges()
Ls = mm_mol.get_lattice_Ls()
qm_center = np.mean(mol.atom_coords(), axis=0)
all_coords = lib.direct_sum('ix+Lx->Lix',
mm_mol.atom_coords(), Ls).reshape(-1,3)
all_charges = np.hstack([mm_mol.atom_charges()] * len(Ls))
all_expnts = np.hstack([np.sqrt(mm_mol.get_zetas())] * len(Ls))
dist2 = all_coords - qm_center
dist2 = lib.einsum('ix,ix->i', dist2, dist2)
mask = dist2 <= mm_mol.rcut_hcore**2
charges = all_charges[mask]
coords = all_coords[mask]
expnts = all_expnts[mask]
g_qm = super().grad_nuc(mol, atmlst)
g_mm = np.zeros_like(all_coords)
for i in range(mol.natm):
q1 = mol.atom_charge(i)
r1 = mol.atom_coord(i)
r = lib.norm(r1-coords, axis=1)
g_mm_ = q1 * lib.einsum('i,ix,i->ix', charges, r1-coords, erf(expnts*r)/r**3)
g_mm_ -= q1 * lib.einsum('i,ix,i->ix', charges * expnts * 2 / np.sqrt(np.pi),
r1-coords, np.exp(-expnts**2 * r**2)/r**2)
g_mm[mask] += g_mm_
g_qm[i] -= np.sum(g_mm_, axis=0)
g_mm = g_mm.reshape(len(Ls), -1, 3)
g_mm = np.sum(g_mm, axis=0)
self.de_nuc_mm = g_mm
logger.timer(self, 'grad_nuc', *cput0)
return g_qm
[docs]
def grad_nuc_mm(self, mol=None):
if self.de_nuc_mm is not None:
return self.de_nuc_mm
cput0 = (logger.process_clock(), logger.perf_counter())
if mol is None:
mol = self.mol
mm_mol = self.base.mm_mol
coords = mm_mol.atom_coords()
charges = mm_mol.atom_charges()
Ls = mm_mol.get_lattice_Ls()
qm_center = np.mean(mol.atom_coords(), axis=0)
all_coords = lib.direct_sum('ix+Lx->Lix',
mm_mol.atom_coords(), Ls).reshape(-1,3)
all_charges = np.hstack([mm_mol.atom_charges()] * len(Ls))
all_expnts = np.hstack([np.sqrt(mm_mol.get_zetas())] * len(Ls))
dist2 = all_coords - qm_center
dist2 = lib.einsum('ix,ix->i', dist2, dist2)
mask = dist2 <= mm_mol.rcut_hcore**2
charges = all_charges[mask]
coords = all_coords[mask]
expnts = all_expnts[mask]
g_mm = np.zeros_like(all_coords)
for i in range(mol.natm):
q1 = mol.atom_charge(i)
r1 = mol.atom_coord(i)
r = lib.norm(r1-coords, axis=1)
g_mm[mask] += q1 * lib.einsum('i,ix,i->ix', charges, r1-coords, erf(expnts*r)/r**3)
g_mm[mask] -= q1 * lib.einsum('i,ix,i->ix', charges * expnts * 2 / np.sqrt(np.pi),
r1-coords, np.exp(-expnts**2 * r**2)/r**2)
g_mm = g_mm.reshape(len(Ls), -1, 3)
g_mm = np.sum(g_mm, axis=0)
logger.timer(self, 'grad_nuc_mm', *cput0)
return g_mm
def _finalize(self):
g_ewald_qm, self.de_ewald_mm = self.grad_ewald(with_mm=True)
self.de += g_ewald_qm
super()._finalize()
