1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#     http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16#
17# Authors: Qiming Sun <osirpt.sun@gmail.com>
18#          Junzi Liu <latrix1247@gmail.com>
19#          Susi Lehtola <susi.lehtola@gmail.com>
20
21import copy
22from functools import reduce
23import numpy
24import scipy.linalg
25from pyscf import lib
26from pyscf.gto import mole
27from pyscf.lib import logger
28from pyscf.lib.scipy_helper import pivoted_cholesky
29from pyscf.scf import hf
30from pyscf import __config__
31
32LINEAR_DEP_THRESHOLD = getattr(__config__, 'scf_addons_remove_linear_dep_threshold', 1e-8)
33CHOLESKY_THRESHOLD = getattr(__config__, 'scf_addons_cholesky_threshold', 1e-10)
34FORCE_PIVOTED_CHOLESKY = getattr(__config__, 'scf_addons_force_cholesky', False)
35LINEAR_DEP_TRIGGER = getattr(__config__, 'scf_addons_remove_linear_dep_trigger', 1e-10)
36
37def smearing_(*args, **kwargs):
38    from pyscf.pbc.scf.addons import smearing_
39    return smearing_(*args, **kwargs)
40
41def frac_occ_(mf, tol=1e-3):
42    '''
43    Addons for SCF methods to assign fractional occupancy for degenerated
44    occpupied HOMOs.
45
46    Examples::
47        >>> mf = gto.M(atom='O 0 0 0; O 0 0 1', verbose=4).RHF()
48        >>> mf = scf.addons.frac_occ(mf)
49        >>> mf.run()
50    '''
51    from pyscf.scf import uhf, rohf
52    old_get_occ = mf.get_occ
53    mol = mf.mol
54
55    def guess_occ(mo_energy, nocc):
56        mo_occ = numpy.zeros_like(mo_energy)
57        if nocc:
58            sorted_idx = numpy.argsort(mo_energy)
59            homo = mo_energy[sorted_idx[nocc-1]]
60            lumo = mo_energy[sorted_idx[nocc]]
61            frac_occ_lst = abs(mo_energy - homo) < tol
62            integer_occ_lst = (mo_energy <= homo) & (~frac_occ_lst)
63            mo_occ[integer_occ_lst] = 1
64            degen = numpy.count_nonzero(frac_occ_lst)
65            frac = nocc - numpy.count_nonzero(integer_occ_lst)
66            mo_occ[frac_occ_lst] = float(frac) / degen
67        else:
68            homo = 0.0
69            lumo = 0.0
70            frac_occ_lst = numpy.zeros_like(mo_energy, dtype=bool)
71        return mo_occ, numpy.where(frac_occ_lst)[0], homo, lumo
72
73    get_grad = None
74
75    if isinstance(mf, uhf.UHF):
76        def get_occ(mo_energy, mo_coeff=None):
77            nocca, noccb = mol.nelec
78            mo_occa, frac_lsta, homoa, lumoa = guess_occ(mo_energy[0], nocca)
79            mo_occb, frac_lstb, homob, lumob = guess_occ(mo_energy[1], noccb)
80
81            if abs(homoa - lumoa) < tol or abs(homob - lumob) < tol:
82                mo_occ = numpy.array([mo_occa, mo_occb])
83                if len(frac_lstb):
84                    logger.warn(mf, 'fraction occ = %6g for alpha orbitals %s  '
85                                '%6g for beta orbitals %s',
86                                mo_occa[frac_lsta[0]], frac_lsta,
87                                mo_occb[frac_lstb[0]], frac_lstb)
88                    logger.info(mf, '  alpha HOMO = %.12g  LUMO = %.12g', homoa, lumoa)
89                    logger.info(mf, '  beta  HOMO = %.12g  LUMO = %.12g', homob, lumob)
90                    logger.debug(mf, '  alpha mo_energy = %s', mo_energy[0])
91                    logger.debug(mf, '  beta  mo_energy = %s', mo_energy[1])
92                else:
93                    logger.warn(mf, 'fraction occ = %6g for alpha orbitals %s  ',
94                                mo_occa[frac_lsta[0]], frac_lsta)
95                    logger.info(mf, '  alpha HOMO = %.12g  LUMO = %.12g', homoa, lumoa)
96                    logger.debug(mf, '  alpha mo_energy = %s', mo_energy[0])
97            else:
98                mo_occ = old_get_occ(mo_energy, mo_coeff)
99            return mo_occ
100
101    elif isinstance(mf, rohf.ROHF):
102        def get_occ(mo_energy, mo_coeff=None):
103            nocca, noccb = mol.nelec
104            mo_occa, frac_lsta, homoa, lumoa = guess_occ(mo_energy, nocca)
105            mo_occb, frac_lstb, homob, lumob = guess_occ(mo_energy, noccb)
106
107            if abs(homoa - lumoa) < tol or abs(homob - lumob) < tol:
108                mo_occ = mo_occa + mo_occb
109                if len(frac_lstb):
110                    logger.warn(mf, 'fraction occ = %6g for alpha orbitals %s  '
111                                '%6g for beta orbitals %s',
112                                mo_occa[frac_lsta[0]], frac_lsta,
113                                mo_occb[frac_lstb[0]], frac_lstb)
114                    logger.info(mf, '  HOMO = %.12g  LUMO = %.12g', homoa, lumoa)
115                    logger.debug(mf, '  mo_energy = %s', mo_energy)
116                else:
117                    logger.warn(mf, 'fraction occ = %6g for alpha orbitals %s  ',
118                                mo_occa[frac_lsta[0]], frac_lsta)
119                    logger.info(mf, '  HOMO = %.12g  LUMO = %.12g', homoa, lumoa)
120                    logger.debug(mf, '  mo_energy = %s', mo_energy)
121            else:
122                mo_occ = old_get_occ(mo_energy, mo_coeff)
123            return mo_occ
124
125        def get_grad(mo_coeff, mo_occ, fock):
126            occidxa = mo_occ > 0
127            occidxb = mo_occ > 1
128            viridxa = ~occidxa
129            viridxb = ~occidxb
130            uniq_var_a = viridxa.reshape(-1,1) & occidxa
131            uniq_var_b = viridxb.reshape(-1,1) & occidxb
132
133            if getattr(fock, 'focka', None) is not None:
134                focka = fock.focka
135                fockb = fock.fockb
136            elif getattr(fock, 'ndim', None) == 3:
137                focka, fockb = fock
138            else:
139                focka = fockb = fock
140            focka = reduce(numpy.dot, (mo_coeff.T.conj(), focka, mo_coeff))
141            fockb = reduce(numpy.dot, (mo_coeff.T.conj(), fockb, mo_coeff))
142
143            g = numpy.zeros_like(focka)
144            g[uniq_var_a]  = focka[uniq_var_a]
145            g[uniq_var_b] += fockb[uniq_var_b]
146            return g[uniq_var_a | uniq_var_b]
147
148    else:  # RHF
149        def get_occ(mo_energy, mo_coeff=None):
150            nocc = (mol.nelectron+1) // 2  # n_docc + n_socc
151            mo_occ, frac_lst, homo, lumo = guess_occ(mo_energy, nocc)
152            n_docc = mol.nelectron // 2
153            n_socc = nocc - n_docc
154            if abs(homo - lumo) < tol or n_socc:
155                mo_occ *= 2
156                degen = len(frac_lst)
157                mo_occ[frac_lst] -= float(n_socc) / degen
158                logger.warn(mf, 'fraction occ = %6g  for orbitals %s',
159                            mo_occ[frac_lst[0]], frac_lst)
160                logger.info(mf, 'HOMO = %.12g  LUMO = %.12g', homo, lumo)
161                logger.debug(mf, '  mo_energy = %s', mo_energy)
162            else:
163                mo_occ = old_get_occ(mo_energy, mo_coeff)
164            return mo_occ
165
166    mf.get_occ = get_occ
167    if get_grad is not None:
168        mf.get_grad = get_grad
169    return mf
170frac_occ = frac_occ_
171
172def dynamic_occ_(mf, tol=1e-3):
173    assert(isinstance(mf, hf.RHF))
174    old_get_occ = mf.get_occ
175    def get_occ(mo_energy, mo_coeff=None):
176        mol = mf.mol
177        nocc = mol.nelectron // 2
178        sort_mo_energy = numpy.sort(mo_energy)
179        lumo = sort_mo_energy[nocc]
180        if abs(sort_mo_energy[nocc-1] - lumo) < tol:
181            mo_occ = numpy.zeros_like(mo_energy)
182            mo_occ[mo_energy<lumo] = 2
183            lst = abs(mo_energy - lumo) < tol
184            mo_occ[lst] = 0
185            logger.warn(mf, 'set charge = %d', mol.charge+int(lst.sum())*2)
186            logger.info(mf, 'HOMO = %.12g  LUMO = %.12g',
187                        sort_mo_energy[nocc-1], sort_mo_energy[nocc])
188            logger.debug(mf, '  mo_energy = %s', sort_mo_energy)
189        else:
190            mo_occ = old_get_occ(mo_energy, mo_coeff)
191        return mo_occ
192    mf.get_occ = get_occ
193    return mf
194dynamic_occ = dynamic_occ_
195
196def dynamic_level_shift_(mf, factor=1.):
197    '''Dynamically change the level shift in each SCF cycle.  The level shift
198    value is set to (HF energy change * factor)
199    '''
200    old_get_fock = mf.get_fock
201    last_e = [None]
202    def get_fock(h1e, s1e, vhf, dm, cycle=-1, diis=None,
203                 diis_start_cycle=None, level_shift_factor=None, damp_factor=None):
204        if cycle >= 0 or diis is not None:
205            ehf =(numpy.einsum('ij,ji', h1e, dm) +
206                  numpy.einsum('ij,ji', vhf, dm) * .5)
207            if last_e[0] is not None:
208                level_shift_factor = abs(ehf-last_e[0]) * factor
209                logger.info(mf, 'Set level shift to %g', level_shift_factor)
210            last_e[0] = ehf
211        return old_get_fock(h1e, s1e, vhf, dm, cycle, diis, diis_start_cycle,
212                            level_shift_factor, damp_factor)
213    mf.get_fock = get_fock
214    return mf
215dynamic_level_shift = dynamic_level_shift_
216
217def float_occ_(mf):
218    '''
219    For UHF, allowing the Sz value being changed during SCF iteration.
220    Determine occupation of alpha and beta electrons based on energy spectrum
221    '''
222    from pyscf.scf import uhf
223    assert(isinstance(mf, uhf.UHF))
224    def get_occ(mo_energy, mo_coeff=None):
225        mol = mf.mol
226        ee = numpy.sort(numpy.hstack(mo_energy))
227        n_a = numpy.count_nonzero(mo_energy[0]<(ee[mol.nelectron-1]+1e-3))
228        n_b = mol.nelectron - n_a
229        if mf.nelec is None:
230            nelec = mf.mol.nelec
231        else:
232            nelec = mf.nelec
233        if n_a != nelec[0]:
234            logger.info(mf, 'change num. alpha/beta electrons '
235                        ' %d / %d -> %d / %d',
236                        nelec[0], nelec[1], n_a, n_b)
237            mf.nelec = (n_a, n_b)
238        return uhf.UHF.get_occ(mf, mo_energy, mo_coeff)
239    mf.get_occ = get_occ
240    return mf
241dynamic_sz_ = float_occ = float_occ_
242
243def follow_state_(mf, occorb=None):
244    occstat = [occorb]
245    old_get_occ = mf.get_occ
246    def get_occ(mo_energy, mo_coeff=None):
247        if occstat[0] is None:
248            mo_occ = old_get_occ(mo_energy, mo_coeff)
249        else:
250            mo_occ = numpy.zeros_like(mo_energy)
251            s = reduce(numpy.dot, (occstat[0].T, mf.get_ovlp(), mo_coeff))
252            nocc = mf.mol.nelectron // 2
253            #choose a subset of mo_coeff, which maximizes <old|now>
254            idx = numpy.argsort(numpy.einsum('ij,ij->j', s, s))
255            mo_occ[idx[-nocc:]] = 2
256            logger.debug(mf, '  mo_occ = %s', mo_occ)
257            logger.debug(mf, '  mo_energy = %s', mo_energy)
258        occstat[0] = mo_coeff[:,mo_occ>0]
259        return mo_occ
260    mf.get_occ = get_occ
261    return mf
262follow_state = follow_state_
263
264def mom_occ_(mf, occorb, setocc):
265    '''Use maximum overlap method to determine occupation number for each orbital in every
266    iteration. It can be applied to unrestricted HF/KS and restricted open-shell
267    HF/KS.'''
268    from pyscf.scf import uhf, rohf
269    if isinstance(mf, uhf.UHF):
270        coef_occ_a = occorb[0][:, setocc[0]>0]
271        coef_occ_b = occorb[1][:, setocc[1]>0]
272    elif isinstance(mf, rohf.ROHF):
273        if mf.mol.spin != (numpy.sum(setocc[0]) - numpy.sum(setocc[1])):
274            raise ValueError('Wrong occupation setting for restricted open-shell calculation.')
275        coef_occ_a = occorb[:, setocc[0]>0]
276        coef_occ_b = occorb[:, setocc[1]>0]
277    else:
278        raise RuntimeError('Cannot support this class of instance %s' % mf)
279    log = logger.Logger(mf.stdout, mf.verbose)
280    def get_occ(mo_energy=None, mo_coeff=None):
281        if mo_energy is None: mo_energy = mf.mo_energy
282        if mo_coeff is None: mo_coeff = mf.mo_coeff
283        if isinstance(mf, rohf.ROHF): mo_coeff = numpy.array([mo_coeff, mo_coeff])
284        mo_occ = numpy.zeros_like(setocc)
285        nocc_a = int(numpy.sum(setocc[0]))
286        nocc_b = int(numpy.sum(setocc[1]))
287        s_a = reduce(numpy.dot, (coef_occ_a.T, mf.get_ovlp(), mo_coeff[0]))
288        s_b = reduce(numpy.dot, (coef_occ_b.T, mf.get_ovlp(), mo_coeff[1]))
289        #choose a subset of mo_coeff, which maximizes <old|now>
290        idx_a = numpy.argsort(numpy.einsum('ij,ij->j', s_a, s_a))[::-1]
291        idx_b = numpy.argsort(numpy.einsum('ij,ij->j', s_b, s_b))[::-1]
292        mo_occ[0][idx_a[:nocc_a]] = 1.
293        mo_occ[1][idx_b[:nocc_b]] = 1.
294
295        log.debug(' New alpha occ pattern: %s', mo_occ[0])
296        log.debug(' New beta occ pattern: %s', mo_occ[1])
297        if isinstance(mf.mo_energy, numpy.ndarray) and mf.mo_energy.ndim == 1:
298            log.debug1(' Current mo_energy(sorted) = %s', mo_energy)
299        else:
300            log.debug1(' Current alpha mo_energy(sorted) = %s', mo_energy[0])
301            log.debug1(' Current beta mo_energy(sorted) = %s', mo_energy[1])
302
303        if (int(numpy.sum(mo_occ[0])) != nocc_a):
304            log.error('mom alpha electron occupation numbers do not match: %d, %d',
305                      nocc_a, int(numpy.sum(mo_occ[0])))
306        if (int(numpy.sum(mo_occ[1])) != nocc_b):
307            log.error('mom beta electron occupation numbers do not match: %d, %d',
308                      nocc_b, int(numpy.sum(mo_occ[1])))
309
310        #output 1-dimension occupation number for restricted open-shell
311        if isinstance(mf, rohf.ROHF): mo_occ = mo_occ[0, :] + mo_occ[1, :]
312        return mo_occ
313    mf.get_occ = get_occ
314    return mf
315mom_occ = mom_occ_
316
317def project_mo_nr2nr(mol1, mo1, mol2):
318    r''' Project orbital coefficients from basis set 1 (C1 for mol1) to basis
319    set 2 (C2 for mol2).
320
321    .. math::
322
323        |\psi1\rangle = |AO1\rangle C1
324
325        |\psi2\rangle = P |\psi1\rangle = |AO2\rangle S^{-1}\langle AO2| AO1\rangle> C1 = |AO2\rangle> C2
326
327        C2 = S^{-1}\langle AO2|AO1\rangle C1
328
329    There are three relevant functions:
330    :func:`project_mo_nr2nr` is the projection for non-relativistic (scalar) basis.
331    :func:`project_mo_nr2r` projects from non-relativistic to relativistic basis.
332    :func:`project_mo_r2r`  is the projection between relativistic (spinor) basis.
333    '''
334    s22 = mol2.intor_symmetric('int1e_ovlp')
335    s21 = mole.intor_cross('int1e_ovlp', mol2, mol1)
336    if isinstance(mo1, numpy.ndarray) and mo1.ndim == 2:
337        return lib.cho_solve(s22, numpy.dot(s21, mo1), strict_sym_pos=False)
338    else:
339        return [lib.cho_solve(s22, numpy.dot(s21, x), strict_sym_pos=False)
340                for x in mo1]
341
342@lib.with_doc(project_mo_nr2nr.__doc__)
343def project_mo_nr2r(mol1, mo1, mol2):
344    assert(not mol1.cart)
345    s22 = mol2.intor_symmetric('int1e_ovlp_spinor')
346    s21 = mole.intor_cross('int1e_ovlp_sph', mol2, mol1)
347
348    ua, ub = mol2.sph2spinor_coeff()
349    s21 = numpy.dot(ua.T.conj(), s21) + numpy.dot(ub.T.conj(), s21) # (*)
350    # mo2: alpha, beta have been summed in Eq. (*)
351    # so DM = mo2[:,:nocc] * 1 * mo2[:,:nocc].H
352    if isinstance(mo1, numpy.ndarray) and mo1.ndim == 2:
353        mo2 = numpy.dot(s21, mo1)
354        return lib.cho_solve(s22, mo2, strict_sym_pos=False)
355    else:
356        return [lib.cho_solve(s22, numpy.dot(s21, x), strict_sym_pos=False)
357                for x in mo1]
358
359@lib.with_doc(project_mo_nr2nr.__doc__)
360def project_mo_r2r(mol1, mo1, mol2):
361    s22 = mol2.intor_symmetric('int1e_ovlp_spinor')
362    t22 = mol2.intor_symmetric('int1e_spsp_spinor')
363    s21 = mole.intor_cross('int1e_ovlp_spinor', mol2, mol1)
364    t21 = mole.intor_cross('int1e_spsp_spinor', mol2, mol1)
365    n2c = s21.shape[1]
366    pl = lib.cho_solve(s22, s21, strict_sym_pos=False)
367    ps = lib.cho_solve(t22, t21, strict_sym_pos=False)
368    if isinstance(mo1, numpy.ndarray) and mo1.ndim == 2:
369        return numpy.vstack((numpy.dot(pl, mo1[:n2c]),
370                             numpy.dot(ps, mo1[n2c:])))
371    else:
372        return [numpy.vstack((numpy.dot(pl, x[:n2c]),
373                              numpy.dot(ps, x[n2c:]))) for x in mo1]
374
375def project_dm_nr2nr(mol1, dm1, mol2):
376    r''' Project density matrix representation from basis set 1 (mol1) to basis
377    set 2 (mol2).
378
379    .. math::
380
381        |AO2\rangle DM_AO2 \langle AO2|
382
383        = |AO2\rangle P DM_AO1 P \langle AO2|
384
385        DM_AO2 = P DM_AO1 P
386
387        P = S_{AO2}^{-1}\langle AO2|AO1\rangle
388
389    There are three relevant functions:
390    :func:`project_dm_nr2nr` is the projection for non-relativistic (scalar) basis.
391    :func:`project_dm_nr2r` projects from non-relativistic to relativistic basis.
392    :func:`project_dm_r2r`  is the projection between relativistic (spinor) basis.
393    '''
394    s22 = mol2.intor_symmetric('int1e_ovlp')
395    s21 = mole.intor_cross('int1e_ovlp', mol2, mol1)
396    p21 = lib.cho_solve(s22, s21, strict_sym_pos=False)
397    if isinstance(dm1, numpy.ndarray) and dm1.ndim == 2:
398        return reduce(numpy.dot, (p21, dm1, p21.conj().T))
399    else:
400        return lib.einsum('pi,nij,qj->npq', p21, dm1, p21.conj())
401
402@lib.with_doc(project_dm_nr2nr.__doc__)
403def project_dm_nr2r(mol1, dm1, mol2):
404    assert(not mol1.cart)
405    s22 = mol2.intor_symmetric('int1e_ovlp_spinor')
406    s21 = mole.intor_cross('int1e_ovlp_sph', mol2, mol1)
407
408    ua, ub = mol2.sph2spinor_coeff()
409    s21 = numpy.dot(ua.T.conj(), s21) + numpy.dot(ub.T.conj(), s21) # (*)
410    # mo2: alpha, beta have been summed in Eq. (*)
411    # so DM = mo2[:,:nocc] * 1 * mo2[:,:nocc].H
412    p21 = lib.cho_solve(s22, s21, strict_sym_pos=False)
413    if isinstance(dm1, numpy.ndarray) and dm1.ndim == 2:
414        return reduce(numpy.dot, (p21, dm1, p21.conj().T))
415    else:
416        return lib.einsum('pi,nij,qj->npq', p21, dm1, p21.conj())
417
418@lib.with_doc(project_dm_nr2nr.__doc__)
419def project_dm_r2r(mol1, dm1, mol2):
420    s22 = mol2.intor_symmetric('int1e_ovlp_spinor')
421    t22 = mol2.intor_symmetric('int1e_spsp_spinor')
422    s21 = mole.intor_cross('int1e_ovlp_spinor', mol2, mol1)
423    t21 = mole.intor_cross('int1e_spsp_spinor', mol2, mol1)
424    pl = lib.cho_solve(s22, s21, strict_sym_pos=False)
425    ps = lib.cho_solve(t22, t21, strict_sym_pos=False)
426    p21 = scipy.linalg.block_diag(pl, ps)
427    if isinstance(dm1, numpy.ndarray) and dm1.ndim == 2:
428        return reduce(numpy.dot, (p21, dm1, p21.conj().T))
429    else:
430        return lib.einsum('pi,nij,qj->npq', p21, dm1, p21.conj())
431
432def canonical_orth_(S, thr=1e-7):
433    '''Löwdin's canonical orthogonalization'''
434    # Ensure the basis functions are normalized (symmetry-adapted ones are not!)
435    normlz = numpy.power(numpy.diag(S), -0.5)
436    Snorm = numpy.dot(numpy.diag(normlz), numpy.dot(S, numpy.diag(normlz)))
437    # Form vectors for normalized overlap matrix
438    Sval, Svec = numpy.linalg.eigh(Snorm)
439    X = Svec[:,Sval>=thr] / numpy.sqrt(Sval[Sval>=thr])
440    # Plug normalization back in
441    X = numpy.dot(numpy.diag(normlz), X)
442    return X
443
444def partial_cholesky_orth_(S, canthr=1e-7, cholthr=1e-9):
445    '''Partial Cholesky orthogonalization for curing overcompleteness.
446
447    References:
448
449    Susi Lehtola, Curing basis set overcompleteness with pivoted
450    Cholesky decompositions, J. Chem. Phys. 151, 241102 (2019),
451    doi:10.1063/1.5139948.
452
453    Susi Lehtola, Accurate reproduction of strongly repulsive
454    interatomic potentials, Phys. Rev. A 101, 032504 (2020),
455    doi:10.1103/PhysRevA.101.032504.
456    '''
457    # Ensure the basis functions are normalized
458    normlz = numpy.power(numpy.diag(S), -0.5)
459    Snorm = numpy.dot(numpy.diag(normlz), numpy.dot(S, numpy.diag(normlz)))
460
461    # Sort the basis functions according to the Gershgorin circle
462    # theorem so that the Cholesky routine is well-initialized
463    odS = numpy.abs(Snorm)
464    numpy.fill_diagonal(odS, 0.0)
465    odSs = numpy.sum(odS, axis=0)
466    sortidx = numpy.argsort(odSs)
467
468    # Run the pivoted Cholesky decomposition
469    Ssort = Snorm[numpy.ix_(sortidx, sortidx)].copy()
470    c, piv, r_c = pivoted_cholesky(Ssort, tol=cholthr)
471    # The functions we're going to use are given by the pivot as
472    idx = sortidx[piv[:r_c]]
473
474    # Get the (un-normalized) sub-basis
475    Ssub = S[numpy.ix_(idx, idx)].copy()
476    # Orthogonalize sub-basis
477    Xsub = canonical_orth_(Ssub, thr=canthr)
478
479    # Full X
480    X = numpy.zeros((S.shape[0], Xsub.shape[1]))
481    X[idx,:] = Xsub
482
483    return X
484
485def remove_linear_dep_(mf, threshold=LINEAR_DEP_THRESHOLD,
486                       lindep=LINEAR_DEP_TRIGGER,
487                       cholesky_threshold=CHOLESKY_THRESHOLD,
488                       force_pivoted_cholesky=FORCE_PIVOTED_CHOLESKY):
489    '''
490    Args:
491        threshold : float
492            The threshold under which the eigenvalues of the overlap matrix are
493            discarded to avoid numerical instability.
494        lindep : float
495            The threshold that triggers the special treatment of the linear
496            dependence issue.
497    '''
498    s = mf.get_ovlp()
499    cond = numpy.max(lib.cond(s))
500    if cond < 1./lindep and not force_pivoted_cholesky:
501        return mf
502
503    logger.info(mf, 'Applying remove_linear_dep_ on SCF object.')
504    logger.debug(mf, 'Overlap condition number %g', cond)
505    if(cond < 1./numpy.finfo(s.dtype).eps and not force_pivoted_cholesky):
506        logger.info(mf, 'Using canonical orthogonalization with threshold {}'.format(threshold))
507        def eigh(h, s):
508            x = canonical_orth_(s, threshold)
509            xhx = reduce(numpy.dot, (x.T.conj(), h, x))
510            e, c = numpy.linalg.eigh(xhx)
511            c = numpy.dot(x, c)
512            return e, c
513        mf._eigh = eigh
514    else:
515        logger.info(mf, 'Using partial Cholesky orthogonalization '
516                    '(doi:10.1063/1.5139948, doi:10.1103/PhysRevA.101.032504)')
517        logger.info(mf, 'Using threshold {} for pivoted Cholesky'.format(cholesky_threshold))
518        logger.info(mf, 'Using threshold {} to orthogonalize the subbasis'.format(threshold))
519        def eigh(h, s):
520            x = partial_cholesky_orth_(s, canthr=threshold, cholthr=cholesky_threshold)
521            xhx = reduce(numpy.dot, (x.T.conj(), h, x))
522            e, c = numpy.linalg.eigh(xhx)
523            c = numpy.dot(x, c)
524            return e, c
525        mf._eigh = eigh
526    return mf
527remove_linear_dep = remove_linear_dep_
528
529def convert_to_uhf(mf, out=None, remove_df=False):
530    '''Convert the given mean-field object to the unrestricted HF/KS object
531
532    Note this conversion only changes the class of the mean-field object.
533    The total energy and wave-function are the same as them in the input
534    mf object. If mf is an second order SCF (SOSCF) object, the SOSCF layer
535    will be discarded. Its underlying SCF object mf._scf will be converted.
536
537    Args:
538        mf : SCF object
539
540    Kwargs
541        remove_df : bool
542            Whether to convert the DF-SCF object to the normal SCF object.
543            This conversion is not applied by default.
544
545    Returns:
546        An unrestricted SCF object
547    '''
548    from pyscf import scf
549    from pyscf import dft
550    assert(isinstance(mf, hf.SCF))
551
552    logger.debug(mf, 'Converting %s to UHF', mf.__class__)
553
554    def update_mo_(mf, mf1):
555        if mf.mo_energy is not None:
556            if isinstance(mf, scf.uhf.UHF):
557                mf1.mo_occ = mf.mo_occ
558                mf1.mo_coeff = mf.mo_coeff
559                mf1.mo_energy = mf.mo_energy
560            elif getattr(mf, 'kpts', None) is None:  # RHF/ROHF
561                mf1.mo_occ = numpy.array((mf.mo_occ>0, mf.mo_occ==2), dtype=numpy.double)
562                # ROHF orbital energies, not canonical UHF orbital energies
563                mo_ea = getattr(mf.mo_energy, 'mo_ea', mf.mo_energy)
564                mo_eb = getattr(mf.mo_energy, 'mo_eb', mf.mo_energy)
565                mf1.mo_energy = (mo_ea, mo_eb)
566                mf1.mo_coeff = (mf.mo_coeff, mf.mo_coeff)
567            else:  # This to handle KRHF object
568                mf1.mo_occ = ([numpy.asarray(occ> 0, dtype=numpy.double)
569                               for occ in mf.mo_occ],
570                              [numpy.asarray(occ==2, dtype=numpy.double)
571                               for occ in mf.mo_occ])
572                mo_ea = getattr(mf.mo_energy, 'mo_ea', mf.mo_energy)
573                mo_eb = getattr(mf.mo_energy, 'mo_eb', mf.mo_energy)
574                mf1.mo_energy = (mo_ea, mo_eb)
575                mf1.mo_coeff = (mf.mo_coeff, mf.mo_coeff)
576        return mf1
577
578    if isinstance(mf, scf.ghf.GHF):
579        raise NotImplementedError
580
581    elif out is not None:
582        assert(isinstance(out, scf.uhf.UHF))
583        out = _update_mf_without_soscf(mf, out, remove_df)
584
585    elif isinstance(mf, scf.uhf.UHF):
586        # Remove with_df for SOSCF method because the post-HF code checks the
587        # attribute .with_df to identify whether an SCF object is DF-SCF method.
588        # with_df in SOSCF is used in orbital hessian approximation only.  For the
589        # returned SCF object, whehter with_df exists in SOSCF has no effects on the
590        # mean-field energy and other properties.
591        if getattr(mf, '_scf', None):
592            return _update_mf_without_soscf(mf, copy.copy(mf._scf), remove_df)
593        else:
594            return copy.copy(mf)
595
596    else:
597        known_cls = {scf.hf.RHF        : scf.uhf.UHF,
598                     scf.rohf.ROHF     : scf.uhf.UHF,
599                     scf.hf_symm.RHF   : scf.uhf_symm.UHF,
600                     scf.hf_symm.ROHF  : scf.uhf_symm.UHF,
601                     dft.rks.RKS       : dft.uks.UKS,
602                     dft.roks.ROKS     : dft.uks.UKS,
603                     dft.rks_symm.RKS  : dft.uks_symm.UKS,
604                     dft.rks_symm.ROKS : dft.uks_symm.UKS}
605        out = _object_without_soscf(mf, known_cls, remove_df)
606
607    return update_mo_(mf, out)
608
609def _object_without_soscf(mf, known_class, remove_df=False):
610    from pyscf.soscf import newton_ah
611    sub_classes = []
612    obj = None
613    for i, cls in enumerate(mf.__class__.__mro__):
614        if cls in known_class:
615            obj = known_class[cls](mf.mol)
616            break
617        else:
618            sub_classes.append(cls)
619
620    if obj is None:
621        raise NotImplementedError(
622            "Incompatible object types. Mean-field `mf` class not found in "
623            "`known_class` type.\n\nmf = '%s'\n\nknown_class = '%s'" %
624            (mf.__class__.__mro__, known_class))
625
626    if isinstance(mf, newton_ah._CIAH_SOSCF):
627        remove_df = (remove_df or
628                     # The main SCF object is not a DFHF object
629                     not getattr(mf._scf, 'with_df', None))
630
631# Mimic the initialization procedure to restore the Hamiltonian
632    for cls in reversed(sub_classes):
633        class_name = cls.__name__
634        if '_DFHF' in class_name:
635            if not remove_df:
636                obj = obj.density_fit()
637        elif '_SGXHF' in class_name:
638            if not remove_df:
639                obj = obj.COSX()
640        elif '_X2C_SCF' in class_name:
641            obj = obj.x2c()
642        elif 'WithSolvent' in class_name:
643            obj = obj.ddCOSMO(mf.with_solvent)
644        elif 'QMMM' in class_name and getattr(mf, 'mm_mol', None):
645            from pyscf.qmmm.itrf import qmmm_for_scf
646            obj = qmmm_for_scf(obj, mf.mm_mol)
647        elif '_DFTD3' in class_name:
648            from pyscf.dftd3.itrf import dftd3
649            obj = dftd3(obj)
650    return _update_mf_without_soscf(mf, obj, remove_df)
651
652def _update_mf_without_soscf(mf, out, remove_df=False):
653    from pyscf.soscf import newton_ah
654    mf_dic = dict(mf.__dict__)
655
656    # if mf is SOSCF object, avoid to overwrite the with_df method
657    # FIXME: it causes bug when converting pbc-SOSCF.
658    if isinstance(mf, newton_ah._CIAH_SOSCF):
659        mf_dic.pop('with_df', None)
660
661    out.__dict__.update(mf_dic)
662
663    if remove_df and getattr(out, 'with_df', None):
664        delattr(out, 'with_df')
665    return out
666
667def convert_to_rhf(mf, out=None, remove_df=False):
668    '''Convert the given mean-field object to the restricted HF/KS object
669
670    Note this conversion only changes the class of the mean-field object.
671    The total energy and wave-function are the same as them in the input
672    mf object. If mf is an second order SCF (SOSCF) object, the SOSCF layer
673    will be discarded. Its underlying SCF object mf._scf will be converted.
674
675    Args:
676        mf : SCF object
677
678    Kwargs
679        remove_df : bool
680            Whether to convert the DF-SCF object to the normal SCF object.
681            This conversion is not applied by default.
682
683    Returns:
684        An unrestricted SCF object
685    '''
686    from pyscf import scf
687    from pyscf import dft
688    assert(isinstance(mf, hf.SCF))
689
690    logger.debug(mf, 'Converting %s to RHF', mf.__class__)
691
692    def update_mo_(mf, mf1):
693        if mf.mo_energy is not None:
694            if isinstance(mf, scf.hf.RHF): # RHF/ROHF/KRHF/KROHF
695                mf1.mo_occ = mf.mo_occ
696                mf1.mo_coeff = mf.mo_coeff
697                mf1.mo_energy = mf.mo_energy
698            elif getattr(mf, 'kpts', None) is None:  # UHF
699                mf1.mo_occ = mf.mo_occ[0] + mf.mo_occ[1]
700                mf1.mo_energy = mf.mo_energy[0]
701                mf1.mo_coeff =  mf.mo_coeff[0]
702                if getattr(mf.mo_coeff[0], 'orbsym', None) is not None:
703                    mf1.mo_coeff = lib.tag_array(mf1.mo_coeff, orbsym=mf.mo_coeff[0].orbsym)
704            else:  # KUHF
705                mf1.mo_occ = [occa+occb for occa, occb in zip(*mf.mo_occ)]
706                mf1.mo_energy = mf.mo_energy[0]
707                mf1.mo_coeff =  mf.mo_coeff[0]
708        return mf1
709
710    if getattr(mf, 'nelec', None) is None:
711        nelec = mf.mol.nelec
712    else:
713        nelec = mf.nelec
714
715    if isinstance(mf, scf.ghf.GHF):
716        raise NotImplementedError
717
718    elif out is not None:
719        assert(isinstance(out, scf.hf.RHF))
720        out = _update_mf_without_soscf(mf, out, remove_df)
721
722    elif (isinstance(mf, scf.hf.RHF) or
723          (nelec[0] != nelec[1] and isinstance(mf, scf.rohf.ROHF))):
724        if getattr(mf, '_scf', None):
725            return _update_mf_without_soscf(mf, copy.copy(mf._scf), remove_df)
726        else:
727            return copy.copy(mf)
728
729    else:
730        if nelec[0] == nelec[1]:
731            known_cls = {scf.uhf.UHF      : scf.hf.RHF      ,
732                         scf.uhf_symm.UHF : scf.hf_symm.RHF ,
733                         dft.uks.UKS      : dft.rks.RKS     ,
734                         dft.uks_symm.UKS : dft.rks_symm.RKS,
735                         scf.rohf.ROHF    : scf.hf.RHF      ,
736                         scf.hf_symm.ROHF : scf.hf_symm.RHF ,
737                         dft.roks.ROKS    : dft.rks.RKS     ,
738                         dft.rks_symm.ROKS: dft.rks_symm.RKS}
739        else:
740            known_cls = {scf.uhf.UHF      : scf.rohf.ROHF    ,
741                         scf.uhf_symm.UHF : scf.hf_symm.ROHF ,
742                         dft.uks.UKS      : dft.roks.ROKS    ,
743                         dft.uks_symm.UKS : dft.rks_symm.ROKS}
744        out = _object_without_soscf(mf, known_cls, remove_df)
745
746    return update_mo_(mf, out)
747
748def convert_to_ghf(mf, out=None, remove_df=False):
749    '''Convert the given mean-field object to the generalized HF/KS object
750
751    Note this conversion only changes the class of the mean-field object.
752    The total energy and wave-function are the same as them in the input
753    mf object. If mf is an second order SCF (SOSCF) object, the SOSCF layer
754    will be discarded. Its underlying SCF object mf._scf will be converted.
755
756    Args:
757        mf : SCF object
758
759    Kwargs
760        remove_df : bool
761            Whether to convert the DF-SCF object to the normal SCF object.
762            This conversion is not applied by default.
763
764    Returns:
765        An generalized SCF object
766    '''
767    from pyscf import scf
768    from pyscf import dft
769    assert(isinstance(mf, hf.SCF))
770
771    logger.debug(mf, 'Converting %s to GHF', mf.__class__)
772
773    def update_mo_(mf, mf1):
774        if mf.mo_energy is not None:
775            if isinstance(mf, scf.hf.RHF): # RHF
776                nao, nmo = mf.mo_coeff.shape
777                orbspin = get_ghf_orbspin(mf.mo_energy, mf.mo_occ, True)
778
779                mf1.mo_energy = numpy.empty(nmo*2)
780                mf1.mo_energy[orbspin==0] = mf.mo_energy
781                mf1.mo_energy[orbspin==1] = mf.mo_energy
782                mf1.mo_occ = numpy.empty(nmo*2)
783                mf1.mo_occ[orbspin==0] = mf.mo_occ > 0
784                mf1.mo_occ[orbspin==1] = mf.mo_occ == 2
785
786                mo_coeff = numpy.zeros((nao*2,nmo*2), dtype=mf.mo_coeff.dtype)
787                mo_coeff[:nao,orbspin==0] = mf.mo_coeff
788                mo_coeff[nao:,orbspin==1] = mf.mo_coeff
789                if getattr(mf.mo_coeff, 'orbsym', None) is not None:
790                    orbsym = numpy.zeros_like(orbspin)
791                    orbsym[orbspin==0] = mf.mo_coeff.orbsym
792                    orbsym[orbspin==1] = mf.mo_coeff.orbsym
793                    mo_coeff = lib.tag_array(mo_coeff, orbsym=orbsym)
794                mf1.mo_coeff = lib.tag_array(mo_coeff, orbspin=orbspin)
795
796            else: # UHF
797                nao, nmo = mf.mo_coeff[0].shape
798                orbspin = get_ghf_orbspin(mf.mo_energy, mf.mo_occ, False)
799
800                mf1.mo_energy = numpy.empty(nmo*2)
801                mf1.mo_energy[orbspin==0] = mf.mo_energy[0]
802                mf1.mo_energy[orbspin==1] = mf.mo_energy[1]
803                mf1.mo_occ = numpy.empty(nmo*2)
804                mf1.mo_occ[orbspin==0] = mf.mo_occ[0]
805                mf1.mo_occ[orbspin==1] = mf.mo_occ[1]
806
807                mo_coeff = numpy.zeros((nao*2,nmo*2), dtype=mf.mo_coeff[0].dtype)
808                mo_coeff[:nao,orbspin==0] = mf.mo_coeff[0]
809                mo_coeff[nao:,orbspin==1] = mf.mo_coeff[1]
810                if getattr(mf.mo_coeff[0], 'orbsym', None) is not None:
811                    orbsym = numpy.zeros_like(orbspin)
812                    orbsym[orbspin==0] = mf.mo_coeff[0].orbsym
813                    orbsym[orbspin==1] = mf.mo_coeff[1].orbsym
814                    mo_coeff = lib.tag_array(mo_coeff, orbsym=orbsym)
815                mf1.mo_coeff = lib.tag_array(mo_coeff, orbspin=orbspin)
816        return mf1
817
818    if out is not None:
819        assert(isinstance(out, scf.ghf.GHF))
820        out = _update_mf_without_soscf(mf, out, remove_df)
821
822    elif isinstance(mf, scf.ghf.GHF):
823        if getattr(mf, '_scf', None):
824            return _update_mf_without_soscf(mf, copy.copy(mf._scf), remove_df)
825        else:
826            return copy.copy(mf)
827
828    else:
829        known_cls = {scf.hf.RHF        : scf.ghf.GHF,
830                     scf.rohf.ROHF     : scf.ghf.GHF,
831                     scf.uhf.UHF       : scf.ghf.GHF,
832                     scf.hf_symm.RHF   : scf.ghf_symm.GHF,
833                     scf.hf_symm.ROHF  : scf.ghf_symm.GHF,
834                     scf.uhf_symm.UHF  : scf.ghf_symm.GHF,
835                     dft.rks.RKS       : dft.gks.GKS,
836                     dft.roks.ROKS     : dft.gks.GKS,
837                     dft.uks.UKS       : dft.gks.GKS,
838                     dft.rks_symm.RKS  : dft.gks_symm.GKS,
839                     dft.rks_symm.ROKS : dft.gks_symm.GKS,
840                     dft.uks_symm.UKS  : dft.gks_symm.GKS}
841        out = _object_without_soscf(mf, known_cls, remove_df)
842
843    return update_mo_(mf, out)
844
845def get_ghf_orbspin(mo_energy, mo_occ, is_rhf=None):
846    '''Spin of each GHF orbital when the GHF orbitals are converted from
847    RHF/UHF orbitals
848
849    For RHF orbitals, the orbspin corresponds to first occupied orbitals then
850    unoccupied orbitals.  In the occupied orbital space, if degenerated, first
851    alpha then beta, last the (open-shell) singly occupied (alpha) orbitals. In
852    the unoccupied orbital space, first the (open-shell) unoccupied (beta)
853    orbitals if applicable, then alpha and beta orbitals
854
855    For UHF orbitals, the orbspin corresponds to first occupied orbitals then
856    unoccupied orbitals.
857    '''
858    if is_rhf is None:  # guess whether the orbitals are RHF orbitals
859        is_rhf = mo_energy[0].ndim == 0
860
861    if is_rhf:
862        nmo = mo_energy.size
863        nocc = numpy.count_nonzero(mo_occ >0)
864        nvir = nmo - nocc
865        ndocc = numpy.count_nonzero(mo_occ==2)
866        nsocc = nocc - ndocc
867        orbspin = numpy.array([0,1]*ndocc + [0]*nsocc + [1]*nsocc + [0,1]*nvir)
868    else:
869        nmo = mo_energy[0].size
870        nocca = numpy.count_nonzero(mo_occ[0]>0)
871        nvira = nmo - nocca
872        noccb = numpy.count_nonzero(mo_occ[1]>0)
873        nvirb = nmo - noccb
874        # round(6) to avoid numerical uncertainty in degeneracy
875        es = numpy.append(mo_energy[0][mo_occ[0] >0],
876                          mo_energy[1][mo_occ[1] >0])
877        oidx = numpy.argsort(es.round(6))
878        es = numpy.append(mo_energy[0][mo_occ[0]==0],
879                          mo_energy[1][mo_occ[1]==0])
880        vidx = numpy.argsort(es.round(6))
881        orbspin = numpy.append(numpy.array([0]*nocca+[1]*noccb)[oidx],
882                               numpy.array([0]*nvira+[1]*nvirb)[vidx])
883    return orbspin
884
885del(LINEAR_DEP_THRESHOLD, LINEAR_DEP_TRIGGER)
886
887def fast_newton(mf, mo_coeff=None, mo_occ=None, dm0=None,
888                auxbasis=None, dual_basis=None, **newton_kwargs):
889    '''This is a wrap function which combines several operations. This
890    function first setup the initial guess
891    from density fitting calculation then use  for
892    Newton solver and call Newton solver.
893    Newton solver attributes [max_cycle_inner, max_stepsize, ah_start_tol,
894    ah_conv_tol, ah_grad_trust_region, ...] can be passed through **newton_kwargs.
895    '''
896    import copy
897    from pyscf.lib import logger
898    from pyscf import df
899    from pyscf.soscf import newton_ah
900    if auxbasis is None:
901        auxbasis = df.addons.aug_etb_for_dfbasis(mf.mol, 'ahlrichs', beta=2.5)
902    if dual_basis:
903        mf1 = mf.newton()
904        pmol = mf1.mol = newton_ah.project_mol(mf.mol, dual_basis)
905        mf1 = mf1.density_fit(auxbasis)
906    else:
907        mf1 = mf.newton().density_fit(auxbasis)
908    mf1.with_df._compatible_format = False
909    mf1.direct_scf_tol = 1e-7
910
911    if getattr(mf, 'grids', None):
912        from pyscf.dft import gen_grid
913        approx_grids = gen_grid.Grids(mf.mol)
914        approx_grids.verbose = 0
915        approx_grids.level = max(0, mf.grids.level-3)
916        mf1.grids = approx_grids
917
918        approx_numint = copy.copy(mf._numint)
919        mf1._numint = approx_numint
920    for key in newton_kwargs:
921        setattr(mf1, key, newton_kwargs[key])
922
923    if mo_coeff is None or mo_occ is None:
924        mo_coeff, mo_occ = mf.mo_coeff, mf.mo_occ
925
926    if dm0 is not None:
927        mo_coeff, mo_occ = mf1.from_dm(dm0)
928    elif mo_coeff is None or mo_occ is None:
929        logger.note(mf, '========================================================')
930        logger.note(mf, 'Generating initial guess with DIIS-SCF for newton solver')
931        logger.note(mf, '========================================================')
932        if dual_basis:
933            mf0 = copy.copy(mf)
934            mf0.mol = pmol
935            mf0 = mf0.density_fit(auxbasis)
936        else:
937            mf0 = mf.density_fit(auxbasis)
938        mf0.direct_scf_tol = 1e-7
939        mf0.conv_tol = 3.
940        mf0.conv_tol_grad = 1.
941        if mf0.level_shift == 0:
942            mf0.level_shift = .2
943        if getattr(mf, 'grids', None):
944            mf0.grids = approx_grids
945            mf0._numint = approx_numint
946# Note: by setting small_rho_cutoff, dft.get_veff function may overwrite
947# approx_grids and approx_numint.  It will further changes the corresponding
948# mf1 grids and _numint.  If inital guess dm0 or mo_coeff/mo_occ were given,
949# dft.get_veff are not executed so that more grid points may be found in
950# approx_grids.
951            mf0.small_rho_cutoff = mf.small_rho_cutoff * 10
952        mf0.kernel()
953        mf1.with_df = mf0.with_df
954        mo_coeff, mo_occ = mf0.mo_coeff, mf0.mo_occ
955
956        if dual_basis:
957            if mo_occ.ndim == 2:
958                mo_coeff =(project_mo_nr2nr(pmol, mo_coeff[0], mf.mol),
959                           project_mo_nr2nr(pmol, mo_coeff[1], mf.mol))
960            else:
961                mo_coeff = project_mo_nr2nr(pmol, mo_coeff, mf.mol)
962            mo_coeff, mo_occ = mf1.from_dm(mf.make_rdm1(mo_coeff,mo_occ))
963        mf0 = None
964
965        logger.note(mf, '============================')
966        logger.note(mf, 'Generating initial guess end')
967        logger.note(mf, '============================')
968
969    mf1.kernel(mo_coeff, mo_occ)
970    mf.converged = mf1.converged
971    mf.e_tot     = mf1.e_tot
972    mf.mo_energy = mf1.mo_energy
973    mf.mo_coeff  = mf1.mo_coeff
974    mf.mo_occ    = mf1.mo_occ
975
976#    mf = copy.copy(mf)
977#    def mf_kernel(*args, **kwargs):
978#        logger.warn(mf, "fast_newton is a wrap function to quickly setup and call Newton solver. "
979#                    "There's no need to call kernel function again for fast_newton.")
980#        del(mf.kernel)  # warn once and remove circular depdence
981#        return mf.e_tot
982#    mf.kernel = mf_kernel
983    return mf
984