1#!/usr/bin/env python
2# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16# Author: Qiming Sun <osirpt.sun@gmail.com>
17#
18
19'''
20Non-relativistic restricted Kohn-Sham
21'''
22
23
24import textwrap
25import numpy
26from pyscf import lib
27from pyscf.lib import logger
28from pyscf import scf
29from pyscf.scf import hf
30from pyscf.scf import _vhf
31from pyscf.scf import jk
32from pyscf.dft import gen_grid
33from pyscf import __config__
34
35
36def get_veff(ks, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
37    '''Coulomb + XC functional
38
39    .. note::
40        This function will modify the input ks object.
41
42    Args:
43        ks : an instance of :class:`RKS`
44            XC functional are controlled by ks.xc attribute.  Attribute
45            ks.grids might be initialized.
46        dm : ndarray or list of ndarrays
47            A density matrix or a list of density matrices
48
49    Kwargs:
50        dm_last : ndarray or a list of ndarrays or 0
51            The density matrix baseline.  If not 0, this function computes the
52            increment of HF potential w.r.t. the reference HF potential matrix.
53        vhf_last : ndarray or a list of ndarrays or 0
54            The reference Vxc potential matrix.
55        hermi : int
56            Whether J, K matrix is hermitian
57
58            | 0 : no hermitian or symmetric
59            | 1 : hermitian
60            | 2 : anti-hermitian
61
62    Returns:
63        matrix Veff = J + Vxc.  Veff can be a list matrices, if the input
64        dm is a list of density matrices.
65    '''
66    if mol is None: mol = ks.mol
67    if dm is None: dm = ks.make_rdm1()
68    t0 = (logger.process_clock(), logger.perf_counter())
69
70    ground_state = (isinstance(dm, numpy.ndarray) and dm.ndim == 2)
71
72    if ks.grids.coords is None:
73        ks.grids.build(with_non0tab=True)
74        if ks.small_rho_cutoff > 1e-20 and ground_state:
75            # Filter grids the first time setup grids
76            ks.grids = prune_small_rho_grids_(ks, mol, dm, ks.grids)
77        t0 = logger.timer(ks, 'setting up grids', *t0)
78    if ks.nlc != '':
79        if ks.nlcgrids.coords is None:
80            ks.nlcgrids.build(with_non0tab=True)
81            if ks.small_rho_cutoff > 1e-20 and ground_state:
82                # Filter grids the first time setup grids
83                ks.nlcgrids = prune_small_rho_grids_(ks, mol, dm, ks.nlcgrids)
84            t0 = logger.timer(ks, 'setting up nlc grids', *t0)
85
86    ni = ks._numint
87    if hermi == 2:  # because rho = 0
88        n, exc, vxc = 0, 0, 0
89    else:
90        max_memory = ks.max_memory - lib.current_memory()[0]
91        n, exc, vxc = ni.nr_rks(mol, ks.grids, ks.xc, dm, max_memory=max_memory)
92        if ks.nlc:
93            assert 'VV10' in ks.nlc.upper()
94            _, enlc, vnlc = ni.nr_rks(mol, ks.nlcgrids, ks.xc+'__'+ks.nlc, dm,
95                                      max_memory=max_memory)
96            exc += enlc
97            vxc += vnlc
98        logger.debug(ks, 'nelec by numeric integration = %s', n)
99        t0 = logger.timer(ks, 'vxc', *t0)
100
101    #enabling range-separated hybrids
102    omega, alpha, hyb = ni.rsh_and_hybrid_coeff(ks.xc, spin=mol.spin)
103
104    if abs(hyb) < 1e-10 and abs(alpha) < 1e-10:
105        vk = None
106        if (ks._eri is None and ks.direct_scf and
107            getattr(vhf_last, 'vj', None) is not None):
108            ddm = numpy.asarray(dm) - numpy.asarray(dm_last)
109            vj = ks.get_j(mol, ddm, hermi)
110            vj += vhf_last.vj
111        else:
112            vj = ks.get_j(mol, dm, hermi)
113        vxc += vj
114    else:
115        if (ks._eri is None and ks.direct_scf and
116            getattr(vhf_last, 'vk', None) is not None):
117            ddm = numpy.asarray(dm) - numpy.asarray(dm_last)
118            vj, vk = ks.get_jk(mol, ddm, hermi)
119            vk *= hyb
120            if abs(omega) > 1e-10:  # For range separated Coulomb operator
121                vklr = ks.get_k(mol, ddm, hermi, omega=omega)
122                vklr *= (alpha - hyb)
123                vk += vklr
124            vj += vhf_last.vj
125            vk += vhf_last.vk
126        else:
127            vj, vk = ks.get_jk(mol, dm, hermi)
128            vk *= hyb
129            if abs(omega) > 1e-10:
130                vklr = ks.get_k(mol, dm, hermi, omega=omega)
131                vklr *= (alpha - hyb)
132                vk += vklr
133        vxc += vj - vk * .5
134
135        if ground_state:
136            exc -= numpy.einsum('ij,ji', dm, vk).real * .5 * .5
137
138    if ground_state:
139        ecoul = numpy.einsum('ij,ji', dm, vj).real * .5
140    else:
141        ecoul = None
142
143    vxc = lib.tag_array(vxc, ecoul=ecoul, exc=exc, vj=vj, vk=vk)
144    return vxc
145
146def get_vsap(ks, mol=None):
147    '''Superposition of atomic potentials
148
149    S. Lehtola, Assessment of initial guesses for self-consistent
150    field calculations. Superposition of Atomic Potentials: simple yet
151    efficient, J. Chem. Theory Comput. 15, 1593 (2019). DOI:
152    10.1021/acs.jctc.8b01089. arXiv:1810.11659.
153
154    This function evaluates the effective charge of a neutral atom,
155    given by exchange-only LDA on top of spherically symmetric
156    unrestricted Hartree-Fock calculations as described in
157
158    S. Lehtola, L. Visscher, E. Engel, Efficient implementation of the
159    superposition of atomic potentials initial guess for electronic
160    structure calculations in Gaussian basis sets, J. Chem. Phys., in
161    press (2020).
162
163    The potentials have been calculated for the ground-states of
164    spherically symmetric atoms at the non-relativistic level of theory
165    as described in
166
167    S. Lehtola, "Fully numerical calculations on atoms with fractional
168    occupations and range-separated exchange functionals", Phys. Rev. A
169    101, 012516 (2020). DOI: 10.1103/PhysRevA.101.012516
170
171    using accurate finite-element calculations as described in
172
173    S. Lehtola, "Fully numerical Hartree-Fock and density functional
174    calculations. I. Atoms", Int. J. Quantum Chem. e25945 (2019).
175    DOI: 10.1002/qua.25945
176
177    .. note::
178        This function will modify the input ks object.
179
180    Args:
181        ks : an instance of :class:`RKS`
182            XC functional are controlled by ks.xc attribute.  Attribute
183            ks.grids might be initialized.
184
185    Returns:
186        matrix Vsap = Vnuc + J + Vxc.
187    '''
188    if mol is None: mol = ks.mol
189    t0 = (logger.process_clock(), logger.perf_counter())
190
191    if ks.grids.coords is None:
192        ks.grids.build(with_non0tab=True)
193        t0 = logger.timer(ks, 'setting up grids', *t0)
194
195    ni = ks._numint
196    max_memory = ks.max_memory - lib.current_memory()[0]
197    vsap = ni.nr_sap(mol, ks.grids, max_memory=max_memory)
198    return vsap
199
200# The vhfopt of standard Coulomb operator can be used here as an approximate
201# opt since long-range part Coulomb is always smaller than standard Coulomb.
202# It's safe to prescreen LR integrals with the integral estimation from
203# standard Coulomb.
204def _get_k_lr(mol, dm, omega=0, hermi=0, vhfopt=None):
205    import sys
206    sys.stderr.write('This function is deprecated. '
207                     'It is replaced by mol.get_k(mol, dm, omege=omega)')
208    dm = numpy.asarray(dm)
209# Note, ks object caches the ERIs for small systems. The cached eris are
210# computed with regular Coulomb operator. ks.get_jk or ks.get_k do not evalute
211# the K matrix with the range separated Coulomb operator.  Here jk.get_jk
212# function computes the K matrix with the modified Coulomb operator.
213    nao = dm.shape[-1]
214    dms = dm.reshape(-1,nao,nao)
215    with mol.with_range_coulomb(omega):
216        # Compute the long range part of ERIs temporarily with omega. Restore
217        # the original omega when the block ends
218        if vhfopt is None:
219            contents = lambda: None # just a place_holder
220        else:
221            contents = vhfopt._this.contents
222        with lib.temporary_env(contents,
223                               fprescreen=_vhf._fpointer('CVHFnrs8_vk_prescreen')):
224            intor = mol._add_suffix('int2e')
225            vklr = jk.get_jk(mol, dms, ['ijkl,jk->il']*len(dms), intor=intor,
226                             vhfopt=vhfopt)
227    return numpy.asarray(vklr).reshape(dm.shape)
228
229
230def energy_elec(ks, dm=None, h1e=None, vhf=None):
231    r'''Electronic part of RKS energy.
232
233    Note this function has side effects which cause mf.scf_summary updated.
234
235    Args:
236        ks : an instance of DFT class
237
238        dm : 2D ndarray
239            one-partical density matrix
240        h1e : 2D ndarray
241            Core hamiltonian
242
243    Returns:
244        RKS electronic energy and the 2-electron contribution
245    '''
246    if dm is None: dm = ks.make_rdm1()
247    if h1e is None: h1e = ks.get_hcore()
248    if vhf is None or getattr(vhf, 'ecoul', None) is None:
249        vhf = ks.get_veff(ks.mol, dm)
250    e1 = numpy.einsum('ij,ji->', h1e, dm)
251    e2 = vhf.ecoul + vhf.exc
252    ks.scf_summary['e1'] = e1.real
253    ks.scf_summary['coul'] = vhf.ecoul.real
254    ks.scf_summary['exc'] = vhf.exc.real
255    logger.debug(ks, 'E1 = %s  Ecoul = %s  Exc = %s', e1, vhf.ecoul, vhf.exc)
256    return (e1+e2).real, e2
257
258
259NELEC_ERROR_TOL = getattr(__config__, 'dft_rks_prune_error_tol', 0.02)
260def prune_small_rho_grids_(ks, mol, dm, grids):
261    rho = ks._numint.get_rho(mol, dm, grids, ks.max_memory)
262    n = numpy.dot(rho, grids.weights)
263    if abs(n-mol.nelectron) < NELEC_ERROR_TOL*n:
264        rho *= grids.weights
265        idx = abs(rho) > ks.small_rho_cutoff / grids.weights.size
266        logger.debug(ks, 'Drop grids %d',
267                     grids.weights.size - numpy.count_nonzero(idx))
268        grids.coords  = numpy.asarray(grids.coords [idx], order='C')
269        grids.weights = numpy.asarray(grids.weights[idx], order='C')
270        grids.non0tab = grids.make_mask(mol, grids.coords)
271    return grids
272
273def define_xc_(ks, description, xctype='LDA', hyb=0, rsh=(0,0,0)):
274    libxc = ks._numint.libxc
275    ks._numint = libxc.define_xc_(ks._numint, description, xctype, hyb, rsh)
276    return ks
277
278
279def _dft_common_init_(mf, xc='LDA,VWN'):
280    from pyscf.dft import numint
281    mf.xc = xc
282    mf.nlc = ''
283    mf.grids = gen_grid.Grids(mf.mol)
284    mf.grids.level = getattr(__config__, 'dft_rks_RKS_grids_level',
285                             mf.grids.level)
286    mf.nlcgrids = gen_grid.Grids(mf.mol)
287    mf.nlcgrids.level = getattr(__config__, 'dft_rks_RKS_nlcgrids_level',
288                                mf.nlcgrids.level)
289    # Use rho to filter grids
290    mf.small_rho_cutoff = getattr(__config__, 'dft_rks_RKS_small_rho_cutoff', 1e-7)
291##################################################
292# don't modify the following attributes, they are not input options
293    mf._numint = numint.NumInt()
294    mf._keys = mf._keys.union(['xc', 'nlc', 'omega', 'grids', 'nlcgrids',
295                               'small_rho_cutoff'])
296
297class KohnShamDFT(object):
298    '''
299    Attributes for Kohn-Sham DFT:
300        xc : str
301            'X_name,C_name' for the XC functional.  Default is 'lda,vwn'
302        nlc : str
303            'NLC_name' for the NLC functional.  Default is '' (i.e., None)
304        omega : float
305            Omega of the range-separated Coulomb operator e^{-omega r_{12}^2} / r_{12}
306        grids : Grids object
307            grids.level (0 - 9)  big number for large mesh grids. Default is 3
308
309            radii_adjust
310                | radi.treutler_atomic_radii_adjust (default)
311                | radi.becke_atomic_radii_adjust
312                | None : to switch off atomic radii adjustment
313
314            grids.atomic_radii
315                | radi.BRAGG_RADII  (default)
316                | radi.COVALENT_RADII
317                | None : to switch off atomic radii adjustment
318
319            grids.radi_method  scheme for radial grids
320                | radi.treutler  (default)
321                | radi.delley
322                | radi.mura_knowles
323                | radi.gauss_chebyshev
324
325            grids.becke_scheme  weight partition function
326                | gen_grid.original_becke  (default)
327                | gen_grid.stratmann
328
329            grids.prune  scheme to reduce number of grids
330                | gen_grid.nwchem_prune  (default)
331                | gen_grid.sg1_prune
332                | gen_grid.treutler_prune
333                | None : to switch off grids pruning
334
335            grids.symmetry  True/False  to symmetrize mesh grids (TODO)
336
337            grids.atom_grid  Set (radial, angular) grids for particular atoms.
338            Eg, grids.atom_grid = {'H': (20,110)} will generate 20 radial
339            grids and 110 angular grids for H atom.
340
341        small_rho_cutoff : float
342            Drop grids if their contribution to total electrons smaller than
343            this cutoff value.  Default is 1e-7.
344
345    Examples:
346
347    >>> mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='ccpvdz', verbose=0)
348    >>> mf = dft.RKS(mol)
349    >>> mf.xc = 'b3lyp'
350    >>> mf.kernel()
351    -76.415443079840458
352    '''
353
354    __init__ = _dft_common_init_
355
356    @property
357    def omega(self):
358        return self._numint.omega
359    @omega.setter
360    def omega(self, v):
361        self._numint.omega = float(v)
362
363    def dump_flags(self, verbose=None):
364        log = logger.new_logger(self, verbose)
365        log.info('XC library %s version %s\n    %s',
366                 self._numint.libxc.__name__,
367                 self._numint.libxc.__version__,
368                 self._numint.libxc.__reference__)
369
370        if log.verbose >= logger.INFO:
371            log.info('XC functionals = %s', self.xc)
372            if hasattr(self._numint.libxc, 'xc_reference'):
373                log.info(textwrap.indent('\n'.join(self._numint.libxc.xc_reference(self.xc)), '    '))
374
375        if self.nlc!='':
376            log.info('NLC functional = %s', self.nlc)
377
378        self.grids.dump_flags(verbose)
379        if self.nlc!='':
380            log.info('** Following is NLC Grids **')
381            self.nlcgrids.dump_flags(verbose)
382
383        log.info('small_rho_cutoff = %g', self.small_rho_cutoff)
384        return self
385
386    define_xc_ = define_xc_
387
388    def to_rhf(self):
389        '''Convert the input mean-field object to a RHF/ROHF object.
390
391        Note this conversion only changes the class of the mean-field object.
392        The total energy and wave-function are the same as them in the input
393        mean-field object.
394        '''
395        mf = scf.RHF(self.mol)
396        mf.__dict__.update(self.to_rks().__dict__)
397        mf.converged = False
398        return mf
399
400    def to_uhf(self):
401        '''Convert the input mean-field object to a UHF object.
402
403        Note this conversion only changes the class of the mean-field object.
404        The total energy and wave-function are the same as them in the input
405        mean-field object.
406        '''
407        mf = scf.UHF(self.mol)
408        mf.__dict__.update(self.to_uks().__dict__)
409        mf.converged = False
410        return mf
411
412    def to_ghf(self):
413        '''Convert the input mean-field object to a GHF object.
414
415        Note this conversion only changes the class of the mean-field object.
416        The total energy and wave-function are the same as them in the input
417        mean-field object.
418        '''
419        mf = scf.GHF(self.mol)
420        mf.__dict__.update(self.to_gks().__dict__)
421        mf.converged = False
422        return mf
423
424    def to_rks(self, xc=None):
425        '''Convert the input mean-field object to a RKS/ROKS object.
426
427        Note this conversion only changes the class of the mean-field object.
428        The total energy and wave-function are the same as them in the input
429        mean-field object.
430        '''
431        mf = scf.addons.convert_to_rhf(self)
432        if xc is not None:
433            mf.xc = xc
434        if xc != self.xc or not isinstance(self, RKS):
435            mf.converged = False
436        return mf
437
438    def to_uks(self, xc=None):
439        '''Convert the input mean-field object to a UKS object.
440
441        Note this conversion only changes the class of the mean-field object.
442        The total energy and wave-function are the same as them in the input
443        mean-field object.
444        '''
445        mf = scf.addons.convert_to_uhf(self)
446        if xc is not None:
447            mf.xc = xc
448        if xc != self.xc:
449            mf.converged = False
450        return mf
451
452    def to_gks(self, xc=None):
453        '''Convert the input mean-field object to a GKS object.
454
455        Note this conversion only changes the class of the mean-field object.
456        The total energy and wave-function are the same as them in the input
457        mean-field object.
458        '''
459        mf = scf.addons.convert_to_ghf(self)
460        if xc is not None:
461            mf.xc = xc
462        if xc != self.xc:
463            mf.converged = False
464        return mf
465
466    def reset(self, mol=None):
467        hf.SCF.reset(self, mol)
468        self.grids.reset(mol)
469        self.nlcgrids.reset(mol)
470        return self
471
472
473def init_guess_by_vsap(mf, mol=None):
474    '''Form SAP guess'''
475    if mol is None: mol = mf.mol
476
477    vsap = mf.get_vsap()
478    t = mol.intor_symmetric('int1e_kin')
479    s = mf.get_ovlp(mol)
480    hsap = t + vsap
481
482    # Form guess orbitals
483    mo_energy, mo_coeff = mf.eig(hsap, s)
484    logger.debug(mf, 'VSAP mo energies\n{}'.format(mo_energy))
485
486    # and guess density
487    mo_occ = mf.get_occ(mo_energy, mo_coeff)
488    return mf.make_rdm1(mo_coeff, mo_occ)
489
490
491class RKS(KohnShamDFT, hf.RHF):
492    __doc__ = '''Restricted Kohn-Sham\n''' + hf.SCF.__doc__ + KohnShamDFT.__doc__
493
494    def __init__(self, mol, xc='LDA,VWN'):
495        hf.RHF.__init__(self, mol)
496        KohnShamDFT.__init__(self, xc)
497
498    def dump_flags(self, verbose=None):
499        hf.RHF.dump_flags(self, verbose)
500        return KohnShamDFT.dump_flags(self, verbose)
501
502    get_veff = get_veff
503    get_vsap = get_vsap
504    energy_elec = energy_elec
505
506    init_guess_by_vsap = init_guess_by_vsap
507
508    def nuc_grad_method(self):
509        from pyscf.grad import rks as rks_grad
510        return rks_grad.Gradients(self)
511