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: Xing Zhang <zhangxing.nju@gmail.com>
17#
18
19'''
20Restricted coupled pertubed Hartree-Fock solver
21Modified from pyscf.scf.cphf
22'''
23
24
25import numpy as np
26from pyscf import lib
27from pyscf.lib import logger
28
29
30def solve(fvind, mo_energy, mo_occ, h1, s1=None,
31          max_cycle=20, tol=1e-9, hermi=False, verbose=logger.WARN):
32    '''
33    Args:
34        fvind : function
35            Given density matrix, compute (ij|kl)D_{lk}*2 - (ij|kl)D_{jk}
36
37    Kwargs:
38        hermi : boolean
39            Whether the matrix defined by fvind is Hermitian or not.
40    '''
41    if s1 is None:
42        return solve_nos1(fvind, mo_energy, mo_occ, h1,
43                          max_cycle, tol, hermi, verbose)
44    else:
45        return solve_withs1(fvind, mo_energy, mo_occ, h1, s1,
46                            max_cycle, tol, hermi, verbose)
47kernel = solve
48
49# h1 shape is (:,nvir,nocc)
50def solve_nos1(fvind, mo_energy, mo_occ, h1,
51               max_cycle=20, tol=1e-9, hermi=False, verbose=logger.WARN):
52    '''For field independent basis. First order overlap matrix is zero'''
53    log = logger.new_logger(verbose=verbose)
54    t0 = (logger.process_clock(), logger.perf_counter())
55
56    nkpt = len(h1)
57    moloc = np.zeros([nkpt+1], dtype=int)
58    for k in range(nkpt):
59        moloc[k+1] = moloc[k] + h1[k].size
60
61    occidx = []
62    viridx = []
63    for k in range(nkpt):
64        occidx.append(mo_occ[k] > 0)
65        viridx.append(mo_occ[k] == 0)
66
67    e_a = [mo_energy[k][viridx[k]] for k in range(nkpt)]
68    e_i = [mo_energy[k][occidx[k]] for k in range(nkpt)]
69    e_ai = [1 / lib.direct_sum('a-i->ai', e_a[k], e_i[k]) for k in range(nkpt)]
70    mo1base = []
71    for k in range(nkpt):
72        mo1base.append((h1[k] * -e_ai[k]).ravel())
73    mo1base = np.hstack(mo1base)
74
75    def vind_vo(mo1):
76        mo1 = mo1.flatten()
77        tmp = []
78        for k in range(nkpt):
79            tmp.append(mo1[moloc[k]:moloc[k+1]].reshape(h1[k].shape))
80        v = fvind(tmp)
81        for k in range(nkpt):
82            v[k] *= e_ai[k]
83            v[k] = v[k].ravel()
84        return np.hstack(v)
85
86    _mo1 = lib.krylov(vind_vo, mo1base,
87                      tol=tol, max_cycle=max_cycle, hermi=hermi, verbose=log).flatten()
88    log.timer('krylov solver in CPHF', *t0)
89    mo1 = []
90    for k in range(nkpt):
91        mo1.append(_mo1[moloc[k]:moloc[k+1]].reshape(h1[k].shape))
92    return mo1, None
93
94# h1 shape is (:,nocc+nvir,nocc)
95def solve_withs1(fvind, mo_energy, mo_occ, h1, s1,
96                 max_cycle=20, tol=1e-9, hermi=False, verbose=logger.WARN):
97    '''For field dependent basis. First order overlap matrix is non-zero.
98    The first order orbitals are set to
99    C^1_{ij} = -1/2 S1
100    e1 = h1 - s1*e0 + (e0_j-e0_i)*c1 + vhf[c1]
101
102    Kwargs:
103        hermi : boolean
104            Whether the matrix defined by fvind is Hermitian or not.
105
106    Returns:
107        First order orbital coefficients (in MO basis) and first order orbital
108        energy matrix
109    '''
110    log = logger.new_logger(verbose=verbose)
111    t0 = (logger.process_clock(), logger.perf_counter())
112
113    nkpt = len(h1)
114    ncomp = h1[0].shape[0]
115    occidx = []
116    viridx = []
117    for k in range(nkpt):
118        occidx.append(mo_occ[k] > 0)
119        viridx.append(mo_occ[k] == 0)
120
121    e_a = [mo_energy[k][viridx[k]] for k in range(nkpt)]
122    e_i = [mo_energy[k][occidx[k]] for k in range(nkpt)]
123    e_ai = [1 / lib.direct_sum('a-i->ai', e_a[k], e_i[k]) for k in range(nkpt)]
124    nocc = np.zeros([nkpt], dtype=int)
125    nvir = np.zeros([nkpt], dtype=int)
126    nmo  = np.zeros([nkpt], dtype=int)
127    moloc = np.zeros([nkpt+1], dtype=int)
128    for k in range(nkpt):
129        nvir_k, nocc_k = e_ai[k].shape
130        nmo_k = nvir_k + nocc_k
131        nvir[k] = nvir_k
132        nocc[k] = nocc_k
133        nmo[k]  = nmo_k
134        moloc[k+1] = moloc[k] + nmo_k * nocc_k * ncomp
135
136    mo1base = []
137    _mo1base = []
138    mo_e1 = []
139    for k in range(nkpt):
140        mo1base.append(h1[k] - s1[k] * e_i[k])
141        mo_e1.append(mo1base[k][:,occidx[k],:].copy())
142        mo1base[k][:,viridx[k]] *= -e_ai[k]
143        mo1base[k][:,occidx[k]] = -s1[k][:,occidx[k]] * .5
144        _mo1base.append(mo1base[k].ravel())
145    _mo1base = np.hstack(_mo1base)
146
147    def vind_vo(mo1):
148        mo1 = mo1.ravel()
149        tmp = []
150        for k in range(nkpt):
151            tmp.append(mo1[moloc[k]:moloc[k+1]].reshape(-1,nmo[k],nocc[k]))
152        v = fvind(tmp)
153        for k in range(nkpt):
154            v[k][:,viridx[k],:] *= e_ai[k]
155            v[k][:,occidx[k],:] = 0
156            v[k] = v[k].ravel()
157        return np.hstack(v)
158    _mo1 = lib.krylov(vind_vo, _mo1base,
159                      tol=tol, max_cycle=max_cycle, hermi=hermi, verbose=log)
160    mo1 = []
161    for k in range(nkpt):
162        mo1.append(_mo1[moloc[k]:moloc[k+1]].reshape(-1,nmo[k],nocc[k]))
163    log.timer('krylov solver in CPHF', *t0)
164
165    v1mo = fvind(mo1)
166    for k in range(nkpt):
167        mo1[k][:,viridx[k]] = mo1base[k][:,viridx[k]] - \
168                             v1mo[k][:,viridx[k]]*e_ai[k]
169
170    # mo_e1 has the same symmetry as the first order Fock matrix (hermitian or
171    # anti-hermitian). mo_e1 = v1mo + u1*lib.direct_sum('i-j->ij',e_i,e_i)
172    for k in range(nkpt):
173        mo_e1[k] += mo1[k][:,occidx[k]] * lib.direct_sum('i-j->ij', e_i[k], e_i[k])
174        mo_e1[k] += v1mo[k][:,occidx[k]]
175
176    return mo1, mo_e1
177