diff --git a/pyscf/qmmm/itrf.py b/pyscf/qmmm/itrf.py index c3650434b7..676ef8bff8 100644 --- a/pyscf/qmmm/itrf.py +++ b/pyscf/qmmm/itrf.py @@ -31,6 +31,8 @@ from pyscf.lib import logger from pyscf.qmmm import mm_mole +from scipy.special import erf + def add_mm_charges(scf_method, atoms_or_coords, charges, radii=None, unit=None): '''Embedding the one-electron (non-relativistic) potential generated by MM @@ -179,14 +181,18 @@ def get_hcore(self, mol=None): def energy_nuc(self): # interactions between QM nuclei and MM particles - assert self.mm_mol.charge_model != 'gaussian' # TODO: support Gaussian charge, same as the one in PCM nuc = self.mol.energy_nuc() coords = self.mm_mol.atom_coords() charges = self.mm_mol.atom_charges() + if self.mm_mol.charge_model == 'gaussian': + expnts = numpy.sqrt(self.mm_mol.get_zetas()) 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/r).sum() + if self.mm_mol.charge_model != 'gaussian': + nuc += q2*(charges/r).sum() + else: + nuc += q2*(charges*erf(expnts*r)/r).sum() return nuc def to_gpu(self): @@ -371,9 +377,11 @@ def grad_hcore_mm(self, dm, mol=None): def grad_nuc(self, mol=None, atmlst=None): if mol is None: mol = self.mol - assert self.base.mm_mol.charge_model != 'gaussian' # TODO: support Gaussian charge, same as the one in PCM coords = self.base.mm_mol.atom_coords() charges = self.base.mm_mol.atom_charges() + if self.base.mm_mol.charge_model == 'gaussian': + expnts = self.base.mm_mol.get_zetas() + radii = 1 / numpy.sqrt(expnts) g_qm = super().grad_nuc(mol, atmlst) # nuclei lattice interaction @@ -382,7 +390,12 @@ def grad_nuc(self, mol=None, atmlst=None): q1 = mol.atom_charge(i) r1 = mol.atom_coord(i) r = lib.norm(r1-coords, axis=1) - g_mm[i] = -q1 * numpy.einsum('i,ix,i->x', charges, r1-coords, 1/r**3) + if self.base.mm_mol.charge_model != 'gaussian': + coulkern = 1/r**3 + else: + coulkern = erf(r/radii)/r - 2/(numpy.sqrt(numpy.pi)*radii) * numpy.exp(-expnts*r**2) + coulkern = coulkern / r**2 + g_mm[i] = -q1 * numpy.einsum('i,ix,i->x', charges, r1-coords, coulkern) if atmlst is not None: g_mm = g_mm[atmlst] return g_qm + g_mm @@ -397,12 +410,20 @@ def grad_nuc_mm(self, mol=None): mm_mol = self.base.mm_mol coords = mm_mol.atom_coords() charges = mm_mol.atom_charges() + if mm_mol.charge_model == 'gaussian': + expnts = mm_mol.get_zetas() + radii = 1 / numpy.sqrt(expnts) g_mm = numpy.zeros_like(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 * numpy.einsum('i,ix,i->ix', charges, r1-coords, 1/r**3) + if mm_mol.charge_model != 'gaussian': + coulkern = 1/r**3 + else: + coulkern = erf(r/radii)/r - 2/(numpy.sqrt(numpy.pi)*radii) * numpy.exp(-expnts*r**2) + coulkern = coulkern / r**2 + g_mm += q1 * numpy.einsum('i,ix,i->ix', charges, r1-coords, coulkern) return g_mm def to_gpu(self): diff --git a/pyscf/qmmm/pbc/__init__.py b/pyscf/qmmm/pbc/__init__.py new file mode 100644 index 0000000000..e91c3aabca --- /dev/null +++ b/pyscf/qmmm/pbc/__init__.py @@ -0,0 +1 @@ +from pyscf.qmmm.pbc import mm_mole, itrf \ No newline at end of file diff --git a/pyscf/qmmm/pbc/itrf.py b/pyscf/qmmm/pbc/itrf.py new file mode 100644 index 0000000000..5d38754c4b --- /dev/null +++ b/pyscf/qmmm/pbc/itrf.py @@ -0,0 +1,1046 @@ +#!/usr/bin/env python +# Copyright 2014-2019 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 +# + +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 + +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) + +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__)) + +class QMMM: + __name_mixin__ = 'QMMM' + +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 + + 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 + + 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()) + + 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 + + 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 + + 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) + + 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 + + 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) + + 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 + + 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) + + 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 + + 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 + + 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) + + 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 + + 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 + + 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 + + def nuc_grad_method(self): + scf_grad = super().nuc_grad_method() + return qmmm_grad_for_scf(scf_grad) + + Gradients = nuc_grad_method + + +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 + +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__))) + +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__) + + 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 + + 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 clasiical 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 clasiical 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 + + 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 + + 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 + + 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 + + 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() diff --git a/pyscf/qmmm/pbc/mm_mole.py b/pyscf/qmmm/pbc/mm_mole.py new file mode 100644 index 0000000000..feb2f35286 --- /dev/null +++ b/pyscf/qmmm/pbc/mm_mole.py @@ -0,0 +1,399 @@ +#!/usr/bin/env python +# Copyright 2014-2019 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 +# + +import numpy as np +from pyscf import gto +from pyscf import pbc +from pyscf import qmmm +from pyscf.gto.mole import is_au +from pyscf.data.elements import charge +from pyscf.lib import param, logger +from scipy.special import erf, erfc, lambertw +from pyscf import lib + +class Cell(qmmm.mm_mole.Mole, pbc.gto.Cell): + '''Cell class for MM particles. + + Args: + atoms : geometry of MM particles (unit Bohr). + + | [[atom1, (x, y, z)], + | [atom2, (x, y, z)], + | ... + | [atomN, (x, y, z)]] + + Kwargs: + charges : 1D array + fractional charges of MM particles + zeta : 1D array + Gaussian charge distribution parameter. + rho(r) = charge * Norm * exp(-zeta * r^2) + + ''' + def __init__(self, atoms, a, + rcut_ewald=None, rcut_hcore=None, + charges=None, zeta=None): + pbc.gto.Cell.__init__(self) + self.atom = self._atom = atoms + self.unit = 'Bohr' + self.charge_model = 'point' + assert np.linalg.norm(a - np.diag(np.diag(a))) < 1e-12 + self.a = a + if rcut_ewald is None: + rcut_ewald = min(np.diag(a)) * 0.5 + logger.warn(self, "Setting rcut_ewald to be half box size") + if rcut_hcore is None: + rcut_hcore = np.linalg.norm(np.diag(a)) / 2 + logger.warn(self, "Setting rcut_hcore to be half box diagonal") + # rcut_ewald has to be < box size cuz my get_lattice_Ls only considers nearest cell + assert rcut_ewald < min(np.diag(a)), "Only rcut_ewald < box size implemented" + self.rcut_ewald = rcut_ewald + self.rcut_hcore = rcut_hcore + + # Initialize ._atm and ._env to save the coordinates and charges and + # other info of MM particles + natm = len(atoms) + _atm = np.zeros((natm,6), dtype=np.int32) + _atm[:,gto.CHARGE_OF] = [charge(a[0]) for a in atoms] + coords = np.asarray([a[1] for a in atoms], dtype=np.double) + if charges is None: + _atm[:,gto.NUC_MOD_OF] = gto.NUC_POINT + charges = _atm[:,gto.CHARGE_OF:gto.CHARGE_OF+1] + else: + _atm[:,gto.NUC_MOD_OF] = gto.NUC_FRAC_CHARGE + charges = np.asarray(charges)[:,np.newaxis] + + self._env = np.append(np.zeros(gto.PTR_ENV_START), + np.hstack((coords, charges)).ravel()) + _atm[:,gto.PTR_COORD] = gto.PTR_ENV_START + np.arange(natm) * 4 + _atm[:,gto.PTR_FRAC_CHARGE] = gto.PTR_ENV_START + np.arange(natm) * 4 + 3 + + if zeta is not None: + self.charge_model = 'gaussian' + zeta = np.asarray(zeta, dtype=float).ravel() + self._env = np.append(self._env, zeta) + _atm[:,gto.PTR_ZETA] = gto.PTR_ENV_START + natm*4 + np.arange(natm) + + self._atm = _atm + + eta, _ = self.get_ewald_params() + e = self.precision + Q = np.sum(self.atom_charges()**2) + L = self.vol**(1/3) + kmax = np.sqrt(3)*eta/2/np.pi * np.sqrt(lambertw( 4*Q**(2/3)/3/np.pi**(2/3)/L**2/eta**(2/3) / e**(4/3) ).real) + self.mesh = np.ceil(np.diag(self.lattice_vectors()) * kmax).astype(int) * 2 + 1 + + self._built = True + + def get_lattice_Ls(self): + Ts = lib.cartesian_prod((np.arange(-1, 2), + np.arange(-1, 2), + np.arange(-1, 2))) + Lall = np.dot(Ts, self.lattice_vectors()) + return Lall + + def get_ewald_params(self, precision=None, rcut=None): + if rcut is None: + ew_cut = self.rcut_ewald + else: + ew_cut = rcut + if precision is None: + precision = self.precision + e = precision + Q = np.sum(self.atom_charges()**2) + ew_eta = 1 / ew_cut * np.sqrt(lambertw(1/e*np.sqrt(Q/2/self.vol)).real) + return ew_eta, ew_cut + + def get_ewald_pot(self, coords1, coords2=None, charges2=None): + assert self.dimension == 3 + assert (coords2 is None and charges2 is None) or (coords2 is not None and charges2 is not None) + + if charges2 is not None: + assert len(charges2) == len(coords2) + else: + coords2 = coords1 + + ew_eta, ew_cut = self.get_ewald_params() + mesh = self.mesh + + logger.debug(self, f"Ewald exponent {ew_eta}") + + # TODO Lall should respect ew_rcut + Lall = self.get_lattice_Ls() + + all_coords2 = lib.direct_sum('jx-Lx->Ljx', coords2, Lall).reshape(-1,3) + if charges2 is not None: + all_charges2 = np.hstack([charges2] * len(Lall)) + else: + all_charges2 = None + dist2 = lib.direct_sum('jx-x->jx', all_coords2, np.mean(coords1, axis=0)) + dist2 = lib.einsum('jx,jx->j', dist2, dist2) + + if all_charges2 is not None: + ewovrl0 = np.zeros(len(coords1)) + ewovrl1 = np.zeros((len(coords1), 3)) + ewovrl2 = np.zeros((len(coords1), 3, 3)) + else: + ewovrl00 = np.zeros((len(coords1), len(coords1))) + ewovrl01 = np.zeros((len(coords1), len(coords1), 3)) + ewovrl11 = np.zeros((len(coords1), len(coords1), 3, 3)) + ewovrl02 = np.zeros((len(coords1), len(coords1), 3, 3)) + ewself00 = np.zeros((len(coords1), len(coords1))) + ewself01 = np.zeros((len(coords1), len(coords1), 3)) + ewself11 = np.zeros((len(coords1), len(coords1), 3, 3)) + ewself02 = np.zeros((len(coords1), len(coords1), 3, 3)) + + mem_avail = self.max_memory - lib.current_memory()[0] # in MB + blksize = int(mem_avail/81/(8e-6*len(all_coords2))) + if blksize == 0: + raise RuntimeError(f"Not enough memory, mem_avail = {mem_avail}, blkszie = {blksize}") + + for i0, i1 in lib.prange(0, len(coords1), blksize): + R = lib.direct_sum('ix-jx->ijx', coords1[i0:i1], all_coords2) + r = np.linalg.norm(R, axis=-1) + r[r<1e-16] = 1e100 + rmax_qm = max(np.linalg.norm(coords1 - np.mean(coords1, axis=0), axis=-1)) + + # substract the real-space Coulomb within rcut_hcore + mask = dist2 <= self.rcut_hcore**2 + Tij = 1 / r[:,mask] + Rij = R[:,mask] + 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)) + if all_charges2 is not None: + charges = all_charges2[mask] + # ew0 = -d^2 E / dQi dqj qj + # ew1 = -d^2 E / dDia dqj qj + # ew2 = -d^2 E / dOiab dqj qj + # qm pc - mm pc + ewovrl0[i0:i1] += -lib.einsum('ij,j->i', Tij, charges) + # qm dip - mm pc + ewovrl1[i0:i1] += -lib.einsum('j,ija->ia', charges, Tija) + # qm quad - mm pc + ewovrl2[i0:i1] += -lib.einsum('j,ijab->iab', charges, Tijab) / 3 + else: + # NOTE a too small rcut_hcore truncates QM atoms, while this correction + # should be applied to all QM pairs regardless of rcut_hcore + # NOTE this is now checked in get_hcore + #assert r[:,mask].shape[0] == r[:,mask].shape[1] # real-space should not see qm images + # ew00 = -d^2 E / dQi dQj + # ew01 = -d^2 E / dQi dDja + # ew11 = -d^2 E / dDia dDjb + # ew02 = -d^2 E / dQi dOjab + ewovrl00[i0:i1] += -Tij + ewovrl01[i0:i1] += Tija + ewovrl11[i0:i1] += Tijab + ewovrl02[i0:i1] += -Tijab / 3 + + # difference between MM gaussain charges and MM point charges + if all_charges2 is not None and self.charge_model == 'gaussian': + mask = dist2 > self.rcut_hcore**2 + min_expnt = min(self.get_zetas()) + max_ewrcut = pbc.gto.cell._estimate_rcut(min_expnt, 0, 1., self.precision) + cut2 = (max_ewrcut + rmax_qm)**2 + mask = mask & (dist2 <= cut2) + expnts = np.hstack([np.sqrt(self.get_zetas())] * len(Lall))[mask] + r_ = r[:,mask] + R_ = R[:,mask] + ekR = np.exp(-lib.einsum('j,ij->ij', expnts**2, r_**2)) + Tij = erfc(lib.einsum('j,ij->ij', expnts, r_)) / r_ + invr3 = (Tij + lib.einsum('j,ij->ij', expnts, 2/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 + lib.einsum('j,ij->ij', expnts**3, 4/3/np.sqrt(np.pi) * ekR) + Tijab = lib.einsum('ijab,ij->ijab', Tijab, invr5) + Tijab += lib.einsum('j,ij,ab->ijab', expnts**3, 4/3/np.sqrt(np.pi)*ekR, np.eye(3)) + ewovrl0[i0:i1] -= lib.einsum('ij,j->i', Tij, all_charges2[mask]) + ewovrl1[i0:i1] -= lib.einsum('j,ija->ia', all_charges2[mask], Tija) + ewovrl2[i0:i1] -= lib.einsum('j,ijab->iab', all_charges2[mask], Tijab) / 3 + + # ewald real-space sum + if all_charges2 is not None: + cut2 = (ew_cut + rmax_qm)**2 + mask = dist2 <= cut2 + r_ = r[:,mask] + R_ = R[:,mask] + all_charges2_ = all_charges2[mask] + else: + # ewald sum will run over all qm images regardless of ew_cut + # this is to ensure r and R will always have the shape of (i1-i0, L*num_qm) + r_ = r + R_ = R + ekR = np.exp(-ew_eta**2 * r_**2) + # Tij = \hat{1/r} = f0 / r = erfc(r) / r + Tij = erfc(ew_eta * r_) / r_ + # Tija = -Rija \hat{1/r^3} = -Rija / r^2 ( \hat{1/r} + 2 eta/sqrt(pi) exp(-eta^2 r^2) ) + invr3 = (Tij + 2*ew_eta/np.sqrt(np.pi) * ekR) / r_**2 + Tija = -lib.einsum('ijx,ij->ijx', R_, invr3) + # Tijab = (3 RijaRijb - Rij^2 delta_ab) \hat{1/r^5} + 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*ew_eta**3/np.sqrt(np.pi) * ekR # NOTE this is invr5 * r**2 + Tijab = lib.einsum('ijab,ij->ijab', Tijab, invr5) + # NOTE the below is present in Eq 8 but missing in Eq 12 + Tijab += 4/3*ew_eta**3/np.sqrt(np.pi)*lib.einsum('ij,ab->ijab', ekR, np.eye(3)) + + if all_charges2 is not None: + ewovrl0[i0:i1] += lib.einsum('ij,j->i', Tij, all_charges2_) + ewovrl1[i0:i1] += lib.einsum('j,ija->ia', all_charges2_, Tija) + ewovrl2[i0:i1] += lib.einsum('j,ijab->iab', all_charges2_, Tijab) / 3 + else: + Tij = np.sum(Tij.reshape(-1, len(Lall), len(coords1)), axis=1) + Tija = np.sum(Tija.reshape(-1, len(Lall), len(coords1), 3), axis=1) + Tijab = np.sum(Tijab.reshape(-1, len(Lall), len(coords1), 3, 3), axis=1) + ewovrl00[i0:i1] += Tij + ewovrl01[i0:i1] -= Tija + ewovrl11[i0:i1] -= Tijab + ewovrl02[i0:i1] += Tijab / 3 + ekR = Tij = invr3 = Tijab = invr5 = None + + if all_charges2 is not None: + pass + else: + ewself01[i0:i1] += 0 + ewself02[i0:i1] += 0 + # -d^2 Eself / dQi dQj + ewself00[i0:i1] += -np.eye(len(coords1))[i0:i1] * 2 * ew_eta / np.sqrt(np.pi) + # -d^2 Eself / dDia dDjb + ewself11[i0:i1] += -lib.einsum('ij,ab->ijab', np.eye(len(coords1)), np.eye(3)) \ + * 4 * ew_eta**3 / 3 / np.sqrt(np.pi) + + r_ = R_ = all_charges2_ = None + + R = r = dist2 = all_charges2 = mask = None + + # g-space sum (using g grid) + logger.debug(self, f"Ewald mesh {mesh}") + + Gv, Gvbase, weights = self.get_Gv_weights(mesh) + absG2 = lib.einsum('gx,gx->g', Gv, Gv) + absG2[absG2==0] = 1e200 + + coulG = 4*np.pi / absG2 + coulG *= weights + # NOTE Gpref is actually Gpref*2 + Gpref = np.exp(-absG2/(4*ew_eta**2)) * coulG + + GvR2 = lib.einsum('gx,ix->ig', Gv, coords2) + cosGvR2 = np.cos(GvR2) + sinGvR2 = np.sin(GvR2) + + if charges2 is not None: + GvR1 = lib.einsum('gx,ix->ig', Gv, coords1) + cosGvR1 = np.cos(GvR1) + sinGvR1 = np.sin(GvR1) + zcosGvR2 = lib.einsum("i,ig->g", charges2, cosGvR2) + zsinGvR2 = lib.einsum("i,ig->g", charges2, sinGvR2) + # qm pc - mm pc + ewg0 = lib.einsum('ig,g,g->i', cosGvR1, zcosGvR2, Gpref) + ewg0 += lib.einsum('ig,g,g->i', sinGvR1, zsinGvR2, Gpref) + # qm dip - mm pc + p = ['einsum_path', (2, 3), (0, 2), (0, 1)] + ewg1 = lib.einsum('gx,ig,g,g->ix', Gv, cosGvR1, zsinGvR2, Gpref, optimize=p) + ewg1 -= lib.einsum('gx,ig,g,g->ix', Gv, sinGvR1, zcosGvR2, Gpref, optimize=p) + # qm quad - mm pc + p = ['einsum_path', (3, 4), (0, 3), (0, 2), (0, 1)] + ewg2 = -lib.einsum('gx,gy,ig,g,g->ixy', Gv, Gv, cosGvR1, zcosGvR2, Gpref, optimize=p) + ewg2 += -lib.einsum('gx,gy,ig,g,g->ixy', Gv, Gv, sinGvR1, zsinGvR2, Gpref, optimize=p) + ewg2 /= 3 + else: + # qm pc - qm pc + ewg00 = lib.einsum('ig,jg,g->ij', cosGvR2, cosGvR2, Gpref) + ewg00 += lib.einsum('ig,jg,g->ij', sinGvR2, sinGvR2, Gpref) + # qm pc - qm dip + ewg01 = lib.einsum('gx,ig,jg,g->ijx', Gv, sinGvR2, cosGvR2, Gpref) + ewg01 -= lib.einsum('gx,ig,jg,g->ijx', Gv, cosGvR2, sinGvR2, Gpref) + # qm dip - qm dip + ewg11 = lib.einsum('gx,gy,ig,jg,g->ijxy', Gv, Gv, cosGvR2, cosGvR2, Gpref) + ewg11 += lib.einsum('gx,gy,ig,jg,g->ijxy', Gv, Gv, sinGvR2, sinGvR2, Gpref) + # qm pc - qm quad + ewg02 = -lib.einsum('gx,gy,ig,jg,g->ijxy', Gv, Gv, cosGvR2, cosGvR2, Gpref) + ewg02 += -lib.einsum('gx,gy,ig,jg,g->ijxy', Gv, Gv, sinGvR2, sinGvR2, Gpref) + ewg02 /= 3 + + if charges2 is not None: + return ewovrl0 + ewg0, ewovrl1 + ewg1, ewovrl2 + ewg2 + else: + return ewovrl00 + ewself00 + ewg00, \ + ewovrl01 + ewself01 + ewg01, \ + ewovrl11 + ewself11 + ewg11, \ + ewovrl02 + ewself02 + ewg02 + +def create_mm_mol(atoms_or_coords, a, charges=None, radii=None, + rcut_ewald=None, rcut_hcore=None, unit='Angstrom'): + '''Create an MM object based on the given coordinates and charges of MM + particles. + + Args: + atoms_or_coords : array-like + Cartesian coordinates of MM atoms, in the form of a 2D array: + [(x1, y1, z1), (x2, y2, z2), ...] + a : (3,3) ndarray + Lattice primitive vectors. Each row represents a lattice vector + Reciprocal lattice vectors are given by b1,b2,b3 = 2 pi inv(a).T + + Kwargs: + charges : 1D array + The charges of MM atoms. + radii : 1D array + The Gaussian charge distribuction radii of MM atoms. + unit : string + The unit of the input. Default is 'Angstrom'. + ''' + if isinstance(atoms_or_coords, np.ndarray): + # atoms_or_coords == np.array([(xx, xx, xx)]) + # Patch ghost atoms + atoms = [(0, c) for c in atoms_or_coords] + elif (isinstance(atoms_or_coords, (list, tuple)) and + atoms_or_coords and + isinstance(atoms_or_coords[0][1], (int, float))): + # atoms_or_coords == [(xx, xx, xx)] + # Patch ghost atoms + atoms = [(0, c) for c in atoms_or_coords] + else: + atoms = atoms_or_coords + atoms = gto.format_atom(atoms, unit=unit) + + if radii is None: + zeta = None + else: + radii = np.asarray(radii, dtype=float).ravel() + if not is_au(unit): + radii = radii / param.BOHR + zeta = 1 / radii**2 + + kwargs = {'charges': charges, 'zeta': zeta} + + if not is_au(unit): + a = a / param.BOHR + if rcut_ewald is not None: + rcut_ewald = rcut_ewald / param.BOHR + if rcut_hcore is not None: + rcut_hcore = rcut_hcore / param.BOHR + + if rcut_ewald is not None: + kwargs['rcut_ewald'] = rcut_ewald + if rcut_hcore is not None: + kwargs['rcut_hcore'] = rcut_hcore + + return Cell(atoms, a, **kwargs) + +create_mm_cell = create_mm_mol diff --git a/pyscf/qmmm/pbc/test/test_qmmm.py b/pyscf/qmmm/pbc/test/test_qmmm.py new file mode 100644 index 0000000000..c4da9831d1 --- /dev/null +++ b/pyscf/qmmm/pbc/test/test_qmmm.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python +# Copyright 2014-2019 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 +# + +import unittest +import numpy as np +import pyscf +from pyscf.dft import rks +from pyscf.qmmm.pbc import itrf + +atom = ''' +O 0.0000000000 -0.0000000000 0.1174000000 +H -0.7570000000 -0.0000000000 -0.4696000000 +H 0.7570000000 0.0000000000 -0.4696000000 +''' + +bas='3-21g' +auxbasis='cc-pvdz-jkfit' +scf_tol = 1e-10 +max_scf_cycles = 50 +grids_level = 3 + +def setUpModule(): + global mol + mol = pyscf.M(atom=atom, basis=bas, max_memory=32000) + mol.output = '/dev/null' + mol.build() + mol.verbose = 1 + +def tearDownModule(): + global mol + mol.stdout.close() + del mol + +def run_dft(xc): + mf = rks.RKS(mol, xc=xc).density_fit(auxbasis=auxbasis) + mf = itrf.add_mm_charges( + mf, [[1,2,-1],[3,4,5]], np.eye(3)*15, [-5,5], [0.8,1.2], rcut_ewald=8, rcut_hcore=6) + mf.conv_tol = scf_tol + mf.max_cycle = max_scf_cycles + mf.grids.level = grids_level + e_dft = mf.kernel() + + g = mf.nuc_grad_method() + g.max_memory = 32000 + g.auxbasis_response = True + g_qm = g.kernel() + + g_mm = g.grad_nuc_mm() + g.grad_hcore_mm(mf.make_rdm1()) + g.de_ewald_mm + return e_dft, g_qm, g_mm + +class KnownValues(unittest.TestCase): + def test_rks_pbe0(self): + print('-------- RKS PBE0 -------------') + e_tot, g_qm, g_mm = run_dft('PBE0') + assert abs(e_tot - -76.00178807) < 1e-7 + assert abs(g_qm - np.array([[ 0.03002572, 0.13947702, -0.09234864], + [-0.00462601, -0.04602809, 0.02750759], + [-0.01821532, -0.18473378, 0.04189843]])).max() < 1e-6 + assert abs(g_mm - np.array([[-0.00914559, 0.08992359, 0.02114633], + [ 0.00196155, 0.00136132, 0.00179565]])).max() < 1e-6 + +if __name__ == "__main__": + print("Full Tests for QMMM PBC") + unittest.main() diff --git a/pyscf/qmmm/test/test_grad_mm.py b/pyscf/qmmm/test/test_grad_mm.py index 2f3fcc2658..fc7597877b 100644 --- a/pyscf/qmmm/test/test_grad_mm.py +++ b/pyscf/qmmm/test/test_grad_mm.py @@ -36,29 +36,28 @@ def setUpModule(): (1.894, 0.486, 0.335), (0.451, 0.165,-0.083)] mm_charges = [-1.040, 0.520, 0.520] - mm_radii = [0.63, 0.32, 0.32] + mm_radii = [6.3, 0.32, 0.32] def tearDownModule(): global mol, mm_coords, mm_charges, mm_radii class KnowValues(unittest.TestCase): - @unittest.skip('gaussian model gradients') def test_grad_mm(self): mf = itrf.mm_charge(scf.RHF(mol), mm_coords, mm_charges, mm_radii) e_hf = mf.kernel() - self.assertAlmostEqual(e_hf, -76.00028338468692, 8) + self.assertAlmostEqual(e_hf, -76.02861628942037, 8) # qm grad g_hf = itrf.mm_charge_grad(grad.RHF(mf), mm_coords, mm_charges, mm_radii) g_hf_qm = g_hf.kernel() - self.assertAlmostEqual(numpy.linalg.norm(g_hf_qm), 0.031003880070795818, 6) + self.assertAlmostEqual(numpy.linalg.norm(g_hf_qm), 0.04561356204333569, 6) # mm grad g_hf_mm_h1 = g_hf.grad_hcore_mm(mf.make_rdm1()) g_hf_mm_nuc = g_hf.grad_nuc_mm() - self.assertAlmostEqual(numpy.linalg.norm(g_hf_mm_h1), 0.5108777013806877, 6) - self.assertAlmostEqual(numpy.linalg.norm(g_hf_mm_nuc), 0.4915404602273757, 6) + self.assertAlmostEqual(numpy.linalg.norm(g_hf_mm_h1), 0.38906742696919, 6) + self.assertAlmostEqual(numpy.linalg.norm(g_hf_mm_nuc), 0.37076081680889555, 6) # finite difference for MM atoms mm_coords1 = [(1.369, 0.147,-0.395), diff --git a/pyscf/qmmm/test/test_itrf.py b/pyscf/qmmm/test/test_itrf.py index 2f52da05d4..8a2da33efc 100644 --- a/pyscf/qmmm/test/test_itrf.py +++ b/pyscf/qmmm/test/test_itrf.py @@ -46,6 +46,19 @@ def test_energy(self): self.assertAlmostEqual(mf.kernel(), 2.0042702433049024, 9) self.assertEqual(mf.undo_qmmm().__class__.__name__, 'RHF') + radii = [1e-30] + mf = itrf.mm_charge(scf.RHF(mol), coords, charges, radii=radii) + self.assertAlmostEqual(mf.kernel(), 2.0042702433049024, 9) + + radii = [1.00] + mf = itrf.mm_charge(scf.RHF(mol), coords, charges, radii=radii) + self.assertAlmostEqual(mf.kernel(), -1.6382613373942227, 9) + + radii = [1e30] + charges = [1e30] + mf = itrf.mm_charge(scf.RHF(mol), coords, charges, radii=radii) + self.assertAlmostEqual(mf.kernel(), scf.RHF(mol).kernel(), 9) + def test_grad(self): coords = [(0.0,0.1,0.0)] charges = [1.00] @@ -76,6 +89,26 @@ def test_grad(self): self.assertAlmostEqual(abs(ref-v).max(), 0, 12) pyscf.DEBUG = bak + radii = [1.00] + mf = itrf.mm_charge(scf.RHF(mol), coords, charges, radii=radii).run(conv_tol=1e-10) + hfg = itrf.mm_charge_grad(grad.RHF(mf), coords, charges, radii=radii).run() + self.assertAlmostEqual(numpy.linalg.norm(hfg.de), 0.287026393860105, 6) + + mfs = mf.as_scanner() + e1 = mfs(''' + H -0.00000000 -0.000 0.001 + H -0.00000000 -0.000 1. + H -0.00000000 -0.82 0. + H -0.91000000 -0.020 0. + ''') + e2 = mfs(''' + H -0.00000000 -0.000 -0.001 + H -0.00000000 -0.000 1. + H -0.00000000 -0.82 0. + H -0.91000000 -0.020 0. + ''') + self.assertAlmostEqual((e1 - e2)/0.002*lib.param.BOHR, hfg.de[0,2], 5) + def test_hcore_cart(self): coords = [(0.0,0.1,0.0)] charges = [1.00]