1#!/usr/bin/env python
2# Copyright 2014-2021 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'''
20General spin-orbital CISD
21'''
22
23import warnings
24
25from functools import reduce
26import numpy
27from pyscf import lib
28from pyscf.lib import logger
29from pyscf.cc import ccsd
30from pyscf.cc import gccsd
31from pyscf.cc import gccsd_rdm
32from pyscf.cc.addons import spatial2spin, spin2spatial
33from pyscf.ci import cisd
34from pyscf.ci import ucisd
35
36
37def make_diagonal(myci, eris):
38    nocc = myci.nocc
39    nmo = myci.nmo
40    nvir = nmo - nocc
41    mo_energy = eris.fock.diagonal()
42    jkdiag = numpy.zeros((nmo,nmo), dtype=mo_energy.dtype)
43    jkdiag[:nocc,:nocc] = numpy.einsum('ijij->ij', eris.oooo)
44    jkdiag[nocc:,nocc:] = numpy.einsum('ijij->ij', eris.vvvv)
45    jkdiag[:nocc,nocc:] = numpy.einsum('ijij->ij', eris.ovov)
46    jksum = jkdiag[:nocc,:nocc].sum()
47    ehf = mo_energy[:nocc].sum() - jksum * .5
48    e1diag = numpy.empty((nocc,nvir), dtype=mo_energy.dtype)
49    e2diag = numpy.empty((nocc,nocc,nvir,nvir), dtype=mo_energy.dtype)
50    for i in range(nocc):
51        for a in range(nocc, nmo):
52            e1diag[i,a-nocc] = ehf - mo_energy[i] + mo_energy[a] - jkdiag[i,a]
53            for j in range(nocc):
54                for b in range(nocc, nmo):
55                    e2diag[i,j,a-nocc,b-nocc] = ehf \
56                            - mo_energy[i] - mo_energy[j] \
57                            + mo_energy[a] + mo_energy[b] \
58                            + jkdiag[i,j] + jkdiag[a,b] \
59                            - jkdiag[i,a] - jkdiag[j,a] \
60                            - jkdiag[i,b] - jkdiag[j,b]
61    return amplitudes_to_cisdvec(ehf, e1diag, e2diag)
62
63def contract(myci, civec, eris):
64    nocc = myci.nocc
65    nmo = myci.nmo
66
67    c0, c1, c2 = cisdvec_to_amplitudes(civec, nmo, nocc)
68
69    fock = eris.fock
70    foo = fock[:nocc,:nocc]
71    fov = fock[:nocc,nocc:]
72    fvo = fock[nocc:,:nocc]
73    fvv = fock[nocc:,nocc:]
74
75    t1  = lib.einsum('ie,ae->ia', c1, fvv)
76    t1 -= lib.einsum('ma,mi->ia', c1, foo)
77    t1 += lib.einsum('imae,me->ia', c2, fov)
78    t1 += lib.einsum('nf,nafi->ia', c1, eris.ovvo)
79    t1 -= 0.5*lib.einsum('imef,maef->ia', c2, eris.ovvv)
80    t1 -= 0.5*lib.einsum('mnae,mnie->ia', c2, eris.ooov)
81
82    tmp = lib.einsum('ijae,be->ijab', c2, fvv)
83    t2  = tmp - tmp.transpose(0,1,3,2)
84    tmp = lib.einsum('imab,mj->ijab', c2, foo)
85    t2 -= tmp - tmp.transpose(1,0,2,3)
86    t2 += 0.5*lib.einsum('mnab,mnij->ijab', c2, eris.oooo)
87    t2 += 0.5*lib.einsum('ijef,abef->ijab', c2, eris.vvvv)
88    tmp = lib.einsum('imae,mbej->ijab', c2, eris.ovvo)
89    tmp+= numpy.einsum('ia,bj->ijab', c1, fvo)
90    tmp = tmp - tmp.transpose(0,1,3,2)
91    t2 += tmp - tmp.transpose(1,0,2,3)
92    tmp = lib.einsum('ie,jeba->ijab', c1, numpy.asarray(eris.ovvv).conj())
93    t2 += tmp - tmp.transpose(1,0,2,3)
94    tmp = lib.einsum('ma,ijmb->ijab', c1, numpy.asarray(eris.ooov).conj())
95    t2 -= tmp - tmp.transpose(0,1,3,2)
96
97    eris_oovv = numpy.asarray(eris.oovv)
98    t1 += fov.conj() * c0
99    t2 += eris_oovv.conj() * c0
100    t0  = numpy.einsum('ia,ia', fov, c1)
101    t0 += numpy.einsum('ijab,ijab', eris_oovv, c2) * .25
102
103    return amplitudes_to_cisdvec(t0, t1, t2)
104
105def amplitudes_to_cisdvec(c0, c1, c2):
106    nocc, nvir = c1.shape
107    ooidx = numpy.tril_indices(nocc, -1)
108    vvidx = numpy.tril_indices(nvir, -1)
109    c2tril = lib.take_2d(c2.reshape(nocc**2,nvir**2),
110                         ooidx[0]*nocc+ooidx[1], vvidx[0]*nvir+vvidx[1])
111    return numpy.hstack((c0, c1.ravel(), c2tril.ravel()))
112
113def cisdvec_to_amplitudes(civec, nmo, nocc):
114    nvir = nmo - nocc
115    c0 = civec[0]
116    c1 = civec[1:nocc*nvir+1].reshape(nocc,nvir)
117    c2 = ccsd._unpack_4fold(civec[nocc*nvir+1:], nocc, nvir)
118    return c0, c1, c2
119
120def from_ucisdvec(civec, nocc, orbspin):
121    '''Convert the (spin-separated) CISD coefficient vector to GCISD
122    coefficient vector'''
123    nmoa = numpy.count_nonzero(orbspin == 0)
124    nmob = numpy.count_nonzero(orbspin == 1)
125    if isinstance(nocc, int):
126        nocca = numpy.count_nonzero(orbspin[:nocc] == 0)
127        noccb = numpy.count_nonzero(orbspin[:nocc] == 1)
128    else:
129        nocca, noccb = nocc
130    nvira = nmoa - nocca
131
132    if civec.size == nocca*nvira + (nocca*nvira)**2 + 1:  # RCISD
133        c0, c1, c2 = cisd.cisdvec_to_amplitudes(civec, nmoa, nocca)
134    else:  # UCISD
135        c0, c1, c2 = ucisd.cisdvec_to_amplitudes(civec, (nmoa,nmob), (nocca,noccb))
136    c1 = spatial2spin(c1, orbspin)
137    c2 = spatial2spin(c2, orbspin)
138    return amplitudes_to_cisdvec(c0, c1, c2)
139from_rcisdvec = from_ucisdvec
140
141def to_ucisdvec(civec, nmo, nocc, orbspin):
142    '''Convert the GCISD coefficient vector to UCISD coefficient vector'''
143    c0, c1, c2 = cisdvec_to_amplitudes(civec, nmo, nocc)
144    c1 = spin2spatial(c1, orbspin)
145    c2 = spin2spatial(c2, orbspin)
146    ucisdvec = ucisd.amplitudes_to_cisdvec(c0, c1, c2)
147    unorm = numpy.linalg.norm(ucisdvec)
148    if unorm < 1e-2:
149        raise RuntimeError('GCISD vector corresponds to spin-flip excitation. '
150                           'It cannot be converted to UCISD wfn.'
151                           'norm(UCISD) = %s' % unorm)
152    elif unorm < 0.99:
153        warnings.warn('GCISD vector has spin-flip excitation. '
154                      'They are ignored when converting to UCISD wfn. '
155                      'norm(UCISD) = %s' % unorm)
156    return ucisdvec
157
158def to_fcivec(cisdvec, nelec, orbspin, frozen=None):
159    assert(numpy.count_nonzero(orbspin == 0) ==
160           numpy.count_nonzero(orbspin == 1))
161    norb = len(orbspin)
162    frozen_mask = numpy.zeros(norb, dtype=bool)
163    if frozen is None:
164        pass
165    elif isinstance(frozen, (int, numpy.integer)):
166        frozen_mask[:frozen] = True
167    else:
168        frozen_mask[frozen] = True
169    frozen = (numpy.where(frozen_mask[orbspin == 0])[0],
170              numpy.where(frozen_mask[orbspin == 1])[0])
171    nelec = (numpy.count_nonzero(orbspin[:nelec] == 0),
172             numpy.count_nonzero(orbspin[:nelec] == 1))
173    orbspin = orbspin[~frozen_mask]
174    nmo = len(orbspin)
175    nocc = numpy.count_nonzero(~frozen_mask[:sum(nelec)])
176    ucisdvec = to_ucisdvec(cisdvec, nmo, nocc, orbspin)
177    return ucisd.to_fcivec(ucisdvec, norb//2, nelec, frozen)
178
179def from_fcivec(ci0, nelec, orbspin, frozen=None):
180    if not (frozen is None or frozen == 0):
181        raise NotImplementedError
182
183    assert(numpy.count_nonzero(orbspin == 0) ==
184           numpy.count_nonzero(orbspin == 1))
185    norb = len(orbspin)
186    frozen_mask = numpy.zeros(norb, dtype=bool)
187    if frozen is None:
188        pass
189    elif isinstance(frozen, (int, numpy.integer)):
190        frozen_mask[:frozen] = True
191    else:
192        frozen_mask[frozen] = True
193    #frozen = (numpy.where(frozen_mask[orbspin == 0])[0],
194    #          numpy.where(frozen_mask[orbspin == 1])[0])
195    nelec = (numpy.count_nonzero(orbspin[:nelec] == 0),
196             numpy.count_nonzero(orbspin[:nelec] == 1))
197    ucisdvec = ucisd.from_fcivec(ci0, norb//2, nelec, frozen)
198    nocc = numpy.count_nonzero(~frozen_mask[:sum(nelec)])
199    return from_ucisdvec(ucisdvec, nocc, orbspin[~frozen_mask])
200
201
202def make_rdm1(myci, civec=None, nmo=None, nocc=None, ao_repr=False):
203    r'''
204    One-particle density matrix in the molecular spin-orbital representation
205    (the occupied-virtual blocks from the orbital response contribution are
206    not included).
207
208    dm1[p,q] = <q^\dagger p>  (p,q are spin-orbitals)
209
210    The convention of 1-pdm is based on McWeeney's book, Eq (5.4.20).
211    The contraction between 1-particle Hamiltonian and rdm1 is
212    E = einsum('pq,qp', h1, rdm1)
213    '''
214    if civec is None: civec = myci.ci
215    if nmo is None: nmo = myci.nmo
216    if nocc is None: nocc = myci.nocc
217    d1 = _gamma1_intermediates(myci, civec, nmo, nocc)
218    return gccsd_rdm._make_rdm1(myci, d1, with_frozen=True, ao_repr=ao_repr)
219
220def make_rdm2(myci, civec=None, nmo=None, nocc=None, ao_repr=False):
221    r'''
222    Two-particle density matrix in the molecular spin-orbital representation
223
224    dm2[p,q,r,s] = <p^\dagger r^\dagger s q>
225
226    where p,q,r,s are spin-orbitals. p,q correspond to one particle and r,s
227    correspond to another particle.  The contraction between ERIs (in
228    Chemist's notation) and rdm2 is
229    E = einsum('pqrs,pqrs', eri, rdm2)
230    '''
231    if civec is None: civec = myci.ci
232    if nmo is None: nmo = myci.nmo
233    if nocc is None: nocc = myci.nocc
234    d1 = _gamma1_intermediates(myci, civec, nmo, nocc)
235    d2 = _gamma2_intermediates(myci, civec, nmo, nocc)
236    return gccsd_rdm._make_rdm2(myci, d1, d2, with_dm1=True, with_frozen=True,
237                                ao_repr=ao_repr)
238
239def _gamma1_intermediates(myci, civec, nmo, nocc):
240    c0, c1, c2 = cisdvec_to_amplitudes(civec, nmo, nocc)
241    dvo = c0.conj() * c1.T
242    dvo += numpy.einsum('jb,ijab->ai', c1.conj(), c2)
243    dov = dvo.T.conj()
244    doo  =-numpy.einsum('ia,ka->ik', c1.conj(), c1)
245    doo -= numpy.einsum('jiab,kiab->jk', c2.conj(), c2) * .5
246    dvv  = numpy.einsum('ia,ic->ac', c1, c1.conj())
247    dvv += numpy.einsum('ijab,ijac->bc', c2, c2.conj()) * .5
248    return doo, dov, dvo, dvv
249
250def _gamma2_intermediates(myci, civec, nmo, nocc):
251    c0, c1, c2 = cisdvec_to_amplitudes(civec, nmo, nocc)
252    goovv = c0 * c2.conj() * .5
253    govvv = numpy.einsum('ia,ikcd->kadc', c1, c2.conj()) * .5
254    gooov = numpy.einsum('ia,klac->klic', c1, c2.conj()) *-.5
255    goooo = numpy.einsum('ijab,klab->ijkl', c2.conj(), c2) * .25
256    gvvvv = numpy.einsum('ijab,ijcd->abcd', c2, c2.conj()) * .25
257    govvo = numpy.einsum('ijab,ikac->jcbk', c2.conj(), c2)
258    govvo+= numpy.einsum('ia,jb->ibaj', c1.conj(), c1)
259
260    dovov = goovv.transpose(0,2,1,3) - goovv.transpose(0,3,1,2)
261    doooo = goooo.transpose(0,2,1,3) - goooo.transpose(0,3,1,2)
262    dvvvv = gvvvv.transpose(0,2,1,3) - gvvvv.transpose(0,3,1,2)
263    dovvo = govvo.transpose(0,2,1,3)
264    dooov = gooov.transpose(0,2,1,3) - gooov.transpose(1,2,0,3)
265    dovvv = govvv.transpose(0,2,1,3) - govvv.transpose(0,3,1,2)
266    doovv = None
267    dvvov = None
268    return dovov, dvvvv, doooo, doovv, dovvo, dvvov, dovvv, dooov
269
270def trans_rdm1(myci, cibra, ciket, nmo=None, nocc=None):
271    r'''
272    One-particle transition density matrix in the molecular spin-orbital
273    representation.
274
275    dm1[p,q] = <q^\dagger p>  (p,q are spin-orbitals)
276
277    The convention of 1-pdm is based on McWeeney's book, Eq (5.4.20).
278    The contraction between 1-particle Hamiltonian and rdm1 is
279    E = einsum('pq,qp', h1, rdm1)
280    '''
281    if nmo is None: nmo = myci.nmo
282    if nocc is None: nocc = myci.nocc
283    c0bra, c1bra, c2bra = myci.cisdvec_to_amplitudes(cibra, nmo, nocc)
284    c0ket, c1ket, c2ket = myci.cisdvec_to_amplitudes(ciket, nmo, nocc)
285
286    dvo = c0bra.conj() * c1ket.T
287    dvo += numpy.einsum('jb,ijab->ai', c1bra.conj(), c2ket)
288
289    dov = c0ket * c1bra.conj()
290    dov += numpy.einsum('jb,ijab->ia', c1ket, c2bra.conj())
291
292    doo  =-numpy.einsum('ia,ka->ik', c1bra.conj(), c1ket)
293    doo -= numpy.einsum('jiab,kiab->jk', c2bra.conj(), c2ket) * .5
294    dvv  = numpy.einsum('ia,ic->ac', c1ket, c1bra.conj())
295    dvv += numpy.einsum('ijab,ijac->bc', c2ket, c2bra.conj()) * .5
296
297    dm1 = numpy.empty((nmo,nmo), dtype=doo.dtype)
298    dm1[:nocc,:nocc] = doo
299    dm1[:nocc,nocc:] = dov
300    dm1[nocc:,:nocc] = dvo
301    dm1[nocc:,nocc:] = dvv
302    norm = numpy.dot(cibra, ciket)
303    dm1[numpy.diag_indices(nocc)] += norm
304
305    if myci.frozen is not None:
306        nmo = myci.mo_occ.size
307        nocc = numpy.count_nonzero(myci.mo_occ > 0)
308        rdm1 = numpy.zeros((nmo,nmo), dtype=dm1.dtype)
309        rdm1[numpy.diag_indices(nocc)] = norm
310        moidx = numpy.where(myci.get_frozen_mask())[0]
311        rdm1[moidx[:,None],moidx] = dm1
312        dm1 = rdm1
313    return dm1
314
315
316class GCISD(cisd.CISD):
317
318    def vector_size(self):
319        nocc = self.nocc
320        nvir = self.nmo - nocc
321        noo = nocc * (nocc-1) // 2
322        nvv = nvir * (nvir-1) // 2
323        return 1 + nocc*nvir + noo*nvv
324
325    def get_init_guess(self, eris=None, nroots=1, diag=None):
326        # MP2 initial guess
327        if eris is None: eris = self.ao2mo(self.mo_coeff)
328        time0 = logger.process_clock(), logger.perf_counter()
329        mo_e = eris.mo_energy
330        nocc = self.nocc
331        eia = mo_e[:nocc,None] - mo_e[None,nocc:]
332        eijab = lib.direct_sum('ia,jb->ijab',eia,eia)
333        ci0 = 1
334        ci1 = eris.fock[:nocc,nocc:] / eia
335        eris_oovv = numpy.array(eris.oovv)
336        ci2 = eris_oovv / eijab
337        self.emp2 = 0.25*numpy.einsum('ijab,ijab', ci2.conj(), eris_oovv).real
338        logger.info(self, 'Init t2, MP2 energy = %.15g', self.emp2)
339        logger.timer(self, 'init mp2', *time0)
340
341        if abs(self.emp2) < 1e-3 and abs(ci1).sum() < 1e-3:
342            ci1 = 1. / eia
343
344        ci_guess = amplitudes_to_cisdvec(ci0, ci1, ci2)
345
346        if nroots > 1:
347            civec_size = ci_guess.size
348            dtype = ci_guess.dtype
349            nroots = min(ci1.size+1, nroots)  # Consider Koopmans' theorem only
350
351            if diag is None:
352                idx = range(1, nroots)
353            else:
354                idx = diag[:ci1.size+1].argsort()[1:nroots]  # exclude HF determinant
355
356            ci_guess = [ci_guess]
357            for i in idx:
358                g = numpy.zeros(civec_size, dtype)
359                g[i] = 1.0
360                ci_guess.append(g)
361
362        return self.emp2, ci_guess
363
364    def ao2mo(self, mo_coeff=None):
365        nmo = self.nmo
366        mem_incore = nmo**4*2 * 8/1e6
367        mem_now = lib.current_memory()[0]
368        if (self._scf._eri is not None and
369            (mem_incore+mem_now < self.max_memory) or self.mol.incore_anyway):
370            return gccsd._make_eris_incore(self, mo_coeff)
371        elif getattr(self._scf, 'with_df', None):
372            raise NotImplementedError
373        else:
374            return gccsd._make_eris_outcore(self, mo_coeff)
375
376    contract = contract
377    make_diagonal = make_diagonal
378    _dot = None
379
380    def to_fcivec(self, cisdvec, nelec, orbspin, frozen=None):
381        return to_fcivec(cisdvec, nelec, orbspin, frozen)
382
383    def from_fcivec(self, fcivec, nelec, orbspin, frozen=None):
384        return from_fcivec(fcivec, nelec, orbspin, frozen)
385
386    make_rdm1 = make_rdm1
387    make_rdm2 = make_rdm2
388
389    trans_rdm1 = trans_rdm1
390
391    def amplitudes_to_cisdvec(self, c0, c1, c2):
392        return amplitudes_to_cisdvec(c0, c1, c2)
393
394    def cisdvec_to_amplitudes(self, civec, nmo=None, nocc=None):
395        if nmo is None: nmo = self.nmo
396        if nocc is None: nocc = self.nocc
397        return cisdvec_to_amplitudes(civec, nmo, nocc)
398
399    @lib.with_doc(from_ucisdvec.__doc__)
400    def from_ucisdvec(self, civec, nocc=None, orbspin=None):
401        if nocc is None: nocc = self.nocc
402        if orbspin is None:
403            orbspin = getattr(self.mo_coeff, 'orbspin', None)
404            if orbspin is not None:
405                orbspin = orbspin[self.get_frozen_mask()]
406        assert(orbspin is not None)
407        return from_ucisdvec(civec, nocc, orbspin=orbspin)
408    from_rcisdvec = from_ucisdvec
409
410    @lib.with_doc(to_ucisdvec.__doc__)
411    def to_ucisdvec(self, civec, orbspin=None):
412        if orbspin is None:
413            orbspin = getattr(self.mo_coeff, 'orbspin', None)
414            if orbspin is not None:
415                orbspin = orbspin[self.get_frozen_mask()]
416        return to_ucisdvec(civec, self.nmo, self.nocc, orbspin)
417
418    def spatial2spin(self, tx, orbspin=None):
419        if orbspin is None:
420            orbspin = getattr(self.mo_coeff, 'orbspin', None)
421            if orbspin is not None:
422                orbspin = orbspin[self.get_frozen_mask()]
423        return spatial2spin(tx, orbspin)
424
425    def spin2spatial(self, tx, orbspin=None):
426        if orbspin is None:
427            orbspin = getattr(self.mo_coeff, 'orbspin', None)
428            if orbspin is not None:
429                orbspin = orbspin[self.get_frozen_mask()]
430        return spin2spatial(tx, orbspin)
431
432CISD = GCISD
433
434from pyscf import scf
435scf.ghf.GHF.CISD = lib.class_as_method(CISD)
436
437
438if __name__ == '__main__':
439    from pyscf import gto
440    from pyscf import scf
441    from pyscf import ao2mo
442
443    mol = gto.Mole()
444    mol.verbose = 0
445    mol.atom = [
446        ['O', ( 0., 0.    , 0.   )],
447        ['H', ( 0., -0.757, 0.587)],
448        ['H', ( 0., 0.757 , 0.587)],]
449    mol.basis = {'H': 'sto-3g',
450                 'O': 'sto-3g',}
451    mol.build()
452    mf = scf.UHF(mol).run(conv_tol=1e-14)
453    gmf = scf.addons.convert_to_ghf(mf)
454    myci = GCISD(gmf)
455    eris = myci.ao2mo()
456    ecisd, civec = myci.kernel(eris=eris)
457    print(ecisd - -0.048878084082066106)
458
459    nmo = eris.mo_coeff.shape[1]
460    rdm1 = myci.make_rdm1(civec, nmo, mol.nelectron)
461    rdm2 = myci.make_rdm2(civec, nmo, mol.nelectron)
462
463    mo = eris.mo_coeff[:7] + eris.mo_coeff[7:]
464    eri = ao2mo.kernel(mf._eri, mo, compact=False).reshape([nmo]*4)
465    eri[eris.orbspin[:,None]!=eris.orbspin,:,:] = 0
466    eri[:,:,eris.orbspin[:,None]!=eris.orbspin] = 0
467    h1a = reduce(numpy.dot, (mf.mo_coeff[0].T, mf.get_hcore(), mf.mo_coeff[0]))
468    h1b = reduce(numpy.dot, (mf.mo_coeff[1].T, mf.get_hcore(), mf.mo_coeff[1]))
469    h1e = numpy.zeros((nmo,nmo))
470    idxa = eris.orbspin == 0
471    idxb = eris.orbspin == 1
472    h1e[idxa[:,None] & idxa] = h1a.ravel()
473    h1e[idxb[:,None] & idxb] = h1b.ravel()
474    e2 = (numpy.einsum('ij,ji', h1e, rdm1) +
475          numpy.einsum('ijkl,ijkl', eri, rdm2) * .5)
476    e2 += mol.energy_nuc()
477    print(myci.e_tot - e2)   # = 0
478
479    print(abs(rdm1 - numpy.einsum('ijkk->ji', rdm2)/(mol.nelectron-1)).sum())
480
481