From 27a8af577912a0eb9b771ba72eff9939686584b4 Mon Sep 17 00:00:00 2001 From: MoleOrbitalHybridAnalyst Date: Wed, 24 Sep 2025 19:16:58 -0700 Subject: [PATCH 1/7] re-enable gaussian MM charges --- pyscf/qmmm/itrf.py | 31 ++++++++++++++++++++++++++----- pyscf/qmmm/test/test_grad_mm.py | 11 +++++------ pyscf/qmmm/test/test_itrf.py | 33 +++++++++++++++++++++++++++++++++ 3 files changed, 64 insertions(+), 11 deletions(-) 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/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] From 8ae05a8a22cc2a93f4e4bff577554d089c800efe Mon Sep 17 00:00:00 2001 From: MoleOrbitalHybridAnalyst Date: Wed, 24 Sep 2025 22:13:55 -0700 Subject: [PATCH 2/7] add periodic QMMM using multipoles --- pyscf/qmmm/pbc/__init__.py | 1 + pyscf/qmmm/pbc/itrf.py | 1046 +++++++++++++++++++++++++++++ pyscf/qmmm/pbc/mm_mole.py | 399 +++++++++++ pyscf/qmmm/pbc/tests/test_qmmm.py | 79 +++ 4 files changed, 1525 insertions(+) create mode 100644 pyscf/qmmm/pbc/__init__.py create mode 100644 pyscf/qmmm/pbc/itrf.py create mode 100644 pyscf/qmmm/pbc/mm_mole.py create mode 100644 pyscf/qmmm/pbc/tests/test_qmmm.py 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..96ace33848 --- /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 pyscf.qmmm.itrf import _QMMM, _QMMMGrad +from pyscf import qmmm as qmmm_gas + +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' + +_QMMM = 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), \ + f"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), \ + f"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 not (type(v) is 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() \ No newline at end of file diff --git a/pyscf/qmmm/pbc/mm_mole.py b/pyscf/qmmm/pbc/mm_mole.py new file mode 100644 index 0000000000..cefc7a7936 --- /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/tests/test_qmmm.py b/pyscf/qmmm/pbc/tests/test_qmmm.py new file mode 100644 index 0000000000..3ac085dfe9 --- /dev/null +++ b/pyscf/qmmm/pbc/tests/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() \ No newline at end of file From 7dd4f164ee1144fdd5e1589949caf1a8d8b853f4 Mon Sep 17 00:00:00 2001 From: MoleOrbitalHybridAnalyst Date: Thu, 25 Sep 2025 10:02:54 -0700 Subject: [PATCH 3/7] minor --- pyscf/qmmm/pbc/itrf.py | 4 ++-- pyscf/qmmm/pbc/mm_mole.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pyscf/qmmm/pbc/itrf.py b/pyscf/qmmm/pbc/itrf.py index 96ace33848..fade531487 100644 --- a/pyscf/qmmm/pbc/itrf.py +++ b/pyscf/qmmm/pbc/itrf.py @@ -783,7 +783,7 @@ def grad_kTij(R, r, eta): # ---------------------------------------------- # # ---------- Ewald k-space gradient ------------ # # ---------------------------------------------- # - + mesh = cell.mesh Gv, Gvbase, weights = cell.get_Gv_weights(mesh) absG2 = lib.einsum('gi,gi->g', Gv, Gv) @@ -858,7 +858,7 @@ def grad_kTij(R, r, eta): 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: diff --git a/pyscf/qmmm/pbc/mm_mole.py b/pyscf/qmmm/pbc/mm_mole.py index cefc7a7936..7b634f681b 100644 --- a/pyscf/qmmm/pbc/mm_mole.py +++ b/pyscf/qmmm/pbc/mm_mole.py @@ -130,7 +130,7 @@ def get_ewald_pot(self, coords1, coords2=None, charges2=None): 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 From 711704abdf168dfd22da02cbd272f1b56dd51943 Mon Sep 17 00:00:00 2001 From: MoleOrbitalHybridAnalyst Date: Thu, 25 Sep 2025 10:05:17 -0700 Subject: [PATCH 4/7] minor --- pyscf/qmmm/pbc/itrf.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/pyscf/qmmm/pbc/itrf.py b/pyscf/qmmm/pbc/itrf.py index fade531487..856050e1e6 100644 --- a/pyscf/qmmm/pbc/itrf.py +++ b/pyscf/qmmm/pbc/itrf.py @@ -25,8 +25,6 @@ from pyscf import grad from pyscf.lib import logger from pyscf.qmmm.pbc import mm_mole -from pyscf.qmmm.itrf import _QMMM, _QMMMGrad -from pyscf import qmmm as qmmm_gas from scipy.special import erfc, erf @@ -107,8 +105,6 @@ def qmmm_for_scf(method, mm_mol): class QMMM: __name_mixin__ = 'QMMM' -_QMMM = QMMM - class QMMMSCF(QMMM): _keys = {'mm_mol', 's1r', 's1rr', 'mm_ewald_pot', 'qm_ewald_hess', 'e_nuc'} @@ -178,14 +174,14 @@ def get_hcore(self, mol=None): 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), \ - f"QM image is within rcut_hcore of QM center. " + \ + "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), \ - f"Not all QM atoms are within rcut_hcore of QM center. " + \ + "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 @@ -219,7 +215,7 @@ def get_hcore(self, mol=None): j3c = df.incore.aux_e2(mol, fakemol, intor=intor, aosym='s2ij', cintopt=cintopt) v += lib.einsum('xk,k->x', j3c, -charges[i0:i1]) - if not (type(v) is int): + if type(v) is not int: v = lib.unpack_tril(v) h1e += v elif mm_mol.charge_model == 'point' and len(coords) != 0: From bd6d4dfc55a701e62d5edbbdfded8cd153573c2f Mon Sep 17 00:00:00 2001 From: MoleOrbitalHybridAnalyst Date: Thu, 25 Sep 2025 10:21:55 -0700 Subject: [PATCH 5/7] pep8 --- pyscf/qmmm/pbc/itrf.py | 48 +++++++++++++++++++++------------------ pyscf/qmmm/pbc/mm_mole.py | 22 +++++++++--------- 2 files changed, 37 insertions(+), 33 deletions(-) diff --git a/pyscf/qmmm/pbc/itrf.py b/pyscf/qmmm/pbc/itrf.py index 856050e1e6..5d38754c4b 100644 --- a/pyscf/qmmm/pbc/itrf.py +++ b/pyscf/qmmm/pbc/itrf.py @@ -28,7 +28,7 @@ from scipy.special import erfc, erf -def add_mm_charges(scf_method, atoms_or_coords, a, charges, radii=None, +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. @@ -38,7 +38,7 @@ def add_mm_charges(scf_method, atoms_or_coords, a, charges, radii=None, 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 + the MM energy, the vdw interaction or other bonding/non-bonding effects between QM region and MM particles. Args: @@ -73,7 +73,7 @@ def add_mm_charges(scf_method, atoms_or_coords, a, charges, radii=None, 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, + 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) @@ -185,7 +185,7 @@ def get_hcore(self, mol=None): 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', + 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 @@ -266,7 +266,7 @@ def get_qm_dipoles(self, dm, s1r=None): def get_s1rr(self): ''' - \int phi_u phi_v [3(r-Rc)\otimes(r-Rc) - |r-Rc|^2] /2 dr + 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()) @@ -302,7 +302,7 @@ def get_qm_quadrupoles(self, dm, s1rr=None): 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 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)) @@ -401,7 +401,7 @@ def energy_nuc(self): 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', + 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)) @@ -436,7 +436,7 @@ def nuc_grad_method(self): Gradients = nuc_grad_method -def add_mm_charges_grad(scf_grad, atoms_or_coords, a, charges, radii=None, +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. @@ -462,7 +462,7 @@ def add_mm_charges_grad(scf_grad, atoms_or_coords, a, charges, radii=None, 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, + 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 @@ -569,7 +569,7 @@ def grad_ewald(self, dm=None, with_mm=False, mm_ewald_pot=None, qm_ewald_pot=Non s1r.append( np.asarray(-mol.intor('int1e_irp', shls_slice=shlslc). reshape(3, 3, -1, mol.nao))) - # s1rr[a,b,x,u,v] = + # 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)) @@ -627,7 +627,7 @@ def grad_Tij(R, r): # 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 = 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) @@ -826,8 +826,10 @@ def grad_kTij(R, r, eta): 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 + 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 @@ -852,8 +854,10 @@ def grad_kTij(R, r, eta): 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 + 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) @@ -875,7 +879,7 @@ def get_hcore(self, mol=None): Ls = mm_mol.get_lattice_Ls() qm_center = np.mean(mol.atom_coords(), axis=0) - all_coords = lib.direct_sum('ix+Lx->Lix', + 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 @@ -928,7 +932,7 @@ def grad_hcore_mm(self, dm, mol=None): Ls = mm_mol.get_lattice_Ls() qm_center = np.mean(mol.atom_coords(), axis=0) - all_coords = lib.direct_sum('ix+Lx->Lix', + 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)) @@ -973,7 +977,7 @@ def grad_nuc(self, mol=None, atmlst=None): 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', + 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)) @@ -991,7 +995,7 @@ def grad_nuc(self, mol=None, atmlst=None): 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), + 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) @@ -1012,7 +1016,7 @@ def grad_nuc_mm(self, mol=None): 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', + 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)) @@ -1029,7 +1033,7 @@ def grad_nuc_mm(self, mol=None): 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), + 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) @@ -1039,4 +1043,4 @@ def grad_nuc_mm(self, mol=None): def _finalize(self): g_ewald_qm, self.de_ewald_mm = self.grad_ewald(with_mm=True) self.de += g_ewald_qm - super()._finalize() \ No newline at end of file + super()._finalize() diff --git a/pyscf/qmmm/pbc/mm_mole.py b/pyscf/qmmm/pbc/mm_mole.py index 7b634f681b..feb2f35286 100644 --- a/pyscf/qmmm/pbc/mm_mole.py +++ b/pyscf/qmmm/pbc/mm_mole.py @@ -45,7 +45,7 @@ class Cell(qmmm.mm_mole.Mole, pbc.gto.Cell): rho(r) = charge * Norm * exp(-zeta * r^2) ''' - def __init__(self, atoms, a, + def __init__(self, atoms, a, rcut_ewald=None, rcut_hcore=None, charges=None, zeta=None): pbc.gto.Cell.__init__(self) @@ -149,14 +149,14 @@ def get_ewald_pot(self, coords1, coords2=None, charges2=None): 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)) + 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))) @@ -174,7 +174,7 @@ def get_ewald_pot(self, coords1, coords2=None, charges2=None): 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 = 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: @@ -337,7 +337,7 @@ def get_ewald_pot(self, coords1, coords2=None, charges2=None): ewovrl11 + ewself11 + ewg11, \ ewovrl02 + ewself02 + ewg02 -def create_mm_mol(atoms_or_coords, a, charges=None, radii=None, +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. From 39f64f1a5a70051192133b3a3ff1aafa9f1b26bf Mon Sep 17 00:00:00 2001 From: MoleOrbitalHybridAnalyst Date: Thu, 25 Sep 2025 10:24:39 -0700 Subject: [PATCH 6/7] pep8 --- pyscf/qmmm/pbc/test/test_qmmm.py | 79 ++++++++++++++++++++++++++++++++ 1 file changed, 79 insertions(+) create mode 100644 pyscf/qmmm/pbc/test/test_qmmm.py 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() From 5d085b714d0b8c746f7b263647f2db72baf45882 Mon Sep 17 00:00:00 2001 From: MoleOrbitalHybridAnalyst Date: Thu, 25 Sep 2025 10:24:54 -0700 Subject: [PATCH 7/7] pep8 --- pyscf/qmmm/pbc/tests/test_qmmm.py | 79 ------------------------------- 1 file changed, 79 deletions(-) delete mode 100644 pyscf/qmmm/pbc/tests/test_qmmm.py diff --git a/pyscf/qmmm/pbc/tests/test_qmmm.py b/pyscf/qmmm/pbc/tests/test_qmmm.py deleted file mode 100644 index 3ac085dfe9..0000000000 --- a/pyscf/qmmm/pbc/tests/test_qmmm.py +++ /dev/null @@ -1,79 +0,0 @@ -#!/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() \ No newline at end of file