#        self.assertAlmostEqual(e[0], 0.16656250872624662, 6)
#        self.assertAlmostEqual(e[1], 0.2394414445283693, 6)
#        self.assertAlmostEqual(e[2], 0.41399434356202935, 6)
#
#    def test_eaccsd_koopmans3(self):
#        e,v = mycc3.eaccsd(nroots=3, koopmans=True)
#        self.assertAlmostEqual(e[0], 0.16737886338859731, 6)
#        self.assertAlmostEqual(e[1], 0.24027613852009164, 6)
#        self.assertAlmostEqual(e[2], 0.73443352557582653, 6)
#
#        e,v = mycc3.eaccsd(nroots=3, guess=v[:3])
#        self.assertAlmostEqual(e[0], 0.16737886338859731, 6)
#        self.assertAlmostEqual(e[1], 0.24027613852009164, 6)
#        self.assertAlmostEqual(e[2], 0.73443352557582653, 6)
#
#    def test_eaccsd_partition3(self):
#        e,v = mycc3.eaccsd(nroots=3, partition='mp')
#        self.assertAlmostEqual(e[0], 0.16947311575051136, 6)
#        self.assertAlmostEqual(e[1], 0.24234326468848749, 6)
#        self.assertAlmostEqual(e[2], 0.7434661346653969 , 6)
#
#        e,v = mycc3.eaccsd(nroots=3, partition='full')
#        self.assertAlmostEqual(e[0], 0.16418276148493574, 6)
#        self.assertAlmostEqual(e[1], 0.23683978491376495, 6)
#        self.assertAlmostEqual(e[2], 0.55640091560545624, 6)


    def test_eeccsd3(self):
        e,v = mycc3.eeccsd(nroots=1)
        self.assertAlmostEqual(e, 0.2757159395886167, 6)

        e,v = mycc3.eeccsd(nroots=4)
        self.assertAlmostEqual(e[0], 0.2757159395886167, 6)
        self.assertAlmostEqual(e[1], 0.2757159395886167, 6)
        self.assertAlmostEqual(e[2], 0.2757159395886167, 6)
        self.assertAlmostEqual(e[3], 0.3005716731825082, 6)

    def test_eeccsd_koopmans3(self):
        e,v = mycc3.eeccsd(nroots=4, koopmans=True)
        self.assertAlmostEqual(e[0], 0.2757159395886167, 6)
        self.assertAlmostEqual(e[1], 0.2757159395886167, 6)
        self.assertAlmostEqual(e[2], 0.2757159395886167, 6)
        self.assertAlmostEqual(e[3], 0.3005716731825082, 6)

        e,v = mycc3.eeccsd(nroots=4, guess=v[:4])
        self.assertAlmostEqual(e[0], 0.2757159395886167, 6)
        self.assertAlmostEqual(e[1], 0.2757159395886167, 6)
        self.assertAlmostEqual(e[2], 0.2757159395886167, 6)
        self.assertAlmostEqual(e[3], 0.3005716731825082, 6)

    def test_eomee_ccsd_matvec_singlet3(self):
        numpy.random.seed(10)
        r1 = numpy.random.random((no,nv)) - .9
        r2 = numpy.random.random((no,no,nv,nv)) - .9
        r2 = r2 + r2.transpose(1,0,3,2)
        myeom = eom_rccsd.EOMEESinglet(mycc31)
        vec = myeom.amplitudes_to_vector(r1,r2)
        imds = myeom.make_imds(eris31)
        vec1 = myeom.matvec(vec, imds)
        r1, r2 = myeom.vector_to_amplitudes(vec1)
        self.assertAlmostEqual(lib.finger(r1), -112883.3791497977, 8)
        self.assertAlmostEqual(lib.finger(r2), -268199.3475813322, 8)

    def test_eomee_ccsd_matvec_triplet3(self):
        numpy.random.seed(10)
        r1 = numpy.random.random((no,nv)) - .9
        r2 = numpy.random.random((2,no,no,nv,nv)) - .9
        r2[0] = r2[0] - r2[0].transpose(0,1,3,2)
        r2[0] = r2[0] - r2[0].transpose(1,0,2,3)
        r2[1] = r2[1] - r2[1].transpose(1,0,3,2)
        myeom = eom_rccsd.EOMEETriplet(mycc31)
        vec = myeom.amplitudes_to_vector(r1, r2)
        imds = myeom.make_imds(eris31)
        vec1 = myeom.matvec(vec, imds)
        r1, r2 = myeom.vector_to_amplitudes(vec1)
        self.assertAlmostEqual(lib.finger(r1   ), 3550.5250670914056, 9)
        self.assertAlmostEqual(lib.finger(r2[0]), -237433.03756895234,8)
        self.assertAlmostEqual(lib.finger(r2[1]), 127680.0182437716 , 8)

    def test_eomsf_ccsd_matvec3(self):
        numpy.random.seed(10)
        r1 = numpy.random.random((no,nv)) - .9
        r2 = numpy.random.random((2,no,no,nv,nv)) - .9
        myeom = eom_rccsd.EOMEESpinFlip(mycc31)
        vec = myeom.amplitudes_to_vector(r1,r2)
        imds = myeom.make_imds(eris31)
        vec1 = myeom.matvec(vec, imds)
        r1, r2 = myeom.vector_to_amplitudes(vec1)
        self.assertAlmostEqual(lib.finger(r1   ), -19368.729268465482, 8)
        self.assertAlmostEqual(lib.finger(r2[0]), 84325.863680611626 , 8)
        self.assertAlmostEqual(lib.finger(r2[1]), 6715.9574457836134 , 8)

    def test_eomee_diag3(self):
        vec1S, vec1T, vec2 = eom_rccsd.EOMEE(mycc31).get_diag()
        self.assertAlmostEqual(lib.finger(vec1S),-2881.6804563818432, 9)
        self.assertAlmostEqual(lib.finger(vec1T), 2039.7385969969259, 9)
        self.assertAlmostEqual(lib.finger(vec2) ,-4271.6230465236358, 9)


    def test_ip_matvec3(self):
        numpy.random.seed(12)
        r1 = numpy.random.random((no)) - .9
        r2 = numpy.random.random((no,no,nv)) - .9
        myeom = eom_rccsd.EOMIP(mycc31)
        vec = myeom.amplitudes_to_vector(r1,r2)
        r1,r2 = myeom.vector_to_amplitudes(vec)
        myeom.partition = 'mp'
        self.assertAlmostEqual(lib.finger(r1), 0.37404344676857076, 12)
        self.assertAlmostEqual(lib.finger(r2), -1.1568913404570922, 12)
        imds = myeom.make_imds(eris31)
        vec1 = myeom.matvec(vec, imds)
        self.assertAlmostEqual(lib.finger(vec1), -14894.669606811192, 9)
        self.assertAlmostEqual(lib.finger(myeom.get_diag()), 1182.3095479451745, 9)

        myeom.partition = 'full'
        imds = myeom.make_imds(eris31)
        diag = myeom.get_diag(imds)
        vec1 = myeom.matvec(vec, imds, diag=diag)
        self.assertAlmostEqual(lib.finger(vec1), -3795.9122245246967, 9)
        self.assertAlmostEqual(lib.finger(diag), 1106.260154202434, 9)

    def test_sort_left_right_eigensystem(self):
        myeom = eom_rccsd.EOMIP(mycc)
        right_evecs = [numpy.ones(10)] * 4
        left_evecs = [numpy.ones(10)] * 5
        right_evecs = [x*i for i, x in enumerate(right_evecs)]
        left_evecs = [x*i for i, x in enumerate(left_evecs)]
        revals, revecs, levecs = eom_rccsd._sort_left_right_eigensystem(
            myeom,
            [True, False, True, True], [-1.1, 0, 1.1, 2.2], right_evecs,
            [True, True, True, False, True], [-2.2, -1.1, 0, 1.1, 2.2], left_evecs)
        self.assertEqual(revals[0], -1.1)
        self.assertEqual(revals[1], 2.2)
        self.assertEqual(revecs[0][0], 0)
        self.assertEqual(revecs[1][0], 3)
        self.assertEqual(levecs[0][0], 1)
        self.assertEqual(levecs[1][0], 4)

        revals, revecs, levecs = eom_rccsd._sort_left_right_eigensystem(
            myeom,
            [True, False, True, True], [-1.1, 0, 1.1, 2.2], right_evecs,
            [True, True, False, True, True], [-2.2, -1.1, 0, 1.1, 2.2], left_evecs)
        self.assertEqual(revals[0], -1.1)
        self.assertEqual(revals[1], 1.1)
        self.assertEqual(revals[2], 2.2)
        self.assertEqual(revecs[0][0], 0)
        self.assertEqual(revecs[1][0], 2)
        self.assertEqual(revecs[2][0], 3)
        self.assertEqual(levecs[0][0], 1)
        self.assertEqual(levecs[1][0], 3)
        self.assertEqual(levecs[2][0], 4)

#    def test_ea_matvec3(self):
#        numpy.random.seed(12)
#        r1 = numpy.random.random((nv)) - .9
#        r2 = numpy.random.random((no,nv,nv)) - .9
#        myeom = eom_rccsd.EOMEA(mycc31)
#        vec = myeom.amplitudes_to_vector(r1,r2)
#        r1,r2 = myeom.vector_to_amplitudes(vec)
#        myeom.partition = 'mp'
#        self.assertAlmostEqual(lib.finger(r1), 1.4488291275539353, 12)
#        self.assertAlmostEqual(lib.finger(r2), 0.97080165032287469, 12)
#        imds = myeom.make_imds(eris31)
#        vec1 = myeom.matvec(vec, imds)
#        self.assertAlmostEqual(lib.finger(vec1), -34426.363943760276, 9)
#        self.assertAlmostEqual(lib.finger(myeom.get_diag()), 2724.8239646679217, 9)
#
#        myeom.partition = 'full'
#        imds = myeom.make_imds(eris31)
#        diag = myeom.get_diag(imds)
#        vec1 = myeom.matvec(vec, imds, diag=diag)
#        self.assertAlmostEqual(lib.finger(vec1), -17030.363405297598, 9)
#        self.assertAlmostEqual(lib.finger(diag), 4688.9122122011922, 9)

    def test_t3p2_intermediates_complex(self):
        '''Although this has not been tested strictly for complex values, it
        was written to be correct for complex values and differences in the complex
        values between versions should be taken into account and corrected.'''
        myt1 = mycc1.t1 + 1j * numpy.sin(mycc1.t1) * mycc1.t1
        myt2 = mycc1.t2 + 1j * numpy.sin(mycc1.t2) * mycc1.t2
        myt2 = myt2 + myt2.transpose(1,0,3,2)
        e, pt1, pt2, Wmcik, Wacek = rintermediates.get_t3p2_imds_slow(mycc1, myt1, myt2, eris=erisi)
        self.assertAlmostEqual(lib.finger(e), 23223.465490572264, 6)
        self.assertAlmostEqual(lib.finger(pt1), (-5.2202836452466705-0.09570164571057749j), 6)
        self.assertAlmostEqual(lib.finger(pt2), (46.188012063609506-1.303867687778909j), 6)
        self.assertAlmostEqual(lib.finger(Wmcik), (-18.438930654297778+1.5734161307568773j), 6)
        self.assertAlmostEqual(lib.finger(Wacek), (-7.187576764072701+0.7399185332889747j), 6)

    def test_t3p2_intermediates_real(self):
        myt1 = mycc1.t1.copy()
        myt2 = mycc1.t2.copy()
        myt2 = myt2 + myt2.transpose(1,0,3,2)
        e, pt1, pt2, Wmcik, Wacek = rintermediates.get_t3p2_imds_slow(mycc1, myt1, myt2)
        self.assertAlmostEqual(lib.finger(e), 23230.479350851536, 6)
        self.assertAlmostEqual(lib.finger(pt1), -5.218888542335442, 6)
        self.assertAlmostEqual(lib.finger(pt2), 46.19512409958347, 6)
        self.assertAlmostEqual(lib.finger(Wmcik), -18.47928005593598, 6)
        self.assertAlmostEqual(lib.finger(Wacek), -7.101360230151883, 6)

    def test_t3p2_intermediates_against_so(self):
        from pyscf.cc.addons import convert_to_gccsd
        myt1 = mycc1.t1.copy()
        myt2 = mycc1.t2.copy()
        e, pt1, pt2, Wmcik, Wacek = rintermediates.get_t3p2_imds_slow(mycc1, myt1, myt2)

        mygcc = convert_to_gccsd(mycc1)
        mygt1 = mygcc.t1.copy()
        mygt2 = mygcc.t2.copy()
        ge, gpt1, gpt2, gWmcik, gWacek = gintermediates.get_t3p2_imds_slow(mygcc, mygt1, mygt2)
        self.assertAlmostEqual(lib.finger(pt1), -2.6094405706617727, 6)
        self.assertAlmostEqual(lib.finger(pt2), 23.097562049844235, 6)
        self.assertAlmostEqual(lib.finger(pt1), lib.finger(gpt1[::2,::2]), 6)
        self.assertAlmostEqual(lib.finger(pt2), lib.finger(gpt2[::2,1::2,::2,1::2]), 6)

    def test_h2o_star(self):
        mol_h2o = gto.Mole()
        mol_h2o.atom = [
                [8, [0.000000000000000, -0.000000000000000, -0.124143731294022]],
                [1, [0.000000000000000, -1.430522735894536,  0.985125550040314]],
                [1, [0.000000000000000,  1.430522735894536,  0.985125550040314]]]
        mol_h2o.unit = 'B'
        mol_h2o.basis = {'H' : [[0,
                               [5.4471780, 0.156285],
                               [0.8245472, 0.904691]],
                               [0, [0.1831916, 1.0]]],
                        'O' : '3-21G'}
        mol_h2o.verbose = 7
        mol_h2o.output = '/dev/null'
        mol_h2o.build()
        mol.conv_tol = 1e-12
        mf_h2o = scf.RHF(mol_h2o)
        mf_h2o.conv_tol_grad = 1e-12
        mf_h2o.kernel()
        mycc_h2o = cc.RCCSD(mf_h2o).run()
        mycc_h2o.conv_tol_normt = 1e-12
        mycc_h2o.conv_tol = 1e-12
        mycc_h2o.kernel()

        myeom = eom_rccsd.EOMIP(mycc_h2o)
        e = myeom.ipccsd_star(nroots=3)
        self.assertAlmostEqual(e[0], 0.410661965883, 6)

        myeom = eom_rccsd.EOMIP_Ta(mycc_h2o)
        e = myeom.ipccsd_star(nroots=3)
        self.assertAlmostEqual(e[0], 0.411695647736, 6)

        myeom = eom_rccsd.EOMEA(mycc_h2o)
        e = myeom.eaccsd_star(nroots=3)
        self.assertAlmostEqual(e[0], 0.250589854185, 6)

        myeom = eom_rccsd.EOMEA_Ta(mycc_h2o)
        e = myeom.eaccsd_star(nroots=3)
        self.assertAlmostEqual(e[0], 0.250720295150, 6)

if __name__ == "__main__":
    print("Tests for EOM RCCSD")
    unittest.main()


# pyscf/cc/test/test_eom_uccsd.py

#!/usr/bin/env python
# Copyright 2014-2018 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.

import copy
import unittest
import numpy
from functools import reduce

from pyscf import gto
from pyscf import scf
from pyscf import cc
from pyscf import lib
from pyscf import ao2mo
from pyscf.cc import gccsd
from pyscf.cc import uccsd
from pyscf.cc import eom_uccsd
from pyscf.cc import eom_gccsd

def setUpModule():
    global mol, mf, mol1, mf0, mf1, gmf, ucc, ucc0, ucc1, eris1
    global ecc, orbspin, nocca, noccb, nvira, nvirb, r1, r2
    mol = gto.Mole()
    mol.verbose = 0
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = '631g'
    mol.spin = 0
    mol.build()
    mf = scf.UHF(mol).run(conv_tol=1e-12)

    mol1 = mol.copy()
    mol1.spin = 2
    mol1.build()
    mf0 = scf.UHF(mol1).run(conv_tol=1e-12)
    mf1 = copy.copy(mf0)

    nocca, noccb = mol1.nelec
    nmo = mol1.nao_nr()
    nvira, nvirb = nmo-nocca, nmo-noccb
    numpy.random.seed(12)
    mf1.mo_coeff = numpy.random.random((2,nmo,nmo)) - .5
    gmf = scf.addons.convert_to_ghf(mf1)
    orbspin = gmf.mo_coeff.orbspin

    ucc1 = cc.UCCSD(mf1)
    eris1 = ucc1.ao2mo()

    numpy.random.seed(11)
    no = nocca + noccb
    nv = nvira + nvirb
    r1 = numpy.random.random((no,nv)) - .9
    r2 = numpy.random.random((no,no,nv,nv)) - .9
    r2 = r2 - r2.transpose(1,0,2,3)
    r2 = r2 - r2.transpose(0,1,3,2)
    r1 = cc.addons.spin2spatial(r1, orbspin)
    r2 = cc.addons.spin2spatial(r2, orbspin)
    r1,r2 = eom_uccsd.vector_to_amplitudes_ee(
        eom_uccsd.amplitudes_to_vector_ee(r1,r2), ucc1.nmo, ucc1.nocc)
    ucc1.t1 = r1
    ucc1.t2 = r2

    ucc = cc.UCCSD(mf)
    ucc.max_space = 0
    ucc.conv_tol = 1e-8
    ecc = ucc.kernel()[0]

    ucc0 = cc.UCCSD(mf0)
    ucc0.conv_tol = 1e-8
    ucc0.direct = True
    ucc0.kernel()

def tearDownModule():
    global mol, mf, mol1, mf0, mf1, gmf, ucc, ucc0, ucc1, eris1
    del mol, mf, mol1, mf0, mf1, gmf, ucc, ucc0, ucc1, eris1

class KnownValues(unittest.TestCase):
    def test_ipccsd(self):
        eom = ucc.eomip_method()
        e,v = eom.kernel(nroots=1, koopmans=False)
        self.assertAlmostEqual(e, 0.42789083399175043, 6)
        e,v = ucc.ipccsd(nroots=8)
        self.assertAlmostEqual(e[0], 0.42789083399175043, 6)
        self.assertAlmostEqual(e[2], 0.50226861340475437, 6)
        self.assertAlmostEqual(e[4], 0.68550641152952585, 6)

        e,v = ucc.ipccsd(nroots=4, guess=v[:4])
        self.assertAlmostEqual(e[2], 0.50226861340475437, 6)

    def test_ipccsd_koopmans(self):
        e,v = ucc.ipccsd(nroots=8, koopmans=True)
        self.assertAlmostEqual(e[0], 0.42789083399175043, 6)
        self.assertAlmostEqual(e[2], 0.50226861340475437, 6)
        self.assertAlmostEqual(e[4], 0.68550641152952585, 6)

    def test_eaccsd(self):
        eom = ucc.eomea_method()
        e,v = eom.kernel(nroots=1, koopmans=False)
        self.assertAlmostEqual(e, 0.19050592137699729, 6)
        e,v = ucc.eaccsd(nroots=8)
        self.assertAlmostEqual(e[0], 0.19050592137699729, 6)
        self.assertAlmostEqual(e[2], 0.28345228891172214, 6)
        self.assertAlmostEqual(e[4], 0.52280673926459342, 6)

        e,v = ucc.eaccsd(nroots=4, guess=v[:4])
        self.assertAlmostEqual(e[2], 0.28345228891172214, 6)

    def test_eaccsd_koopmans(self):
        e,v = ucc.eaccsd(nroots=6, koopmans=True)
        self.assertAlmostEqual(e[0], 0.19050592137699729, 6)
        self.assertAlmostEqual(e[2], 0.28345228891172214, 6)
        self.assertAlmostEqual(e[4], 1.02136493172648370, 6)

        gcc1 = gccsd.GCCSD(scf.addons.convert_to_ghf(mf)).run()
        e1 = gcc1.eaccsd(nroots=6, koopmans=True)[0]
        self.assertAlmostEqual(abs(e1-e).max(), 0, 6)


    def test_eomee(self):
        self.assertAlmostEqual(ecc, -0.13539788719099638, 6)
        eom = ucc.eomee_method()
        e,v = eom.kernel(nroots=1, koopmans=False)
        self.assertAlmostEqual(e, 0.28114509667240556, 6)

        e,v = ucc.eeccsd(nroots=4)
        self.assertAlmostEqual(e[0], 0.28114509667240556, 6)
        self.assertAlmostEqual(e[1], 0.28114509667240556, 6)
        self.assertAlmostEqual(e[2], 0.28114509667240556, 6)
        self.assertAlmostEqual(e[3], 0.30819728420902842, 6)

        e,v = ucc.eeccsd(nroots=4, guess=v[:4])
        self.assertAlmostEqual(e[3], 0.30819728420902842, 6)

    def test_eomee_ccsd_spin_keep(self):
        e, v = ucc.eomee_ccsd(nroots=2, koopmans=False)
        self.assertAlmostEqual(e[0], 0.28114509667240556, 6)
        self.assertAlmostEqual(e[1], 0.30819728420902842, 6)

        e, v = ucc.eomee_ccsd(nroots=2, koopmans=True)
        self.assertAlmostEqual(e[0], 0.28114509667240556, 6)
        self.assertAlmostEqual(e[1], 0.30819728420902842, 6)

    def test_eomsf_ccsd(self):
        e, v = ucc.eomsf_ccsd(nroots=2, koopmans=False)
        self.assertAlmostEqual(e[0], 0.28114509667240556, 6)
        self.assertAlmostEqual(e[1], 0.28114509667240556, 6)

        e, v = ucc.eomsf_ccsd(nroots=2, koopmans=True)
        self.assertAlmostEqual(e[0], 0.28114509667240556, 6)
        self.assertAlmostEqual(e[1], 0.28114509667240556, 6)

    def test_ucc_update_amps(self):
        gcc1 = gccsd.GCCSD(gmf)
        r1g = gcc1.spatial2spin(ucc1.t1, orbspin)
        r2g = gcc1.spatial2spin(ucc1.t2, orbspin)
        r1g, r2g = gcc1.update_amps(r1g, r2g, gcc1.ao2mo())
        u1g = gcc1.spin2spatial(r1g, orbspin)
        u2g = gcc1.spin2spatial(r2g, orbspin)
        t1, t2 = ucc1.update_amps(ucc1.t1, ucc1.t2, eris1)
        self.assertAlmostEqual(abs(u1g[0]-t1[0]).max(), 0, 7)
        self.assertAlmostEqual(abs(u1g[1]-t1[1]).max(), 0, 7)
        self.assertAlmostEqual(abs(u2g[0]-t2[0]).max(), 0, 6)
        self.assertAlmostEqual(abs(u2g[1]-t2[1]).max(), 0, 6)
        self.assertAlmostEqual(abs(u2g[2]-t2[2]).max(), 0, 6)
        self.assertAlmostEqual(float(abs(r1g-gcc1.spatial2spin(t1, orbspin)).max()), 0, 6)
        self.assertAlmostEqual(float(abs(r2g-gcc1.spatial2spin(t2, orbspin)).max()), 0, 6)
        self.assertAlmostEqual(uccsd.energy(ucc1, r1, r2, eris1), -7.2775115532675771, 8)
        e0, t1, t2 = ucc1.init_amps(eris1)
        self.assertAlmostEqual(lib.finger(cc.addons.spatial2spin(t1, orbspin)), 148.57054876656397, 8)
        self.assertAlmostEqual(lib.finger(cc.addons.spatial2spin(t2, orbspin)),-349.94207953071475, 8)
        self.assertAlmostEqual(e0, 30.640616265644827, 2)

    def test_ucc_eomee_ccsd_matvec(self):
        numpy.random.seed(10)
        r1 = [numpy.random.random((nocca,nvira))-.9,
              numpy.random.random((noccb,nvirb))-.9]
        r2 = [numpy.random.random((nocca,nocca,nvira,nvira))-.9,
              numpy.random.random((nocca,noccb,nvira,nvirb))-.9,
              numpy.random.random((noccb,noccb,nvirb,nvirb))-.9]
        r2[0] = r2[0] - r2[0].transpose(1,0,2,3)
        r2[0] = r2[0] - r2[0].transpose(0,1,3,2)
        r2[2] = r2[2] - r2[2].transpose(1,0,2,3)
        r2[2] = r2[2] - r2[2].transpose(0,1,3,2)

        gcc1 = cc.addons.convert_to_gccsd(ucc1)
        gr1 = gcc1.spatial2spin(r1)
        gr2 = gcc1.spatial2spin(r2)
        gee1 = eom_gccsd.EOMEE(gcc1)
        gvec = eom_gccsd.amplitudes_to_vector_ee(gr1, gr2)
        vecref = eom_gccsd.eeccsd_matvec(gee1, gvec)

        vec = eom_uccsd.amplitudes_to_vector_ee(r1,r2)
        uee1 = eom_uccsd.EOMEESpinKeep(ucc1)
        vec1 = eom_uccsd.eomee_ccsd_matvec(uee1, vec)

        uv = eom_uccsd.amplitudes_to_vector_ee(r1, r2)
        gv = eom_gccsd.amplitudes_to_vector_ee(gr1, gr2)
        r1, r2 = uee1.vector_to_amplitudes(uee1.matvec(uv))
        gr1, gr2 = gee1.vector_to_amplitudes(gee1.matvec(gv))

        r1, r2 = uee1.vector_to_amplitudes(vec1)
        gr1, gr2 = gee1.vector_to_amplitudes(vecref)
        self.assertAlmostEqual(float(abs(gr1-gcc1.spatial2spin(r1)).max()), 0, 9)
        self.assertAlmostEqual(float(abs(gr2-gcc1.spatial2spin(r2)).max()), 0, 9)
        self.assertAlmostEqual(lib.finger(vec1), 49.499911123484523, 9)

    def test_ucc_eomee_ccsd_diag(self):  # FIXME: compare to EOMEE-GCCSD diag
        vec1, vec2 = eom_uccsd.EOMEE(ucc1).get_diag()
        self.assertAlmostEqual(lib.finger(vec1), 62.767648620751018, 9)
        self.assertAlmostEqual(lib.finger(vec2), 156.2976365433517, 9)

    def test_ucc_eomee_init_guess(self):
        uee = eom_uccsd.EOMEESpinKeep(ucc1)
        diag = uee.get_diag()[0]
        guess = uee.get_init_guess(nroots=1, koopmans=False, diag=diag)
        self.assertAlmostEqual(lib.finger(guess[0]), -0.99525784369029358, 9)

        guess = uee.get_init_guess(nroots=1, koopmans=True, diag=diag)
        self.assertAlmostEqual(lib.finger(guess[0]), -0.84387013299273794, 9)

        guess = uee.get_init_guess(nroots=4, koopmans=False, diag=diag)
        self.assertAlmostEqual(lib.finger(guess), -0.98261980006133565, 9)

        guess = uee.get_init_guess(nroots=4, koopmans=True, diag=diag)
        self.assertAlmostEqual(lib.finger(guess), -0.38124032366955651, 9)

    def test_ucc_eomsf_ccsd_matvec(self):
        numpy.random.seed(10)
        myeom = eom_uccsd.EOMEESpinFlip(ucc1)
        vec = numpy.random.random(myeom.vector_size()) - .9
        vec1 = eom_uccsd.eomsf_ccsd_matvec(myeom, vec)
        self.assertAlmostEqual(lib.finger(vec1), -1655.5564756993756, 8)

        r1, r2 = myeom.vector_to_amplitudes(vec)
        gr1 = eom_uccsd.spatial2spin_eomsf(r1, orbspin)
        gr2 = eom_uccsd.spatial2spin_eomsf(r2, orbspin)
        gvec = eom_gccsd.amplitudes_to_vector_ee(gr1, gr2)

        gcc1 = cc.addons.convert_to_gccsd(ucc1)
        gee1 = eom_gccsd.EOMEE(gcc1)
        vecref = eom_gccsd.eeccsd_matvec(gee1, gvec)
        gr1, gr2 = gee1.vector_to_amplitudes(vecref)
        v1, v2 = myeom.vector_to_amplitudes(vec1)
        self.assertAlmostEqual(float(abs(gr1-eom_uccsd.spatial2spin_eomsf(v1, orbspin)).max()), 0, 9)
        self.assertAlmostEqual(float(abs(gr2-eom_uccsd.spatial2spin_eomsf(v2, orbspin)).max()), 0, 9)

#    def test_ucc_eomip_matvec(self):
#
#    def test_ucc_eomea_matvec(self):

########################################
# With 4-fold symmetry in integrals
# max_memory = 0
# direct = True
    def test_eomee1(self):
        self.assertAlmostEqual(ucc0.e_corr, -0.10805861805688141, 6)
        e,v = ucc0.eeccsd(nroots=4)
        self.assertAlmostEqual(e[0],-0.28757438579564343, 6)
        self.assertAlmostEqual(e[1], 7.0932490003970672e-05, 6)
        self.assertAlmostEqual(e[2], 0.026861582690761672, 6)
        self.assertAlmostEqual(e[3], 0.091111388761653589, 6)

        e,v = ucc0.eeccsd(nroots=4, guess=v[:4])
        self.assertAlmostEqual(e[3], 0.091111388761653589, 6)

    def test_vector_to_amplitudes_eomsf(self):
        eomsf = eom_uccsd.EOMEESpinFlip(ucc0)
        size = eomsf.vector_size()
        v = numpy.random.random(size)
        r1, r2 = eomsf.vector_to_amplitudes(v)
        v1 = eomsf.amplitudes_to_vector(r1, r2)
        self.assertAlmostEqual(abs(v-v1).max(), 0, 12)

    def test_spatial2spin_eomsf(self):
        eomsf = eom_uccsd.EOMEESpinFlip(ucc0)
        size = eomsf.vector_size()
        v = numpy.random.random(size)
        r1, r2 = eomsf.vector_to_amplitudes(v)
        v1 = eom_uccsd.spin2spatial_eomsf(eom_uccsd.spatial2spin_eomsf(r1, orbspin), orbspin)
        v2 = eom_uccsd.spin2spatial_eomsf(eom_uccsd.spatial2spin_eomsf(r2, orbspin), orbspin)
        self.assertAlmostEqual(abs(r1[0]-v1[0]).max(), 0, 12)
        self.assertAlmostEqual(abs(r1[1]-v1[1]).max(), 0, 12)
        self.assertAlmostEqual(abs(r2[0]-v2[0]).max(), 0, 12)
        self.assertAlmostEqual(abs(r2[1]-v2[1]).max(), 0, 12)
        self.assertAlmostEqual(abs(r2[2]-v2[2]).max(), 0, 12)
        self.assertAlmostEqual(abs(r2[3]-v2[3]).max(), 0, 12)

    def test_vector_to_amplitudes_eom_spin_keep(self):
        eomsf = eom_uccsd.EOMEESpinKeep(ucc0)
        size = eomsf.vector_size()
        v = numpy.random.random(size)
        r1, r2 = eomsf.vector_to_amplitudes(v)
        v1 = eomsf.amplitudes_to_vector(r1, r2)
        self.assertAlmostEqual(abs(v-v1).max(), 0, 12)

    def test_vector_to_amplitudes_overwritten(self):
        mol = gto.M()
        mycc = scf.UHF(mol).apply(cc.UCCSD)
        nelec = (3,3)
        nocc = nelec
        nmo = (5, 5)
        mycc.nocc = nocc
        mycc.nmo = nmo
        def check_overwritten(method):
            vec = numpy.zeros(method.vector_size())
            vec_orig = vec.copy()
            ts = method.vector_to_amplitudes(vec)
            for ti in ts:
                if isinstance(ti, numpy.ndarray):
                    ti[:] = 1
                else:
                    for t in ti:
                        t[:] = 1
            self.assertAlmostEqual(abs(vec - vec_orig).max(), 0, 15)

        check_overwritten(mycc)
        check_overwritten(mycc.EOMIP())
        check_overwritten(mycc.EOMEA())
        check_overwritten(mycc.EOMEESpinKeep())
        check_overwritten(mycc.EOMEESpinFlip())


if __name__ == "__main__":
    print("Tests for UCCSD")
    unittest.main()


# pyscf/cc/test/test_gccsd.py

#!/usr/bin/env python
# Copyright 2014-2020 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.

import unittest
import numpy
from functools import reduce

from pyscf import lib
from pyscf import scf
from pyscf import gto
from pyscf import cc
from pyscf import ao2mo
from pyscf import mp
from pyscf.cc import gccsd
from pyscf.cc import gccsd_rdm
from pyscf.cc import ccsd
from pyscf.cc import uccsd

def setUpModule():
    global mol, mf, gcc1
    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.verbose = 5
    mol.output = '/dev/null'
    mol.basis = '631g'
    mol.spin = 2
    mol.build()
    mf = scf.UHF(mol).run(conv_tol=1e-12)
    mf = scf.addons.convert_to_ghf(mf)

    gcc1 = gccsd.GCCSD(mf).run(conv_tol=1e-9)

def tearDownModule():
    global mol, mf, gcc1
    mol.stdout.close()
    del mol, mf, gcc1

class KnownValues(unittest.TestCase):
    def test_gccsd(self):
        self.assertAlmostEqual(gcc1.e_corr, -0.10805861695870976, 7)

    def test_ERIS(self):
        gcc = gccsd.GCCSD(mf, frozen=4)
        numpy.random.seed(9)
        mo_coeff0 = numpy.random.random(mf.mo_coeff.shape) - .9
        nao = mo_coeff0.shape[0]//2
        orbspin = numpy.array([0,1,0,1,0,1,0,0,1,0,1,0,1,1,0,0,1,0,1,0,1,1,0,1,0,1])
        mo_coeff0[nao:,orbspin==0] = 0
        mo_coeff0[:nao,orbspin==1] = 0
        mo_coeff1 = mo_coeff0.copy()
        mo_coeff1[-1,0] = 1e-12

        eris = gccsd._make_eris_incore(gcc, mo_coeff0)
        self.assertAlmostEqual(lib.finger(eris.oooo),  15.97533838570434, 9)
        self.assertAlmostEqual(lib.finger(eris.ooov), -80.97666019169982, 9)
        self.assertAlmostEqual(lib.finger(eris.oovv), 278.00028168381675, 9)
        self.assertAlmostEqual(lib.finger(eris.ovov),   2.34326750142844, 9)
        self.assertAlmostEqual(lib.finger(eris.ovvv), 908.61659731634768, 9)
        self.assertAlmostEqual(lib.finger(eris.vvvv), 756.77383112217694, 9)

        eris = gccsd._make_eris_outcore(gcc, mo_coeff0)
        self.assertAlmostEqual(lib.finger(eris.oooo),  15.97533838570434, 9)
        self.assertAlmostEqual(lib.finger(eris.ooov), -80.97666019169982, 9)
        self.assertAlmostEqual(lib.finger(eris.oovv), 278.00028168381675, 9)
        self.assertAlmostEqual(lib.finger(eris.ovov),   2.34326750142844, 9)
        self.assertAlmostEqual(lib.finger(eris.ovvv), 908.61659731634768, 9)
        self.assertAlmostEqual(lib.finger(eris.vvvv), 756.77383112217694, 9)

        eris = gccsd._make_eris_incore(gcc, mo_coeff1)
        self.assertAlmostEqual(lib.finger(eris.oooo),  15.97533838570434, 9)
        self.assertAlmostEqual(lib.finger(eris.ooov), -80.97666019169982, 9)
        self.assertAlmostEqual(lib.finger(eris.oovv), 278.00028168381675, 9)
        self.assertAlmostEqual(lib.finger(eris.ovov),   2.34326750142844, 9)
        self.assertAlmostEqual(lib.finger(eris.ovvv), 908.61659731634768, 9)
        self.assertAlmostEqual(lib.finger(eris.vvvv), 756.77383112217694, 9)

        gcc.max_memory = 0
        eris = gcc.ao2mo(mo_coeff1)
        self.assertAlmostEqual(lib.finger(eris.oooo),  15.97533838570434, 9)
        self.assertAlmostEqual(lib.finger(eris.ooov), -80.97666019169982, 9)
        self.assertAlmostEqual(lib.finger(eris.oovv), 278.00028168381675, 9)
        self.assertAlmostEqual(lib.finger(eris.ovov),   2.34326750142844, 9)
        self.assertAlmostEqual(lib.finger(eris.ovvv), 908.61659731634768, 9)
        self.assertAlmostEqual(lib.finger(eris.vvvv), 756.77383112217694, 9)

    def test_spin2spatial(self):
        nocca, noccb = mol.nelec
        nvira = mol.nao_nr() - nocca
        nvirb = mol.nao_nr() - noccb
        numpy.random.seed(1)
        t1 = [numpy.random.random((nocca,nvira))*.1 - .1,
              numpy.random.random((noccb,nvirb))*.1 - .1]
        t2 = [numpy.random.random((nocca,nocca,nvira,nvira))*.1 - .1,
              numpy.random.random((nocca,noccb,nvira,nvirb))*.1 - .1,
              numpy.random.random((noccb,noccb,nvirb,nvirb))*.1 - .1]
        t2[0] = t2[0] - t2[0].transpose(1,0,2,3)
        t2[0] = t2[0] - t2[0].transpose(0,1,3,2)
        t2[2] = t2[2] - t2[2].transpose(1,0,2,3)
        t2[2] = t2[2] - t2[2].transpose(0,1,3,2)
        t1u = gcc1.spin2spatial(gcc1.spatial2spin(t1))
        t2u = gcc1.spin2spatial(gcc1.spatial2spin(t2))
        self.assertAlmostEqual(abs(t1[0] - t1u[0]).max(), 0, 12)
        self.assertAlmostEqual(abs(t1[1] - t1u[1]).max(), 0, 12)
        self.assertAlmostEqual(abs(t2[0] - t2u[0]).max(), 0, 12)
        self.assertAlmostEqual(abs(t2[1] - t2u[1]).max(), 0, 12)
        self.assertAlmostEqual(abs(t2[2] - t2u[2]).max(), 0, 12)

    def test_amplitudes_from_rccsd_or_uccsd(self):
        t1u = gcc1.spin2spatial(gcc1.t1)
        t2u = gcc1.spin2spatial(gcc1.t2)
        t1, t2 = gcc1.amplitudes_from_rccsd(t1u, t2u, mf.mo_coeff.orbspin)
        self.assertAlmostEqual(abs(t1[0] - gcc1.t1[0]).max(), 0, 12)
        self.assertAlmostEqual(abs(t1[1] - gcc1.t1[1]).max(), 0, 12)
        self.assertAlmostEqual(abs(t2[0] - gcc1.t2[0]).max(), 0, 12)
        self.assertAlmostEqual(abs(t2[1] - gcc1.t2[1]).max(), 0, 12)
        self.assertAlmostEqual(abs(t2[2] - gcc1.t2[2]).max(), 0, 12)

    def test_vector_size(self):
        self.assertEqual(gcc1.vector_size(), 5560)

    def test_update_amps(self):
        mol = gto.M()
        nocc, nvir = 8, 14
        nmo = nocc + nvir
        nmo_pair = nmo*(nmo+1)//2
        mf = scf.GHF(mol)
        numpy.random.seed(12)
        mf._eri = numpy.random.random(nmo_pair*(nmo_pair+1)//2)
        mf.mo_coeff = numpy.random.random((nmo,nmo))
        mf.mo_energy = numpy.arange(0., nmo)
        mf.mo_occ = numpy.zeros(nmo)
        mf.mo_occ[:nocc] = 1
        vhf = numpy.random.random((nmo,nmo)) + numpy.random.random((nmo,nmo))+1j
        vhf = vhf + vhf.conj().T
        mf.get_veff = lambda *args: vhf
        cinv = numpy.linalg.inv(mf.mo_coeff)
        mf.get_hcore = lambda *args: (reduce(numpy.dot, (cinv.T*mf.mo_energy, cinv)) - vhf)
        nmo_pair = nmo*(nmo//2+1)//4
        mf._eri = numpy.random.random(nmo_pair*(nmo_pair+1)//2)
        mycc = gccsd.GCCSD(mf)
        eris = mycc.ao2mo()
        eris.oooo = eris.oooo + numpy.sin(eris.oooo)*1j
        eris.oooo = eris.oooo + eris.oooo.conj().transpose(2,3,0,1)
        eris.ooov = eris.ooov + numpy.sin(eris.ooov)*1j
        eris.oovv = eris.oovv + numpy.sin(eris.oovv)*1j
        eris.ovov = eris.ovov + numpy.sin(eris.ovov)*1j
        eris.ovov = eris.ovov + eris.ovov.conj().transpose(2,3,0,1)
        eris.ovvv = eris.ovvv + numpy.sin(eris.ovvv)*1j
        eris.vvvv = eris.vvvv + numpy.sin(eris.vvvv)*1j
        eris.vvvv = eris.vvvv + eris.vvvv.conj().transpose(2,3,0,1)
        a = numpy.random.random((nmo,nmo)) * .1
        eris.fock += a + a.T
        t1 = numpy.random.random((nocc,nvir))*.1 + numpy.random.random((nocc,nvir))*.1j
        t2 = (numpy.random.random((nocc,nocc,nvir,nvir))*.1 +
              numpy.random.random((nocc,nocc,nvir,nvir))*.1j)
        t2 = t2 - t2.transpose(0,1,3,2)
        t2 = t2 - t2.transpose(1,0,2,3)
        r1, r2 = mycc.vector_to_amplitudes(mycc.amplitudes_to_vector(t1, t2))
        self.assertAlmostEqual(abs(t1-r1).max(), 0, 14)
        self.assertAlmostEqual(abs(t2-r2).max(), 0, 14)

        t1a, t2a = mycc.update_amps(t1, t2, eris)
        self.assertAlmostEqual(lib.finger(t1a), 20.805393111419136-300.1138026015621j, 9)
        self.assertAlmostEqual(lib.finger(t2a),-313.54117398035567+8.3700078645035205j, 9)

    def test_rdm_real(self):
        nocc = 6
        nvir = 10
        mf = scf.GHF(mol)
        nmo = nocc + nvir
        npair = nmo*(nmo//2+1)//4
        numpy.random.seed(12)
        mf._eri = numpy.random.random(npair*(npair+1)//2)*.3
        hcore = numpy.random.random((nmo,nmo)) * .5
        hcore = hcore + hcore.T + numpy.diag(range(nmo))*2
        mf.get_hcore = lambda *args: hcore
        mf.get_ovlp = lambda *args: numpy.eye(nmo)
        mf.mo_coeff = numpy.eye(nmo)
        mf.mo_occ = numpy.zeros(nmo)
        mf.mo_occ[:nocc] = 1
        mycc = gccsd.GCCSD(mf)
        ecc, t1, t2 = mycc.kernel()
        l1, l2 = mycc.solve_lambda()
        dm1 = gccsd_rdm.make_rdm1(mycc, t1, t2, l1, l2)
        dm2 = gccsd_rdm.make_rdm2(mycc, t1, t2, l1, l2)
        nao = nmo // 2
        mo_a = mf.mo_coeff[:nao]
        mo_b = mf.mo_coeff[nao:]
        eri  = ao2mo.kernel(mf._eri, mo_a)
        eri += ao2mo.kernel(mf._eri, mo_b)
        eri1 = ao2mo.kernel(mf._eri, (mo_a,mo_a,mo_b,mo_b))
        eri += eri1
        eri += eri1.T
        eri = ao2mo.restore(1, eri, nmo)
        h1 = reduce(numpy.dot, (mf.mo_coeff.T.conj(), hcore, mf.mo_coeff))
        e1 = numpy.einsum('ij,ji', h1, dm1)
        e1+= numpy.einsum('ijkl,ijkl', eri, dm2) * .5
        e1+= mol.energy_nuc()
        self.assertAlmostEqual(e1, mycc.e_tot, 7)

        self.assertAlmostEqual(abs(dm2-dm2.transpose(1,0,3,2).conj()).max(), 0, 9)
        self.assertAlmostEqual(abs(dm2-dm2.transpose(2,3,0,1)       ).max(), 0, 9)
        self.assertAlmostEqual(abs(dm2+dm2.transpose(2,1,0,3)       ).max(), 0, 9)
        self.assertAlmostEqual(abs(dm2+dm2.transpose(0,3,2,1)       ).max(), 0, 9)

    def test_rdm_real1(self):
        numpy.random.seed(1)
        mol = gto.M(
            atom = 'H 0 0 0; F 0 0 1.1',
            basis = '321g',
            spin = 2)
        myhf = mol.GHF()
        hcoreX = myhf.get_hcore()
        # a small random potential to break the Sz symmetry:
        pot = (numpy.random.random(hcoreX.shape) - 0.5) * 3e-2
        pot = pot + pot.T
        hcoreX += pot
        myhf.get_hcore = lambda *args: hcoreX
        myhf.kernel()
        mycc = myhf.CCSD().run()
        mycc.solve_lambda()
        rdm1 = mycc.make_rdm1(ao_repr=False)
        rdm2 = mycc.make_rdm2(ao_repr=False)

        # integrals in MO basis 
        nao = mol.nao
        C_ao_mo = myhf.mo_coeff
        hcore = reduce(numpy.dot, (C_ao_mo.conj().T, myhf.get_hcore(), C_ao_mo))
        mo_a = C_ao_mo[:nao]
        mo_b = C_ao_mo[nao:]
        eri  = ao2mo.kernel(myhf._eri, mo_a)
        eri += ao2mo.kernel(myhf._eri, mo_b)
        eri1 = ao2mo.kernel(myhf._eri, (mo_a, mo_a, mo_b, mo_b))
        eri += eri1
        eri += eri1.T
        eri1 = None
        eri_s1 = ao2mo.restore(1, eri, nao*2)
        E_ref = myhf.energy_nuc()
        E_ref += numpy.einsum('pq, qp ->', hcore, rdm1)
        E_ref += 0.5 * numpy.einsum('pqrs, pqrs ->', eri_s1, rdm2)
        self.assertAlmostEqual(E_ref, mycc.e_tot, 7)

    def test_rdm_complex(self):
        mol = gto.M()
        mol.verbose = 0
        nocc = 6
        nvir = 8
        mf = scf.GHF(mol)
        nmo = nocc + nvir
        numpy.random.seed(1)
        eri = (numpy.random.random((nmo,nmo,nmo,nmo)) +
               numpy.random.random((nmo,nmo,nmo,nmo))* 1j - (.5+.5j))
        eri = eri + eri.transpose(1,0,3,2).conj()
        eri = eri + eri.transpose(2,3,0,1)
        eri *= .1

        def get_jk(mol, dm, *args,**kwargs):
            vj = numpy.einsum('ijkl,lk->ij', eri, dm)
            vk = numpy.einsum('ijkl,jk->il', eri, dm)
            return vj, vk
        def get_veff(mol, dm, *args, **kwargs):
            vj, vk = get_jk(mol, dm)
            return vj - vk
        def ao2mofn(mos):
            return eri

        mf.get_jk = get_jk
        mf.get_veff = get_veff
        hcore = numpy.random.random((nmo,nmo)) * .2 + numpy.random.random((nmo,nmo))* .2j
        hcore = hcore + hcore.T.conj() + numpy.diag(range(nmo))*2
        mf.get_hcore = lambda *args: hcore
        mf.get_ovlp = lambda *args: numpy.eye(nmo)
        orbspin = numpy.zeros(nmo, dtype=int)
        orbspin[1::2] = 1
        mf.mo_coeff = lib.tag_array(numpy.eye(nmo) + 0j, orbspin=orbspin)
        mf.mo_occ = numpy.zeros(nmo)
        mf.mo_occ[:nocc] = 1

        mycc = cc.GCCSD(mf)
        eris = gccsd._make_eris_incore(mycc, mf.mo_coeff, ao2mofn)
        mycc.ao2mo = lambda *args, **kwargs: eris
        mycc.kernel(eris=eris)
        mycc.solve_lambda(eris=eris)
        dm1 = mycc.make_rdm1()
        dm2 = mycc.make_rdm2()

        e1 = numpy.einsum('ij,ji', hcore, dm1)
        e1+= numpy.einsum('ijkl,ijkl', eri, dm2) * .5
        self.assertAlmostEqual(e1, mycc.e_tot, 7)

        self.assertAlmostEqual(abs(dm2-dm2.transpose(1,0,3,2).conj()).max(), 0, 9)
        self.assertAlmostEqual(abs(dm2-dm2.transpose(2,3,0,1)       ).max(), 0, 9)
        self.assertAlmostEqual(abs(dm2+dm2.transpose(2,1,0,3)       ).max(), 0, 9)
        self.assertAlmostEqual(abs(dm2+dm2.transpose(0,3,2,1)       ).max(), 0, 9)

    def test_rdm_vs_uccsd(self):
        mol = gto.Mole()
        mol.atom = [
            [8 , (0. , 0.     , 0.)],
            [1 , (0. , -0.757 , 0.587)],
            [1 , (0. , 0.757  , 0.587)]]
        mol.verbose = 5
        mol.output = '/dev/null'
        mol.basis = '631g'
        mol.spin = 2
        mol.build()
        mf = scf.UHF(mol).run()
        myucc = uccsd.UCCSD(mf)
        myucc.frozen = 1
        myucc.kernel()
        udm1 = myucc.make_rdm1()
        udm2 = myucc.make_rdm2()

        mf = scf.addons.convert_to_ghf(mf)
        mygcc = gccsd.GCCSD(mf)
        mygcc.frozen = 2
        ecc, t1, t2 = mygcc.kernel()
        l1, l2 = mygcc.solve_lambda()
        dm1 = gccsd_rdm.make_rdm1(mygcc, t1, t2, l1, l2)
        dm2 = gccsd_rdm.make_rdm2(mygcc, t1, t2, l1, l2)

        nao = mol.nao_nr()
        mo_a = mf.mo_coeff[:nao]
        mo_b = mf.mo_coeff[nao:]
        nmo = mo_a.shape[1]
        eri = ao2mo.kernel(mf._eri, mo_a+mo_b, compact=False).reshape([nmo]*4)
        orbspin = mf.mo_coeff.orbspin
        sym_forbid = (orbspin[:,None] != orbspin)
        eri[sym_forbid,:,:] = 0
        eri[:,:,sym_forbid] = 0
        hcore = scf.RHF(mol).get_hcore()
        h1 = reduce(numpy.dot, (mo_a.T.conj(), hcore, mo_a))
        h1+= reduce(numpy.dot, (mo_b.T.conj(), hcore, mo_b))
        e1 = numpy.einsum('ij,ji', h1, dm1)
        e1+= numpy.einsum('ijkl,ijkl', eri, dm2) * .5
        e1+= mol.energy_nuc()
        self.assertAlmostEqual(e1, mygcc.e_tot, 7)

        idxa = numpy.where(orbspin == 0)[0]
        idxb = numpy.where(orbspin == 1)[0]
        self.assertAlmostEqual(abs(dm1[idxa[:,None],idxa] - udm1[0]).max(), 0, 7)
        self.assertAlmostEqual(abs(dm1[idxb[:,None],idxb] - udm1[1]).max(), 0, 7)
        self.assertAlmostEqual(abs(dm2[idxa[:,None,None,None],idxa[:,None,None],idxa[:,None],idxa] - udm2[0]).max(), 0, 7)
        self.assertAlmostEqual(abs(dm2[idxa[:,None,None,None],idxa[:,None,None],idxb[:,None],idxb] - udm2[1]).max(), 0, 7)
        self.assertAlmostEqual(abs(dm2[idxb[:,None,None,None],idxb[:,None,None],idxb[:,None],idxb] - udm2[2]).max(), 0, 7)

        ut1 = [0] * 2
        ul1 = [0] * 2
        ut2 = [0] * 3
        ul2 = [0] * 3
        ut1[0] = myucc.t1[0] + numpy.cos(myucc.t1[0]) * .2j
        ut1[1] = myucc.t1[1] + numpy.cos(myucc.t1[1]) * .2j
        ul1[0] = myucc.l1[0] + numpy.cos(myucc.l1[0]) * .2j
        ul1[1] = myucc.l1[1] + numpy.cos(myucc.l1[1]) * .2j
        ut2[0] = myucc.t2[0] + numpy.sin(myucc.t2[0]) * .8j
        ut2[1] = myucc.t2[1] + numpy.sin(myucc.t2[1]) * .8j
        ut2[2] = myucc.t2[2] + numpy.sin(myucc.t2[2]) * .8j
        ul2[0] = myucc.l2[0] + numpy.sin(myucc.l2[0]) * .8j
        ul2[1] = myucc.l2[1] + numpy.sin(myucc.l2[1]) * .8j
        ul2[2] = myucc.l2[2] + numpy.sin(myucc.l2[2]) * .8j
        udm1 = myucc.make_rdm1(ut1, ut2, ul1, ul2)
        udm2 = myucc.make_rdm2(ut1, ut2, ul1, ul2)

        gt1 = mygcc.spatial2spin(ut1)
        gt2 = mygcc.spatial2spin(ut2)
        gl1 = mygcc.spatial2spin(ul1)
        gl2 = mygcc.spatial2spin(ul2)
        gdm1 = mygcc.make_rdm1(gt1, gt2, gl1, gl2)
        gdm2 = mygcc.make_rdm2(gt1, gt2, gl1, gl2)

        self.assertAlmostEqual(abs(gdm1[idxa[:,None],idxa] - udm1[0]).max(), 0, 9)
        self.assertAlmostEqual(abs(gdm1[idxb[:,None],idxb] - udm1[1]).max(), 0, 9)
        self.assertAlmostEqual(abs(gdm2[idxa[:,None,None,None],idxa[:,None,None],idxa[:,None],idxa] - udm2[0]).max(), 0, 9)
        self.assertAlmostEqual(abs(gdm2[idxa[:,None,None,None],idxa[:,None,None],idxb[:,None],idxb] - udm2[1]).max(), 0, 9)
        self.assertAlmostEqual(abs(gdm2[idxb[:,None,None,None],idxb[:,None,None],idxb[:,None],idxb] - udm2[2]).max(), 0, 9)

    def test_rdm_vs_rccsd(self):
        mol = gto.Mole()
        mol.atom = [
            [8 , (0. , 0.     , 0.)],
            [1 , (0. , -0.757 , 0.587)],
            [1 , (0. , 0.757  , 0.587)]]
        mol.verbose = 5
        mol.output = '/dev/null'
        mol.basis = '631g'
        mol.build()
        mf = scf.RHF(mol).run()
        myrcc = ccsd.CCSD(mf).set(diis_start_cycle=1).run()
        rdm1 = myrcc.make_rdm1()
        rdm2 = myrcc.make_rdm2()

        mf = scf.addons.convert_to_ghf(mf)
        mygcc = gccsd.GCCSD(mf).set(diis_start_cycle=1).run()
        dm1 = mygcc.make_rdm1()
        dm2 = mygcc.make_rdm2()

        nao = mol.nao_nr()
        mo_a = mf.mo_coeff[:nao]
        mo_b = mf.mo_coeff[nao:]
        nmo = mo_a.shape[1]
        eri = ao2mo.kernel(mf._eri, mo_a+mo_b, compact=False).reshape([nmo]*4)
        orbspin = mf.mo_coeff.orbspin
        sym_forbid = (orbspin[:,None] != orbspin)
        eri[sym_forbid,:,:] = 0
        eri[:,:,sym_forbid] = 0
        hcore = scf.RHF(mol).get_hcore()
        h1 = reduce(numpy.dot, (mo_a.T.conj(), hcore, mo_a))
        h1+= reduce(numpy.dot, (mo_b.T.conj(), hcore, mo_b))
        e1 = numpy.einsum('ij,ji', h1, dm1)
        e1+= numpy.einsum('ijkl,ijkl', eri, dm2) * .5
        e1+= mol.energy_nuc()
        self.assertAlmostEqual(e1, mygcc.e_tot, 7)

        idxa = numpy.where(orbspin == 0)[0]
        idxb = numpy.where(orbspin == 1)[0]
        trdm1 = dm1[idxa[:,None],idxa]
        trdm1+= dm1[idxb[:,None],idxb]
        trdm2 = dm2[idxa[:,None,None,None],idxa[:,None,None],idxa[:,None],idxa]
        trdm2+= dm2[idxb[:,None,None,None],idxb[:,None,None],idxb[:,None],idxb]
        dm2ab = dm2[idxa[:,None,None,None],idxa[:,None,None],idxb[:,None],idxb]
        trdm2+= dm2ab
        trdm2+= dm2ab.transpose(2,3,0,1)
        self.assertAlmostEqual(abs(trdm1 - rdm1).max(), 0, 6)
        self.assertAlmostEqual(abs(trdm2 - rdm2).max(), 0, 6)

        rt1 = myrcc.t1 + numpy.cos(myrcc.t1) * .2j
        rl1 = myrcc.l1 + numpy.cos(myrcc.l1) * .2j
        rt2 = myrcc.t2 + numpy.sin(myrcc.t2) * .8j
        rl2 = myrcc.l2 + numpy.sin(myrcc.l2) * .8j
        rdm1 = myrcc.make_rdm1(rt1, rt2, rl1, rl2)
        rdm2 = myrcc.make_rdm2(rt1, rt2, rl1, rl2)

        gt1 = mygcc.spatial2spin(rt1)
        gt2 = mygcc.spatial2spin(rt2)
        gl1 = mygcc.spatial2spin(rl1)
        gl2 = mygcc.spatial2spin(rl2)
        gdm1 = mygcc.make_rdm1(gt1, gt2, gl1, gl2)
        gdm2 = mygcc.make_rdm2(gt1, gt2, gl1, gl2)

        trdm1 = gdm1[idxa[:,None],idxa]
        trdm1+= gdm1[idxb[:,None],idxb]
        trdm2 = gdm2[idxa[:,None,None,None],idxa[:,None,None],idxa[:,None],idxa]
        trdm2+= gdm2[idxb[:,None,None,None],idxb[:,None,None],idxb[:,None],idxb]
        dm2ab = gdm2[idxa[:,None,None,None],idxa[:,None,None],idxb[:,None],idxb]
        trdm2+= dm2ab
        trdm2+= dm2ab.transpose(2,3,0,1)
        self.assertAlmostEqual(abs(trdm1 - rdm1).max(), 0, 9)
        self.assertAlmostEqual(abs(trdm2 - rdm2).max(), 0, 9)

    def test_mbpt2(self):
        mygcc = gccsd.GCCSD(mf)
        e = mygcc.kernel(mbpt2=True)[0]
        self.assertAlmostEqual(e, -0.096257842171487293, 9)
        emp2 = mp.MP2(mf).kernel()[0]
        self.assertAlmostEqual(e, emp2, 10)


if __name__ == "__main__":
    print("Tests for GCCSD")
    unittest.main()


# pyscf/cc/test/test_gccsd_lambda.py

#!/usr/bin/env python
# Copyright 2014-2018 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.

import unittest
import numpy
import numpy as np
from functools import reduce

from pyscf import ao2mo
from pyscf import lib
from pyscf import gto
from pyscf import scf
from pyscf import cc
from pyscf.cc import addons
from pyscf.cc import gccsd_lambda

def setUpModule():
    global mol
    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = '631g'
    mol.spin = 2
    mol.build()

def tearDownModule():
    global mol
    del mol

class ccsd:
    def __init__(self,nso,nelec,h1e,int2e,h1e_is_fock=False):
        self.nso = nso
        self.nocc = nelec
        self.nvir = nso-nelec
        self.t1 = None
        self.t2 = None
        self.tau = None
        self.tau_tilde = None
        self.L1 = None
        self.L2 = None
        self.rspdm1 = None

        self.fock = h1e.copy()

        if h1e_is_fock is False:
            for p in range(self.nso):
                for q in range(self.nso):
                    for i in range(self.nocc):
                        self.fock[p,q] += int2e[p,i,q,i]

        self.f_oo = self.fock[:self.nocc,:self.nocc]
        self.f_ov = self.fock[:self.nocc,self.nocc:]
        self.f_vo = self.fock[self.nocc:,:self.nocc]
        self.f_vv = self.fock[self.nocc:,self.nocc:]
        self.w_oooo = int2e[:self.nocc,:self.nocc,:self.nocc,:self.nocc]
        self.w_ooov = int2e[:self.nocc,:self.nocc,:self.nocc,self.nocc:]
        self.w_oovo = int2e[:self.nocc,:self.nocc,self.nocc:,:self.nocc]
        self.w_ovoo = int2e[:self.nocc,self.nocc:,:self.nocc,:self.nocc]
        self.w_vooo = int2e[self.nocc:,:self.nocc,:self.nocc,:self.nocc]
        self.w_oovv = int2e[:self.nocc,:self.nocc,self.nocc:,self.nocc:]
        self.w_vvoo = int2e[self.nocc:,self.nocc:,:self.nocc,:self.nocc]
        self.w_ovov = int2e[:self.nocc,self.nocc:,:self.nocc,self.nocc:]
        self.w_ovvo = int2e[:self.nocc,self.nocc:,self.nocc:,:self.nocc]
        self.w_vovo = int2e[self.nocc:,:self.nocc,self.nocc:,:self.nocc]
        self.w_ovvv = int2e[:self.nocc,self.nocc:,self.nocc:,self.nocc:]
        self.w_vovv = int2e[self.nocc:,:self.nocc,self.nocc:,self.nocc:]
        self.w_vvov = int2e[self.nocc:,self.nocc:,:self.nocc,self.nocc:]
        self.w_vvvo = int2e[self.nocc:,self.nocc:,self.nocc:,:self.nocc]
        self.w_vvvv = int2e[self.nocc:,self.nocc:,self.nocc:,self.nocc:]

def update_l1l2_sub(mycc, t1, t2, l1, l2):
    l1new  = mycc.f_ov.copy()
    l1new += np.einsum('jiba,jb->ia', mycc.w_oovv, t1)
    l1new += np.einsum('ib,ba->ia', l1, mycc.f_vv)
    l1new -= np.einsum('ja,ij->ia', l1, mycc.f_oo)
    l1new += np.einsum('jb,ibaj->ia', l1, mycc.w_ovvo)
    l1new -= np.einsum('ib,ja,jb->ia', l1, mycc.f_ov, t1)
    l1new -= np.einsum('ja,ib,jb->ia', l1, mycc.f_ov, t1)
    l1new += np.einsum('jc,ciba,jb->ia', l1, mycc.w_vovv, t1)
    l1new -= np.einsum('ic,cjba,jb->ia', l1, mycc.w_vovv, t1)
    l1new -= np.einsum('kb,jika,jb->ia', l1, mycc.w_ooov, t1)
    l1new += np.einsum('ka,jikb,jb->ia', l1, mycc.w_ooov, t1)
    l1new += np.einsum('kc  ,kica->ia', np.einsum('jb,jkbc->kc  ', l1, t2), mycc.w_oovv)
    l1new += np.einsum('jkic,jkca->ia', np.einsum('ib,jkbc->jkic', l1, t2), mycc.w_oovv) * .5
    l1new += np.einsum('akbc,kibc->ia', np.einsum('ja,jkbc->akbc', l1, t2), mycc.w_oovv) * .5
    l1new -= np.einsum('kb  ,kiba->ia', np.einsum('jc,kc,jb->kb  ', l1, t1, t1), mycc.w_oovv)
    l1new += np.einsum('ikjb,kjba->ia', np.einsum('ic,kc,jb->ikjb', l1, t1, t1), mycc.w_oovv)
    l1new += np.einsum('acjb,jicb->ia', np.einsum('ka,kc,jb->acjb', l1, t1, t1), mycc.w_oovv)
    l1new += np.einsum('jica,cj->ia', l2, np.einsum('cb,jb->cj', mycc.f_vv, t1))
    l1new -= np.einsum('kiba,kb->ia', l2, np.einsum('jk,jb->kb', mycc.f_oo, t1))
    l1new += np.einsum('jicd,cdja->ia', l2, np.einsum('cdba,jb->cdja', mycc.w_vvvv, t1)) * .5
    l1new += np.einsum('klba,bikl->ia', l2, np.einsum('jikl,jb->bikl', mycc.w_oooo, t1)) * .5
    l1new -= np.einsum('kicb,bcak->ia', l2, np.einsum('jcak,jb->bcak', mycc.w_ovvo, t1))

    l1new -= np.einsum('kjca,icjk->ia', l2, np.einsum('icbk,jb->icjk', mycc.w_ovvo, t1))
    l1new += np.einsum('kica,ck  ->ia', l2, np.einsum('jcbk,jb->ck  ', mycc.w_ovvo, t1))
    l1new += np.einsum('jiba,bj->ia', l2, mycc.f_vo)
    l1new -= np.einsum('jibc,bcaj->ia', l2, mycc.w_vvvo) * .5
    l1new += np.einsum('jkba,ibjk->ia', l2, mycc.w_ovoo) * .5
    l1new -= np.einsum('jibc,jabc->ia', l2, np.einsum('ka,jkbc->jabc', mycc.f_ov, t2)) * .5
    l1new -= np.einsum('jkba,jkbi->ia', l2, np.einsum('ic,jkbc->jkbi', mycc.f_ov, t2)) * .5
    l1new += np.einsum('jiba,jb  ->ia', l2, np.einsum('kc,jkbc->jb  ', mycc.f_ov, t2))
    l1new -= np.einsum('jica,jc->ia', l2, np.einsum('kb,kc,jb->jc', mycc.f_ov, t1, t1))
    l1new -= np.einsum('dc,dica->ia', np.einsum('jkdb,jkbc->dc', l2, t2), mycc.w_vovv) * .5
    l1new += np.einsum('jidb,dajb->ia', l2, np.einsum('dkca,jkbc->dajb', mycc.w_vovv, t2))
    l1new -= np.einsum('jkda,dijk->ia', l2, np.einsum('dibc,jkbc->dijk', mycc.w_vovv, t2)) * .25
    l1new += np.einsum('jida,jd->ia', l2, np.einsum('dkbc,jkbc->jd', mycc.w_vovv, t2)) * .5
    l1new += np.einsum('kila,lk->ia', mycc.w_ooov, np.einsum('ljbc,jkbc->lk', l2, t2)) * .5
    l1new += np.einsum('libc,labc->ia', l2, np.einsum('jkla,jkbc->labc', mycc.w_ooov, t2)) * .25
    l1new -= np.einsum('ljba,iljb->ia', l2, np.einsum('kilc,jkbc->iljb', mycc.w_ooov, t2))
    l1new -= np.einsum('liba,lb  ->ia', l2, np.einsum('jklc,jkbc->lb  ', mycc.w_ooov, t2)) * .5
    l1new -= np.einsum('jidc,dcja->ia', l2, np.einsum('dkba,kc,jb->dcja', mycc.w_vovv, t1, t1))
    l1new -= np.einsum('kjda,dikj->ia', l2, np.einsum('dicb,kc,jb->dikj', mycc.w_vovv, t1, t1)) * .5
    l1new += np.einsum('kida,dk  ->ia', l2, np.einsum('djcb,kc,jb->dk  ', mycc.w_vovv, t1, t1))
    l1new += np.einsum('licb,cbla->ia', l2, np.einsum('kjla,kc,jb->cbla', mycc.w_ooov, t1, t1)) * .5

    l1new += np.einsum('ljca,cilj->ia', l2, np.einsum('kilb,kc,jb->cilj', mycc.w_ooov, t1, t1))
    l1new -= np.einsum('lica,cl  ->ia', l2, np.einsum('kjlb,kc,jb->cl  ', mycc.w_ooov, t1, t1))
    l1new -= np.einsum('jl,lija->ia  ', np.einsum('kjcd,klcd->jl  ', l2, t2), np.einsum('liba,jb->lija', mycc.w_oovv, t1)) * .5
    l1new += np.einsum('il,la->ia    ', np.einsum('kicd,klcd->il  ', l2, t2), np.einsum('ljba,jb->la  ', mycc.w_oovv, t1)) * .5
    l1new += np.einsum('jikl,klja->ia', np.einsum('jicd,klcd->jikl', l2, t2), np.einsum('klba,jb->klja', mycc.w_oovv, t1)) * .25
    l1new -= np.einsum('bd,bida->ia  ', np.einsum('klcb,klcd->bd', l2, t2), np.einsum('jida,jb->bida', mycc.w_oovv, t1)) * .5
    l1new += np.einsum('ad,id->ia    ', np.einsum('klca,klcd->ad', l2, t2), np.einsum('jidb,jb->id  ', mycc.w_oovv, t1)) * .5
    l1new -= np.einsum('kicj,jakc->ia', np.einsum('kicb,jb->kicj', l2, t1), np.einsum('ljda,klcd->jakc', mycc.w_oovv, t2))
    l1new -= np.einsum('kbca,ibkc->ia', np.einsum('kjca,jb->kbca', l2, t1), np.einsum('lidb,klcd->ibkc', mycc.w_oovv, t2))
    l1new += np.einsum('iald,ld->ia  ', np.einsum('kica,klcd->iald', l2, t2), np.einsum('ljdb,jb->ld', mycc.w_oovv, t1))
    l1new += np.einsum('bica,bc->ia  ', np.einsum('jica,jb->bica', l2, t1), np.einsum('kldb,klcd->bc', mycc.w_oovv, t2)) * .5
    l1new += np.einsum('klja,jikl->ia', np.einsum('klba,jb->klja', l2, t1), np.einsum('jicd,klcd->jikl', mycc.w_oovv, t2)) * .25
    l1new += np.einsum('kija,jk  ->ia', np.einsum('kiba,jb->kija', l2, t1), np.einsum('ljcd,klcd->jk  ', mycc.w_oovv, t2)) * .5
    l1new += np.einsum('jilk,lkja->ia', np.einsum('jidc,ld,kc->jilk', l2, t1, t1), np.einsum('lkba,jb->lkja', mycc.w_oovv, t1)) * .5
    l1new += np.einsum('kjla,likj->ia', np.einsum('kjda,ld->kjla', l2, t1), np.einsum('licb,kc,jb->likj', mycc.w_oovv, t1, t1)) * .5
    l1new -= np.einsum('kila,lk  ->ia', np.einsum('kida,ld->kila', l2, t1), np.einsum('ljcb,kc,jb->lk', mycc.w_oovv, t1, t1))

    def Pij(v2):
        return v2 - v2.transpose(1,0,2,3)
    def Pab(v2):
        return v2 - v2.transpose(0,1,3,2)
    def Pijab(v2):
        return Pij(Pab(v2))
    Pabij = Pijab

    l2new  =       mycc.w_oovv.copy()
    l2new += Pijab(np.einsum('ia,jb->ijab', l1, mycc.f_ov))
    l2new += Pijab(np.einsum('ia,kjcb,kc->ijab', l1, mycc.w_oovv, t1))
    l2new += Pij  (np.einsum('ic,cjab->ijab', l1, mycc.w_vovv))
    l2new -= Pab  (np.einsum('ka,ijkb->ijab', l1, mycc.w_ooov))
    l2new -= Pij  (np.einsum('ic,kc,kjab->ijab', l1, t1, mycc.w_oovv))
    l2new -= Pab  (np.einsum('ka,kc,ijcb->ijab', l1, t1, mycc.w_oovv))
    l2new += Pij  (np.einsum('kiab,jk->ijab', l2, mycc.f_oo))
    l2new -= Pab  (np.einsum('ijca,cb->ijab', l2, mycc.f_vv))
    l2new += np.einsum('ijcd,cdab->ijab', l2, mycc.w_vvvv) * .5
    l2new += np.einsum('klab,ijkl->ijab', l2, mycc.w_oooo) * .5
    l2new += Pabij(np.einsum('kica,jcbk->ijab', l2, mycc.w_ovvo))
    l2new += Pab  (np.einsum('ijca,kb,kc->ijab', l2, mycc.f_ov, t1))
    l2new += Pij  (np.einsum('kiab,jc,kc->ijab', l2, mycc.f_ov, t1))
    l2new -=      (np.einsum('ijdc,dcab->ijab', l2, np.einsum('dkab,kc->dcab', mycc.w_vovv, t1)))
    l2new += Pabij(np.einsum('kida,djkb->ijab', l2, np.einsum('djcb,kc->djkb', mycc.w_vovv, t1)))
    l2new += Pab  (np.einsum('ijda,db  ->ijab', l2, np.einsum('dkcb,kc->db  ', mycc.w_vovv, t1)))
    l2new -= Pabij(np.einsum('lica,cjlb->ijab', l2, np.einsum('kjlb,kc->cjlb', mycc.w_ooov, t1)))
    l2new +=      (np.einsum('lkab,ijlk->ijab', l2, np.einsum('ijlc,kc->ijlk', mycc.w_ooov, t1)))
    l2new -= Pij  (np.einsum('liab,jl  ->ijab', l2, np.einsum('kjlc,kc->jl  ', mycc.w_ooov, t1)))
    l2new -= Pij  (np.einsum('il  ,ljab->ijab', np.einsum('kicd,klcd->il  ', l2, t2), mycc.w_oovv)) * .5
    l2new +=      (np.einsum('ijkl,klab->ijab', np.einsum('ijcd,klcd->ijkl', l2, t2), mycc.w_oovv)) * .25
    l2new -= Pab  (np.einsum('ad  ,ijdb->ijab', np.einsum('klca,klcd->ad  ', l2, t2), mycc.w_oovv)) * .5
    l2new += Pabij(np.einsum('iald,ljdb->ijab', np.einsum('kica,klcd->iald', l2, t2), mycc.w_oovv))
    l2new -= Pab  (np.einsum('ijca,bc->ijab', l2, np.einsum('kldb,klcd->bc', mycc.w_oovv, t2))) * .5
    l2new +=      (np.einsum('klab,ijkl->ijab', l2, np.einsum('ijcd,klcd->ijkl', mycc.w_oovv, t2))) * .25
    l2new -= Pij  (np.einsum('kiab,jk  ->ijab', l2, np.einsum('ljcd,klcd->jk  ', mycc.w_oovv, t2))) * .5
    l2new +=      (np.einsum('ijlk,lkab->ijab', np.einsum('ijdc,ld,kc->ijlk', l2, t1, t1), mycc.w_oovv)) * .5
    l2new -= Pabij(np.einsum('kida,djkb->ijab', l2, np.einsum('ljcb,ld,kc->djkb', mycc.w_oovv, t1, t1)))
    l2new -= Pab  (np.einsum('ijda,db  ->ijab', l2, np.einsum('lkcb,kc,ld->db', mycc.w_oovv, t1, t1)))
    l2new +=      (np.einsum('lkab,ijlk->ijab', l2, np.einsum('ijdc,ld,kc->ijlk', mycc.w_oovv, t1, t1))) * .5
    l2new -= Pij  (np.einsum('liab,jl  ->ijab', l2, np.einsum('kjdc,ld,kc->jl', mycc.w_oovv, t1, t1)))

    eia = mycc.f_oo.diagonal()[:,None] - mycc.f_vv.diagonal()
    l1new = l1 + l1new / eia
    eijab = lib.direct_sum('ia+jb->ijab', eia, eia)
    l2new = l2 + l2new / eijab
    return l1new, l2new

def update_l1l2(mf, t1, t2, l1, l2, orbspin):
    mol = mf.mol
    nao,nmo = mf.mo_coeff[0].shape
    nelec = mol.nelectron
    nso = nmo * 2
    hcore = mf.get_hcore()
    mo = np.zeros((nao,nso))
    mo[:,orbspin==0] = mf.mo_coeff[0]
    mo[:,orbspin==1] = mf.mo_coeff[1]
    h1e = reduce(np.dot, (mo.T, hcore, mo))
    h1e[orbspin[:,None]!=orbspin] = 0
    int2e = ao2mo.full(mf._eri, mo)
    sym_forbid = (orbspin[:,None] != orbspin)[np.tril_indices(nso)]
    int2e[sym_forbid] = 0
    int2e[:,sym_forbid] = 0
    int2e = ao2mo.restore(1, int2e, nso)
    int2e = int2e.transpose(0,2,1,3)
    int2e = int2e - int2e.transpose(0,1,3,2)
    mycc = ccsd(nso,nelec,h1e,int2e,h1e_is_fock=False)
    l1, l2 = update_l1l2_sub(mycc, t1, t2, l1, l2)

    return l1, l2


class KnownValues(unittest.TestCase):
    def test_update_amps(self):
        mf = scf.UHF(mol).run()
        numpy.random.seed(21)
        mycc = cc.GCCSD(mf)
        eris = mycc.ao2mo()
        orbspin = eris.orbspin

        nocc = mol.nelectron
        nvir = mol.nao_nr()*2 - nocc

        t1 = numpy.random.random((nocc,nvir))*.1
        t2 = numpy.random.random((nocc,nocc,nvir,nvir))*.1
        t2 = t2 - t2.transpose(1,0,2,3)
        t2 = t2 - t2.transpose(0,1,3,2)
        l1 = numpy.random.random((nocc,nvir))*.1
        l2 = numpy.random.random((nocc,nocc,nvir,nvir))*.1
        l2 = l2 - l2.transpose(1,0,2,3)
        l2 = l2 - l2.transpose(0,1,3,2)
        l1ref, l2ref = update_l1l2(mf, t1, t2, l1, l2, orbspin)

        imds = gccsd_lambda.make_intermediates(mycc, t1, t2, eris)
        l1, l2 = gccsd_lambda.update_lambda(mycc, t1, t2, l1, l2, eris, imds)
        self.assertAlmostEqual(abs(l1-l1ref).max(), 0, 8)
        self.assertAlmostEqual(abs(l2-l2ref).max(), 0, 8)


if __name__ == "__main__":
    print("Full Tests for GCCSD lambda")
    unittest.main()


# pyscf/cc/test/test_gccsd_t.py

#!/usr/bin/env python
# Copyright 2014-2018 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.

import unittest
import numpy
import copy
from functools import reduce

from pyscf import gto, scf, lib, symm
from pyscf import cc
from pyscf.cc import uccsd_t
from pyscf.cc import gccsd_t

def setUpModule():
    global mol, mol1, mf, myucc, mygcc
    mol = gto.Mole()
    mol.verbose = 5
    mol.output = '/dev/null'
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -.757 , .587)],
        [1 , (0. ,  .757 , .587)]]
    mol.spin = 2
    mol.basis = '3-21g'
    mol.symmetry = 'C2v'
    mol.build()
    mol1 = copy.copy(mol)
    mol1.symmetry = False

    mf = scf.UHF(mol1).run(conv_tol=1e-14)
    myucc = cc.UCCSD(mf).run()
    mygcc = cc.GCCSD(mf).run()

def tearDownModule():
    global mol, mol1, mf, myucc, mygcc
    mol.stdout.close()
    del mol, mol1, mf, myucc, mygcc

class KnownValues(unittest.TestCase):
    def test_gccsd_t_compare_uccsd_t(self):
        self.assertAlmostEqual(myucc.ccsd_t(), mygcc.ccsd_t(t1=None), 7)

    def test_gccsd_t(self):
        mf1 = copy.copy(mf)
        nao, nmo = mf.mo_coeff[0].shape
        numpy.random.seed(10)
        mf1.mo_coeff = numpy.random.random((2,nao,nmo))

        numpy.random.seed(12)
        nocca, noccb = mol.nelec
        nmo = mf1.mo_occ[0].size
        nvira = nmo - nocca
        nvirb = nmo - noccb
        t1a  = .1 * numpy.random.random((nocca,nvira))
        t1b  = .1 * numpy.random.random((noccb,nvirb))
        t2aa = .1 * numpy.random.random((nocca,nocca,nvira,nvira))
        t2aa = t2aa - t2aa.transpose(0,1,3,2)
        t2aa = t2aa - t2aa.transpose(1,0,2,3)
        t2bb = .1 * numpy.random.random((noccb,noccb,nvirb,nvirb))
        t2bb = t2bb - t2bb.transpose(0,1,3,2)
        t2bb = t2bb - t2bb.transpose(1,0,2,3)
        t2ab = .1 * numpy.random.random((nocca,noccb,nvira,nvirb))

        mycc = cc.GCCSD(mf1)
        t1 = mycc.spatial2spin((t1a, t1b        ))
        t2 = mycc.spatial2spin((t2aa, t2ab, t2bb))
        eris = mycc.ao2mo()
        e3a = gccsd_t.kernel(mycc, eris, t1, t2)
        self.assertAlmostEqual(e3a, 9877.2780859693339, 6)

    def test_gccsd_t_complex(self):
        nocc, nvir = 4, 6
        nmo = nocc + nvir
        numpy.random.seed(1)
        eris = cc.gccsd._PhysicistsERIs()
        h = (numpy.random.random((nmo,nmo)) +
             numpy.random.random((nmo,nmo)) * .6j - .5-.3j)
        eris.fock = h + h.T.conj() + numpy.diag(numpy.arange(nmo)) * 2
        eri1 = (numpy.random.random((nmo,nmo,nmo,nmo)) +
                numpy.random.random((nmo,nmo,nmo,nmo))*.8j - .5-.4j)
        eri1 = eri1 - eri1.transpose(0,1,3,2)
        eri1 = eri1 - eri1.transpose(1,0,2,3)
        eri1 = eri1 + eri1.transpose(2,3,0,1).conj()
        eris.ovvv = eri1[:nocc,nocc:,nocc:,nocc:]
        eris.oovv = eri1[:nocc,:nocc,nocc:,nocc:]
        eris.ooov = eri1[:nocc,:nocc,:nocc,nocc:]
        t2 = (numpy.random.random((nocc,nocc,nvir,nvir)) +
              numpy.random.random((nocc,nocc,nvir,nvir))*.8j - .5-.4j)
        t2 = t2 - t2.transpose(0,1,3,2)
        t2 = t2 - t2.transpose(1,0,2,3)
        t1 = (numpy.random.random((nocc,nvir)) +
              numpy.random.random((nocc,nvir))*.8j - .5-.4j)
        eris.mo_energy = eris.fock.diagonal().real

        gcc = cc.gccsd.GCCSD(scf.GHF(gto.M()))
        self.assertAlmostEqual(gccsd_t.kernel(gcc, eris, t1, t2),
                               (-104.15886718888137+0.30739952563327672j), 9)


if __name__ == "__main__":
    print("Full Tests for GCCSD(T)")
    unittest.main()



# pyscf/cc/test/test_h2o.py

#!/usr/bin/env python
# Copyright 2014-2018 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.

import unittest
import numpy
from functools import reduce

from pyscf import gto, lib
from pyscf import scf
from pyscf import cc
from pyscf import ao2mo
from pyscf.cc import ccsd
from pyscf.cc import uccsd
from pyscf.cc import gccsd
from pyscf.cc import rccsd
from pyscf.cc import dfccsd

def setUpModule():
    global mol, mf
    mol = gto.Mole()
    mol.verbose = 7
    mol.output = '/dev/null'
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]

    mol.basis = {'H': 'cc-pvdz',
                 'O': 'cc-pvdz',}
    mol.build()
    mf = scf.RHF(mol)
    mf.conv_tol_grad = 1e-8
    ehf = mf.kernel()

def tearDownModule():
    global mol, mf
    mol.stdout.close()
    del mol, mf


class KnownValues(unittest.TestCase):
    def test_ccsd(self):
        mcc = cc.ccsd.CC(mf)
        mcc.conv_tol = 1e-9
        mcc.conv_tol_normt = 1e-7
        eris = mcc.ao2mo()
        emp2, t1, t2 = mcc.init_amps(eris)
        self.assertAlmostEqual(abs(t2).sum(), 4.9556571211255909, 6)
        self.assertAlmostEqual(emp2, -0.2040199672883385, 10)
        t1, t2 = cc.ccsd.update_amps(mcc, t1, t2, eris)
        self.assertAlmostEqual(abs(t1).sum(), 0.0475038989126, 8)
        self.assertAlmostEqual(abs(t2).sum(), 5.4018238455030, 6)
        self.assertAlmostEqual(cc.ccsd.energy(mcc, t1, t2, eris),
                               -0.208967840546667, 10)
        t1, t2 = cc.ccsd.update_amps(mcc, t1, t2, eris)
        self.assertAlmostEqual(cc.ccsd.energy(mcc, t1, t2, eris),
                               -0.212173678670510, 10)
        self.assertAlmostEqual(abs(t1).sum(), 0.05470123093500083, 8)
        self.assertAlmostEqual(abs(t2).sum(), 5.5605208386554716, 6)

        mcc.kernel()
        self.assertTrue(numpy.allclose(mcc.t2,mcc.t2.transpose(1,0,3,2)))
        self.assertAlmostEqual(mcc.ecc, -0.2133432312951, 8)
        self.assertAlmostEqual(abs(mcc.t2).sum(), 5.63970279799556984, 6)

    def test_ccsd_frozen(self):
        mcc = cc.ccsd.CC(mf, frozen=range(1))
        mcc.conv_tol = 1e-10
        mcc.kernel()
        self.assertAlmostEqual(mcc.ecc, -0.21124878189922872, 8)
        self.assertAlmostEqual(abs(mcc.t2).sum(), 5.4996425901189347, 6)

    def test_ccsd_cart(self):
        pmol = mol.copy()
        pmol.cart = True
        pmol.build()
        mf = scf.RHF(pmol).set(conv_tol_grad=1e-8).run()
        mcc = cc.ccsd.CC(mf, frozen=range(1))
        mcc.conv_tol = 1e-10
        mcc.kernel()
        self.assertAlmostEqual(mcc.ecc, -0.21303885376969361, 8)

    def test_h2o_non_hf_orbital_high_cost(self):
        nmo = mf.mo_energy.size
        nocc = mol.nelectron // 2
        nvir = nmo - nocc
        u = numpy.eye(nmo)
        numpy.random.seed(1)
        u[:nocc,:nocc] = numpy.linalg.svd(numpy.random.random((nocc,nocc)))[0]
        u[nocc:,nocc:] = numpy.linalg.svd(numpy.random.random((nvir,nvir)))[0]
        mycc = cc.ccsd.CCSD(mf)
        mycc.conv_tol = 1e-12
        mycc.diis_start_energy_diff = 1e2
        mycc.max_cycle = 400
        mycc.mo_coeff = mo_coeff=numpy.dot(mf.mo_coeff,u)
        self.assertAlmostEqual(mycc.kernel()[0], -0.21334323320620596, 8)

## FIXME
#    def test_h2o_without_scf(self):
#        mycc = cc.ccsd.CCSD(mf)
#        nmo = mf.mo_energy.size
#        nocc = mol.nelectron // 2
#        nvir = nmo - nocc
#        numpy.random.seed(1)
#        u = numpy.eye(nmo) + numpy.random.random((nmo,nmo))*.2
#        u, w, vh = numpy.linalg.svd(u)
#        u = numpy.dot(u, vh)
#
#        mo1 = numpy.dot(mf.mo_coeff, u)
#        dm1 = mf.make_rdm1(mo1, mf.mo_occ)
#
#        mycc.diis_start_energy_diff = 1e2
#        mycc.max_cycle = 1000
#        mycc.conv_tol = 1e-12
#        mycc.mo_coeff = mo1
#        self.assertAlmostEqual(mf.energy_tot(dm1)+mycc.kernel()[0],
#                               ehf-0.21334323320620596, 8)

    def test_ccsd_lambda(self):
        mcc = cc.ccsd.CC(mf)
        mcc.conv_tol = 1e-9
        mcc.conv_tol_normt = 1e-7
        mcc.kernel()
        mcc.solve_lambda()
        self.assertAlmostEqual(numpy.linalg.norm(mcc.l1), 0.01326267012100099, 7)
        self.assertAlmostEqual(numpy.linalg.norm(mcc.l2), 0.21257559872380857, 7)

    def test_ccsd_rdm(self):
        mcc = cc.ccsd.CC(mf)
        mcc.max_memory = 0
        mcc.conv_tol = 1e-9
        mcc.conv_tol_normt = 1e-7
        mcc.kernel()
        mcc.solve_lambda()
        dm1 = mcc.make_rdm1()
        dm2 = mcc.make_rdm2()
        self.assertAlmostEqual(numpy.linalg.norm(dm1), 4.4227785269355078, 3)
        self.assertAlmostEqual(numpy.linalg.norm(dm2), 20.074587448789089, 3)

        nmo = mf.mo_coeff.shape[1]
        eri = ao2mo.restore(1, ao2mo.kernel(mf._eri, mf.mo_coeff), nmo)
        h1 = reduce(numpy.dot, (mf.mo_coeff.T.conj(), mf.get_hcore(), mf.mo_coeff))
        e1 = numpy.einsum('ij,ji', h1, dm1)
        e1+= numpy.einsum('ijkl,ijkl', eri, dm2) * .5
        e1+= mf.energy_nuc()
        self.assertAlmostEqual(e1, mcc.e_tot, 7)

    def test_scanner(self):
        cc_scanner = scf.RHF(mol).apply(cc.CCSD).as_scanner()
        cc_scanner.conv_tol = 1e-8
        self.assertAlmostEqual(cc_scanner(mol), -76.240108935038691, 6)
        geom = '''
        O   0.   0.       .1
        H   0.   -0.757   0.587
        H   0.   0.757    0.587'''
        cc_scanner = cc_scanner.as_scanner()
        self.assertAlmostEqual(cc_scanner(geom), -76.228972886940639, 6)

    def test_init(self):
        self.assertTrue(isinstance(cc.CCSD(mf), ccsd.CCSD))
        self.assertTrue(isinstance(cc.CCSD(mf.density_fit()), dfccsd.RCCSD))
        self.assertTrue(isinstance(cc.CCSD(mf.newton()), ccsd.CCSD))
        self.assertTrue(isinstance(cc.CCSD(mf.density_fit().newton()), dfccsd.RCCSD))
        self.assertTrue(isinstance(cc.CCSD(mf.newton().density_fit()), ccsd.CCSD))
        self.assertTrue(not isinstance(cc.CCSD(mf.newton().density_fit()), dfccsd.RCCSD))
        self.assertTrue(isinstance(cc.CCSD(mf.density_fit().newton().density_fit()), dfccsd.RCCSD))

        self.assertTrue(isinstance(cc.UCCSD(mf), uccsd.UCCSD))
#        self.assertTrue(isinstance(cc.UCCSD(mf.density_fit()), dfccsd.UCCSD))
        self.assertTrue(isinstance(cc.UCCSD(mf.newton()), uccsd.UCCSD))
#        self.assertTrue(isinstance(cc.UCCSD(mf.density_fit().newton()), dfccsd.UCCSD))
        self.assertTrue(isinstance(cc.UCCSD(mf.newton().density_fit()), uccsd.UCCSD))
#        self.assertTrue(not isinstance(cc.UCCSD(mf.newton().density_fit()), dfccsd.UCCSD))
#        self.assertTrue(isinstance(cc.UCCSD(mf.density_fit().newton().density_fit()), dfccsd.UCCSD))
        self.assertTrue(isinstance(cc.CCSD(scf.GHF(mol)), gccsd.GCCSD))

        umf = scf.convert_to_uhf(mf, scf.UHF(mol))
        self.assertTrue(isinstance(cc.CCSD(umf), uccsd.UCCSD))
#        self.assertTrue(isinstance(cc.CCSD(umf.density_fit()), dfccsd.UCCSD))
        self.assertTrue(isinstance(cc.CCSD(umf.newton()), uccsd.UCCSD))
#        self.assertTrue(isinstance(cc.CCSD(umf.density_fit().newton()), dfccsd.UCCSD))
        self.assertTrue(isinstance(cc.CCSD(umf.newton().density_fit()), uccsd.UCCSD))
#        self.assertTrue(not isinstance(cc.CCSD(umf.newton().density_fit()), dfccsd.UCCSD))
#        self.assertTrue(isinstance(cc.CCSD(umf.density_fit().newton().density_fit()), dfccsd.UCCSD))

        self.assertTrue(isinstance(cc.CCSD(mf, mo_coeff=mf.mo_coeff*1j), rccsd.RCCSD))


    # May be related to issue #509
    def test_race_condition_skip(self):
        current_nthreads = lib.num_threads()
        lib.num_threads(32)

        mycc = cc.ccsd.CCSD(mf)
        mycc.max_memory = 1
        mycc.conv_tol = 1e-10
        ecc = mycc.kernel()[0]
        self.assertAlmostEqual(ecc, -0.2133432312951, 8)

        lib.num_threads(current_nthreads)

if __name__ == "__main__":
    print("Full Tests for H2O")
    unittest.main()



# pyscf/cc/test/test_rccsd.py

#!/usr/bin/env python
# Copyright 2014-2018 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.

import tempfile
from functools import reduce
import unittest
import copy
import numpy
import numpy as np
import h5py

from pyscf import gto, lib
from pyscf import scf, dft
from pyscf import df
from pyscf import cc
from pyscf import ao2mo
from pyscf import mp
from pyscf.cc import ccsd
from pyscf.cc import rccsd

def setUpModule():
    global mol, mf, eris, mycc
    mol = gto.Mole()
    mol.verbose = 7
    mol.output = '/dev/null'
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]

    mol.basis = '631g'
    mol.build()
    mf = scf.RHF(mol)
    mf.chkfile = tempfile.NamedTemporaryFile().name
    mf.conv_tol_grad = 1e-8
    mf.kernel()

    mycc = rccsd.RCCSD(mf)
    mycc.conv_tol = 1e-10
    eris = mycc.ao2mo()
    mycc.kernel(eris=eris)

def tearDownModule():
    global mol, mf, eris, mycc
    mol.stdout.close()
    del mol, mf, eris, mycc


class KnownValues(unittest.TestCase):
    def test_roccsd(self):
        mf = scf.ROHF(mol).run()
        mycc = cc.RCCSD(mf).run()
        self.assertAlmostEqual(mycc.e_tot, -76.119346385357446, 7)

    def test_denisty_fit_interface(self):
        mydf = df.DF(mol)
        mycc1 = ccsd.CCSD(mf).density_fit(auxbasis='ccpvdz-ri', with_df=mydf).run()
        self.assertAlmostEqual(mycc1.e_tot, -76.119348934346789, 7)

    def test_ERIS(self):
        mycc = rccsd.RCCSD(mf)
        numpy.random.seed(1)
        mo_coeff = numpy.random.random(mf.mo_coeff.shape)
        eris = rccsd._make_eris_incore(mycc, mo_coeff)

        self.assertAlmostEqual(lib.finger(eris.oooo),  4.963884938282539, 11)
        self.assertAlmostEqual(lib.finger(eris.ovoo), -1.362368189698315, 11)
        self.assertAlmostEqual(lib.finger(eris.ovov),125.815506844421580, 11)
        self.assertAlmostEqual(lib.finger(eris.oovv), 55.123681017639463, 11)
        self.assertAlmostEqual(lib.finger(eris.ovvo),133.480835278982620, 11)
        self.assertAlmostEqual(lib.finger(eris.ovvv), 95.756230114113222, 11)
        self.assertAlmostEqual(lib.finger(eris.vvvv),-10.450387490987071, 11)

        ccsd.MEMORYMIN, bak = 0, ccsd.MEMORYMIN
        mycc.max_memory = 0
        eris1 = mycc.ao2mo(mo_coeff)
        ccsd.MEMORYMIN = bak
        self.assertAlmostEqual(abs(numpy.array(eris1.oooo)-eris.oooo).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.ovoo)-eris.ovoo).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.ovov)-eris.ovov).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.oovv)-eris.oovv).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.ovvo)-eris.ovvo).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.ovvv)-eris.ovvv).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.vvvv)-eris.vvvv).max(), 0, 11)

        # Testing the complex MO integrals
        def ao2mofn(mos):
            if isinstance(mos, numpy.ndarray) and mos.ndim == 2:
                mos = [mos]*4
            nmos = [mo.shape[1] for mo in mos]
            eri_mo = ao2mo.kernel(mf._eri, mos, compact=False).reshape(nmos)
            return eri_mo * 1j
        eris1 = rccsd._make_eris_incore(mycc, mo_coeff, ao2mofn=ao2mofn)
        self.assertAlmostEqual(abs(eris1.oooo.imag-eris.oooo).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.ovoo.imag-eris.ovoo).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.ovov.imag-eris.ovov).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.oovv.imag-eris.oovv).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.ovvo.imag-eris.ovvo).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.ovvv.imag-eris.ovvv).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.vvvv.imag-eris.vvvv).max(), 0, 11)

    def test_dump_chk(self):
        cc1 = copy.copy(mycc)
        cc1.nmo = mf.mo_energy.size
        cc1.nocc = mol.nelectron // 2
        cc1.dump_chk()
        cc1 = cc.CCSD(mf)
        cc1.__dict__.update(lib.chkfile.load(cc1.chkfile, 'ccsd'))
        e = cc1.energy(cc1.t1, cc1.t2, eris)
        self.assertAlmostEqual(e, -0.13539788638119823, 8)

    def test_ccsd_t(self):
        e = mycc.ccsd_t()
        self.assertAlmostEqual(e, -0.0009964234049929792, 10)

    def test_mbpt2(self):
        e = mycc.kernel(mbpt2=True)[0]
        #emp2 = mp.MP2(mf).kernel()[0]
        self.assertAlmostEqual(e, -0.12886859466216125, 10)

    def test_ao_direct(self):
        cc1 = cc.CCSD(mf)
        cc1.direct = True
        cc1.conv_tol = 1e-10
        cc1.kernel(t1=numpy.zeros_like(mycc.t1))
        self.assertAlmostEqual(cc1.e_corr, -0.13539788638119823, 8)

    def test_incore_complete(self):
        cc1 = cc.CCSD(mf)
        cc1.incore_complete = True
        cc1.conv_tol = 1e-10
        cc1.kernel()
        self.assertAlmostEqual(cc1.e_corr, -0.13539788638119823, 8)

    def test_no_diis(self):
        cc1 = cc.CCSD(mf)
        cc1.diis = False
        cc1.max_cycle = 4
        cc1.kernel()
        self.assertAlmostEqual(cc1.e_corr, -0.13516622806104395, 8)

    def test_restart(self):
        ftmp = tempfile.NamedTemporaryFile()
        cc1 = cc.CCSD(mf)
        cc1.max_cycle = 5
        cc1.kernel()
        ref = cc1.e_corr

        adiis = lib.diis.DIIS(mol)
        adiis.filename = ftmp.name
        cc1.diis = adiis
        cc1.max_cycle = 3
        cc1.kernel(t1=None, t2=None)
        self.assertAlmostEqual(cc1.e_corr, -0.13529291367331436, 8)

        t1, t2 = cc1.vector_to_amplitudes(adiis.extrapolate())
        self.assertAlmostEqual(abs(t1-cc1.t1).max(), 0, 9)
        self.assertAlmostEqual(abs(t2-cc1.t2).max(), 0, 9)
        cc1.diis = None
        cc1.max_cycle = 1
        cc1.kernel(t1, t2)
        self.assertAlmostEqual(cc1.e_corr, -0.13535690694539226, 8)

        cc1.diis = adiis
        cc1.max_cycle = 2
        cc1.kernel(t1, t2)
        self.assertAlmostEqual(cc1.e_corr, ref, 8)

        cc2 = cc.CCSD(mf)
        cc2.restore_from_diis_(ftmp.name)
        self.assertAlmostEqual(abs(cc1.t1 - cc2.t1).max(), 0, 9)
        self.assertAlmostEqual(abs(cc1.t2 - cc2.t2).max(), 0, 9)

    def test_iterative_dampling(self):
        cc1 = cc.CCSD(mf)
        cc1.max_cycle = 3
        cc1.iterative_damping = 0.7
        cc1.kernel()
        self.assertAlmostEqual(cc1.e_corr, -0.13508743605375528, 8)

    def test_amplitudes_to_vector(self):
        vec = mycc.amplitudes_to_vector(mycc.t1, mycc.t2)
        #self.assertAlmostEqual(lib.finger(vec), -0.056992042448099592, 6)
        r1, r2 = mycc.vector_to_amplitudes(vec)
        self.assertAlmostEqual(abs(r1-mycc.t1).max(), 0, 14)
        self.assertAlmostEqual(abs(r2-mycc.t2).max(), 0, 14)

        vec = numpy.random.random(vec.size)
        r1, r2 = mycc.vector_to_amplitudes(vec)
        vec1 = mycc.amplitudes_to_vector(r1, r2)
        self.assertAlmostEqual(abs(vec-vec1).max(), 0, 14)

    def test_vector_to_amplitudes_overwritten(self):
        mol = gto.M()
        mycc = scf.RHF(mol).apply(cc.CCSD)
        nelec = (3,3)
        nocc, nvir = nelec[0], 4
        nmo = nocc + nvir
        mycc.nocc = nocc
        mycc.nmo = nmo
        vec = numpy.zeros(mycc.vector_size())
        vec_orig = vec.copy()
        t1, t2 = mycc.vector_to_amplitudes(vec)
        t1[:] = 1
        t2[:] = 1
        self.assertAlmostEqual(abs(vec - vec_orig).max(), 0, 15)

    def test_vector_size(self):
        self.assertEqual(mycc.vector_size(), 860)

    def test_rccsd_frozen(self):
        cc1 = copy.copy(mycc)
        cc1.frozen = 1
        self.assertEqual(cc1.nmo, 12)
        self.assertEqual(cc1.nocc, 4)
        cc1.set_frozen()
        self.assertEqual(cc1.nmo, 12)
        self.assertEqual(cc1.nocc, 4)
        cc1.frozen = [0,1]
        self.assertEqual(cc1.nmo, 11)
        self.assertEqual(cc1.nocc, 3)
        cc1.frozen = [1,9]
        self.assertEqual(cc1.nmo, 11)
        self.assertEqual(cc1.nocc, 4)
        cc1.frozen = [9,10,12]
        self.assertEqual(cc1.nmo, 10)
        self.assertEqual(cc1.nocc, 5)
        cc1.nmo = 10
        cc1.nocc = 6
        self.assertEqual(cc1.nmo, 10)
        self.assertEqual(cc1.nocc, 6)

    def test_update_amps(self):
        mol = gto.M()
        nocc, nvir = 5, 12
        nmo = nocc + nvir
        nmo_pair = nmo*(nmo+1)//2
        mf = scf.RHF(mol)
        np.random.seed(12)
        mf._eri = np.random.random(nmo_pair*(nmo_pair+1)//2)
        mf.mo_coeff = np.random.random((nmo,nmo))
        mf.mo_energy = np.arange(0., nmo)
        mf.mo_occ = np.zeros(nmo)
        mf.mo_occ[:nocc] = 2
        vhf = mf.get_veff(mol, mf.make_rdm1())
        cinv = np.linalg.inv(mf.mo_coeff)
        mf.get_hcore = lambda *args: (reduce(np.dot, (cinv.T*mf.mo_energy, cinv)) - vhf)

        mycc1 = rccsd.RCCSD(mf)
        eris1 = mycc1.ao2mo()
        mycc2 = ccsd.CCSD(mf)
        eris2 = mycc2.ao2mo()
        a = np.random.random((nmo,nmo)) * .1
        eris1.fock += a + a.T.conj()
        eris2.fock += a + a.T
        t1 = np.random.random((nocc,nvir)) * .1
        t2 = np.random.random((nocc,nocc,nvir,nvir)) * .1
        t2 = t2 + t2.transpose(1,0,3,2)

        t1b, t2b = ccsd.update_amps(mycc2, t1, t2, eris2)
        self.assertAlmostEqual(lib.finger(t1b), -106360.5276951083, 6)
        self.assertAlmostEqual(lib.finger(t2b), 66540.100267798145, 6)

        mycc2.max_memory = 0
        t1a, t2a = ccsd.update_amps(mycc2, t1, t2, eris2)
        self.assertAlmostEqual(abs(t1a-t1b).max(), 0, 9)
        self.assertAlmostEqual(abs(t2a-t2b).max(), 0, 9)

        t2tril = ccsd._add_vvvv_tril(mycc2, t1, t2, eris2)
        self.assertAlmostEqual(lib.finger(t2tril), 13306.139402693696, 8)

        Ht2 = ccsd._add_vvvv_full(mycc2, t1, t2, eris2)
        self.assertAlmostEqual(lib.finger(Ht2), 760.50164232208408, 9)

        mycc1.cc2 = False
        t1a, t2a = rccsd.update_amps(mycc1, t1, t2, eris1)
        self.assertAlmostEqual(lib.finger(t1a), -106360.5276951083, 7)
        self.assertAlmostEqual(lib.finger(t2a), 66540.100267798145, 6)
        self.assertAlmostEqual(abs(t1a-t1b).max(), 0, 6)
        self.assertAlmostEqual(abs(t2a-t2b).max(), 0, 6)
        mycc1.cc2 = True
        t1a, t2a = rccsd.update_amps(mycc1, t1, t2, eris1)
        self.assertAlmostEqual(lib.finger(t1a), -106360.5276951083, 7)
        self.assertAlmostEqual(lib.finger(t2a), -1517.9391800662809, 7)

        mol = gto.Mole()
        mol.verbose = 0
        mol.atom = [
            [8 , (0. , 0.     , 0.)],
            [1 , (1. , 0.     , 0.)],
            [1 , (0. , -0.757 , 0.587)],
            [1 , (0. , 0.757  , 0.587)]]
        mol.basis = {'H': 'sto3g',
                     'O': 'cc-pvdz',}
        mol.charge = 1
        mol.build(0, 0)
        mycc2.direct = True
        eris2.vvvv = None
        eris2.mol = mol
        mycc2.mo_coeff, eris2.mo_coeff = eris2.mo_coeff, None
        t2tril = ccsd._add_vvvv_tril(mycc2, t1, t2, eris2, with_ovvv=True)
        self.assertAlmostEqual(lib.finger(t2tril), 680.07199094501584, 9)
        t2tril = ccsd._add_vvvv_tril(mycc2, t1, t2, eris2, with_ovvv=False)
        self.assertAlmostEqual(lib.finger(t2tril), 446.56702664171348, 9)
        Ht2 = ccsd._add_vvvv_full(mycc2, t1, t2, eris2)
        self.assertAlmostEqual(lib.finger(Ht2), 48.122317842230686, 9)

        eri1 = np.random.random((nmo,nmo,nmo,nmo)) + np.random.random((nmo,nmo,nmo,nmo))*1j
        eri1 = eri1.transpose(0,2,1,3)
        eri1 = eri1 + eri1.transpose(1,0,3,2).conj()
        eri1 = eri1 + eri1.transpose(2,3,0,1)
        eri1 *= .1
        eris1.oooo = eri1[:nocc,:nocc,:nocc,:nocc].copy()
        eris1.ovoo = eri1[:nocc,nocc:,:nocc,:nocc].copy()
        eris1.ovov = eri1[:nocc,nocc:,:nocc,nocc:].copy()
        eris1.oovv = eri1[:nocc,:nocc,nocc:,nocc:].copy()
        eris1.ovvo = eri1[:nocc,nocc:,nocc:,:nocc].copy()
        eris1.ovvv = eri1[:nocc,nocc:,nocc:,nocc:].copy()
        eris1.vvvv = eri1[nocc:,nocc:,nocc:,nocc:].copy()
        a = np.random.random((nmo,nmo)) * .1j
        eris1.fock = eris1.fock + a + a.T.conj()

        t1 = t1 + np.random.random((nocc,nvir)) * .1j
        t2 = t2 + np.random.random((nocc,nocc,nvir,nvir)) * .1j
        t2 = t2 + t2.transpose(1,0,3,2)
        mycc1.cc2 = False
        t1a, t2a = rccsd.update_amps(mycc1, t1, t2, eris1)
        self.assertAlmostEqual(lib.finger(t1a), -13.32050019680894-1.8825765910430254j, 9)
        self.assertAlmostEqual(lib.finger(t2a), 9.2521062044785189+29.999480274811873j, 9)
        mycc1.cc2 = True
        t1a, t2a = rccsd.update_amps(mycc1, t1, t2, eris1)
        self.assertAlmostEqual(lib.finger(t1a), -13.32050019680894-1.8825765910430254j, 9)
        self.assertAlmostEqual(lib.finger(t2a), -0.056223856104895858+0.025472249329733986j, 9)

    def test_eris_contract_vvvv_t2(self):
        mol = gto.Mole()
        nocc, nvir = 5, 12
        nvir_pair = nvir*(nvir+1)//2
        numpy.random.seed(9)
        t2 = numpy.random.random((nocc,nocc,nvir,nvir)) - .5
        t2 = t2 + t2.transpose(1,0,3,2)
        eris = ccsd._ChemistsERIs()
        vvvv = numpy.random.random((nvir_pair,nvir_pair)) - .5
        eris.vvvv = vvvv + vvvv.T
        eris.mol = mol
        mycc.max_memory, bak = 0, mycc.max_memory
        vt2 = eris._contract_vvvv_t2(mycc, t2, eris.vvvv)
        mycc.max_memory = bak
        self.assertAlmostEqual(lib.finger(vt2), -39.572579908080087, 11)
        vvvv = ao2mo.restore(1, eris.vvvv, nvir)
        ref = lib.einsum('acbd,ijcd->ijab', vvvv, t2)
        self.assertAlmostEqual(abs(vt2 - ref).max(), 0, 11)

        # _contract_s1vvvv_t2, testing complex and real mixed contraction
        vvvv =(numpy.random.random((nvir,nvir,nvir,nvir)) +
               numpy.random.random((nvir,nvir,nvir,nvir))*1j - (.5+.5j))
        vvvv = vvvv + vvvv.transpose(1,0,3,2).conj()
        vvvv = vvvv + vvvv.transpose(2,3,0,1)
        eris.vvvv = vvvv
        eris.mol = mol
        mycc.max_memory, bak = 0, mycc.max_memory
        vt2 = eris._contract_vvvv_t2(mycc, t2, eris.vvvv)
        mycc.max_memory = bak
        self.assertAlmostEqual(lib.finger(vt2), 23.502736435296871+113.90422480013488j, 11)
        ref = lib.einsum('acbd,ijcd->ijab', eris.vvvv, t2)
        self.assertAlmostEqual(abs(vt2 - ref).max(), 0, 11)

    def test_add_vvvv(self):
        t1 = mycc.t1
        t2 = mycc.t2
        nocc, nvir = t1.shape
        tau = t2 + numpy.einsum('ia,jb->ijab', t1, t1)
        eris1 = copy.copy(eris)
        mycc1 = copy.copy(mycc)
        ovvv = eris1.get_ovvv()
        tmp = -numpy.einsum('ijcd,ka,kdcb->ijba', tau, t1, ovvv)
        t2a = tmp + tmp.transpose(1,0,3,2)
        t2a += mycc1._add_vvvv(t1, t2, eris1)
        mycc1.direct = True
        eris1.vvvv = None  # == with_ovvv=True in the call below
        t2b = mycc1._add_vvvv(t1, t2, eris1, t2sym='jiba')
        self.assertAlmostEqual(abs(t2a-t2b).max(), 0, 12)

    def test_diagnostic(self):
        t1_diag = mycc.get_t1_diagnostic()
        d1_diag = mycc.get_d1_diagnostic()
        d2_diag = mycc.get_d2_diagnostic()
        self.assertAlmostEqual(t1_diag, 0.006002754773812036, 6)
        self.assertAlmostEqual(d1_diag, 0.012738043220198926, 6)
        self.assertAlmostEqual(d2_diag, 0.1169239107130769, 6)

    def test_ao2mo(self):
        mycc = ccsd.CCSD(mf)
        numpy.random.seed(2)
        mo = numpy.random.random(mf.mo_coeff.shape)
        mycc.max_memory = 2000
        eri_incore = mycc.ao2mo(mo)
        mycc.max_memory = 0
        eri_outcore = mycc.ao2mo(mo)
        self.assertTrue(isinstance(eri_outcore.oovv, h5py.Dataset))
        self.assertAlmostEqual(abs(eri_incore.oooo - eri_outcore.oooo).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.oovv - eri_outcore.oovv).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.ovoo - eri_outcore.ovoo).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.ovvo - eri_outcore.ovvo).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.ovov - eri_outcore.ovov).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.ovvv - eri_outcore.ovvv).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.vvvv - eri_outcore.vvvv).max(), 0, 12)

        mycc1 = ccsd.CCSD(mf.density_fit(auxbasis='ccpvdz-ri'))
        mycc1.max_memory = 0
        eri_df = mycc1.ao2mo(mo)
        self.assertAlmostEqual(lib.fp(eri_df.oooo), -493.98003157749906, 9)
        self.assertAlmostEqual(lib.fp(eri_df.oovv), -91.84858398271658 , 9)
        self.assertAlmostEqual(lib.fp(eri_df.ovoo), -203.89515661847437, 9)
        self.assertAlmostEqual(lib.fp(eri_df.ovvo), -14.883877359169205, 9)
        self.assertAlmostEqual(lib.fp(eri_df.ovov), -57.62195194777554 , 9)
        self.assertAlmostEqual(lib.fp(eri_df.ovvv), -24.359418953533535, 9)
        self.assertAlmostEqual(lib.fp(eri_df.vvvv),  76.9017539373456  , 9)

if __name__ == "__main__":
    print("Full Tests for RCCSD")
    unittest.main()


# pyscf/cc/test/test_rccsd_lambda.py

#!/usr/bin/env python
# Copyright 2014-2018 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.

import copy
import unittest
import numpy
import numpy as np
from functools import reduce

from pyscf import lib
from pyscf import gto
from pyscf import scf
from pyscf import ao2mo
from pyscf.cc import rccsd
from pyscf.cc import addons
from pyscf.cc import rccsd_lambda
from pyscf.cc import ccsd_rdm
from pyscf.cc import gccsd, gccsd_lambda

def setUpModule():
    global mol, mf, mycc
    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = '631g'
    mol.verbose = 5
    mol.output = '/dev/null'
    mol.build()
    mf = scf.RHF(mol).run()
    mycc = rccsd.RCCSD(mf)

def tearDownModule():
    global mol, mf, mycc
    mol.stdout.close()
    del mol, mf, mycc

class KnownValues(unittest.TestCase):
    def test_update_lambda_real(self):
        mycc = rccsd.RCCSD(mf)
        np.random.seed(12)
        nocc = 5
        nmo = 12
        nvir = nmo - nocc
        eri0 = np.random.random((nmo,nmo,nmo,nmo))
        eri0 = ao2mo.restore(1, ao2mo.restore(8, eri0, nmo), nmo)
        fock0 = np.random.random((nmo,nmo))
        fock0 = fock0 + fock0.T + np.diag(range(nmo))*2
        t1 = np.random.random((nocc,nvir))
        t2 = np.random.random((nocc,nocc,nvir,nvir))
        t2 = t2 + t2.transpose(1,0,3,2)
        l1 = np.random.random((nocc,nvir))
        l2 = np.random.random((nocc,nocc,nvir,nvir))
        l2 = l2 + l2.transpose(1,0,3,2)

        eris = rccsd._ChemistsERIs(mol)
        eris.oooo = eri0[:nocc,:nocc,:nocc,:nocc].copy()
        eris.ovoo = eri0[:nocc,nocc:,:nocc,:nocc].copy()
        eris.oovv = eri0[:nocc,:nocc,nocc:,nocc:].copy()
        eris.ovvo = eri0[:nocc,nocc:,nocc:,:nocc].copy()
        eris.ovov = eri0[:nocc,nocc:,:nocc,nocc:].copy()
        idx = np.tril_indices(nvir)
        eris.ovvv = eri0[:nocc,nocc:,nocc:,nocc:].copy()
        eris.vvvv = eri0[nocc:,nocc:,nocc:,nocc:].copy()
        eris.fock = fock0

        imds = rccsd_lambda.make_intermediates(mycc, t1, t2, eris)
        l1new, l2new = rccsd_lambda.update_lambda(mycc, t1, t2, l1, l2, eris, imds)
        self.assertAlmostEqual(lib.finger(l1new), -6699.5335665027187, 9)
        self.assertAlmostEqual(lib.finger(l2new), -514.7001243502192 , 9)
        self.assertAlmostEqual(abs(l2new-l2new.transpose(1,0,3,2)).max(), 0, 12)

        mycc.max_memory = 0
        imds = rccsd_lambda.make_intermediates(mycc, t1, t2, eris)
        l1new, l2new = rccsd_lambda.update_lambda(mycc, t1, t2, l1, l2, eris, imds)
        self.assertAlmostEqual(lib.finger(l1new), -6699.5335665027187, 9)
        self.assertAlmostEqual(lib.finger(l2new), -514.7001243502192 , 9)
        self.assertAlmostEqual(abs(l2new-l2new.transpose(1,0,3,2)).max(), 0, 12)

    def test_update_lambda_complex(self):
        mo_coeff = mf.mo_coeff + np.sin(mf.mo_coeff) * .01j
        nao = mo_coeff.shape[0]
        eri = ao2mo.restore(1, mf._eri, nao)
        eri0 = lib.einsum('pqrs,pi,qj,rk,sl->ijkl', eri, mo_coeff.conj(), mo_coeff,
                          mo_coeff.conj(), mo_coeff)

        nocc, nvir = 5, nao-5
        eris = rccsd._ChemistsERIs(mol)
        eris.oooo = eri0[:nocc,:nocc,:nocc,:nocc].copy()
        eris.ovoo = eri0[:nocc,nocc:,:nocc,:nocc].copy()
        eris.oovv = eri0[:nocc,:nocc,nocc:,nocc:].copy()
        eris.ovvo = eri0[:nocc,nocc:,nocc:,:nocc].copy()
        eris.ovov = eri0[:nocc,nocc:,:nocc,nocc:].copy()
        eris.ovvv = eri0[:nocc,nocc:,nocc:,nocc:].copy()
        eris.vvvv = eri0[nocc:,nocc:,nocc:,nocc:].copy()
        eris.fock = np.diag(mf.mo_energy)

        np.random.seed(1)
        t1 = np.random.random((nocc,nvir)) + np.random.random((nocc,nvir))*.1j - .5
        t2 = np.random.random((nocc,nocc,nvir,nvir)) - .5
        t2 = t2 + np.sin(t2) * .1j
        t2 = t2 + t2.transpose(1,0,3,2)

        l1 = np.random.random((nocc,nvir)) + np.random.random((nocc,nvir))*.1j - .5
        l2 = np.random.random((nocc,nocc,nvir,nvir)) - .5
        l2 = l2 + np.sin(l2) * .1j
        l2 = l2 + l2.transpose(1,0,3,2)
        mycc = rccsd.RCCSD(mf)
        imds = rccsd_lambda.make_intermediates(mycc, t1, t2, eris)
        l1new_ref, l2new_ref = rccsd_lambda.update_lambda(mycc, t1, t2, l1, l2, eris, imds)

        orbspin = np.zeros(nao*2, dtype=int)
        orbspin[1::2] = 1
        eri1 = np.zeros([nao*2]*4, dtype=np.complex128)
        eri1[0::2,0::2,0::2,0::2] = \
        eri1[0::2,0::2,1::2,1::2] = \
        eri1[1::2,1::2,0::2,0::2] = \
        eri1[1::2,1::2,1::2,1::2] = eri0
        eri1 = eri1.transpose(0,2,1,3) - eri1.transpose(0,2,3,1)
        erig = gccsd._PhysicistsERIs(mol)
        nocc *= 2
        nvir *= 2
        erig.oooo = eri1[:nocc,:nocc,:nocc,:nocc].copy()
        erig.ooov = eri1[:nocc,:nocc,:nocc,nocc:].copy()
        erig.ovov = eri1[:nocc,nocc:,:nocc,nocc:].copy()
        erig.ovvo = eri1[:nocc,nocc:,nocc:,:nocc].copy()
        erig.oovv = eri1[:nocc,:nocc,nocc:,nocc:].copy()
        erig.ovvv = eri1[:nocc,nocc:,nocc:,nocc:].copy()
        erig.vvvv = eri1[nocc:,nocc:,nocc:,nocc:].copy()
        mo_e = np.array([mf.mo_energy]*2)
        erig.fock = np.diag(mo_e.T.ravel())
        erig.mo_energy = erig.fock.diagonal()

        myccg = gccsd.GCCSD(scf.addons.convert_to_ghf(mf))
        t1, t2 = myccg.amplitudes_from_ccsd(t1, t2)
        l1, l2 = myccg.amplitudes_from_ccsd(l1, l2)
        imds = gccsd_lambda.make_intermediates(myccg, t1, t2, erig)
        l1new, l2new = gccsd_lambda.update_lambda(myccg, t1, t2, l1, l2, erig, imds)
        self.assertAlmostEqual(float(abs(l1new[0::2,0::2]-l1new_ref).max()), 0, 9)
        l2aa = l2new[0::2,0::2,0::2,0::2]
        l2ab = l2new[0::2,1::2,0::2,1::2]
        self.assertAlmostEqual(float(abs(l2ab-l2new_ref).max()), 0, 9)
        self.assertAlmostEqual(float(abs(l2ab-l2ab.transpose(1,0,2,3) - l2aa).max()), 0, 9)

    def test_rdm(self):
        mycc = rccsd.RCCSD(mf)
        mycc.frozen = 1
        mycc.kernel()
        dm1 = mycc.make_rdm1()
        dm2 = mycc.make_rdm2()
        h1 = reduce(numpy.dot, (mf.mo_coeff.T, mf.get_hcore(), mf.mo_coeff))
        nmo = mf.mo_coeff.shape[1]
        eri = ao2mo.restore(1, ao2mo.kernel(mf._eri, mf.mo_coeff), nmo)
        e1 = numpy.einsum('ij,ji', h1, dm1)
        e1+= numpy.einsum('ijkl,ijkl', eri, dm2) * .5
        e1+= mol.energy_nuc()
        self.assertAlmostEqual(e1, mycc.e_tot, 7)

        d1 = ccsd_rdm._gamma1_intermediates(mycc, mycc.t1, mycc.t2, mycc.l1, mycc.l2)
        mycc1 = copy.copy(mycc)
        mycc1.max_memory = 0
        d2 = ccsd_rdm._gamma2_intermediates(mycc1, mycc.t1, mycc.t2, mycc.l1, mycc.l2, True)
        dm2 = ccsd_rdm._make_rdm2(mycc, d1, d2, with_dm1=True, with_frozen=True)
        e1 = numpy.einsum('ij,ji', h1, dm1)
        e1+= numpy.einsum('ijkl,ijkl', eri, dm2) * .5
        e1+= mol.energy_nuc()
        self.assertAlmostEqual(e1, mycc.e_tot, 7)

    def test_rdm_trace(self):
        mycc = rccsd.RCCSD(mf)
        numpy.random.seed(2)
        nocc = 5
        nmo = 12
        nvir = nmo - nocc
        eri0 = numpy.random.random((nmo,nmo,nmo,nmo))
        eri0 = ao2mo.restore(1, ao2mo.restore(8, eri0, nmo), nmo)
        fock0 = numpy.random.random((nmo,nmo))
        fock0 = fock0 + fock0.T + numpy.diag(range(nmo))*2
        t1 = numpy.random.random((nocc,nvir))
        t2 = numpy.random.random((nocc,nocc,nvir,nvir))
        t2 = t2 + t2.transpose(1,0,3,2)
        l1 = numpy.random.random((nocc,nvir))
        l2 = numpy.random.random((nocc,nocc,nvir,nvir))
        l2 = l2 + l2.transpose(1,0,3,2)
        h1 = fock0 - (numpy.einsum('kkpq->pq', eri0[:nocc,:nocc])*2
                    - numpy.einsum('pkkq->pq', eri0[:,:nocc,:nocc]))

        eris = lambda:None
        eris.oooo = eri0[:nocc,:nocc,:nocc,:nocc].copy()
        eris.ooov = eri0[:nocc,:nocc,:nocc,nocc:].copy()
        eris.ovoo = eri0[:nocc,nocc:,:nocc,:nocc].copy()
        eris.oovv = eri0[:nocc,:nocc,nocc:,nocc:].copy()
        eris.ovov = eri0[:nocc,nocc:,:nocc,nocc:].copy()
        eris.ovvo = eri0[:nocc,nocc:,nocc:,:nocc].copy()
        eris.ovvv = eri0[:nocc,nocc:,nocc:,nocc:].copy()
        eris.vvvv = eri0[nocc:,nocc:,nocc:,nocc:].copy()
        eris.fock = fock0

        doo, dov, dvo, dvv = ccsd_rdm._gamma1_intermediates(mycc, t1, t2, l1, l2)
        self.assertAlmostEqual((numpy.einsum('ij,ij', doo, fock0[:nocc,:nocc]))*2, -20166.329861034799, 8)
        self.assertAlmostEqual((numpy.einsum('ab,ab', dvv, fock0[nocc:,nocc:]))*2,  58078.964019246778, 8)
        self.assertAlmostEqual((numpy.einsum('ai,ia', dvo, fock0[:nocc,nocc:]))*2, -74994.356886784764, 8)
        self.assertAlmostEqual((numpy.einsum('ia,ai', dov, fock0[nocc:,:nocc]))*2,  34.010188025702391, 9)

        fdm2 = lib.H5TmpFile()
        dovov, dvvvv, doooo, doovv, dovvo, dvvov, dovvv, dooov = \
                ccsd_rdm._gamma2_outcore(mycc, t1, t2, l1, l2, fdm2, True)
        self.assertAlmostEqual(lib.finger(numpy.array(dovov)), -14384.907042073517, 9)
        self.assertAlmostEqual(lib.finger(numpy.array(dvvvv)), -25.374007033024839, 9)
        self.assertAlmostEqual(lib.finger(numpy.array(doooo)),  60.114594698129963, 9)
        self.assertAlmostEqual(lib.finger(numpy.array(doovv)), -79.176348067958401, 9)
        self.assertAlmostEqual(lib.finger(numpy.array(dovvo)),   9.864134457251815, 9)
        self.assertAlmostEqual(lib.finger(numpy.array(dovvv)), -421.90333700061342, 9)
        self.assertAlmostEqual(lib.finger(numpy.array(dooov)), -592.66863759586136, 9)
        fdm2 = None

        dovov, dvvvv, doooo, doovv, dovvo, dvvov, dovvv, dooov = \
                ccsd_rdm._gamma2_intermediates(mycc, t1, t2, l1, l2)
        self.assertAlmostEqual(lib.finger(numpy.array(dovov)), -14384.907042073517, 9)
        self.assertAlmostEqual(lib.finger(numpy.array(dvvvv)),  45.872344902116758, 9)
        self.assertAlmostEqual(lib.finger(numpy.array(doooo)),  60.114594698129963, 9)
        self.assertAlmostEqual(lib.finger(numpy.array(doovv)), -79.176348067958401, 9)
        self.assertAlmostEqual(lib.finger(numpy.array(dovvo)),   9.864134457251815, 9)
        self.assertAlmostEqual(lib.finger(numpy.array(dovvv)), -421.90333700061342, 9)
        self.assertAlmostEqual(lib.finger(numpy.array(dooov)), -592.66863759586136, 9)

        self.assertAlmostEqual(numpy.einsum('kilj,kilj', doooo, eris.oooo)*2, 15939.9007625418, 7)
        self.assertAlmostEqual(numpy.einsum('acbd,acbd', dvvvv, eris.vvvv)*2, 37581.823919588 , 7)
        self.assertAlmostEqual(numpy.einsum('jkia,jkia', dooov, eris.ooov)*2, 128470.009687716, 7)
        self.assertAlmostEqual(numpy.einsum('icba,icba', dovvv, eris.ovvv)*2,-166794.225195056, 7)
        self.assertAlmostEqual(numpy.einsum('iajb,iajb', dovov, eris.ovov)*2,-719279.812916893, 7)
        self.assertAlmostEqual(numpy.einsum('jbai,jbia', dovvo, eris.ovov)*2+
                               numpy.einsum('jiab,jiba', doovv, eris.oovv)*2,-53634.0012286654, 7)

        dm1 = ccsd_rdm.make_rdm1(mycc, t1, t2, l1, l2)
        dm2 = ccsd_rdm.make_rdm2(mycc, t1, t2, l1, l2)
        e2 =(numpy.einsum('ijkl,ijkl', doooo, eris.oooo)*2
            +numpy.einsum('acbd,acbd', dvvvv, eris.vvvv)*2
            +numpy.einsum('jkia,jkia', dooov, eris.ooov)*2
            +numpy.einsum('icba,icba', dovvv, eris.ovvv)*2
            +numpy.einsum('iajb,iajb', dovov, eris.ovov)*2
            +numpy.einsum('jbai,jbia', dovvo, eris.ovov)*2
            +numpy.einsum('ijab,ijab', doovv, eris.oovv)*2
            +numpy.einsum('ij,ij', doo, fock0[:nocc,:nocc])*2
            +numpy.einsum('ai,ia', dvo, fock0[:nocc,nocc:])*2
            +numpy.einsum('ia,ai', dov, fock0[nocc:,:nocc])*2
            +numpy.einsum('ab,ab', dvv, fock0[nocc:,nocc:])*2
            +fock0[:nocc].trace()*2
            -numpy.einsum('kkpq->pq', eri0[:nocc,:nocc,:nocc,:nocc]).trace()*2
            +numpy.einsum('pkkq->pq', eri0[:nocc,:nocc,:nocc,:nocc]).trace())
        self.assertAlmostEqual(e2, -794721.197459942, 8)
        self.assertAlmostEqual(numpy.einsum('pqrs,pqrs', dm2, eri0)*.5 +
                               numpy.einsum('pq,qp', dm1, h1), e2, 9)

        self.assertAlmostEqual(abs(dm2-dm2.transpose(1,0,3,2)).max(), 0, 9)
        self.assertAlmostEqual(abs(dm2-dm2.transpose(2,3,0,1)).max(), 0, 9)

        d1 = numpy.einsum('kkpq->qp', dm2) / 9
        self.assertAlmostEqual(abs(d1-dm1).max(), 0, 9)


if __name__ == "__main__":
    print("Full Tests for RCCSD lambda")
    unittest.main()


# pyscf/cc/test/test_uccsd.py

#!/usr/bin/env python
# Copyright 2014-2020 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.

import unittest
import copy
import numpy
import h5py
from functools import reduce

from pyscf import gto, lib
from pyscf import scf, dft
from pyscf import mp
from pyscf import cc
from pyscf import ao2mo
from pyscf.cc import uccsd
from pyscf.cc import gccsd
from pyscf.cc import addons
from pyscf.cc import uccsd_rdm
from pyscf.fci import direct_uhf

def setUpModule():
    global mol, rhf, mf, myucc, mol_s2, mf_s2, eris
    mol = gto.Mole()
    mol.verbose = 7
    mol.output = '/dev/null'
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]

    mol.basis = '631g'
    mol.build()
    rhf = scf.RHF(mol)
    rhf.conv_tol_grad = 1e-8
    rhf.kernel()
    mf = scf.addons.convert_to_uhf(rhf)

    myucc = cc.UCCSD(mf).run(conv_tol=1e-10)

    mol_s2 = gto.Mole()
    mol_s2.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol_s2.basis = '631g'
    mol_s2.spin = 2
    mol_s2.verbose = 5
    mol_s2.output = '/dev/null'
    mol_s2.build()
    mf_s2 = scf.UHF(mol_s2).run()
    eris = uccsd.UCCSD(mf_s2).ao2mo()

def tearDownModule():
    global mol, rhf, mf, myucc, mol_s2, mf_s2, eris
    mol.stdout.close()
    mol_s2.stdout.close()
    del mol, rhf, mf, myucc, mol_s2, mf_s2, eris

class KnownValues(unittest.TestCase):

    def test_with_df_s0(self):
        mf = scf.UHF(mol).density_fit(auxbasis='weigend').run()
        mycc = cc.UCCSD(mf).run()
        self.assertAlmostEqual(mycc.e_tot, -76.118403942938741, 7)

    def test_with_df_s2(self):
        mf = scf.UHF(mol_s2).density_fit(auxbasis='weigend').run()
        mycc = cc.UCCSD(mf).run()
        self.assertAlmostEqual(mycc.e_tot, -75.83360033370676, 7)

    def test_ERIS(self):
        ucc1 = cc.UCCSD(mf)
        nao,nmo = mf.mo_coeff[0].shape
        numpy.random.seed(1)
        mo_coeff = numpy.random.random((2,nao,nmo))
        eris = cc.uccsd._make_eris_incore(ucc1, mo_coeff)

        self.assertAlmostEqual(lib.fp(eris.oooo), 4.9638849382825754, 11)
        self.assertAlmostEqual(lib.fp(eris.ovoo),-1.3623681896983584, 11)
        self.assertAlmostEqual(lib.fp(eris.ovov), 125.81550684442163, 11)
        self.assertAlmostEqual(lib.fp(eris.oovv), 55.123681017639598, 11)
        self.assertAlmostEqual(lib.fp(eris.ovvo), 133.48083527898248, 11)
        self.assertAlmostEqual(lib.fp(eris.ovvv), 59.421927525288183, 11)
        self.assertAlmostEqual(lib.fp(eris.vvvv), 43.556602622204778, 11)
        self.assertAlmostEqual(lib.fp(eris.OOOO),-407.05319440524585, 11)
        self.assertAlmostEqual(lib.fp(eris.OVOO), 56.284299937160796, 11)
        self.assertAlmostEqual(lib.fp(eris.OVOV),-287.72899895597448, 11)
        self.assertAlmostEqual(lib.fp(eris.OOVV),-85.484299959144522, 11)
        self.assertAlmostEqual(lib.fp(eris.OVVO),-228.18996145476956, 11)
        self.assertAlmostEqual(lib.fp(eris.OVVV),-10.715902258877399, 11)
        self.assertAlmostEqual(lib.fp(eris.VVVV),-89.908425473958303, 11)
        self.assertAlmostEqual(lib.fp(eris.ooOO),-336.65979260175226, 11)
        self.assertAlmostEqual(lib.fp(eris.ovOO),-16.405125847288176, 11)
        self.assertAlmostEqual(lib.fp(eris.ovOV), 231.59042209500075, 11)
        self.assertAlmostEqual(lib.fp(eris.ooVV), 20.338077193028354, 11)
        self.assertAlmostEqual(lib.fp(eris.ovVO), 206.48662856981386, 11)
        self.assertAlmostEqual(lib.fp(eris.ovVV),-71.273249852220516, 11)
        self.assertAlmostEqual(lib.fp(eris.vvVV), 172.47130671068496, 11)
        self.assertAlmostEqual(lib.fp(eris.OVoo),-19.927660309103977, 11)
        self.assertAlmostEqual(lib.fp(eris.OOvv),-27.761433381797019, 11)
        self.assertAlmostEqual(lib.fp(eris.OVvo),-140.09648311337384, 11)
        self.assertAlmostEqual(lib.fp(eris.OVvv), 40.700983950220547, 11)

        uccsd.MEMORYMIN, bak = 0, uccsd.MEMORYMIN
        ucc1.max_memory = 0
        eris1 = ucc1.ao2mo(mo_coeff)
        uccsd.MEMORYMIN = bak
        self.assertAlmostEqual(abs(numpy.array(eris1.oooo)-eris.oooo).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.ovoo)-eris.ovoo).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.ovov)-eris.ovov).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.oovv)-eris.oovv).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.ovvo)-eris.ovvo).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.ovvv)-eris.ovvv).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.vvvv)-eris.vvvv).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.OOOO)-eris.OOOO).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.OVOO)-eris.OVOO).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.OVOV)-eris.OVOV).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.OOVV)-eris.OOVV).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.OVVO)-eris.OVVO).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.OVVV)-eris.OVVV).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.VVVV)-eris.VVVV).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.ooOO)-eris.ooOO).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.ovOO)-eris.ovOO).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.ovOV)-eris.ovOV).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.ooVV)-eris.ooVV).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.ovVO)-eris.ovVO).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.ovVV)-eris.ovVV).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.vvVV)-eris.vvVV).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.OVoo)-eris.OVoo).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.OOvv)-eris.OOvv).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.OVvo)-eris.OVvo).max(), 0, 11)
        self.assertAlmostEqual(abs(numpy.array(eris1.OVvv)-eris.OVvv).max(), 0, 11)

        # Testing the complex MO integrals
        def ao2mofn(mos):
            if isinstance(mos, numpy.ndarray) and mos.ndim == 2:
                mos = [mos]*4
            nmos = [mo.shape[1] for mo in mos]
            eri_mo = ao2mo.kernel(mf._eri, mos, compact=False).reshape(nmos)
            return eri_mo * 1j
        eris1 = cc.uccsd._make_eris_incore(ucc1, mo_coeff, ao2mofn=ao2mofn)
        self.assertAlmostEqual(abs(eris1.oooo.imag-eris.oooo).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.ovoo.imag-eris.ovoo).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.ovov.imag-eris.ovov).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.oovv.imag-eris.oovv).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.ovvo.imag-eris.ovvo).max(), 0, 11)
        #self.assertAlmostEqual(abs(eris1.ovvv.imag-eris.ovvv).max(), 0, 11)
        #self.assertAlmostEqual(abs(eris1.vvvv.imag-eris.vvvv).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.OOOO.imag-eris.OOOO).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.OVOO.imag-eris.OVOO).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.OVOV.imag-eris.OVOV).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.OOVV.imag-eris.OOVV).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.OVVO.imag-eris.OVVO).max(), 0, 11)
        #self.assertAlmostEqual(abs(eris1.OVVV.imag-eris.OVVV).max(), 0, 11)
        #self.assertAlmostEqual(abs(eris1.VVVV.imag-eris.VVVV).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.ooOO.imag-eris.ooOO).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.ovOO.imag-eris.ovOO).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.ovOV.imag-eris.ovOV).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.ooVV.imag-eris.ooVV).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.ovVO.imag-eris.ovVO).max(), 0, 11)
        #self.assertAlmostEqual(abs(eris1.ovVV.imag-eris.ovVV).max(), 0, 11)
        #self.assertAlmostEqual(abs(eris1.vvVV.imag-eris.vvVV).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.OVoo.imag-eris.OVoo).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.OOvv.imag-eris.OOvv).max(), 0, 11)
        self.assertAlmostEqual(abs(eris1.OVvo.imag-eris.OVvo).max(), 0, 11)
        #self.assertAlmostEqual(abs(eris1.OVvv.imag-eris.OVvv).max(), 0, 11)

    def test_amplitudes_from_rccsd(self):
        e, t1, t2 = cc.RCCSD(rhf).set(conv_tol=1e-10).kernel()
        t1, t2 = myucc.amplitudes_from_rccsd(t1, t2)
        self.assertAlmostEqual(abs(t1[0]-myucc.t1[0]).max(), 0, 6)
        self.assertAlmostEqual(abs(t1[1]-myucc.t1[1]).max(), 0, 6)
        self.assertAlmostEqual(abs(t2[0]-myucc.t2[0]).max(), 0, 6)
        self.assertAlmostEqual(abs(t2[1]-myucc.t2[1]).max(), 0, 6)
        self.assertAlmostEqual(abs(t2[2]-myucc.t2[2]).max(), 0, 6)

    def test_uccsd_frozen(self):
        ucc1 = copy.copy(myucc)
        ucc1.frozen = 1
        self.assertEqual(ucc1.nmo, (12,12))
        self.assertEqual(ucc1.nocc, (4,4))
        ucc1.set_frozen()
        self.assertEqual(ucc1.nmo, (12,12))
        self.assertEqual(ucc1.nocc, (4,4))
        ucc1.frozen = [0,1]
        self.assertEqual(ucc1.nmo, (11,11))
        self.assertEqual(ucc1.nocc, (3,3))
        ucc1.frozen = [[0,1], [0,1]]
        self.assertEqual(ucc1.nmo, (11,11))
        self.assertEqual(ucc1.nocc, (3,3))
        ucc1.frozen = [1,9]
        self.assertEqual(ucc1.nmo, (11,11))
        self.assertEqual(ucc1.nocc, (4,4))
        ucc1.frozen = [[1,9], [1,9]]
        self.assertEqual(ucc1.nmo, (11,11))
        self.assertEqual(ucc1.nocc, (4,4))
        ucc1.frozen = [9,10,12]
        self.assertEqual(ucc1.nmo, (10,10))
        self.assertEqual(ucc1.nocc, (5,5))
        ucc1.nmo = (13,12)
        ucc1.nocc = (5,4)
        self.assertEqual(ucc1.nmo, (13,12))
        self.assertEqual(ucc1.nocc, (5,4))

    def test_uccsd_frozen(self):
        # Freeze 1s electrons
        frozen = [[0,1], [0,1]]
        ucc = cc.UCCSD(mf_s2, frozen=frozen)
        ucc.diis_start_cycle = 1
        ecc, t1, t2 = ucc.kernel()
        self.assertAlmostEqual(ecc, -0.07414978284611283, 8)

    def test_rdm(self):
        nocc = 5
        nvir = 7
        mol = gto.M()
        mf = scf.UHF(mol)
        mf.mo_occ = numpy.zeros((2,nocc+nvir))
        mf.mo_occ[:,:nocc] = 1
        mycc = uccsd.UCCSD(mf)

        def antisym(t2):
            t2 = t2 - t2.transpose(0,1,3,2)
            t2 = t2 - t2.transpose(1,0,2,3)
            return t2
        orbspin = numpy.zeros((nocc+nvir)*2, dtype=int)
        orbspin[1::2] = 1
        numpy.random.seed(1)
        t1 = numpy.random.random((2,nocc,nvir))*.1 - .1
        t2ab = numpy.random.random((nocc,nocc,nvir,nvir))*.1 - .1
        t2aa = antisym(numpy.random.random((nocc,nocc,nvir,nvir))*.1 - .1)
        t2bb = antisym(numpy.random.random((nocc,nocc,nvir,nvir))*.1 - .1)
        t2 = (t2aa,t2ab,t2bb)
        l1 = numpy.random.random((2,nocc,nvir))*.1 - .1
        l2ab = numpy.random.random((nocc,nocc,nvir,nvir))*.1 - .1
        l2aa = antisym(numpy.random.random((nocc,nocc,nvir,nvir))*.1 - .1)
        l2bb = antisym(numpy.random.random((nocc,nocc,nvir,nvir))*.1 - .1)
        l2 = (l2aa,l2ab,l2bb)

        dm1a, dm1b = mycc.make_rdm1(t1, t2, l1, l2)
        dm2aa, dm2ab, dm2bb = mycc.make_rdm2(t1, t2, l1, l2)
        ia = orbspin == 0
        ib = orbspin == 1
        oa = orbspin[:nocc*2] == 0
        ob = orbspin[:nocc*2] == 1
        va = orbspin[nocc*2:] == 0
        vb = orbspin[nocc*2:] == 1

        t1 = addons.spatial2spin(t1, orbspin)
        t2 = addons.spatial2spin(t2, orbspin)
        l1 = addons.spatial2spin(l1, orbspin)
        l2 = addons.spatial2spin(l2, orbspin)
        mf1 = scf.GHF(mol)
        mf1.mo_occ = numpy.zeros((nocc+nvir)*2)
        mf.mo_occ[:,:nocc*2] = 1
        mycc1 = gccsd.GCCSD(mf1)
        dm1 = mycc1.make_rdm1(t1, t2, l1, l2)
        dm2 = mycc1.make_rdm2(t1, t2, l1, l2)
        self.assertAlmostEqual(abs(dm1[ia][:,ia]-dm1a).max(), 0, 9)
        self.assertAlmostEqual(abs(dm1[ib][:,ib]-dm1b).max(), 0, 9)
        self.assertAlmostEqual(abs(dm2[ia][:,ia][:,:,ia][:,:,:,ia]-dm2aa).max(), 0, 9)
        self.assertAlmostEqual(abs(dm2[ia][:,ia][:,:,ib][:,:,:,ib]-dm2ab).max(), 0, 9)
        self.assertAlmostEqual(abs(dm2[ib][:,ib][:,:,ib][:,:,:,ib]-dm2bb).max(), 0, 9)

    def test_h2o_rdm(self):
        mol = mol_s2
        mf = mf_s2
        mycc = uccsd.UCCSD(mf)
        mycc.frozen = 2
        ecc, t1, t2 = mycc.kernel()
        l1, l2 = mycc.solve_lambda()
        dm1a,dm1b = mycc.make_rdm1(t1, t2, l1, l2)
        dm2aa,dm2ab,dm2bb = mycc.make_rdm2(t1, t2, l1, l2)
        mo_a = mf.mo_coeff[0]
        mo_b = mf.mo_coeff[1]
        nmoa = mo_a.shape[1]
        nmob = mo_b.shape[1]
        eriaa = ao2mo.kernel(mf._eri, mo_a, compact=False).reshape([nmoa]*4)
        eribb = ao2mo.kernel(mf._eri, mo_b, compact=False).reshape([nmob]*4)
        eriab = ao2mo.kernel(mf._eri, (mo_a,mo_a,mo_b,mo_b), compact=False)
        eriab = eriab.reshape([nmoa,nmoa,nmob,nmob])
        hcore = mf.get_hcore()
        h1a = reduce(numpy.dot, (mo_a.T.conj(), hcore, mo_a))
        h1b = reduce(numpy.dot, (mo_b.T.conj(), hcore, mo_b))
        e1 = numpy.einsum('ij,ji', h1a, dm1a)
        e1+= numpy.einsum('ij,ji', h1b, dm1b)
        e1+= numpy.einsum('ijkl,ijkl', eriaa, dm2aa) * .5
        e1+= numpy.einsum('ijkl,ijkl', eriab, dm2ab)
        e1+= numpy.einsum('ijkl,ijkl', eribb, dm2bb) * .5
        e1+= mol.energy_nuc()
        self.assertAlmostEqual(e1, mycc.e_tot, 7)

        d1 = uccsd_rdm._gamma1_intermediates(mycc, mycc.t1, mycc.t2, mycc.l1, mycc.l2)
        mycc.max_memory = 0
        d2 = uccsd_rdm._gamma2_intermediates(mycc, mycc.t1, mycc.t2, mycc.l1, mycc.l2, True)
        dm2 = uccsd_rdm._make_rdm2(mycc, d1, d2, with_dm1=True, with_frozen=True)
        e1 = numpy.einsum('ij,ji', h1a, dm1a)
        e1+= numpy.einsum('ij,ji', h1b, dm1b)
        e1+= numpy.einsum('ijkl,ijkl', eriaa, dm2[0]) * .5
        e1+= numpy.einsum('ijkl,ijkl', eriab, dm2[1])
        e1+= numpy.einsum('ijkl,ijkl', eribb, dm2[2]) * .5
        e1+= mol.energy_nuc()
        self.assertAlmostEqual(e1, mycc.e_tot, 7)

    def test_h4_rdm(self):
        mol = gto.Mole()
        mol.verbose = 0
        mol.atom = [
            ['H', ( 1.,-1.    , 0.   )],
            ['H', ( 0.,-1.    ,-1.   )],
            ['H', ( 1.,-0.5   , 0.   )],
            ['H', ( 0.,-1.    , 1.   )],
        ]
        mol.charge = 2
        mol.spin = 2
        mol.basis = '6-31g'
        mol.build()
        mf = scf.UHF(mol).set(init_guess='1e').run(conv_tol=1e-14)
        ehf0 = mf.e_tot - mol.energy_nuc()
        mycc = uccsd.UCCSD(mf).run()
        mycc.solve_lambda()
        eri_aa = ao2mo.kernel(mf._eri, mf.mo_coeff[0])
        eri_bb = ao2mo.kernel(mf._eri, mf.mo_coeff[1])
        eri_ab = ao2mo.kernel(mf._eri, [mf.mo_coeff[0], mf.mo_coeff[0],
                                        mf.mo_coeff[1], mf.mo_coeff[1]])
        h1a = reduce(numpy.dot, (mf.mo_coeff[0].T, mf.get_hcore(), mf.mo_coeff[0]))
        h1b = reduce(numpy.dot, (mf.mo_coeff[1].T, mf.get_hcore(), mf.mo_coeff[1]))
        efci, fcivec = direct_uhf.kernel((h1a,h1b), (eri_aa,eri_ab,eri_bb),
                                         h1a.shape[0], mol.nelec)
        dm1ref, dm2ref = direct_uhf.make_rdm12s(fcivec, h1a.shape[0], mol.nelec)
        t1, t2 = mycc.t1, mycc.t2
        l1, l2 = mycc.l1, mycc.l2
        rdm1 = mycc.make_rdm1(t1, t2, l1, l2)
        rdm2 = mycc.make_rdm2(t1, t2, l1, l2)
        self.assertAlmostEqual(abs(dm1ref[0] - rdm1[0]).max(), 0, 6)
        self.assertAlmostEqual(abs(dm1ref[1] - rdm1[1]).max(), 0, 6)
        self.assertAlmostEqual(abs(dm2ref[0] - rdm2[0]).max(), 0, 6)
        self.assertAlmostEqual(abs(dm2ref[1] - rdm2[1]).max(), 0, 6)
        self.assertAlmostEqual(abs(dm2ref[2] - rdm2[2]).max(), 0, 6)

    def test_eris_contract_vvvv_t2(self):
        mol = gto.Mole()
        nocca, noccb, nvira, nvirb = 5, 4, 12, 13
        nvira_pair = nvira*(nvira+1)//2
        nvirb_pair = nvirb*(nvirb+1)//2
        numpy.random.seed(9)
        t2 = numpy.random.random((nocca,noccb,nvira,nvirb))
        eris = uccsd._ChemistsERIs()
        eris.vvVV = numpy.random.random((nvira_pair,nvirb_pair))
        eris.mol = mol
        myucc.max_memory, bak = 0, myucc.max_memory
        vt2 = eris._contract_vvVV_t2(myucc, t2, eris.vvVV)
        myucc.max_memory = bak
        self.assertAlmostEqual(lib.fp(vt2), 12.00904827896089, 11)
        idxa = lib.square_mat_in_trilu_indices(nvira)
        idxb = lib.square_mat_in_trilu_indices(nvirb)
        vvVV = eris.vvVV[:,idxb][idxa]
        ref = lib.einsum('acbd,ijcd->ijab', vvVV, t2)
        self.assertAlmostEqual(abs(vt2 - ref).max(), 0, 11)

        # _contract_VVVV_t2, testing complex and real mixed contraction
        VVVV =(numpy.random.random((nvirb,nvirb,nvirb,nvirb)) +
               numpy.random.random((nvirb,nvirb,nvirb,nvirb))*1j - (.5+.5j))
        VVVV = VVVV + VVVV.transpose(1,0,3,2).conj()
        VVVV = VVVV + VVVV.transpose(2,3,0,1)
        eris.VVVV = VVVV
        t2 = numpy.random.random((noccb,noccb,nvirb,nvirb))
        t2 = t2 - t2.transpose(0,1,3,2)
        t2 = t2 - t2.transpose(1,0,3,2)
        myucc.max_memory, bak = 0, myucc.max_memory
        vt2 = eris._contract_VVVV_t2(myucc, t2, eris.VVVV)
        myucc.max_memory = bak
        self.assertAlmostEqual(lib.fp(vt2), 47.903883794299404-50.501573400833429j, 11)
        ref = lib.einsum('acbd,ijcd->ijab', eris.VVVV, t2)
        self.assertAlmostEqual(abs(vt2 - ref).max(), 0, 11)

    def test_update_amps1(self):
        mf = scf.UHF(mol_s2)
        numpy.random.seed(9)
        nmo = mf_s2.mo_occ[0].size
        mf.mo_coeff = numpy.random.random((2,nmo,nmo)) - 0.5
        mf.mo_occ = numpy.zeros((2,nmo))
        mf.mo_occ[0,:6] = 1
        mf.mo_occ[1,:5] = 1
        mycc = uccsd.UCCSD(mf)
        nocca, noccb = 6, 5
        nvira, nvirb = nmo-nocca, nmo-noccb
        nvira_pair = nvira*(nvira+1)//2
        nvirb_pair = nvirb*(nvirb+1)//2

        eris = mycc.ao2mo()
        fakeris = uccsd._ChemistsERIs()
        fakeris.mo_coeff = eris.mo_coeff
        fakeris.vvVV = eris.vvVV
        fakeris.mol = mol_s2
        t2ab = numpy.random.random((nocca,noccb,nvira,nvirb))
        t1a = numpy.zeros((nocca,nvira))
        t1b = numpy.zeros((noccb,nvirb))
        self.assertAlmostEqual(lib.fp(mycc._add_vvVV(None, t2ab, fakeris)), 21.652482203108928, 9)
        fakeris.vvVV = None
        mycc.direct = True
        mycc.max_memory = 0
        self.assertAlmostEqual(lib.fp(mycc._add_vvVV(None, t2ab, fakeris)), 21.652482203108928, 9)

        t1 = (numpy.random.random((nocca,nvira)), numpy.random.random((noccb,nvirb)))
        t2 = (numpy.random.random((nocca,nocca,nvira,nvira)),
              numpy.random.random((nocca,noccb,nvira,nvirb)),
              numpy.random.random((noccb,noccb,nvirb,nvirb)))
        t1, t2 = mycc.vector_to_amplitudes(mycc.amplitudes_to_vector(t1, t2))
        t1, t2 = mycc.update_amps(t1, t2, eris)
        self.assertAlmostEqual(lib.fp(t1[0]),  49.912690337392938, 10)
        self.assertAlmostEqual(lib.fp(t1[1]),  74.596097348134776, 10)
        self.assertAlmostEqual(lib.fp(t2[0]), -41.784696524955393, 10)
        self.assertAlmostEqual(lib.fp(t2[1]), -9675.767769478574, 6)
        self.assertAlmostEqual(lib.fp(t2[2]),  270.75447826471577, 7)
        self.assertAlmostEqual(lib.fp(mycc.amplitudes_to_vector(t1, t2)), 4341.9623137256776, 6)

    def test_vector_to_amplitudes(self):
        t1, t2 = myucc.vector_to_amplitudes(myucc.amplitudes_to_vector(myucc.t1, myucc.t2))
        self.assertAlmostEqual(abs(t1[0]-myucc.t1[0]).max(), 0, 12)
        self.assertAlmostEqual(abs(t1[1]-myucc.t1[1]).max(), 0, 12)
        self.assertAlmostEqual(abs(t2[0]-myucc.t2[0]).max(), 0, 12)
        self.assertAlmostEqual(abs(t2[1]-myucc.t2[1]).max(), 0, 12)
        self.assertAlmostEqual(abs(t2[2]-myucc.t2[2]).max(), 0, 12)

    def test_vector_to_amplitudes_overwritten(self):
        mol = gto.M()
        mycc = scf.UHF(mol).apply(cc.UCCSD)
        nelec = (3, 3)
        nocc = nelec
        nmo = (5, 5)
        mycc.nocc = nocc
        mycc.nmo = nmo
        vec = numpy.zeros(mycc.vector_size())
        vec_orig = vec.copy()
        t1, t2 = mycc.vector_to_amplitudes(vec)
        t1a, t1b = t1
        t2aa, t2ab, t2bb = t2
        t1a[:] = 1
        t1b[:] = 1
        t2aa[:] = 1
        t2ab[:] = 1
        t2bb[:] = 1
        self.assertAlmostEqual(abs(vec - vec_orig).max(), 0, 15)

    def test_vector_size(self):
        self.assertEqual(myucc.vector_size(), 2240)

    def test_update_amps2(self):  # compare to gccsd.update_amps
        mol = mol_s2
        mf = mf_s2
        myucc = uccsd.UCCSD(mf)
        nocca, noccb = 6,4
        nmo = mol.nao_nr()
        nvira,nvirb = nmo-nocca, nmo-noccb
        numpy.random.seed(9)
        t1 = [numpy.random.random((nocca,nvira))-.9,
              numpy.random.random((noccb,nvirb))-.9]
        t2 = [numpy.random.random((nocca,nocca,nvira,nvira))-.9,
              numpy.random.random((nocca,noccb,nvira,nvirb))-.9,
              numpy.random.random((noccb,noccb,nvirb,nvirb))-.9]
        t2[0] = t2[0] - t2[0].transpose(1,0,2,3)
        t2[0] = t2[0] - t2[0].transpose(0,1,3,2)
        t2[2] = t2[2] - t2[2].transpose(1,0,2,3)
        t2[2] = t2[2] - t2[2].transpose(0,1,3,2)

        mo_a = mf.mo_coeff[0] + numpy.sin(mf.mo_coeff[0]) * .01j
        mo_b = mf.mo_coeff[1] + numpy.sin(mf.mo_coeff[1]) * .01j
        nao = mo_a.shape[0]
        eri = ao2mo.restore(1, mf._eri, nao)
        eri0aa = lib.einsum('pqrs,pi,qj,rk,sl->ijkl', eri, mo_a.conj(), mo_a, mo_a.conj(), mo_a)
        eri0ab = lib.einsum('pqrs,pi,qj,rk,sl->ijkl', eri, mo_a.conj(), mo_a, mo_b.conj(), mo_b)
        eri0bb = lib.einsum('pqrs,pi,qj,rk,sl->ijkl', eri, mo_b.conj(), mo_b, mo_b.conj(), mo_b)
        eri0ba = eri0ab.transpose(2,3,0,1)

        nvira = nao - nocca
        nvirb = nao - noccb
        eris = uccsd._ChemistsERIs(mol)
        eris.oooo = eri0aa[:nocca,:nocca,:nocca,:nocca].copy()
        eris.ovoo = eri0aa[:nocca,nocca:,:nocca,:nocca].copy()
        eris.oovv = eri0aa[:nocca,:nocca,nocca:,nocca:].copy()
        eris.ovvo = eri0aa[:nocca,nocca:,nocca:,:nocca].copy()
        eris.ovov = eri0aa[:nocca,nocca:,:nocca,nocca:].copy()
        eris.ovvv = eri0aa[:nocca,nocca:,nocca:,nocca:].copy()
        eris.vvvv = eri0aa[nocca:,nocca:,nocca:,nocca:].copy()

        eris.OOOO = eri0bb[:noccb,:noccb,:noccb,:noccb].copy()
        eris.OVOO = eri0bb[:noccb,noccb:,:noccb,:noccb].copy()
        eris.OOVV = eri0bb[:noccb,:noccb,noccb:,noccb:].copy()
        eris.OVVO = eri0bb[:noccb,noccb:,noccb:,:noccb].copy()
        eris.OVOV = eri0bb[:noccb,noccb:,:noccb,noccb:].copy()
        eris.OVVV = eri0bb[:noccb,noccb:,noccb:,noccb:].copy()
        eris.VVVV = eri0bb[noccb:,noccb:,noccb:,noccb:].copy()

        eris.ooOO = eri0ab[:nocca,:nocca,:noccb,:noccb].copy()
        eris.ovOO = eri0ab[:nocca,nocca:,:noccb,:noccb].copy()
        eris.ooVV = eri0ab[:nocca,:nocca,noccb:,noccb:].copy()
        eris.ovVO = eri0ab[:nocca,nocca:,noccb:,:noccb].copy()
        eris.ovOV = eri0ab[:nocca,nocca:,:noccb,noccb:].copy()
        eris.ovVV = eri0ab[:nocca,nocca:,noccb:,noccb:].copy()
        eris.vvVV = eri0ab[nocca:,nocca:,noccb:,noccb:].copy()

        eris.OOoo = eri0ba[:noccb,:noccb,:nocca,:nocca].copy()
        eris.OVoo = eri0ba[:noccb,noccb:,:nocca,:nocca].copy()
        eris.OOvv = eri0ba[:noccb,:noccb,nocca:,nocca:].copy()
        eris.OVvo = eri0ba[:noccb,noccb:,nocca:,:nocca].copy()
        eris.OVov = eri0ba[:noccb,noccb:,:nocca,nocca:].copy()
        eris.OVvv = eri0ba[:noccb,noccb:,nocca:,nocca:].copy()
        eris.VVvv = eri0ba[noccb:,noccb:,nocca:,nocca:].copy()

        eris.focka = numpy.diag(mf.mo_energy[0])
        eris.fockb = numpy.diag(mf.mo_energy[1])
        eris.mo_energy = mf.mo_energy

        t1[0] = t1[0] + numpy.sin(t1[0]) * .05j
        t1[1] = t1[1] + numpy.sin(t1[1]) * .05j
        t2[0] = t2[0] + numpy.sin(t2[0]) * .05j
        t2[1] = t2[1] + numpy.sin(t2[1]) * .05j
        t2[2] = t2[2] + numpy.sin(t2[2]) * .05j
        t1new_ref, t2new_ref = uccsd.update_amps(myucc, t1, t2, eris)

        nocc = nocca + noccb
        orbspin = numpy.zeros(nao*2, dtype=int)
        orbspin[1::2] = 1
        orbspin[nocc-1] = 0
        orbspin[nocc  ] = 1
        eri1 = numpy.zeros([nao*2]*4, dtype=numpy.complex128)
        idxa = numpy.where(orbspin == 0)[0]
        idxb = numpy.where(orbspin == 1)[0]
        eri1[idxa[:,None,None,None],idxa[:,None,None],idxa[:,None],idxa] = eri0aa
        eri1[idxa[:,None,None,None],idxa[:,None,None],idxb[:,None],idxb] = eri0ab
        eri1[idxb[:,None,None,None],idxb[:,None,None],idxa[:,None],idxa] = eri0ba
        eri1[idxb[:,None,None,None],idxb[:,None,None],idxb[:,None],idxb] = eri0bb
        eri1 = eri1.transpose(0,2,1,3) - eri1.transpose(0,2,3,1)
        erig = gccsd._PhysicistsERIs()
        erig.oooo = eri1[:nocc,:nocc,:nocc,:nocc].copy()
        erig.ooov = eri1[:nocc,:nocc,:nocc,nocc:].copy()
        erig.ovov = eri1[:nocc,nocc:,:nocc,nocc:].copy()
        erig.ovvo = eri1[:nocc,nocc:,nocc:,:nocc].copy()
        erig.oovv = eri1[:nocc,:nocc,nocc:,nocc:].copy()
        erig.ovvv = eri1[:nocc,nocc:,nocc:,nocc:].copy()
        erig.vvvv = eri1[nocc:,nocc:,nocc:,nocc:].copy()
        mo_e = numpy.empty(nao*2)
        mo_e[orbspin==0] = mf.mo_energy[0]
        mo_e[orbspin==1] = mf.mo_energy[1]
        erig.fock = numpy.diag(mo_e)
        erig.mo_energy = mo_e.real

        myccg = gccsd.GCCSD(scf.addons.convert_to_ghf(mf))
        t1 = myccg.spatial2spin(t1, orbspin)
        t2 = myccg.spatial2spin(t2, orbspin)
        t1new, t2new = gccsd.update_amps(myccg, t1, t2, erig)
        t1new = myccg.spin2spatial(t1new, orbspin)
        t2new = myccg.spin2spatial(t2new, orbspin)
        self.assertAlmostEqual(abs(t1new[0] - t1new_ref[0]).max(), 0, 12)
        self.assertAlmostEqual(abs(t1new[1] - t1new_ref[1]).max(), 0, 12)
        self.assertAlmostEqual(abs(t2new[0] - t2new_ref[0]).max(), 0, 12)
        self.assertAlmostEqual(abs(t2new[1] - t2new_ref[1]).max(), 0, 12)
        self.assertAlmostEqual(abs(t2new[2] - t2new_ref[2]).max(), 0, 12)

    def test_mbpt2(self):
        myucc = uccsd.UCCSD(mf)
        e = myucc.kernel(mbpt2=True)[0]
        self.assertAlmostEqual(e, -0.12886859466216125, 10)
        emp2 = mp.MP2(mf).kernel()[0]
        self.assertAlmostEqual(e, emp2, 10)

        myucc = uccsd.UCCSD(mf_s2)
        e = myucc.kernel(mbpt2=True)[0]
        self.assertAlmostEqual(e, -0.096257842171487293, 10)
        emp2 = mp.MP2(mf_s2).kernel()[0]
        self.assertAlmostEqual(e, emp2, 10)

    def test_uintermediats(self):
        from pyscf.cc import uintermediates
        self.assertTrue(eris.get_ovvv().ndim == 4)
        self.assertTrue(eris.get_ovVV().ndim == 4)
        self.assertTrue(eris.get_OVvv().ndim == 4)
        self.assertTrue(eris.get_OVVV().ndim == 4)
        self.assertTrue(eris.get_ovvv(slice(None), slice(2,4)).ndim == 4)
        self.assertTrue(eris.get_ovVV(slice(None), slice(2,4)).ndim == 4)
        self.assertTrue(eris.get_OVvv(slice(None), slice(2,4)).ndim == 4)
        self.assertTrue(eris.get_OVVV(slice(None), slice(2,4)).ndim == 4)
        self.assertTrue(uintermediates._get_vvvv(eris).ndim == 4)
        self.assertTrue(uintermediates._get_vvVV(eris).ndim == 4)
        self.assertTrue(uintermediates._get_VVVV(eris).ndim == 4)

    def test_add_vvvv(self):
        myucc = uccsd.UCCSD(mf_s2)
        nocca, noccb = 6,4
        nmo = mf_s2.mo_occ[0].size
        nvira, nvirb = nmo-nocca, nmo-noccb
        numpy.random.seed(9)
        t1 = [numpy.zeros((nocca,nvira)),
              numpy.zeros((noccb,nvirb))]
        t2 = [numpy.random.random((nocca,nocca,nvira,nvira))-.9,
              numpy.random.random((nocca,noccb,nvira,nvirb))-.9,
              numpy.random.random((noccb,noccb,nvirb,nvirb))-.9]
        t2[0] = t2[0] - t2[0].transpose(1,0,2,3)
        t2[0] = t2[0] - t2[0].transpose(0,1,3,2)
        t2[2] = t2[2] - t2[2].transpose(1,0,2,3)
        t2[2] = t2[2] - t2[2].transpose(0,1,3,2)

        eris1 = copy.copy(eris)
        idxa = lib.square_mat_in_trilu_indices(nvira)
        idxb = lib.square_mat_in_trilu_indices(nvirb)
        ref =(lib.einsum('acbd,ijcd->ijab', eris1.vvvv[:,idxa][idxa], t2[0]),
              lib.einsum('acbd,ijcd->ijab', eris1.vvVV[:,idxb][idxa], t2[1]),
              lib.einsum('acbd,ijcd->ijab', eris1.VVVV[:,idxb][idxb], t2[2]))

        t2a = myucc._add_vvvv((t1[0]*0,t1[1]*0), t2, eris, t2sym=False)
        self.assertAlmostEqual(abs(ref[0]-t2a[0]).max(), 0, 12)
        self.assertAlmostEqual(abs(ref[1]-t2a[1]).max(), 0, 12)
        self.assertAlmostEqual(abs(ref[2]-t2a[2]).max(), 0, 12)

        myucc.direct = True
        eris1.vvvv = None  # == with_ovvv=True in the call below
        eris1.VVVV = None
        eris1.vvVV = None
        t1 = None
        myucc.mo_coeff, eris1.mo_coeff = eris1.mo_coeff, None
        t2b = myucc._add_vvvv(t1, t2, eris1)
        self.assertAlmostEqual(abs(ref[0]-t2b[0]).max(), 0, 12)
        self.assertAlmostEqual(abs(ref[1]-t2b[1]).max(), 0, 12)
        self.assertAlmostEqual(abs(ref[2]-t2b[2]).max(), 0, 12)

    def test_add_vvVV(self):
        myucc = uccsd.UCCSD(mf_s2)
        nocca, noccb = 6,4
        nmo = mf_s2.mo_occ[0].size
        nvira, nvirb = nmo-nocca, nmo-noccb
        numpy.random.seed(9)
        t1 = [numpy.zeros((nocca,nvira)),
              numpy.zeros((noccb,nvirb))]
        t2 = [numpy.random.random((nocca,nocca,nvira,nvira))-.9,
              numpy.random.random((nocca,noccb,nvira,nvirb))-.9,
              numpy.random.random((noccb,noccb,nvirb,nvirb))-.9]
        t2[0] = t2[0] - t2[0].transpose(1,0,2,3)
        t2[0] = t2[0] - t2[0].transpose(0,1,3,2)
        t2[2] = t2[2] - t2[2].transpose(1,0,2,3)
        t2[2] = t2[2] - t2[2].transpose(0,1,3,2)

        eris1 = copy.copy(eris)
        idxa = lib.square_mat_in_trilu_indices(nvira)
        idxb = lib.square_mat_in_trilu_indices(nvirb)
        vvVV = eris1.vvVV[:,idxb][idxa]
        ref = lib.einsum('acbd,ijcd->ijab', vvVV, t2[1])

        t2a = myucc._add_vvVV((t1[0]*0,t1[1]*0), t2[1], eris)
        self.assertAlmostEqual(abs(ref-t2a).max(), 0, 12)

        myucc.direct = True
        eris1.vvvv = None  # == with_ovvv=True in the call below
        eris1.VVVV = None
        eris1.vvVV = None
        t1 = None
        myucc.mo_coeff, eris1.mo_coeff = eris1.mo_coeff, None
        t2b = myucc._add_vvVV(t1, t2[1], eris1)
        self.assertAlmostEqual(abs(ref-t2b).max(), 0, 12)

    def test_zero_beta_electrons(self):
        mol = gto.M(atom='H', basis=('631g', [[0, (.2, 1)], [0, (.5, 1)]]),
                    spin=1, verbose=0)
        mf = scf.UHF(mol).run()
        mycc = uccsd.UCCSD(mf).run()
        self.assertAlmostEqual(mycc.e_corr, 0, 9)

        mol = gto.M(atom='He', basis=('631g', [[0, (.2, 1)], [0, (.5, 1)]]),
                    spin=2, verbose=0)
        mf = scf.UHF(mol).run()
        mycc = uccsd.UCCSD(mf).run()
        self.assertAlmostEqual(mycc.e_corr, -2.6906675843462455e-05, 9)
        self.assertEqual(mycc.t1[1].size, 0)
        self.assertEqual(mycc.t2[1].size, 0)
        self.assertEqual(mycc.t2[2].size, 0)

    def test_reset(self):
        mycc = cc.CCSD(scf.UHF(mol).newton())
        mycc.reset(mol_s2)
        self.assertTrue(mycc.mol is mol_s2)
        self.assertTrue(mycc._scf.mol is mol_s2)

    def test_ao2mo(self):
        mycc = uccsd.UCCSD(mf)
        numpy.random.seed(2)
        nao = mol.nao
        mo = numpy.random.random((2, nao, nao))
        eri_incore = mycc.ao2mo(mo)
        mycc.max_memory = 0
        eri_outcore = mycc.ao2mo(mo)
        self.assertTrue(isinstance(eri_outcore.oovv, h5py.Dataset))
        self.assertAlmostEqual(abs(eri_incore.oooo - eri_outcore.oooo).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.oovv - eri_outcore.oovv).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.ovoo - eri_outcore.ovoo).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.ovvo - eri_outcore.ovvo).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.ovov - eri_outcore.ovov).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.ovvv - eri_outcore.ovvv).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.vvvv - eri_outcore.vvvv).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.OOOO - eri_outcore.OOOO).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.OOVV - eri_outcore.OOVV).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.OVOO - eri_outcore.OVOO).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.OVVO - eri_outcore.OVVO).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.OVOV - eri_outcore.OVOV).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.OVVV - eri_outcore.OVVV).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.VVVV - eri_outcore.VVVV).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.ooOO - eri_outcore.ooOO).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.ooVV - eri_outcore.ooVV).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.ovOO - eri_outcore.ovOO).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.ovVO - eri_outcore.ovVO).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.ovOV - eri_outcore.ovOV).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.ovVV - eri_outcore.ovVV).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.vvVV - eri_outcore.vvVV).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.OOvv - eri_outcore.OOvv).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.OVoo - eri_outcore.OVoo).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.OVvo - eri_outcore.OVvo).max(), 0, 12)
        self.assertAlmostEqual(abs(eri_incore.OVvv - eri_outcore.OVvv).max(), 0, 12)


if __name__ == "__main__":
    print("Full Tests for UCCSD")
    unittest.main()


# pyscf/cc/test/test_uccsd_lambda.py

#!/usr/bin/env python
# Copyright 2014-2018 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.

import unittest
import numpy
import numpy as np
from functools import reduce

from pyscf import lib
from pyscf import gto
from pyscf import scf
from pyscf import ao2mo
from pyscf.cc import uccsd
from pyscf.cc import addons
from pyscf.cc import uccsd_lambda
from pyscf.cc import gccsd, gccsd_lambda

def setUpModule():
    global mol, mf, mycc
    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = '631g'
    mol.spin = 2
    mol.build()
    mf = scf.UHF(mol).run()
    mycc = uccsd.UCCSD(mf)

def tearDownModule():
    global mol, mf, mycc
    del mol, mf, mycc

class KnownValues(unittest.TestCase):
    def test_update_lambda_real(self):
        numpy.random.seed(21)
        eris = mycc.ao2mo()
        gcc1 = gccsd.GCCSD(scf.addons.convert_to_ghf(mf))
        eri1 = gcc1.ao2mo()
        orbspin = eri1.orbspin

        nocc = mol.nelectron
        nvir = mol.nao_nr()*2 - nocc

        t1r = numpy.random.random((nocc,nvir))*.1
        t2r = numpy.random.random((nocc,nocc,nvir,nvir))*.1
        t2r = t2r - t2r.transpose(1,0,2,3)
        t2r = t2r - t2r.transpose(0,1,3,2)
        l1r = numpy.random.random((nocc,nvir))*.1
        l2r = numpy.random.random((nocc,nocc,nvir,nvir))*.1
        l2r = l2r - l2r.transpose(1,0,2,3)
        l2r = l2r - l2r.transpose(0,1,3,2)
        t1r = addons.spin2spatial(t1r, orbspin)
        t2r = addons.spin2spatial(t2r, orbspin)
        t1r = addons.spatial2spin(t1r, orbspin)
        t2r = addons.spatial2spin(t2r, orbspin)
        l1r = addons.spin2spatial(l1r, orbspin)
        l2r = addons.spin2spatial(l2r, orbspin)
        l1r = addons.spatial2spin(l1r, orbspin)
        l2r = addons.spatial2spin(l2r, orbspin)
        imds = gccsd_lambda.make_intermediates(gcc1, t1r, t2r, eri1)
        l1ref, l2ref = gccsd_lambda.update_lambda(gcc1, t1r, t2r, l1r, l2r, eri1, imds)

        t1 = addons.spin2spatial(t1r, orbspin)
        t2 = addons.spin2spatial(t2r, orbspin)
        l1 = addons.spin2spatial(l1r, orbspin)
        l2 = addons.spin2spatial(l2r, orbspin)
        imds = uccsd_lambda.make_intermediates(mycc, t1, t2, eris)
        l1, l2 = uccsd_lambda.update_lambda(mycc, t1, t2, l1, l2, eris, imds)
        self.assertAlmostEqual(float(abs(addons.spatial2spin(l1, orbspin)-l1ref).max()), 0, 8)
        self.assertAlmostEqual(float(abs(addons.spatial2spin(l2, orbspin)-l2ref).max()), 0, 8)

        l1ref = addons.spin2spatial(l1ref, orbspin)
        l2ref = addons.spin2spatial(l2ref, orbspin)
        self.assertAlmostEqual(abs(l1[0]-l1ref[0]).max(), 0, 8)
        self.assertAlmostEqual(abs(l1[1]-l1ref[1]).max(), 0, 8)
        self.assertAlmostEqual(abs(l2[0]-l2ref[0]).max(), 0, 8)
        self.assertAlmostEqual(abs(l2[1]-l2ref[1]).max(), 0, 8)
        self.assertAlmostEqual(abs(l2[2]-l2ref[2]).max(), 0, 8)

    def test_update_lambda_complex(self):
        nocca, noccb = mol.nelec
        nmo = mol.nao_nr()
        nvira,nvirb = nmo-nocca, nmo-noccb
        numpy.random.seed(9)
        t1 = [numpy.random.random((nocca,nvira))-.9,
              numpy.random.random((noccb,nvirb))-.9]
        l1 = [numpy.random.random((nocca,nvira))-.9,
              numpy.random.random((noccb,nvirb))-.9]
        t2 = [numpy.random.random((nocca,nocca,nvira,nvira))-.9,
              numpy.random.random((nocca,noccb,nvira,nvirb))-.9,
              numpy.random.random((noccb,noccb,nvirb,nvirb))-.9]
        t2[0] = t2[0] - t2[0].transpose(1,0,2,3)
        t2[0] = t2[0] - t2[0].transpose(0,1,3,2)
        t2[2] = t2[2] - t2[2].transpose(1,0,2,3)
        t2[2] = t2[2] - t2[2].transpose(0,1,3,2)
        l2 = [numpy.random.random((nocca,nocca,nvira,nvira))-.9,
              numpy.random.random((nocca,noccb,nvira,nvirb))-.9,
              numpy.random.random((noccb,noccb,nvirb,nvirb))-.9]
        l2[0] = l2[0] - l2[0].transpose(1,0,2,3)
        l2[0] = l2[0] - l2[0].transpose(0,1,3,2)
        l2[2] = l2[2] - l2[2].transpose(1,0,2,3)
        l2[2] = l2[2] - l2[2].transpose(0,1,3,2)

#        eris = mycc.ao2mo()
#        imds = make_intermediates(mycc, t1, t2, eris)
#        l1new, l2new = update_lambda(mycc, t1, t2, l1, l2, eris, imds)
#        print(lib.finger(l1new[0]) --104.55975252585894)
#        print(lib.finger(l1new[1]) --241.12677819375281)
#        print(lib.finger(l2new[0]) --0.4957533529669417)
#        print(lib.finger(l2new[1]) - 15.46423057451851 )
#        print(lib.finger(l2new[2]) - 5.8430776663704407)

        nocca, noccb = mol.nelec
        mo_a = mf.mo_coeff[0] + numpy.sin(mf.mo_coeff[0]) * .01j
        mo_b = mf.mo_coeff[1] + numpy.sin(mf.mo_coeff[1]) * .01j
        nao = mo_a.shape[0]
        eri = ao2mo.restore(1, mf._eri, nao)
        eri0aa = lib.einsum('pqrs,pi,qj,rk,sl->ijkl', eri, mo_a.conj(), mo_a, mo_a.conj(), mo_a)
        eri0ab = lib.einsum('pqrs,pi,qj,rk,sl->ijkl', eri, mo_a.conj(), mo_a, mo_b.conj(), mo_b)
        eri0bb = lib.einsum('pqrs,pi,qj,rk,sl->ijkl', eri, mo_b.conj(), mo_b, mo_b.conj(), mo_b)
        eri0ba = eri0ab.transpose(2,3,0,1)

        nvira = nao - nocca
        nvirb = nao - noccb
        eris = uccsd._ChemistsERIs(mol)
        eris.oooo = eri0aa[:nocca,:nocca,:nocca,:nocca].copy()
        eris.ovoo = eri0aa[:nocca,nocca:,:nocca,:nocca].copy()
        eris.oovv = eri0aa[:nocca,:nocca,nocca:,nocca:].copy()
        eris.ovvo = eri0aa[:nocca,nocca:,nocca:,:nocca].copy()
        eris.ovov = eri0aa[:nocca,nocca:,:nocca,nocca:].copy()
        eris.ovvv = eri0aa[:nocca,nocca:,nocca:,nocca:].copy()
        eris.vvvv = eri0aa[nocca:,nocca:,nocca:,nocca:].copy()

        eris.OOOO = eri0bb[:noccb,:noccb,:noccb,:noccb].copy()
        eris.OVOO = eri0bb[:noccb,noccb:,:noccb,:noccb].copy()
        eris.OOVV = eri0bb[:noccb,:noccb,noccb:,noccb:].copy()
        eris.OVVO = eri0bb[:noccb,noccb:,noccb:,:noccb].copy()
        eris.OVOV = eri0bb[:noccb,noccb:,:noccb,noccb:].copy()
        eris.OVVV = eri0bb[:noccb,noccb:,noccb:,noccb:].copy()
        eris.VVVV = eri0bb[noccb:,noccb:,noccb:,noccb:].copy()

        eris.ooOO = eri0ab[:nocca,:nocca,:noccb,:noccb].copy()
        eris.ovOO = eri0ab[:nocca,nocca:,:noccb,:noccb].copy()
        eris.ooVV = eri0ab[:nocca,:nocca,noccb:,noccb:].copy()
        eris.ovVO = eri0ab[:nocca,nocca:,noccb:,:noccb].copy()
        eris.ovOV = eri0ab[:nocca,nocca:,:noccb,noccb:].copy()
        eris.ovVV = eri0ab[:nocca,nocca:,noccb:,noccb:].copy()
        eris.vvVV = eri0ab[nocca:,nocca:,noccb:,noccb:].copy()

        eris.OOoo = eri0ba[:noccb,:noccb,:nocca,:nocca].copy()
        eris.OVoo = eri0ba[:noccb,noccb:,:nocca,:nocca].copy()
        eris.OOvv = eri0ba[:noccb,:noccb,nocca:,nocca:].copy()
        eris.OVvo = eri0ba[:noccb,noccb:,nocca:,:nocca].copy()
        eris.OVov = eri0ba[:noccb,noccb:,:nocca,nocca:].copy()
        eris.OVvv = eri0ba[:noccb,noccb:,nocca:,nocca:].copy()
        eris.VVvv = eri0ba[noccb:,noccb:,nocca:,nocca:].copy()

        eris.focka = numpy.diag(mf.mo_energy[0])
        eris.fockb = numpy.diag(mf.mo_energy[1])
        eris.mo_energy = mf.mo_energy

        t1[0] = t1[0] + numpy.sin(t1[0]) * .05j
        t1[1] = t1[1] + numpy.sin(t1[1]) * .05j
        t2[0] = t2[0] + numpy.sin(t2[0]) * .05j
        t2[1] = t2[1] + numpy.sin(t2[1]) * .05j
        t2[2] = t2[2] + numpy.sin(t2[2]) * .05j
        l1[0] = l1[0] + numpy.sin(l1[0]) * .05j
        l1[1] = l1[1] + numpy.sin(l1[1]) * .05j
        l2[0] = l2[0] + numpy.sin(l2[0]) * .05j
        l2[1] = l2[1] + numpy.sin(l2[1]) * .05j
        l2[2] = l2[2] + numpy.sin(l2[2]) * .05j
        imds = uccsd_lambda.make_intermediates(mycc, t1, t2, eris)
        l1new_ref, l2new_ref = uccsd_lambda.update_lambda(mycc, t1, t2, l1, l2, eris, imds)

        nocc = nocca + noccb
        orbspin = numpy.zeros(nao*2, dtype=int)
        orbspin[1::2] = 1
        orbspin[nocc-1] = 0
        orbspin[nocc  ] = 1
        eri1 = numpy.zeros([nao*2]*4, dtype=numpy.complex128)
        idxa = numpy.where(orbspin == 0)[0]
        idxb = numpy.where(orbspin == 1)[0]
        eri1[idxa[:,None,None,None],idxa[:,None,None],idxa[:,None],idxa] = eri0aa
        eri1[idxa[:,None,None,None],idxa[:,None,None],idxb[:,None],idxb] = eri0ab
        eri1[idxb[:,None,None,None],idxb[:,None,None],idxa[:,None],idxa] = eri0ba
        eri1[idxb[:,None,None,None],idxb[:,None,None],idxb[:,None],idxb] = eri0bb
        eri1 = eri1.transpose(0,2,1,3) - eri1.transpose(0,2,3,1)
        erig = gccsd._PhysicistsERIs()
        erig.oooo = eri1[:nocc,:nocc,:nocc,:nocc].copy()
        erig.ooov = eri1[:nocc,:nocc,:nocc,nocc:].copy()
        erig.ovov = eri1[:nocc,nocc:,:nocc,nocc:].copy()
        erig.ovvo = eri1[:nocc,nocc:,nocc:,:nocc].copy()
        erig.oovv = eri1[:nocc,:nocc,nocc:,nocc:].copy()
        erig.ovvv = eri1[:nocc,nocc:,nocc:,nocc:].copy()
        erig.vvvv = eri1[nocc:,nocc:,nocc:,nocc:].copy()
        mo_e = numpy.empty(nao*2)
        mo_e[orbspin==0] = mf.mo_energy[0]
        mo_e[orbspin==1] = mf.mo_energy[1]
        erig.fock = numpy.diag(mo_e)
        erig.mo_energy = mo_e.real

        myccg = gccsd.GCCSD(scf.addons.convert_to_ghf(mf))
        t1 = myccg.spatial2spin(t1, orbspin)
        t2 = myccg.spatial2spin(t2, orbspin)
        l1 = myccg.spatial2spin(l1, orbspin)
        l2 = myccg.spatial2spin(l2, orbspin)
        imds = gccsd_lambda.make_intermediates(myccg, t1, t2, erig)
        l1new, l2new = gccsd_lambda.update_lambda(myccg, t1, t2, l1, l2, erig, imds)
        l1new = myccg.spin2spatial(l1new, orbspin)
        l2new = myccg.spin2spatial(l2new, orbspin)
        self.assertAlmostEqual(abs(l1new[0] - l1new_ref[0]).max(), 0, 11)
        self.assertAlmostEqual(abs(l1new[1] - l1new_ref[1]).max(), 0, 11)
        self.assertAlmostEqual(abs(l2new[0] - l2new_ref[0]).max(), 0, 11)
        self.assertAlmostEqual(abs(l2new[1] - l2new_ref[1]).max(), 0, 11)
        self.assertAlmostEqual(abs(l2new[2] - l2new_ref[2]).max(), 0, 11)


if __name__ == "__main__":
    print("Full Tests for UCCSD lambda")
    unittest.main()


# pyscf/cc/test/test_uccsd_t.py

#!/usr/bin/env python
# Copyright 2014-2018 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.

import unittest
import numpy
import copy
from functools import reduce

from pyscf import gto, scf, lib, symm
from pyscf import cc
from pyscf.cc import uccsd_t
from pyscf.cc import gccsd, gccsd_t

def setUpModule():
    global mol, mol1, mf, mcc
    mol = gto.Mole()
    mol.verbose = 5
    mol.output = '/dev/null'
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -.757 , .587)],
        [1 , (0. ,  .757 , .587)]]
    mol.spin = 2
    mol.basis = '3-21g'
    mol.symmetry = 'C2v'
    mol.build()
    mol1 = copy.copy(mol)
    mol1.symmetry = False

    mf = scf.UHF(mol1).run(conv_tol=1e-14)
    mcc = cc.UCCSD(mf)
    mcc.conv_tol = 1e-14
    mcc.kernel()

def tearDownModule():
    global mol, mol1, mf, mcc
    mol.stdout.close()
    del mol, mol1, mf, mcc

class KnownValues(unittest.TestCase):
    def test_uccsd_t(self):
        mf1 = copy.copy(mf)
        nao, nmo = mf.mo_coeff[0].shape
        numpy.random.seed(10)
        mf1.mo_coeff = numpy.random.random((2,nao,nmo)) - .5
        numpy.random.seed(12)
        nocca, noccb = mol.nelec
        nmo = mf1.mo_occ[0].size
        nvira = nmo - nocca
        nvirb = nmo - noccb
        t1a  = .1 * numpy.random.random((nocca,nvira))
        t1b  = .1 * numpy.random.random((noccb,nvirb))
        t2aa = .1 * numpy.random.random((nocca,nocca,nvira,nvira))
        t2aa = t2aa - t2aa.transpose(0,1,3,2)
        t2aa = t2aa - t2aa.transpose(1,0,2,3)
        t2bb = .1 * numpy.random.random((noccb,noccb,nvirb,nvirb))
        t2bb = t2bb - t2bb.transpose(0,1,3,2)
        t2bb = t2bb - t2bb.transpose(1,0,2,3)
        t2ab = .1 * numpy.random.random((nocca,noccb,nvira,nvirb))
        t1 = t1a, t1b
        t2 = t2aa, t2ab, t2bb
        mycc = cc.UCCSD(mf1)
        eris = mycc.ao2mo(mf1.mo_coeff)
        mycc.incore_complete = True
        e3a = mycc.ccsd_t([t1a,t1b], [t2aa, t2ab, t2bb], eris)
        self.assertAlmostEqual(e3a, 15.582860941071505, 8)

        mycc.incore_complete = False
        mycc.max_memory = 0
        e3a = uccsd_t.kernel(mycc, eris, [t1a,t1b], [t2aa, t2ab, t2bb])
        self.assertAlmostEqual(e3a, 15.582860941071505, 8)

        e3a = mcc.ccsd_t()
        self.assertAlmostEqual(e3a, -0.0009857042572475674, 11)

    #def test_uccsd_t_symm(self):
    #    mf = scf.UHF(mol).run(conv_tol=1e-14)
    #    mcc = cc.UCCSD(mf)
    #    mcc.conv_tol = 1e-14
    #    e3a = mcc.run().ccsd_t()
    #    self.assertAlmostEqual(e3a, -0.0030600226107389866, 11)

    def test_uccsd_t_complex(self):
        mol = gto.M()
        numpy.random.seed(12)
        nocca, noccb, nvira, nvirb = 3, 2, 4, 5
        nmo = nocca + nvira
        eris = cc.uccsd._ChemistsERIs()
        eris.nocca = nocca
        eris.noccb = noccb
        eris.nocc = (nocca, noccb)
        eri1 = (numpy.random.random((3,nmo,nmo,nmo,nmo)) +
                numpy.random.random((3,nmo,nmo,nmo,nmo)) * .8j - .5-.4j)
        eri1 = eri1 + eri1.transpose(0,2,1,4,3).conj()
        eri1[0] = eri1[0] + eri1[0].transpose(2,3,0,1)
        eri1[2] = eri1[2] + eri1[2].transpose(2,3,0,1)
        eri1 *= .1
        eris.ovvv = eri1[0,:nocca,nocca:,nocca:,nocca:]
        eris.ovov = eri1[0,:nocca,nocca:,:nocca,nocca:]
        eris.ovoo = eri1[0,:nocca,nocca:,:nocca,:nocca]
        eris.OVVV = eri1[2,:noccb,noccb:,noccb:,noccb:]
        eris.OVOV = eri1[2,:noccb,noccb:,:noccb,noccb:]
        eris.OVOO = eri1[2,:noccb,noccb:,:noccb,:noccb]
        eris.voVP = eri1[1,nocca:,:nocca,noccb:,:     ]
        eris.ovVV = eri1[1,:nocca,nocca:,noccb:,noccb:]
        eris.ovOV = eri1[1,:nocca,nocca:,:noccb,noccb:]
        eris.ovOO = eri1[1,:nocca,nocca:,:noccb,:noccb]
        eris.OVvv = eri1[1,nocca:,nocca:,:noccb,noccb:].transpose(2,3,0,1)
        eris.OVoo = eri1[1,:nocca,:nocca,:noccb,noccb:].transpose(2,3,0,1)
        t1a  = .1 * numpy.random.random((nocca,nvira)) + numpy.random.random((nocca,nvira))*.1j
        t1b  = .1 * numpy.random.random((noccb,nvirb)) + numpy.random.random((noccb,nvirb))*.1j
        t2aa = .1 * numpy.random.random((nocca,nocca,nvira,nvira)) + numpy.random.random((nocca,nocca,nvira,nvira))*.1j
        t2aa = t2aa - t2aa.transpose(0,1,3,2)
        t2aa = t2aa - t2aa.transpose(1,0,2,3)
        t2bb = .1 * numpy.random.random((noccb,noccb,nvirb,nvirb)) + numpy.random.random((noccb,noccb,nvirb,nvirb))*.1j
        t2bb = t2bb - t2bb.transpose(0,1,3,2)
        t2bb = t2bb - t2bb.transpose(1,0,2,3)
        t2ab = .1 * numpy.random.random((nocca,noccb,nvira,nvirb)) + numpy.random.random((nocca,noccb,nvira,nvirb))*.1j
        f = (numpy.random.random((2,nmo,nmo)) * .4 +
             numpy.random.random((2,nmo,nmo)) * .4j)
        eris.focka = f[0]+f[0].T.conj() + numpy.diag(numpy.arange(nmo))
        eris.fockb = f[1]+f[1].T.conj() + numpy.diag(numpy.arange(nmo))
        eris.mo_energy = (eris.focka.diagonal().real,
                          eris.fockb.diagonal().real)
        t1 = t1a, t1b
        t2 = t2aa, t2ab, t2bb
        mcc = cc.UCCSD(scf.UHF(mol))
        mcc.nocc = eris.nocc
        e0 = uccsd_t.kernel(mcc, eris, t1, t2)

        eri2 = numpy.zeros((nmo*2,nmo*2,nmo*2,nmo*2), dtype=eri1.dtype)
        orbspin = numpy.zeros(nmo*2,dtype=int)
        orbspin[1::2] = 1
        eri2[0::2,0::2,0::2,0::2] = eri1[0]
        eri2[1::2,1::2,0::2,0::2] = eri1[1].transpose(2,3,0,1)
        eri2[0::2,0::2,1::2,1::2] = eri1[1]
        eri2[1::2,1::2,1::2,1::2] = eri1[2]
        eri2 = eri2.transpose(0,2,1,3) - eri2.transpose(0,2,3,1)
        fock = numpy.zeros((nmo*2,nmo*2), dtype=eris.focka.dtype)
        fock[0::2,0::2] = eris.focka
        fock[1::2,1::2] = eris.fockb
        eris1 = gccsd._PhysicistsERIs()
        nocc = nocca + noccb
        eris1.ovvv = eri2[:nocc,nocc:,nocc:,nocc:]
        eris1.oovv = eri2[:nocc,:nocc,nocc:,nocc:]
        eris1.ooov = eri2[:nocc,:nocc,:nocc,nocc:]
        eris1.fock = fock
        eris1.mo_energy = fock.diagonal().real
        t1 = gccsd.spatial2spin(t1, orbspin)
        t2 = gccsd.spatial2spin(t2, orbspin)
        gcc = gccsd.GCCSD(scf.GHF(gto.M()))
        e1 = gccsd_t.kernel(gcc, eris1, t1, t2)
        self.assertAlmostEqual(e0, e1.real, 9)
        self.assertAlmostEqual(e1, -0.056092415718338388-0.011390417704868244j, 9)


if __name__ == "__main__":
    print("Full Tests for UCCSD(T)")
    unittest.main()



# pyscf/__all__.py

from . import lib, gto, scf, ao2mo
from .post_scf import *  # noqa: F403
from . import grad
from . import gw
from . import hessian
from . import lo
from . import qmmm
#from . import semiempirical
from . import sgx
from . import solvent
from . import tools
from . import x2c

from . import geomopt

#try:
#    from . import mrpt
#except ImportError:
#    pass
#try:
#    from . import prop
#except ImportError:
#    pass
#try:
#    from . import dftd3
#except ImportError:
#    pass
#try:
#    from . import dmrgscf
#except ImportError:
#    pass
#try:
#    from . import icmpspt
#except ImportError:
#    pass
#try:
#    from . import shciscf
#except ImportError:
#    pass
#try:
#    from . import cornell_shci
#except ImportError:
#    pass

from . import pbc
from .pbc import __all__
del __all__



# pyscf/__config__.py

import os, sys

#
# All parameters initialized before loading pyscf_conf.py will be overwritten
# by the dynamic importing procedure.
#

DEBUG = False

MAX_MEMORY = int(os.environ.get('PYSCF_MAX_MEMORY', 4000)) # MB
TMPDIR = os.environ.get('TMPDIR', '.')
TMPDIR = os.environ.get('PYSCF_TMPDIR', TMPDIR)

VERBOSE = 3  # default logger level (logger.NOTE)
UNIT = 'angstrom'

#
# Loading pyscf_conf.py and overwriting above parameters
#
for conf_file in (os.environ.get('PYSCF_CONFIG_FILE', None),
                  os.path.join(os.path.abspath('.'), '.pyscf_conf.py'),
                  os.path.join(os.environ.get('HOME', '.'), '.pyscf_conf.py')):
    if conf_file is not None and os.path.isfile(conf_file):
        break
else:
    conf_file = None

if conf_file is not None:
    if sys.version_info < (3,0):
        import imp
        imp.load_source('pyscf.__config__', conf_file)
        del(imp)
    else:
        from importlib import machinery
        machinery.SourceFileLoader('pyscf.__config__', conf_file).load_module()
        del(machinery)
del(os, sys)

#
# All parameters initialized after loading pyscf_conf.py will be kept in the
# program.
#


# pyscf/__init__.py

# Copyright 2014-2020 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.

'''
*****************************************************
PySCF Python-based simulations of chemistry framework
*****************************************************

How to use
----------
There are two ways to access the documentation: the docstrings come with
the code, and an online program reference, available from
http://www.sunqm.net/pyscf/index.html

We recommend the enhanced Python interpreter `IPython <http://ipython.org>`_
and the web-based Python IDE `Ipython notebook <http://ipython.org/notebook.html>`_
to try out the package::

    >>> from pyscf import gto, scf
    >>> mol = gto.M(atom='H 0 0 0; H 0 0 1.2', basis='cc-pvdz')
    >>> mol.apply(scf.RHF).run()
    converged SCF energy = -1.06111199785749
    -1.06111199786

'''

__version__ = '2.0.1'

import os
import sys
# Avoid too many threads being created in OMP loops.
# See issue https://github.com/pyscf/pyscf/issues/317
if 'OPENBLAS_NUM_THREADS' not in os.environ:
    os.environ['OPENBLAS_NUM_THREADS'] = '1'
if 'MKL_NUM_THREADS' not in os.environ:
    os.environ['MKL_NUM_THREADS'] = '1'

# Load modules which are developed as plugins of the namespace package
PYSCF_EXT_PATH = os.getenv('PYSCF_EXT_PATH')
if PYSCF_EXT_PATH:
    for p in PYSCF_EXT_PATH.split(':'):
        if os.path.isdir(p):
            submodules = os.listdir(p)
            if 'pyscf' in submodules:
                # PYSCF_EXT_PATH points to the root directory of a submodule
                __path__.append(os.path.join(p, 'pyscf'))
            else:
                # Load all modules in PYSCF_EXT_PATH if it's a folder that
                # contains all extended modules
                for f in submodules:
                    __path__.append(os.path.join(p, f, 'pyscf'))
                del f
        elif os.path.exists(p):
            # Load all moduels defined inside the file PYSCF_EXT_PATH
            with open(p, 'r') as f:
                __path__.extend(f.read().splitlines())
            del f
    del p
elif '/site-packages/' in __file__ or '/dist-packages/' in __file__:
    # If pyscf has been installed in the standard runtime path (typically
    # under site-packages), and plugins are installed with the pip editable
    # mode, load namespace plugins. In this case, it is likely all modules are
    # managed by pip/conda or their virtual environments. It is safe to search
    # namespace plugins according to the old style of PEP 420.
    __path__ = __import__('pkgutil').extend_path(__path__, __name__)
else:
    # We need a better way to load plugins if pyscf is imported by the
    # environment variable PYTHONPATH. Current treatment may mix installed
    # plugins (e.g.  through pip install) with the developing plugins which
    # were accidentally placed under PYTHONPATH. When PYTHONPATH mechanism is
    # taken, an explicit list of extended paths (using environment
    # PYSCF_EXT_PATH) is recommended.
    __path__ = __import__('pkgutil').extend_path(__path__, __name__)
    if not all('/site-packages/' in p for p in __path__[1:]):
        sys.stderr.write('pyscf plugins found in \n%s\n'
                         'When PYTHONPATH is set, it is recommended to load '
                         'these plugins through the environment variable '
                         'PYSCF_EXT_PATH' % '\n'.join(__path__[1:]))

from distutils.version import LooseVersion
import numpy
if LooseVersion(numpy.__version__) <= '1.8.0':
    raise SystemError("You're using an old version of Numpy (%s). "
                      "It is recommended to upgrade numpy to 1.8.0 or newer. \n"
                      "You still can use all features of PySCF with the old numpy by removing this warning msg. "
                      "Some modules (DFT, CC, MRPT) might be affected because of the bug in old numpy." %
                      numpy.__version__)
elif '1.16.2' <= LooseVersion(numpy.__version__) < '1.18':
    #sys.stderr.write('Numpy 1.16 has memory leak bug  '
    #                 'https://github.com/numpy/numpy/issues/13808\n'
    #                 'It is recommended to downgrade to numpy 1.15 or older\n')
    import ctypes
    from numpy.core import _internal
    def _get_void_ptr(arr):
        simple_arr = numpy.asarray(_internal._unsafe_first_element_pointer(arr))
        c_arr = (ctypes.c_char * 0).from_buffer(simple_arr)
        return ctypes.cast(ctypes.byref(c_arr), ctypes.c_void_p)
    # patch _get_void_ptr as a workaround to numpy issue #13808
    _internal._get_void_ptr = _get_void_ptr

from pyscf import __config__
from pyscf import lib
from pyscf import gto
from pyscf import scf
from pyscf import ao2mo

# Whether to enable debug mode. When this flag is set, some modules may run
# extra debug code.
DEBUG = __config__.DEBUG

def M(**kwargs):
    '''Main driver to create Molecule object (mol) or Material crystal object (cell)'''
    from pyscf import __all__
    if kwargs.get('a') is not None:  # a is crystal lattice parameter
        return __all__.pbc.gto.M(**kwargs)
    else:  # Molecule
        return gto.M(**kwargs)

del os, sys, LooseVersion


# pyscf/post_scf.py

from . import mp, ci, cc, mcscf, fci, tdscf
try:
    from . import doci
except ImportError:
    pass


# pyscf/cc/__init__.py

# Copyright 2014-2020 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.

'''
Coupled Cluster
===============

Simple usage::

    >>> from pyscf import gto, scf, cc
    >>> mol = gto.M(atom='H 0 0 0; H 0 0 1')
    >>> mf = scf.RHF(mol).run()
    >>> cc.CCSD(mf).run()

:func:`cc.CCSD` returns an instance of CCSD class.  Followings are parameters
to control CCSD calculation.

    verbose : int
        Print level.  Default value equals to :class:`Mole.verbose`
    max_memory : float or int
        Allowed memory in MB.  Default value equals to :class:`Mole.max_memory`
    conv_tol : float
        converge threshold.  Default is 1e-7.
    conv_tol_normt : float
        converge threshold for norm(t1,t2).  Default is 1e-5.
    max_cycle : int
        max number of iterations.  Default is 50.
    diis_space : int
        DIIS space size.  Default is 6.
    diis_start_cycle : int
        The step to start DIIS.  Default is 0.
    direct : bool
        AO-direct CCSD. Default is False.
    async_io : bool
        Allow for asynchronous function execution. Default is True.
    incore_complete : bool
        Avoid all I/O. Default is False.
    frozen : int or list
        If integer is given, the inner-most orbitals are frozen from CC
        amplitudes.  Given the orbital indices (0-based) in a list, both
        occupied and virtual orbitals can be frozen in CC calculation.


Saved results

    converged : bool
        CCSD converged or not
    e_tot : float
        Total CCSD energy (HF + correlation)
    t1, t2 :
        t1[i,a], t2[i,j,a,b]  (i,j in occ, a,b in virt)
    l1, l2 :
        Lambda amplitudes l1[i,a], l2[i,j,a,b]  (i,j in occ, a,b in virt)
'''

from pyscf.cc import ccsd
from pyscf.cc import ccsd_lambda
from pyscf.cc import ccsd_rdm
from pyscf.cc import addons
from pyscf.cc import rccsd
from pyscf.cc import uccsd
from pyscf.cc import gccsd
from pyscf.cc import eom_rccsd
from pyscf.cc import eom_uccsd
from pyscf.cc import eom_gccsd
from pyscf.cc import qcisd
from pyscf import scf

def CCSD(mf, frozen=None, mo_coeff=None, mo_occ=None):
    if isinstance(mf, scf.uhf.UHF):
        return UCCSD(mf, frozen, mo_coeff, mo_occ)
    elif isinstance(mf, scf.ghf.GHF):
        return GCCSD(mf, frozen, mo_coeff, mo_occ)
    else:
        return RCCSD(mf, frozen, mo_coeff, mo_occ)
CCSD.__doc__ = ccsd.CCSD.__doc__

scf.hf.SCF.CCSD = CCSD


def RCCSD(mf, frozen=None, mo_coeff=None, mo_occ=None):
    import numpy
    from pyscf import lib
    from pyscf.soscf import newton_ah
    from pyscf.cc import dfccsd

    if isinstance(mf, scf.uhf.UHF):
        raise RuntimeError('RCCSD cannot be used with UHF method.')
    elif isinstance(mf, scf.rohf.ROHF):
        lib.logger.warn(mf, 'RCCSD method does not support ROHF method. ROHF object '
                        'is converted to UHF object and UCCSD method is called.')
        mf = scf.addons.convert_to_uhf(mf)
        return UCCSD(mf, frozen, mo_coeff, mo_occ)

    if isinstance(mf, newton_ah._CIAH_SOSCF) or not isinstance(mf, scf.hf.RHF):
        mf = scf.addons.convert_to_rhf(mf)

    if getattr(mf, 'with_df', None):
        return dfccsd.RCCSD(mf, frozen, mo_coeff, mo_occ)

    elif numpy.iscomplexobj(mo_coeff) or numpy.iscomplexobj(mf.mo_coeff):
        return rccsd.RCCSD(mf, frozen, mo_coeff, mo_occ)

    else:
        return ccsd.CCSD(mf, frozen, mo_coeff, mo_occ)
RCCSD.__doc__ = ccsd.CCSD.__doc__


def UCCSD(mf, frozen=None, mo_coeff=None, mo_occ=None):
    from pyscf.soscf import newton_ah

    if isinstance(mf, newton_ah._CIAH_SOSCF) or not isinstance(mf, scf.uhf.UHF):
        mf = scf.addons.convert_to_uhf(mf)

    if getattr(mf, 'with_df', None):
        # TODO: DF-UCCSD with memory-efficient particle-particle ladder,
        # similar to dfccsd.RCCSD
        return uccsd.UCCSD(mf, frozen, mo_coeff, mo_occ)
    else:
        return uccsd.UCCSD(mf, frozen, mo_coeff, mo_occ)
UCCSD.__doc__ = uccsd.UCCSD.__doc__


def GCCSD(mf, frozen=None, mo_coeff=None, mo_occ=None):
    from pyscf.soscf import newton_ah

    if isinstance(mf, newton_ah._CIAH_SOSCF) or not isinstance(mf, scf.ghf.GHF):
        mf = scf.addons.convert_to_ghf(mf)

    if getattr(mf, 'with_df', None):
        raise NotImplementedError('DF-GCCSD')
    else:
        return gccsd.GCCSD(mf, frozen, mo_coeff, mo_occ)
GCCSD.__doc__ = gccsd.GCCSD.__doc__


def QCISD(mf, frozen=None, mo_coeff=None, mo_occ=None):
    if isinstance(mf, scf.uhf.UHF):
        raise NotImplementedError
    elif isinstance(mf, scf.ghf.GHF):
        raise NotImplementedError
    else:
        return RQCISD(mf, frozen, mo_coeff, mo_occ)
QCISD.__doc__ = qcisd.QCISD.__doc__

scf.hf.SCF.QCISD = QCISD

def RQCISD(mf, frozen=None, mo_coeff=None, mo_occ=None):
    import numpy
    from pyscf import lib
    from pyscf.soscf import newton_ah

    if isinstance(mf, scf.uhf.UHF):
        raise RuntimeError('RQCISD cannot be used with UHF method.')
    elif isinstance(mf, scf.rohf.ROHF):
        lib.logger.warn(mf, 'RQCISD method does not support ROHF method. ROHF object '
                        'is converted to UHF object and UQCISD method is called.')
        mf = scf.addons.convert_to_uhf(mf)
        raise NotImplementedError

    if isinstance(mf, newton_ah._CIAH_SOSCF) or not isinstance(mf, scf.hf.RHF):
        mf = scf.addons.convert_to_rhf(mf)

    elif numpy.iscomplexobj(mo_coeff) or numpy.iscomplexobj(mf.mo_coeff):
        raise NotImplementedError

    else:
        return qcisd.QCISD(mf, frozen, mo_coeff, mo_occ)
RQCISD.__doc__ = qcisd.QCISD.__doc__


def FNOCCSD(mf, thresh=1e-6, pct_occ=None, nvir_act=None):
    """Frozen natural orbital CCSD

    Attributes:
        thresh : float
            Threshold on NO occupation numbers. Default is 1e-6 (very conservative).
        pct_occ : float
            Percentage of total occupation number. Default is None. If present, overrides `thresh`.
        nvir_act : int
            Number of virtual NOs to keep. Default is None. If present, overrides `thresh` and `pct_occ`.
    """
    from pyscf import mp
    pt = mp.MP2(mf).set(verbose=0).run()
    frozen, no_coeff = pt.make_fno(thresh=thresh, pct_occ=pct_occ, nvir_act=nvir_act)
    pt_no = mp.MP2(mf, frozen=frozen, mo_coeff=no_coeff).set(verbose=0).run()
    mycc = ccsd.CCSD(mf, frozen=frozen, mo_coeff=no_coeff)
    mycc.delta_emp2 = pt.e_corr - pt_no.e_corr
    from pyscf.lib import logger
    def _finalize(self):
        '''Hook for dumping results and clearing up the object.'''
        if self.converged:
            logger.info(self, 'FNO-%s converged', self.__class__.__name__)
        else:
            logger.note(self, 'FNO-%s not converged', self.__class__.__name__)
        logger.note(self, 'E(FNO-%s) = %.16g  E_corr = %.16g',
                    self.__class__.__name__, self.e_tot, self.e_corr)
        logger.note(self, 'E(FNO-%s+delta-MP2) = %.16g  E_corr = %.16g',
                    self.__class__.__name__, self.e_tot+self.delta_emp2,
                    self.e_corr+self.delta_emp2)
        return self
    mycc._finalize = _finalize.__get__(mycc, mycc.__class__)
    return mycc


# pyscf/cc/_ccsd.py

#!/usr/bin/env python
# Copyright 2014-2020 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.

import numpy
from pyscf import lib

libcc = lib.load_library('libcc')  # noqa

# t2 + numpy.einsum('ia,jb->ijab', t1a, t1b)
def make_tau(t2, t1a, t1b, fac=1, out=None):
    if out is None:
        out = numpy.empty_like(t2)
    out = numpy.einsum('ia,jb->ijab', t1a*fac, t1b, out=out)
    out += t2
    return out



# pyscf/cc/addons.py

#!/usr/bin/env python
# Copyright 2014-2018 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.

import numpy
from pyscf import lib

def spatial2spin(tx, orbspin=None):
    '''Convert T1/T2 of spatial orbital representation to T1/T2 of
    spin-orbital representation
    '''
    if isinstance(tx, numpy.ndarray) and tx.ndim == 2:
        # RCCSD t1 amplitudes
        return spatial2spin((tx,tx), orbspin)
    elif isinstance(tx, numpy.ndarray) and tx.ndim == 4:
        # RCCSD t2 amplitudes
        t2aa = tx - tx.transpose(0,1,3,2)
        return spatial2spin((t2aa,tx,t2aa), orbspin)
    elif len(tx) == 2:  # t1
        t1a, t1b = tx
        nocc_a, nvir_a = t1a.shape
        nocc_b, nvir_b = t1b.shape
    else:
        t2aa, t2ab, t2bb = tx
        nocc_a, nocc_b, nvir_a, nvir_b = t2ab.shape

    if orbspin is None:
        orbspin = numpy.zeros((nocc_a+nvir_a)*2, dtype=int)
        orbspin[1::2] = 1

    nocc = nocc_a + nocc_b
    nvir = nvir_a + nvir_b
    idxoa = numpy.where(orbspin[:nocc] == 0)[0]
    idxob = numpy.where(orbspin[:nocc] == 1)[0]
    idxva = numpy.where(orbspin[nocc:] == 0)[0]
    idxvb = numpy.where(orbspin[nocc:] == 1)[0]

    if len(tx) == 2:  # t1
        t1 = numpy.zeros((nocc,nvir), dtype=t1a.dtype)
        lib.takebak_2d(t1, t1a, idxoa, idxva)
        lib.takebak_2d(t1, t1b, idxob, idxvb)
        t1 = lib.tag_array(t1, orbspin=orbspin)
        return t1

    else:
        t2 = numpy.zeros((nocc**2,nvir**2), dtype=t2aa.dtype)
        idxoaa = idxoa[:,None] * nocc + idxoa
        idxoab = idxoa[:,None] * nocc + idxob
        idxoba = idxob[:,None] * nocc + idxoa
        idxobb = idxob[:,None] * nocc + idxob
        idxvaa = idxva[:,None] * nvir + idxva
        idxvab = idxva[:,None] * nvir + idxvb
        idxvba = idxvb[:,None] * nvir + idxva
        idxvbb = idxvb[:,None] * nvir + idxvb
        t2aa = t2aa.reshape(nocc_a*nocc_a,nvir_a*nvir_a)
        t2ab = t2ab.reshape(nocc_a*nocc_b,nvir_a*nvir_b)
        t2bb = t2bb.reshape(nocc_b*nocc_b,nvir_b*nvir_b)
        lib.takebak_2d(t2, t2aa, idxoaa.ravel()  , idxvaa.ravel()  )
        lib.takebak_2d(t2, t2bb, idxobb.ravel()  , idxvbb.ravel()  )
        lib.takebak_2d(t2, t2ab, idxoab.ravel()  , idxvab.ravel()  )
        lib.takebak_2d(t2, t2ab, idxoba.T.ravel(), idxvba.T.ravel())
        abba = -t2ab
        lib.takebak_2d(t2, abba, idxoab.ravel()  , idxvba.T.ravel())
        lib.takebak_2d(t2, abba, idxoba.T.ravel(), idxvab.ravel()  )
        t2 = lib.tag_array(t2, orbspin=orbspin)
        return t2.reshape(nocc,nocc,nvir,nvir)

spatial2spinorb = spatial2spin

def spin2spatial(tx, orbspin):
    if tx.ndim == 2:  # t1
        nocc, nvir = tx.shape
    else:
        nocc, nvir = tx.shape[1:3]

    idxoa = numpy.where(orbspin[:nocc] == 0)[0]
    idxob = numpy.where(orbspin[:nocc] == 1)[0]
    idxva = numpy.where(orbspin[nocc:] == 0)[0]
    idxvb = numpy.where(orbspin[nocc:] == 1)[0]
    nocc_a = len(idxoa)
    nocc_b = len(idxob)
    nvir_a = len(idxva)
    nvir_b = len(idxvb)

    if tx.ndim == 2:  # t1
        t1a = lib.take_2d(tx, idxoa, idxva)
        t1b = lib.take_2d(tx, idxob, idxvb)
        return t1a, t1b
    else:
        idxoaa = idxoa[:,None] * nocc + idxoa
        idxoab = idxoa[:,None] * nocc + idxob
        idxobb = idxob[:,None] * nocc + idxob
        idxvaa = idxva[:,None] * nvir + idxva
        idxvab = idxva[:,None] * nvir + idxvb
        idxvbb = idxvb[:,None] * nvir + idxvb
        t2 = tx.reshape(nocc**2,nvir**2)
        t2aa = lib.take_2d(t2, idxoaa.ravel(), idxvaa.ravel())
        t2bb = lib.take_2d(t2, idxobb.ravel(), idxvbb.ravel())
        t2ab = lib.take_2d(t2, idxoab.ravel(), idxvab.ravel())
        t2aa = t2aa.reshape(nocc_a,nocc_a,nvir_a,nvir_a)
        t2bb = t2bb.reshape(nocc_b,nocc_b,nvir_b,nvir_b)
        t2ab = t2ab.reshape(nocc_a,nocc_b,nvir_a,nvir_b)
        return t2aa,t2ab,t2bb

def convert_to_uccsd(mycc):
    from pyscf import scf
    from pyscf.cc import uccsd, gccsd
    if isinstance(mycc, uccsd.UCCSD):
        return mycc
    elif isinstance(mycc, gccsd.GCCSD):
        raise NotImplementedError

    mf = scf.addons.convert_to_uhf(mycc._scf)
    ucc = uccsd.UCCSD(mf)
    assert(mycc._nocc is None)
    assert(mycc._nmo is None)
    ucc.__dict__.update(mycc.__dict__)
    ucc._scf = mf
    ucc.mo_coeff = mf.mo_coeff
    ucc.mo_occ = mf.mo_occ
    if not (mycc.frozen is None or isinstance(mycc.frozen, (int, numpy.integer))):
        raise NotImplementedError
    ucc.t1, ucc.t2 = uccsd.amplitudes_from_rccsd(mycc.t1, mycc.t2)
    return ucc

def convert_to_gccsd(mycc):
    from pyscf import scf
    from pyscf.cc import gccsd
    if isinstance(mycc, gccsd.GCCSD):
        return mycc

    mf = scf.addons.convert_to_ghf(mycc._scf)
    gcc = gccsd.GCCSD(mf)
    assert(mycc._nocc is None)
    assert(mycc._nmo is None)
    gcc.__dict__.update(mycc.__dict__)
    gcc._scf = mf
    gcc.mo_coeff = mf.mo_coeff
    gcc.mo_occ = mf.mo_occ
    if isinstance(mycc.frozen, (int, numpy.integer)):
        gcc.frozen = mycc.frozen * 2
    elif not (mycc.frozen is None or mycc.frozen == 0):
        raise NotImplementedError
    gcc.t1 = spatial2spin(mycc.t1, mf.mo_coeff.orbspin)
    gcc.t2 = spatial2spin(mycc.t2, mf.mo_coeff.orbspin)
    return gcc


# pyscf/cc/ccd.py

#!/usr/bin/env python
# Copyright 2014-2020 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: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Coupled cluster doubles
'''

import numpy
from pyscf.cc import ccsd
from pyscf.cc import ccsd_lambda
from pyscf.cc import ccsd_rdm

class CCD(ccsd.CCSD):
    def update_amps(self, t1, t2, eris):
        t1, t2 = ccsd.update_amps(self, t1, t2, eris)
        return numpy.zeros_like(t1), t2

    def kernel(self, t2=None, eris=None):
        nocc = self.nocc
        nvir = self.nmo - nocc
        t1 = numpy.zeros((nocc, nvir))
        ccsd.CCSD.kernel(self, t1, t2, eris)
        return self.e_corr, self.t2

    def solve_lambda(self, t2=None, l2=None, eris=None):
        if t2 is None: t2 = self.t2
        if eris is None: eris = self.ao2mo(self.mo_coeff)

        nocc = self.nocc
        nvir = self.nmo - nocc
        l1 = t1 = numpy.zeros((nocc, nvir))

        def update_lambda(mycc, t1, t2, l1, l2, eris=None, imds=None):
            l1, l2 = ccsd_lambda.update_lambda(mycc, t1, t2, l1, l2, eris, imds)
            return numpy.zeros_like(l1), l2

        self.converged_lambda, self.l1, self.l2 = \
                ccsd_lambda.kernel(self, eris, t1, t2, l1, l2,
                                   max_cycle=self.max_cycle,
                                   tol=self.conv_tol_normt,
                                   verbose=self.verbose, fupdate=update_lambda)
        return self.l2

    def make_rdm1(self, t2=None, l2=None, ao_repr=False):
        '''Un-relaxed 1-particle density matrix in MO space'''
        if t2 is None: t2 = self.t2
        if l2 is None: l2 = self.l2
        if l2 is None: l2 = self.solve_lambda(t2)

        nocc = self.nocc
        nvir = self.nmo - nocc
        l1 = t1 = numpy.zeros((nocc, nvir))

        return ccsd_rdm.make_rdm1(self, t1, t2, l1, l2, ao_repr=ao_repr)

    def make_rdm2(self, t1=None, t2=None, l1=None, l2=None, ao_repr=False):
        '''2-particle density matrix in MO space.  The density matrix is
        stored as

        dm2[p,r,q,s] = <p^+ q^+ s r>
        '''
        if t2 is None: t2 = self.t2
        if l2 is None: l2 = self.l2
        if l2 is None: l2 = self.solve_lambda(t2)

        nocc = self.nocc
        nvir = self.nmo - nocc
        l1 = t1 = numpy.zeros((nocc, nvir))

        return ccsd_rdm.make_rdm2(self, t1, t2, l1, l2, ao_repr=ao_repr)


# pyscf/cc/ccsd.py

#!/usr/bin/env python
# Copyright 2014-2021 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: Qiming Sun <osirpt.sun@gmail.com>
#

'''
RCCSD for real integrals
8-fold permutation symmetry has been used
(ij|kl) = (ji|kl) = (kl|ij) = ...
'''


import ctypes
from functools import reduce
import numpy
from pyscf import gto
from pyscf import lib
from pyscf.lib import logger
from pyscf import ao2mo
from pyscf.ao2mo import _ao2mo
from pyscf.cc import _ccsd
from pyscf.mp.mp2 import get_nocc, get_nmo, get_frozen_mask, _mo_without_core
from pyscf import __config__

BLKMIN = getattr(__config__, 'cc_ccsd_blkmin', 4)
MEMORYMIN = getattr(__config__, 'cc_ccsd_memorymin', 2000)


# t1: ia
# t2: ijab
def kernel(mycc, eris=None, t1=None, t2=None, max_cycle=50, tol=1e-8,
           tolnormt=1e-6, verbose=None, callback=None):
    log = logger.new_logger(mycc, verbose)
    if eris is None:
        eris = mycc.ao2mo(mycc.mo_coeff)
    if t1 is None and t2 is None:
        t1, t2 = mycc.get_init_guess(eris)
    elif t2 is None:
        t2 = mycc.get_init_guess(eris)[1]

    cput1 = cput0 = (logger.process_clock(), logger.perf_counter())
    eold = 0
    eccsd = mycc.energy(t1, t2, eris)
    log.info('Init E_corr(CCSD) = %.15g', eccsd)

    if isinstance(mycc.diis, lib.diis.DIIS):
        adiis = mycc.diis
    elif mycc.diis:
        adiis = lib.diis.DIIS(mycc, mycc.diis_file, incore=mycc.incore_complete)
        adiis.space = mycc.diis_space
    else:
        adiis = None

    conv = False
    for istep in range(max_cycle):
        t1new, t2new = mycc.update_amps(t1, t2, eris)
        if callback is not None:
            callback(locals())
        tmpvec = mycc.amplitudes_to_vector(t1new, t2new)
        tmpvec -= mycc.amplitudes_to_vector(t1, t2)
        normt = numpy.linalg.norm(tmpvec)
        tmpvec = None
        if mycc.iterative_damping < 1.0:
            alpha = mycc.iterative_damping
            t1new = (1-alpha) * t1 + alpha * t1new
            t2new *= alpha
            t2new += (1-alpha) * t2
        t1, t2 = t1new, t2new
        t1new = t2new = None
        t1, t2 = mycc.run_diis(t1, t2, istep, normt, eccsd-eold, adiis)
        eold, eccsd = eccsd, mycc.energy(t1, t2, eris)
        log.info('cycle = %d  E_corr(CCSD) = %.15g  dE = %.9g  norm(t1,t2) = %.6g',
                 istep+1, eccsd, eccsd - eold, normt)
        cput1 = log.timer('CCSD iter', *cput1)
        if abs(eccsd-eold) < tol and normt < tolnormt:
            conv = True
            break
    log.timer('CCSD', *cput0)
    return conv, eccsd, t1, t2


def update_amps(mycc, t1, t2, eris):
    if mycc.cc2:
        raise NotImplementedError
    assert(isinstance(eris, _ChemistsERIs))

    time0 = logger.process_clock(), logger.perf_counter()
    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocc, nvir = t1.shape
    fock = eris.fock
    mo_e_o = eris.mo_energy[:nocc]
    mo_e_v = eris.mo_energy[nocc:] + mycc.level_shift

    t1new = numpy.zeros_like(t1)
    t2new = mycc._add_vvvv(t1, t2, eris, t2sym='jiba')
    t2new *= .5  # *.5 because t2+t2.transpose(1,0,3,2) in the end
    time1 = log.timer_debug1('vvvv', *time0)

#** make_inter_F
    fov = fock[:nocc,nocc:].copy()
    t1new += fov

    foo = fock[:nocc,:nocc] - numpy.diag(mo_e_o)
    foo += .5 * numpy.einsum('ia,ja->ij', fock[:nocc,nocc:], t1)

    fvv = fock[nocc:,nocc:] - numpy.diag(mo_e_v)
    fvv -= .5 * numpy.einsum('ia,ib->ab', t1, fock[:nocc,nocc:])

    if mycc.incore_complete:
        fswap = None
    else:
        fswap = lib.H5TmpFile()
    fwVOov, fwVooV = _add_ovvv_(mycc, t1, t2, eris, fvv, t1new, t2new, fswap)
    time1 = log.timer_debug1('ovvv', *time1)

    woooo = numpy.asarray(eris.oooo).transpose(0,2,1,3).copy()

    unit = nocc**2*nvir*7 + nocc**3 + nocc*nvir**2
    mem_now = lib.current_memory()[0]
    max_memory = max(0, mycc.max_memory - mem_now)
    blksize = min(nvir, max(BLKMIN, int((max_memory*.9e6/8-nocc**4)/unit)))
    log.debug1('max_memory %d MB,  nocc,nvir = %d,%d  blksize = %d',
               max_memory, nocc, nvir, blksize)

    for p0, p1 in lib.prange(0, nvir, blksize):
        wVOov = fwVOov[p0:p1]
        wVooV = fwVooV[p0:p1]
        eris_ovoo = eris.ovoo[:,p0:p1]
        eris_oovv = numpy.empty((nocc,nocc,p1-p0,nvir))
        def load_oovv(p0, p1):
            eris_oovv[:] = eris.oovv[:,:,p0:p1]
        with lib.call_in_background(load_oovv, sync=not mycc.async_io) as prefetch_oovv:
            #:eris_oovv = eris.oovv[:,:,p0:p1]
            prefetch_oovv(p0, p1)
            foo += numpy.einsum('kc,kcji->ij', 2*t1[:,p0:p1], eris_ovoo)
            foo += numpy.einsum('kc,icjk->ij',  -t1[:,p0:p1], eris_ovoo)
            tmp = lib.einsum('la,jaik->lkji', t1[:,p0:p1], eris_ovoo)
            woooo += tmp + tmp.transpose(1,0,3,2)
            tmp = None

            wVOov -= lib.einsum('jbik,ka->bjia', eris_ovoo, t1)
            t2new[:,:,p0:p1] += wVOov.transpose(1,2,0,3)

            wVooV += lib.einsum('kbij,ka->bija', eris_ovoo, t1)
            eris_ovoo = None
        load_oovv = prefetch_oovv = None

        eris_ovvo = numpy.empty((nocc,p1-p0,nvir,nocc))
        def load_ovvo(p0, p1):
            eris_ovvo[:] = eris.ovvo[:,p0:p1]
        with lib.call_in_background(load_ovvo, sync=not mycc.async_io) as prefetch_ovvo:
            #:eris_ovvo = eris.ovvo[:,p0:p1]
            prefetch_ovvo(p0, p1)
            t1new[:,p0:p1] -= numpy.einsum('jb,jiab->ia', t1, eris_oovv)
            wVooV -= eris_oovv.transpose(2,0,1,3)
            wVOov += wVooV*.5  #: bjia + bija*.5
        load_ovvo = prefetch_ovvo = None

        t2new[:,:,p0:p1] += (eris_ovvo*0.5).transpose(0,3,1,2)
        eris_voov = eris_ovvo.conj().transpose(1,0,3,2)
        t1new[:,p0:p1] += 2*numpy.einsum('jb,aijb->ia', t1, eris_voov)
        eris_ovvo = None

        tmp  = lib.einsum('ic,kjbc->ibkj', t1, eris_oovv)
        tmp += lib.einsum('bjkc,ic->jbki', eris_voov, t1)
        t2new[:,:,p0:p1] -= lib.einsum('ka,jbki->jiba', t1, tmp)
        eris_oovv = tmp = None

        fov[:,p0:p1] += numpy.einsum('kc,aikc->ia', t1, eris_voov) * 2
        fov[:,p0:p1] -= numpy.einsum('kc,akic->ia', t1, eris_voov)

        tau  = numpy.einsum('ia,jb->ijab', t1[:,p0:p1]*.5, t1)
        tau += t2[:,:,p0:p1]
        theta  = tau.transpose(1,0,2,3) * 2
        theta -= tau
        fvv -= lib.einsum('cjia,cjib->ab', theta.transpose(2,1,0,3), eris_voov)
        foo += lib.einsum('aikb,kjab->ij', eris_voov, theta)
        tau = theta = None

        tau = t2[:,:,p0:p1] + numpy.einsum('ia,jb->ijab', t1[:,p0:p1], t1)
        woooo += lib.einsum('ijab,aklb->ijkl', tau, eris_voov)
        tau = None

        def update_wVooV(q0, q1, tau):
            wVooV[:] += lib.einsum('bkic,jkca->bija', eris_voov[:,:,:,q0:q1], tau)
        with lib.call_in_background(update_wVooV, sync=not mycc.async_io) as update_wVooV:
            for q0, q1 in lib.prange(0, nvir, blksize):
                tau  = t2[:,:,q0:q1] * .5
                tau += numpy.einsum('ia,jb->ijab', t1[:,q0:q1], t1)
                #:wVooV += lib.einsum('bkic,jkca->bija', eris_voov[:,:,:,q0:q1], tau)
                update_wVooV(q0, q1, tau)
        tau = update_wVooV = None
        def update_t2(q0, q1, tmp):
            t2new[:,:,q0:q1] += tmp.transpose(2,0,1,3)
            tmp *= .5
            t2new[:,:,q0:q1] += tmp.transpose(0,2,1,3)
        with lib.call_in_background(update_t2, sync=not mycc.async_io) as update_t2:
            for q0, q1 in lib.prange(0, nvir, blksize):
                tmp = lib.einsum('jkca,ckib->jaib', t2[:,:,p0:p1,q0:q1], wVooV)
                #:t2new[:,:,q0:q1] += tmp.transpose(2,0,1,3)
                #:tmp *= .5
                #:t2new[:,:,q0:q1] += tmp.transpose(0,2,1,3)
                update_t2(q0, q1, tmp)
                tmp = None

        wVOov += eris_voov
        eris_VOov = -.5 * eris_voov.transpose(0,2,1,3)
        eris_VOov += eris_voov
        eris_voov = None
        def update_wVOov(q0, q1, tau):
            wVOov[:,:,:,q0:q1] += .5 * lib.einsum('aikc,kcjb->aijb', eris_VOov, tau)
        with lib.call_in_background(update_wVOov, sync=not mycc.async_io) as update_wVOov:
            for q0, q1 in lib.prange(0, nvir, blksize):
                tau  = t2[:,:,q0:q1].transpose(1,3,0,2) * 2
                tau -= t2[:,:,q0:q1].transpose(0,3,1,2)
                tau -= numpy.einsum('ia,jb->ibja', t1[:,q0:q1]*2, t1)
                #:wVOov[:,:,:,q0:q1] += .5 * lib.einsum('aikc,kcjb->aijb', eris_VOov, tau)
                update_wVOov(q0, q1, tau)
                tau = None
        def update_t2(q0, q1, theta):
            t2new[:,:,q0:q1] += lib.einsum('kica,ckjb->ijab', theta, wVOov)
        with lib.call_in_background(update_t2, sync=not mycc.async_io) as update_t2:
            for q0, q1 in lib.prange(0, nvir, blksize):
                theta  = t2[:,:,p0:p1,q0:q1] * 2
                theta -= t2[:,:,p0:p1,q0:q1].transpose(1,0,2,3)
                #:t2new[:,:,q0:q1] += lib.einsum('kica,ckjb->ijab', theta, wVOov)
                update_t2(q0, q1, theta)
                theta = None
        eris_VOov = wVOov = wVooV = update_wVOov = None
        time1 = log.timer_debug1('voov [%d:%d]'%(p0, p1), *time1)
    fwVOov = fwVooV = fswap = None

    for p0, p1 in lib.prange(0, nvir, blksize):
        theta = t2[:,:,p0:p1].transpose(1,0,2,3) * 2 - t2[:,:,p0:p1]
        t1new += numpy.einsum('jb,ijba->ia', fov[:,p0:p1], theta)
        t1new -= lib.einsum('jbki,kjba->ia', eris.ovoo[:,p0:p1], theta)

        tau = numpy.einsum('ia,jb->ijab', t1[:,p0:p1], t1)
        tau += t2[:,:,p0:p1]
        t2new[:,:,p0:p1] += .5 * lib.einsum('ijkl,klab->ijab', woooo, tau)
        theta = tau = None

    ft_ij = foo + numpy.einsum('ja,ia->ij', .5*t1, fov)
    ft_ab = fvv - numpy.einsum('ia,ib->ab', .5*t1, fov)
    t2new += lib.einsum('ijac,bc->ijab', t2, ft_ab)
    t2new -= lib.einsum('ki,kjab->ijab', ft_ij, t2)

    eia = mo_e_o[:,None] - mo_e_v
    t1new += numpy.einsum('ib,ab->ia', t1, fvv)
    t1new -= numpy.einsum('ja,ji->ia', t1, foo)
    t1new /= eia

    #: t2new = t2new + t2new.transpose(1,0,3,2)
    for i in range(nocc):
        if i > 0:
            t2new[i,:i] += t2new[:i,i].transpose(0,2,1)
            t2new[i,:i] /= lib.direct_sum('a,jb->jab', eia[i], eia[:i])
            t2new[:i,i] = t2new[i,:i].transpose(0,2,1)
        t2new[i,i] = t2new[i,i] + t2new[i,i].T
        t2new[i,i] /= lib.direct_sum('a,b->ab', eia[i], eia[i])

    time0 = log.timer_debug1('update t1 t2', *time0)
    return t1new, t2new


def _add_ovvv_(mycc, t1, t2, eris, fvv, t1new, t2new, fswap):
    time1 = logger.process_clock(), logger.perf_counter()
    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocc, nvir = t1.shape
    nvir_pair = nvir * (nvir+1) // 2

    if fswap is None:
        wVOov = numpy.zeros((nvir,nocc,nocc,nvir))
    else:
        wVOov = fswap.create_dataset('wVOov', (nvir,nocc,nocc,nvir), 'f8')
    wooVV = numpy.zeros((nocc,nocc*nvir_pair))

    max_memory = mycc.max_memory - lib.current_memory()[0]
    unit = nocc*nvir**2*3 + nocc**2*nvir + 2
    blksize = min(nvir, max(BLKMIN, int((max_memory*.95e6/8-wooVV.size)/unit)))
    if not mycc.direct:
        unit = nocc*nvir**2*3 + nocc**2*nvir + 2 + nocc*nvir**2 + nocc*nvir
        blksize = min(nvir, max(BLKMIN, int((max_memory*.95e6/8-wooVV.size-nocc**2*nvir)/unit)))
    log.debug1('max_memory %d MB,  nocc,nvir = %d,%d  blksize = %d',
               max_memory, nocc, nvir, blksize)

    def load_ovvv(buf, p0):
        if p0 < nvir:
            p1 = min(nvir, p0+blksize)
            buf[:p1-p0] = eris.ovvv[:,p0:p1].transpose(1,0,2)

    with lib.call_in_background(load_ovvv, sync=not mycc.async_io) as prefetch:
        buf = numpy.empty((blksize,nocc,nvir_pair))
        buf_prefetch = numpy.empty((blksize,nocc,nvir_pair))

        load_ovvv(buf_prefetch, 0)
        for p0, p1 in lib.prange(0, nvir, blksize):
            buf, buf_prefetch = buf_prefetch, buf
            prefetch(buf_prefetch, p1)

            eris_vovv = buf[:p1-p0]

            #:wooVV -= numpy.einsum('jc,ciba->jiba', t1[:,p0:p1], eris_vovv)
            lib.ddot(numpy.asarray(t1[:,p0:p1], order='C'),
                     eris_vovv.reshape(p1-p0,-1), -1, wooVV, 1)

            eris_vovv = lib.unpack_tril(eris_vovv.reshape((p1-p0)*nocc,nvir_pair))
            eris_vovv = eris_vovv.reshape(p1-p0,nocc,nvir,nvir)

            fvv += 2*numpy.einsum('kc,ckab->ab', t1[:,p0:p1], eris_vovv)
            fvv[:,p0:p1] -= numpy.einsum('kc,bkca->ab', t1, eris_vovv)

            if not mycc.direct:
                vvvo = eris_vovv.transpose(0,2,3,1).copy()
                for i in range(nocc):
                    tau = t2[i,:,p0:p1] + numpy.einsum('a,jb->jab', t1[i,p0:p1], t1)
                    tmp = lib.einsum('jcd,cdbk->jbk', tau, vvvo)
                    t2new[i] -= lib.einsum('ka,jbk->jab', t1, tmp)
                    tau = tmp = None

            wVOov[p0:p1] = lib.einsum('biac,jc->bija', eris_vovv, t1)

            theta = t2[:,:,p0:p1].transpose(1,2,0,3) * 2
            theta -= t2[:,:,p0:p1].transpose(0,2,1,3)
            t1new += lib.einsum('icjb,cjba->ia', theta, eris_vovv)
            theta = None
            time1 = log.timer_debug1('vovv [%d:%d]'%(p0, p1), *time1)

    if fswap is None:
        wooVV = lib.unpack_tril(wooVV.reshape(nocc**2,nvir_pair))
        return wVOov, wooVV.reshape(nocc,nocc,nvir,nvir).transpose(2,1,0,3)
    else:
        fswap.create_dataset('wVooV', (nvir,nocc,nocc,nvir), 'f8')
        wooVV = wooVV.reshape(nocc,nocc,nvir_pair)
        tril2sq = lib.square_mat_in_trilu_indices(nvir)
        for p0, p1 in lib.prange(0, nvir, blksize):
            fswap['wVooV'][p0:p1] = wooVV[:,:,tril2sq[p0:p1]].transpose(2,1,0,3)
        return fswap['wVOov'], fswap['wVooV']

def _add_vvvv(mycc, t1, t2, eris, out=None, with_ovvv=None, t2sym=None):
    '''t2sym: whether t2 has the symmetry t2[ijab]==t2[jiba] or
    t2[ijab]==-t2[jiab] or t2[ijab]==-t2[jiba]
    '''
    #TODO: Guess the symmetry of t2 amplitudes
    #if t2sym is None:
    #    if t2.shape[0] != t2.shape[1]:
    #        t2sym = ''
    #    elif abs(t2-t2.transpose(1,0,3,2)).max() < 1e-12:
    #        t2sym = 'jiba'
    #    elif abs(t2+t2.transpose(1,0,2,3)).max() < 1e-12:
    #        t2sym = '-jiab'
    #    elif abs(t2+t2.transpose(1,0,3,2)).max() < 1e-12:
    #        t2sym = '-jiba'

    if t2sym in ('jiba', '-jiba', '-jiab'):
        Ht2tril = _add_vvvv_tril(mycc, t1, t2, eris, with_ovvv=with_ovvv)
        nocc, nvir = t2.shape[1:3]
        Ht2 = _unpack_t2_tril(Ht2tril, nocc, nvir, out, t2sym)
    else:
        Ht2 = _add_vvvv_full(mycc, t1, t2, eris, out, with_ovvv)
    return Ht2

def _add_vvvv_tril(mycc, t1, t2, eris, out=None, with_ovvv=None):
    '''Ht2 = numpy.einsum('ijcd,acdb->ijab', t2, vvvv)
    Using symmetry t2[ijab] = t2[jiba] and Ht2[ijab] = Ht2[jiba], compute the
    lower triangular part of  Ht2
    '''
    time0 = logger.process_clock(), logger.perf_counter()
    log = logger.Logger(mycc.stdout, mycc.verbose)
    if with_ovvv is None:
        with_ovvv = mycc.direct
    nocc, nvir = t2.shape[1:3]
    nocc2 = nocc*(nocc+1)//2
    if t1 is None:
        tau = t2[numpy.tril_indices(nocc)]
    else:
        tau = numpy.empty((nocc2,nvir,nvir), dtype=t2.dtype)
        p1 = 0
        for i in range(nocc):
            p0, p1 = p1, p1 + i+1
            tau[p0:p1] = numpy.einsum('a,jb->jab', t1[i], t1[:i+1])
            tau[p0:p1] += t2[i,:i+1]

    if mycc.direct:   # AO-direct CCSD
        mo = getattr(eris, 'mo_coeff', None)
        if mo is None:  # If eris does not have the attribute mo_coeff
            mo = _mo_without_core(mycc, mycc.mo_coeff)
        nao, nmo = mo.shape
        aos = numpy.asarray(mo[:,nocc:].T, order='F')
        tau = _ao2mo.nr_e2(tau.reshape(nocc2,nvir**2), aos, (0,nao,0,nao), 's1', 's1')
        tau = tau.reshape(nocc2,nao,nao)
        time0 = log.timer_debug1('vvvv-tau', *time0)

        buf = eris._contract_vvvv_t2(mycc, tau, mycc.direct, out, log)
        buf = buf.reshape(nocc2,nao,nao)
        Ht2tril = _ao2mo.nr_e2(buf, mo.conj(), (nocc,nmo,nocc,nmo), 's1', 's1')
        Ht2tril = Ht2tril.reshape(nocc2,nvir,nvir)

        if with_ovvv:
            #: tmp = numpy.einsum('ijcd,ka,kdcb->ijba', tau, t1, eris.ovvv)
            #: t2new -= tmp + tmp.transpose(1,0,3,2)
            tmp = _ao2mo.nr_e2(buf, mo.conj(), (nocc,nmo,0,nocc), 's1', 's1')
            Ht2tril -= lib.ddot(tmp.reshape(nocc2*nvir,nocc), t1).reshape(nocc2,nvir,nvir)
            tmp = _ao2mo.nr_e2(buf, mo.conj(), (0,nocc,nocc,nmo), 's1', 's1')
            #: Ht2tril -= numpy.einsum('xkb,ka->xab', tmp.reshape(-1,nocc,nvir), t1)
            tmp = lib.transpose(tmp.reshape(nocc2,nocc,nvir), axes=(0,2,1), out=buf)
            tmp = lib.ddot(tmp.reshape(nocc2*nvir,nocc), t1, 1,
                           numpy.ndarray((nocc2*nvir,nvir), buffer=tau), 0)
            tmp = lib.transpose(tmp.reshape(nocc2,nvir,nvir), axes=(0,2,1), out=buf)
            Ht2tril -= tmp.reshape(nocc2,nvir,nvir)
    else:
        assert(not with_ovvv)
        Ht2tril = eris._contract_vvvv_t2(mycc, tau, mycc.direct, out, log)
    return Ht2tril

def _add_vvvv_full(mycc, t1, t2, eris, out=None, with_ovvv=False):
    '''Ht2 = numpy.einsum('ijcd,acdb->ijab', t2, vvvv)
    without using symmetry t2[ijab] = t2[jiba] in t2 or Ht2
    '''
    time0 = logger.process_clock(), logger.perf_counter()
    log = logger.Logger(mycc.stdout, mycc.verbose)
    if t1 is None:
        tau = t2
    else:
        tau = numpy.einsum('ia,jb->ijab', t1, t1)
        tau += t2

    if mycc.direct:   # AO-direct CCSD
        if with_ovvv:
            raise NotImplementedError
        mo = getattr(eris, 'mo_coeff', None)
        if mo is None:  # If eris does not have the attribute mo_coeff
            mo = _mo_without_core(mycc, mycc.mo_coeff)
        nocc, nvir = t2.shape[1:3]
        nao, nmo = mo.shape
        aos = numpy.asarray(mo[:,nocc:].T, order='F')
        tau = _ao2mo.nr_e2(tau.reshape(nocc**2,nvir,nvir), aos, (0,nao,0,nao), 's1', 's1')
        tau = tau.reshape(nocc,nocc,nao,nao)
        time0 = log.timer_debug1('vvvv-tau mo2ao', *time0)

        buf = eris._contract_vvvv_t2(mycc, tau, mycc.direct, out, log)
        buf = buf.reshape(nocc**2,nao,nao)
        Ht2 = _ao2mo.nr_e2(buf, mo.conj(), (nocc,nmo,nocc,nmo), 's1', 's1')
    else:
        assert(not with_ovvv)
        Ht2 = eris._contract_vvvv_t2(mycc, tau, mycc.direct, out, log)

    return Ht2.reshape(t2.shape)


def _contract_vvvv_t2(mycc, mol, vvvv, t2, out=None, verbose=None):
    '''Ht2 = numpy.einsum('ijcd,acbd->ijab', t2, vvvv)

    Args:
        vvvv : None or integral object
            if vvvv is None, contract t2 to AO-integrals using AO-direct algorithm
    '''
    if vvvv is None or len(vvvv.shape) == 2:
        # AO-direct or vvvv in 4-fold symmetry
        return _contract_s4vvvv_t2(mycc, mol, vvvv, t2, out, verbose)
    else:
        return _contract_s1vvvv_t2(mycc, mol, vvvv, t2, out, verbose)


def _contract_s4vvvv_t2(mycc, mol, vvvv, t2, out=None, verbose=None):
    '''Ht2 = numpy.einsum('ijcd,acbd->ijab', t2, vvvv)
    where vvvv has to be real and has the 4-fold permutation symmetry

    Args:
        vvvv : None or integral object
            if vvvv is None, contract t2 to AO-integrals using AO-direct algorithm
    '''
    assert(t2.dtype == numpy.double)
    if t2.size == 0:
        return numpy.zeros_like(t2)

    _dgemm = lib.numpy_helper._dgemm
    time0 = logger.process_clock(), logger.perf_counter()
    log = logger.new_logger(mycc, verbose)

    nvira, nvirb = t2.shape[-2:]
    x2 = t2.reshape(-1,nvira,nvirb)
    nocc2 = x2.shape[0]
    nvir2 = nvira * nvirb
    Ht2 = numpy.ndarray(x2.shape, dtype=x2.dtype, buffer=out)
    Ht2[:] = 0

    def contract_blk_(eri, i0, i1, j0, j1):
        ic = i1 - i0
        jc = j1 - j0
        #:Ht2[:,j0:j1] += numpy.einsum('xef,efab->xab', x2[:,i0:i1], eri)
        _dgemm('N', 'N', nocc2, jc*nvirb, ic*nvirb,
               x2.reshape(-1,nvir2), eri.reshape(-1,jc*nvirb),
               Ht2.reshape(-1,nvir2), 1, 1, i0*nvirb, 0, j0*nvirb)

        if i0 > j0:
            #:Ht2[:,i0:i1] += numpy.einsum('xef,abef->xab', x2[:,j0:j1], eri)
            _dgemm('N', 'T', nocc2, ic*nvirb, jc*nvirb,
                   x2.reshape(-1,nvir2), eri.reshape(-1,jc*nvirb),
                   Ht2.reshape(-1,nvir2), 1, 1, j0*nvirb, 0, i0*nvirb)

    max_memory = max(MEMORYMIN, mycc.max_memory - lib.current_memory()[0])
    if vvvv is None:   # AO-direct CCSD
        ao_loc = mol.ao_loc_nr()
        assert(nvira == nvirb == ao_loc[-1])

        intor = mol._add_suffix('int2e')
        ao2mopt = _ao2mo.AO2MOpt(mol, intor, 'CVHFnr_schwarz_cond',
                                 'CVHFsetnr_direct_scf')
        blksize = max(BLKMIN, numpy.sqrt(max_memory*.9e6/8/nvirb**2/2.5))
        blksize = int(min((nvira+3)/4, blksize))
        sh_ranges = ao2mo.outcore.balance_partition(ao_loc, blksize)
        blksize = max(x[2] for x in sh_ranges)
        eribuf = numpy.empty((blksize,blksize,nvirb,nvirb))
        loadbuf = numpy.empty((blksize,blksize,nvirb,nvirb))
        fint = gto.moleintor.getints4c

        for ip, (ish0, ish1, ni) in enumerate(sh_ranges):
            for jsh0, jsh1, nj in sh_ranges[:ip]:
                eri = fint(intor, mol._atm, mol._bas, mol._env,
                           shls_slice=(ish0,ish1,jsh0,jsh1), aosym='s2kl',
                           ao_loc=ao_loc, cintopt=ao2mopt._cintopt, out=eribuf)
                i0, i1 = ao_loc[ish0], ao_loc[ish1]
                j0, j1 = ao_loc[jsh0], ao_loc[jsh1]
                tmp = numpy.ndarray((i1-i0,nvirb,j1-j0,nvirb), buffer=loadbuf)
                _ccsd.libcc.CCload_eri(tmp.ctypes.data_as(ctypes.c_void_p),
                                       eri.ctypes.data_as(ctypes.c_void_p),
                                       (ctypes.c_int*4)(i0, i1, j0, j1),
                                       ctypes.c_int(nvirb))
                contract_blk_(tmp, i0, i1, j0, j1)
                time0 = log.timer_debug1('AO-vvvv [%d:%d,%d:%d]' %
                                         (ish0,ish1,jsh0,jsh1), *time0)

            eri = fint(intor, mol._atm, mol._bas, mol._env,
                       shls_slice=(ish0,ish1,ish0,ish1), aosym='s4',
                       ao_loc=ao_loc, cintopt=ao2mopt._cintopt, out=eribuf)
            i0, i1 = ao_loc[ish0], ao_loc[ish1]
            eri = lib.unpack_tril(eri, axis=0)
            tmp = numpy.ndarray((i1-i0,nvirb,i1-i0,nvirb), buffer=loadbuf)
            _ccsd.libcc.CCload_eri(tmp.ctypes.data_as(ctypes.c_void_p),
                                   eri.ctypes.data_as(ctypes.c_void_p),
                                   (ctypes.c_int*4)(i0, i1, i0, i1),
                                   ctypes.c_int(nvirb))
            eri = None
            contract_blk_(tmp, i0, i1, i0, i1)
            time0 = log.timer_debug1('AO-vvvv [%d:%d,%d:%d]' %
                                     (ish0,ish1,ish0,ish1), *time0)

    else:
        nvir_pair = nvirb * (nvirb+1) // 2
        unit = nvira*nvir_pair*2 + nvirb**2*nvira/4 + 1

        if mycc.async_io:
            fmap = lib.map_with_prefetch
            unit += nvira*nvir_pair
        else:
            fmap = map

        blksize = numpy.sqrt(max(BLKMIN**2, max_memory*.95e6/8/unit))
        blksize = int(min((nvira+3)/4, blksize))

        def load(v_slice):
            i0, i1 = v_slice
            off0 = i0*(i0+1)//2
            off1 = i1*(i1+1)//2
            return numpy.asarray(vvvv[off0:off1], order='C')

        tril2sq = lib.square_mat_in_trilu_indices(nvira)
        loadbuf = numpy.empty((blksize,blksize,nvirb,nvirb))

        slices = [(i0, i1) for i0, i1 in lib.prange(0, nvira, blksize)]
        for istep, wwbuf in enumerate(fmap(load, lib.prange(0, nvira, blksize))):
            i0, i1 = slices[istep]
            off0 = i0*(i0+1)//2
            for j0, j1 in lib.prange(0, i1, blksize):
                eri = wwbuf[tril2sq[i0:i1,j0:j1]-off0]
                tmp = numpy.ndarray((i1-i0,nvirb,j1-j0,nvirb), buffer=loadbuf)
                _ccsd.libcc.CCload_eri(tmp.ctypes.data_as(ctypes.c_void_p),
                                       eri.ctypes.data_as(ctypes.c_void_p),
                                       (ctypes.c_int*4)(i0, i1, j0, j1),
                                       ctypes.c_int(nvirb))
                contract_blk_(tmp, i0, i1, j0, j1)
            wwbuf = None
            time0 = log.timer_debug1('vvvv [%d:%d]'%(i0,i1), *time0)
    return Ht2.reshape(t2.shape)

def _contract_s1vvvv_t2(mycc, mol, vvvv, t2, out=None, verbose=None):
    '''Ht2 = numpy.einsum('ijcd,acdb->ijab', t2, vvvv)
    where vvvv can be real or complex and no permutation symmetry is available in vvvv.

    Args:
        vvvv : None or integral object
            if vvvv is None, contract t2 to AO-integrals using AO-direct algorithm
    '''
    # vvvv == None means AO-direct CCSD. It should redirect to
    # _contract_s4vvvv_t2(mycc, mol, vvvv, t2, out, verbose)
    assert(vvvv is not None)

    time0 = logger.process_clock(), logger.perf_counter()
    log = logger.new_logger(mycc, verbose)

    nvira, nvirb = t2.shape[-2:]
    x2 = t2.reshape(-1,nvira,nvirb)
    nocc2 = x2.shape[0]
    dtype = numpy.result_type(t2, vvvv)
    Ht2 = numpy.ndarray(x2.shape, dtype=dtype, buffer=out)

    max_memory = mycc.max_memory - lib.current_memory()[0]
    unit = nvirb**2*nvira*2 + nocc2*nvirb + 1
    blksize = min(nvira, max(BLKMIN, int(max_memory*1e6/8/unit)))

    for p0,p1 in lib.prange(0, nvira, blksize):
        Ht2[:,p0:p1] = lib.einsum('xcd,acbd->xab', x2, vvvv[p0:p1])
        time0 = log.timer_debug1('vvvv [%d:%d]' % (p0,p1), *time0)
    return Ht2.reshape(t2.shape)

def _unpack_t2_tril(t2tril, nocc, nvir, out=None, t2sym='jiba'):
    t2 = numpy.ndarray((nocc,nocc,nvir,nvir), dtype=t2tril.dtype, buffer=out)
    idx,idy = numpy.tril_indices(nocc)
    if t2sym == 'jiba':
        t2[idy,idx] = t2tril.transpose(0,2,1)
        t2[idx,idy] = t2tril
    elif t2sym == '-jiba':
        t2[idy,idx] = -t2tril.transpose(0,2,1)
        t2[idx,idy] = t2tril
    elif t2sym == '-jiab':
        t2[idy,idx] =-t2tril
        t2[idx,idy] = t2tril
        t2[numpy.diag_indices(nocc)] = 0
    return t2

def _unpack_4fold(c2vec, nocc, nvir, anti_symm=True):
    t2 = numpy.zeros((nocc**2,nvir**2), dtype=c2vec.dtype)
    if nocc > 1 and nvir > 1:
        t2tril = c2vec.reshape(nocc*(nocc-1)//2,nvir*(nvir-1)//2)
        otril = numpy.tril_indices(nocc, k=-1)
        vtril = numpy.tril_indices(nvir, k=-1)
        lib.takebak_2d(t2, t2tril, otril[0]*nocc+otril[1], vtril[0]*nvir+vtril[1])
        lib.takebak_2d(t2, t2tril, otril[1]*nocc+otril[0], vtril[1]*nvir+vtril[0])
        if anti_symm:  # anti-symmetry when exchanging two particle indices
            t2tril = -t2tril
        lib.takebak_2d(t2, t2tril, otril[0]*nocc+otril[1], vtril[1]*nvir+vtril[0])
        lib.takebak_2d(t2, t2tril, otril[1]*nocc+otril[0], vtril[0]*nvir+vtril[1])
    return t2.reshape(nocc,nocc,nvir,nvir)

def amplitudes_to_vector(t1, t2, out=None):
    nocc, nvir = t1.shape
    nov = nocc * nvir
    size = nov + nov*(nov+1)//2
    vector = numpy.ndarray(size, t1.dtype, buffer=out)
    vector[:nov] = t1.ravel()
    lib.pack_tril(t2.transpose(0,2,1,3).reshape(nov,nov), out=vector[nov:])
    return vector

def vector_to_amplitudes(vector, nmo, nocc):
    nvir = nmo - nocc
    nov = nocc * nvir
    t1 = vector[:nov].copy().reshape((nocc,nvir))
    # filltriu=lib.SYMMETRIC because t2[iajb] == t2[jbia]
    t2 = lib.unpack_tril(vector[nov:], filltriu=lib.SYMMETRIC)
    t2 = t2.reshape(nocc,nvir,nocc,nvir).transpose(0,2,1,3)
    return t1, numpy.asarray(t2, order='C')

def amplitudes_to_vector_s4(t1, t2, out=None):
    nocc, nvir = t1.shape
    nov = nocc * nvir
    size = nov + nocc*(nocc-1)//2*nvir*(nvir-1)//2
    vector = numpy.ndarray(size, t1.dtype, buffer=out)
    vector[:nov] = t1.ravel()
    otril = numpy.tril_indices(nocc, k=-1)
    vtril = numpy.tril_indices(nvir, k=-1)
    lib.take_2d(t2.reshape(nocc**2,nvir**2), otril[0]*nocc+otril[1],
                vtril[0]*nvir+vtril[1], out=vector[nov:])
    return vector

def vector_to_amplitudes_s4(vector, nmo, nocc):
    nvir = nmo - nocc
    nov = nocc * nvir
    size = nov + nocc*(nocc-1)//2*nvir*(nvir-1)//2
    t1 = vector[:nov].copy().reshape(nocc,nvir)
    t2 = numpy.zeros((nocc,nocc,nvir,nvir), dtype=vector.dtype)
    t2 = _unpack_4fold(vector[nov:size], nocc, nvir)
    return t1, t2


def energy(mycc, t1=None, t2=None, eris=None):
    '''CCSD correlation energy'''
    if t1 is None: t1 = mycc.t1
    if t2 is None: t2 = mycc.t2
    if eris is None: eris = mycc.ao2mo()

    nocc, nvir = t1.shape
    fock = eris.fock
    e = numpy.einsum('ia,ia', fock[:nocc,nocc:], t1) * 2
    max_memory = mycc.max_memory - lib.current_memory()[0]
    blksize = int(min(nvir, max(BLKMIN, max_memory*.3e6/8/(nocc**2*nvir+1))))
    for p0, p1 in lib.prange(0, nvir, blksize):
        eris_ovvo = eris.ovvo[:,p0:p1]
        tau = t2[:,:,p0:p1] + numpy.einsum('ia,jb->ijab', t1[:,p0:p1], t1)
        e += 2 * numpy.einsum('ijab,iabj', tau, eris_ovvo)
        e -=     numpy.einsum('jiab,iabj', tau, eris_ovvo)
    if abs(e.imag) > 1e-4:
        logger.warn(mycc, 'Non-zero imaginary part found in CCSD energy %s', e)
    return e.real

def restore_from_diis_(mycc, diis_file, inplace=True):
    '''Reuse an existed DIIS object in the CCSD calculation.

    The CCSD amplitudes will be restored from the DIIS object to generate t1
    and t2 amplitudes. The t1/t2 amplitudes of the CCSD object will be
    overwritten by the generated t1 and t2 amplitudes. The amplitudes vector
    and error vector will be reused in the CCSD calculation.
    '''
    adiis = lib.diis.DIIS(mycc, mycc.diis_file, incore=mycc.incore_complete)
    adiis.restore(diis_file, inplace=inplace)

    ccvec = adiis.extrapolate()
    mycc.t1, mycc.t2 = mycc.vector_to_amplitudes(ccvec)
    if inplace:
        mycc.diis = adiis
    return mycc

def get_t1_diagnostic(t1):
    '''Returns the t1 amplitude norm, normalized by number of correlated electrons.'''
    nelectron = 2 * t1.shape[0]
    return numpy.sqrt(numpy.linalg.norm(t1)**2 / nelectron)

def get_d1_diagnostic(t1):
    '''D1 diagnostic given in

        Janssen, et. al Chem. Phys. Lett. 290 (1998) 423
    '''
    f = lambda x: numpy.sqrt(numpy.sort(numpy.abs(x[0])))[-1]
    d1norm_ij = f(numpy.linalg.eigh(numpy.einsum('ia,ja->ij',t1,t1)))
    d1norm_ab = f(numpy.linalg.eigh(numpy.einsum('ia,ib->ab',t1,t1)))
    d1norm = max(d1norm_ij, d1norm_ab)
    return d1norm

def get_d2_diagnostic(t2):
    '''D2 diagnostic given in

        Nielsen, et. al Chem. Phys. Lett. 310 (1999) 568

    Note: This is currently only defined in the literature for restricted
    closed-shell systems.
    '''
    f = lambda x: numpy.sqrt(numpy.sort(numpy.abs(x[0])))[-1]
    d2norm_ij = f(numpy.linalg.eigh(numpy.einsum('ikab,jkab->ij',t2,t2)))
    d2norm_ab = f(numpy.linalg.eigh(numpy.einsum('ijac,ijbc->ab',t2,t2)))
    d2norm = max(d2norm_ij, d2norm_ab)
    return d2norm

def set_frozen(mycc, method='auto', window=(-1000.0, 1000.0), is_gcc=False):
    if method == 'auto':
        from pyscf.data import elements
        mycc.frozen = elements.chemcore(mycc.mol, spinorb=is_gcc)
    elif method == 'window':
        emin, emax = window
        mo_e = numpy.asarray(mycc._scf.mo_energy)
        if mo_e.ndim == 1:
            fr1 = list(numpy.flatnonzero(mo_e < emin))
            fr2 = list(numpy.flatnonzero(mo_e > emax))
            frozen = fr1 + fr2
        elif mo_e.ndim == 2:
            fr1a = list(numpy.flatnonzero(mo_e[0] < emin))
            fr2a = list(numpy.flatnonzero(mo_e[0] > emax))
            fr1b = list(numpy.flatnonzero(mo_e[1] < emin))
            fr2b = list(numpy.flatnonzero(mo_e[1] > emax))
            frozen = [fr1a+fr2a, fr1b+fr2b]
        mycc.frozen = frozen
    return mycc


def as_scanner(cc):
    '''Generating a scanner/solver for CCSD PES.

    The returned solver is a function. This function requires one argument
    "mol" as input and returns total CCSD energy.

    The solver will automatically use the results of last calculation as the
    initial guess of the new calculation.  All parameters assigned in the
    CCSD and the underlying SCF objects (conv_tol, max_memory etc) are
    automatically applied in the solver.

    Note scanner has side effects.  It may change many underlying objects
    (_scf, with_df, with_x2c, ...) during calculation.

    Examples::

        >>> from pyscf import gto, scf, cc
        >>> mol = gto.M(atom='H 0 0 0; F 0 0 1')
        >>> cc_scanner = cc.CCSD(scf.RHF(mol)).as_scanner()
        >>> e_tot = cc_scanner(gto.M(atom='H 0 0 0; F 0 0 1.1'))
        >>> e_tot = cc_scanner(gto.M(atom='H 0 0 0; F 0 0 1.5'))
    '''
    if isinstance(cc, lib.SinglePointScanner):
        return cc

    logger.info(cc, 'Set %s as a scanner', cc.__class__)

    class CCSD_Scanner(cc.__class__, lib.SinglePointScanner):
        def __init__(self, cc):
            self.__dict__.update(cc.__dict__)
            self._scf = cc._scf.as_scanner()
        def __call__(self, mol_or_geom, **kwargs):
            if isinstance(mol_or_geom, gto.Mole):
                mol = mol_or_geom
            else:
                mol = self.mol.set_geom_(mol_or_geom, inplace=False)

            if self.t2 is not None:
                last_size = self.vector_size()
            else:
                last_size = 0

            self.reset(mol)

            mf_scanner = self._scf
            mf_scanner(mol)
            self.mo_coeff = mf_scanner.mo_coeff
            self.mo_occ = mf_scanner.mo_occ
            if last_size != self.vector_size():
                self.t1 = self.t2 = None
            self.kernel(self.t1, self.t2, **kwargs)
            return self.e_tot
    return CCSD_Scanner(cc)


class CCSD(lib.StreamObject):
    '''restricted CCSD

    Attributes:
        verbose : int
            Print level.  Default value equals to :class:`Mole.verbose`
        max_memory : float or int
            Allowed memory in MB.  Default value equals to :class:`Mole.max_memory`
        conv_tol : float
            converge threshold.  Default is 1e-7.
        conv_tol_normt : float
            converge threshold for norm(t1,t2).  Default is 1e-5.
        max_cycle : int
            max number of iterations.  Default is 50.
        diis_space : int
            DIIS space size.  Default is 6.
        diis_start_cycle : int
            The step to start DIIS.  Default is 0.
        iterative_damping : float
            The self consistent damping parameter.
        direct : bool
            AO-direct CCSD. Default is False.
        async_io : bool
            Allow for asynchronous function execution. Default is True.
        incore_complete : bool
            Avoid all I/O (also for DIIS). Default is False.
        level_shift : float
            A shift on virtual orbital energies to stablize the CCSD iteration
        frozen : int or list
            If integer is given, the inner-most orbitals are frozen from CC
            amplitudes.  Given the orbital indices (0-based) in a list, both
            occupied and virtual orbitals can be frozen in CC calculation.

            >>> mol = gto.M(atom = 'H 0 0 0; F 0 0 1.1', basis = 'ccpvdz')
            >>> mf = scf.RHF(mol).run()
            >>> # freeze 2 core orbitals
            >>> mycc = cc.CCSD(mf).set(frozen = 2).run()
            >>> # auto-generate the number of core orbitals to be frozen (1 in this case)
            >>> mycc = cc.CCSD(mf).set_frozen().run()
            >>> # freeze 2 core orbitals and 3 high lying unoccupied orbitals
            >>> mycc.set(frozen = [0,1,16,17,18]).run()

        callback : function(envs_dict) => None
            callback function takes one dict as the argument which is
            generated by the builtin function :func:`locals`, so that the
            callback function can access all local variables in the current
            environment.

    Saved results

        converged : bool
            CCSD converged or not
        e_corr : float
            CCSD correlation correction
        e_tot : float
            Total CCSD energy (HF + correlation)
        t1, t2 :
            T amplitudes t1[i,a], t2[i,j,a,b]  (i,j in occ, a,b in virt)
        l1, l2 :
            Lambda amplitudes l1[i,a], l2[i,j,a,b]  (i,j in occ, a,b in virt)
    '''

    max_cycle = getattr(__config__, 'cc_ccsd_CCSD_max_cycle', 50)
    conv_tol = getattr(__config__, 'cc_ccsd_CCSD_conv_tol', 1e-7)
    iterative_damping = getattr(__config__, 'cc_ccsd_CCSD_iterative_damping', 1.0)
    conv_tol_normt = getattr(__config__, 'cc_ccsd_CCSD_conv_tol_normt', 1e-5)

    diis = getattr(__config__, 'cc_ccsd_CCSD_diis', True)
    diis_space = getattr(__config__, 'cc_ccsd_CCSD_diis_space', 6)
    diis_file = None
    diis_start_cycle = getattr(__config__, 'cc_ccsd_CCSD_diis_start_cycle', 0)
    # FIXME: Should we avoid DIIS starting early?
    diis_start_energy_diff = getattr(__config__, 'cc_ccsd_CCSD_diis_start_energy_diff', 1e9)

    direct = getattr(__config__, 'cc_ccsd_CCSD_direct', False)
    async_io = getattr(__config__, 'cc_ccsd_CCSD_async_io', True)
    incore_complete = getattr(__config__, 'cc_ccsd_CCSD_incore_complete', False)
    cc2 = getattr(__config__, 'cc_ccsd_CCSD_cc2', False)

    def __init__(self, mf, frozen=None, mo_coeff=None, mo_occ=None):
        if isinstance(mf, gto.Mole):
            raise RuntimeError('''
You see this error message because of the API updates in pyscf v0.10.
In the new API, the first argument of CC class is HF objects.  Please see
http://sunqm.net/pyscf/code-rule.html#api-rules for the details of API conventions''')

        from pyscf.scf import hf
        if isinstance(mf, hf.KohnShamDFT):
            raise RuntimeError('CCSD Warning: The first argument mf is a DFT object. '
                               'CCSD calculation should be initialized with HF object.\n'
                               'DFT object can be converted to HF object with '
                               'the code:\n'
                               '    mf_hf = mol.HF()\n'
                               '    if getattr(mf_dft, "with_x2c", False):\n'
                               '        mf_hf = mf_hf.x2c()\n'
                               '    mf_hf.__dict__.update(mf_dft.__dict__)\n')

        if mo_coeff is None: mo_coeff = mf.mo_coeff
        if mo_occ is None: mo_occ = mf.mo_occ

        self.mol = mf.mol
        self._scf = mf
        self.verbose = self.mol.verbose
        self.stdout = self.mol.stdout
        self.max_memory = mf.max_memory

        self.frozen = frozen
        self.incore_complete = self.incore_complete or self.mol.incore_anyway
        self.level_shift = 0

##################################################
# don't modify the following attributes, they are not input options
        self.mo_coeff = mo_coeff
        self.mo_occ = mo_occ
        self.converged = False
        self.converged_lambda = False
        self.emp2 = None
        self.e_hf = None
        self.e_corr = None
        self.t1 = None
        self.t2 = None
        self.l1 = None
        self.l2 = None
        self._nocc = None
        self._nmo = None
        self.chkfile = mf.chkfile
        self.callback = None

        keys = set(('max_cycle', 'conv_tol', 'iterative_damping',
                    'conv_tol_normt', 'diis', 'diis_space', 'diis_file',
                    'diis_start_cycle', 'diis_start_energy_diff', 'direct',
                    'async_io', 'incore_complete', 'cc2'))
        self._keys = set(self.__dict__.keys()).union(keys)

    @property
    def ecc(self):
        return self.e_corr

    @property
    def e_tot(self):
        return (self.e_hf or self._scf.e_tot) + self.e_corr

    @property
    def nocc(self):
        return self.get_nocc()
    @nocc.setter
    def nocc(self, n):
        self._nocc = n

    @property
    def nmo(self):
        return self.get_nmo()
    @nmo.setter
    def nmo(self, n):
        self._nmo = n

    def reset(self, mol=None):
        if mol is not None:
            self.mol = mol
        self._scf.reset(mol)
        return self

    get_nocc = get_nocc
    get_nmo = get_nmo
    get_frozen_mask = get_frozen_mask

    def set_frozen(self, method='auto', window=(-1000.0, 1000.0)):
        from pyscf import cc
        is_gcc = isinstance(self, cc.gccsd.GCCSD)
        return set_frozen(self, method=method, window=window, is_gcc=is_gcc)

    def dump_flags(self, verbose=None):
        log = logger.new_logger(self, verbose)
        log.info('')
        log.info('******** %s ********', self.__class__)
        log.info('CC2 = %g', self.cc2)
        log.info('CCSD nocc = %s, nmo = %s', self.nocc, self.nmo)
        if self.frozen is not None:
            log.info('frozen orbitals %s', self.frozen)
        log.info('max_cycle = %d', self.max_cycle)
        log.info('direct = %d', self.direct)
        log.info('conv_tol = %g', self.conv_tol)
        log.info('conv_tol_normt = %s', self.conv_tol_normt)
        log.info('diis_space = %d', self.diis_space)
        #log.info('diis_file = %s', self.diis_file)
        log.info('diis_start_cycle = %d', self.diis_start_cycle)
        log.info('diis_start_energy_diff = %g', self.diis_start_energy_diff)
        log.info('max_memory %d MB (current use %d MB)',
                 self.max_memory, lib.current_memory()[0])
        if (log.verbose >= logger.DEBUG1 and
            self.__class__ == CCSD):
            nocc = self.nocc
            nvir = self.nmo - self.nocc
            flops = _flops(nocc, nvir)
            log.debug1('total FLOPs %s', flops)
        return self

    def get_init_guess(self, eris=None):
        return self.init_amps(eris)[1:]
    def init_amps(self, eris=None):
        time0 = logger.process_clock(), logger.perf_counter()
        if eris is None:
            eris = self.ao2mo(self.mo_coeff)
        mo_e = eris.mo_energy
        nocc = self.nocc
        nvir = mo_e.size - nocc
        eia = mo_e[:nocc,None] - mo_e[None,nocc:]

        t1 = eris.fock[:nocc,nocc:] / eia
        t2 = numpy.empty((nocc,nocc,nvir,nvir), dtype=eris.ovov.dtype)
        max_memory = self.max_memory - lib.current_memory()[0]
        blksize = int(min(nvir, max(BLKMIN, max_memory*.3e6/8/(nocc**2*nvir+1))))
        emp2 = 0
        for p0, p1 in lib.prange(0, nvir, blksize):
            eris_ovov = eris.ovov[:,p0:p1]
            t2[:,:,p0:p1] = (eris_ovov.transpose(0,2,1,3).conj()
                             / lib.direct_sum('ia,jb->ijab', eia[:,p0:p1], eia))
            emp2 += 2 * numpy.einsum('ijab,iajb', t2[:,:,p0:p1], eris_ovov)
            emp2 -=     numpy.einsum('jiab,iajb', t2[:,:,p0:p1], eris_ovov)
        self.emp2 = emp2.real

        e_hf = self.e_hf or eris.e_hf

        logger.info(self, 'Init t2, MP2 energy = %.15g  E_corr(MP2) %.15g',
                    e_hf + self.emp2, self.emp2)
        logger.timer(self, 'init mp2', *time0)
        return self.emp2, t1, t2

    energy = energy
    _add_vvvv = _add_vvvv
    update_amps = update_amps

    def kernel(self, t1=None, t2=None, eris=None):
        return self.ccsd(t1, t2, eris)
    def ccsd(self, t1=None, t2=None, eris=None):
        assert(self.mo_coeff is not None)
        assert(self.mo_occ is not None)

        if self.verbose >= logger.WARN:
            self.check_sanity()
        self.dump_flags()

        if eris is None:
            eris = self.ao2mo(self.mo_coeff)

        self.e_hf = getattr(eris, 'e_hf', None)
        if self.e_hf is None:
            self.e_hf = self._scf.e_tot

        self.converged, self.e_corr, self.t1, self.t2 = \
                kernel(self, eris, t1, t2, max_cycle=self.max_cycle,
                       tol=self.conv_tol, tolnormt=self.conv_tol_normt,
                       verbose=self.verbose, callback=self.callback)
        self._finalize()
        return self.e_corr, self.t1, self.t2

    def _finalize(self):
        '''Hook for dumping results and clearing up the object.'''
        if self.converged:
            logger.info(self, '%s converged', self.__class__.__name__)
        else:
            logger.note(self, '%s not converged', self.__class__.__name__)
        logger.note(self, 'E(%s) = %.16g  E_corr = %.16g',
                    self.__class__.__name__, self.e_tot, self.e_corr)
        return self

    as_scanner = as_scanner
    restore_from_diis_ = restore_from_diis_


    def solve_lambda(self, t1=None, t2=None, l1=None, l2=None,
                     eris=None):
        from pyscf.cc import ccsd_lambda
        if t1 is None: t1 = self.t1
        if t2 is None: t2 = self.t2
        if eris is None: eris = self.ao2mo(self.mo_coeff)
        self.converged_lambda, self.l1, self.l2 = \
                ccsd_lambda.kernel(self, eris, t1, t2, l1, l2,
                                   max_cycle=self.max_cycle,
                                   tol=self.conv_tol_normt,
                                   verbose=self.verbose)
        return self.l1, self.l2

    def ccsd_t(self, t1=None, t2=None, eris=None):
        from pyscf.cc import ccsd_t
        if t1 is None: t1 = self.t1
        if t2 is None: t2 = self.t2
        if eris is None: eris = self.ao2mo(self.mo_coeff)
        return ccsd_t.kernel(self, eris, t1, t2, self.verbose)

    def ipccsd(self, nroots=1, left=False, koopmans=False, guess=None,
               partition=None, eris=None):
        from pyscf.cc import eom_rccsd
        return eom_rccsd.EOMIP(self).kernel(nroots, left, koopmans, guess,
                                            partition, eris)

    def eaccsd(self, nroots=1, left=False, koopmans=False, guess=None,
               partition=None, eris=None):
        from pyscf.cc import eom_rccsd
        return eom_rccsd.EOMEA(self).kernel(nroots, left, koopmans, guess,
                                            partition, eris)

    def eeccsd(self, nroots=1, koopmans=False, guess=None, eris=None):
        from pyscf.cc import eom_rccsd
        return eom_rccsd.EOMEE(self).kernel(nroots, koopmans, guess, eris)

    def eomee_ccsd_singlet(self, nroots=1, koopmans=False, guess=None, eris=None):
        from pyscf.cc import eom_rccsd
        return eom_rccsd.EOMEESinglet(self).kernel(nroots, koopmans, guess, eris)

    def eomee_ccsd_triplet(self, nroots=1, koopmans=False, guess=None, eris=None):
        from pyscf.cc import eom_rccsd
        return eom_rccsd.EOMEETriplet(self).kernel(nroots, koopmans, guess, eris)

    def eomsf_ccsd(self, nroots=1, koopmans=False, guess=None, eris=None):
        from pyscf.cc import eom_rccsd
        return eom_rccsd.EOMEESpinFlip(self).kernel(nroots, koopmans, guess, eris)

    def eomip_method(self):
        from pyscf.cc import eom_rccsd
        return eom_rccsd.EOMIP(self)

    def eomea_method(self):
        from pyscf.cc import eom_rccsd
        return eom_rccsd.EOMEA(self)

    def eomee_method(self):
        from pyscf.cc import eom_rccsd
        return eom_rccsd.EOMEE(self)

    def make_rdm1(self, t1=None, t2=None, l1=None, l2=None, ao_repr=False,
                  with_frozen=True, with_mf=True):
        '''Un-relaxed 1-particle density matrix in MO space'''
        from pyscf.cc import ccsd_rdm
        if t1 is None: t1 = self.t1
        if t2 is None: t2 = self.t2
        if l1 is None: l1 = self.l1
        if l2 is None: l2 = self.l2
        if l1 is None: l1, l2 = self.solve_lambda(t1, t2)
        return ccsd_rdm.make_rdm1(self, t1, t2, l1, l2, ao_repr=ao_repr,
                                  with_frozen=with_frozen, with_mf=with_mf)

    def make_rdm2(self, t1=None, t2=None, l1=None, l2=None, ao_repr=False,
                  with_frozen=True, with_dm1=True):
        '''2-particle density matrix in MO space.  The density matrix is
        stored as

        dm2[p,r,q,s] = <p^+ q^+ s r>
        '''
        from pyscf.cc import ccsd_rdm
        if t1 is None: t1 = self.t1
        if t2 is None: t2 = self.t2
        if l1 is None: l1 = self.l1
        if l2 is None: l2 = self.l2
        if l1 is None: l1, l2 = self.solve_lambda(t1, t2)
        return ccsd_rdm.make_rdm2(self, t1, t2, l1, l2, ao_repr=ao_repr,
                                  with_frozen=with_frozen, with_dm1=with_dm1)

    def ao2mo(self, mo_coeff=None):
        # Pseudo code how eris are implemented:
        # nocc = self.nocc
        # nmo = self.nmo
        # nvir = nmo - nocc
        # eris = _ChemistsERIs()
        # eri = ao2mo.incore.full(self._scf._eri, mo_coeff)
        # eri = ao2mo.restore(1, eri, nmo)
        # eris.oooo = eri[:nocc,:nocc,:nocc,:nocc].copy()
        # eris.ovoo = eri[:nocc,nocc:,:nocc,:nocc].copy()
        # eris.ovvo = eri[nocc:,:nocc,nocc:,:nocc].copy()
        # eris.ovov = eri[nocc:,:nocc,:nocc,nocc:].copy()
        # eris.oovv = eri[:nocc,:nocc,nocc:,nocc:].copy()
        # ovvv = eri[:nocc,nocc:,nocc:,nocc:].copy()
        # eris.ovvv = lib.pack_tril(ovvv.reshape(-1,nvir,nvir))
        # eris.vvvv = ao2mo.restore(4, eri[nocc:,nocc:,nocc:,nocc:], nvir)
        # eris.fock = numpy.diag(self._scf.mo_energy)
        # return eris

        nmo = self.nmo
        nao = self.mo_coeff.shape[0]
        nmo_pair = nmo * (nmo+1) // 2
        nao_pair = nao * (nao+1) // 2
        mem_incore = (max(nao_pair**2, nmo**4) + nmo_pair**2) * 8/1e6
        mem_now = lib.current_memory()[0]
        if (self._scf._eri is not None and
            (mem_incore+mem_now < self.max_memory or self.incore_complete)):
            return _make_eris_incore(self, mo_coeff)

        elif getattr(self._scf, 'with_df', None):
            logger.warn(self, 'CCSD detected DF being used in the HF object. '
                        'MO integrals are computed based on the DF 3-index tensors.\n'
                        'It\'s recommended to use dfccsd.CCSD for the '
                        'DF-CCSD calculations')
            return _make_df_eris_outcore(self, mo_coeff)

        else:
            return _make_eris_outcore(self, mo_coeff)

    def run_diis(self, t1, t2, istep, normt, de, adiis):
        if (adiis and
            istep >= self.diis_start_cycle and
            abs(de) < self.diis_start_energy_diff):
            vec = self.amplitudes_to_vector(t1, t2)
            t1, t2 = self.vector_to_amplitudes(adiis.update(vec))
            logger.debug1(self, 'DIIS for step %d', istep)
        return t1, t2

    def amplitudes_to_vector(self, t1, t2, out=None):
        return amplitudes_to_vector(t1, t2, out)

    def vector_to_amplitudes(self, vec, nmo=None, nocc=None):
        if nocc is None: nocc = self.nocc
        if nmo is None: nmo = self.nmo
        return vector_to_amplitudes(vec, nmo, nocc)

    def vector_size(self, nmo=None, nocc=None):
        if nocc is None: nocc = self.nocc
        if nmo is None: nmo = self.nmo
        nvir = nmo - nocc
        nov = nocc * nvir
        return nov + nov*(nov+1)//2

    def dump_chk(self, t1_t2=None, frozen=None, mo_coeff=None, mo_occ=None):
        if not self.chkfile:
            return self
        if t1_t2 is None: t1_t2 = self.t1, self.t2
        t1, t2 = t1_t2

        if frozen is None: frozen = self.frozen
        # "None" cannot be serialized by the chkfile module
        if frozen is None:
            frozen = 0

        cc_chk = {'e_corr': self.e_corr,
                  't1': t1,
                  't2': t2,
                  'frozen': frozen}

        if mo_coeff is not None: cc_chk['mo_coeff'] = mo_coeff
        if mo_occ is not None: cc_chk['mo_occ'] = mo_occ
        if self._nmo is not None: cc_chk['_nmo'] = self._nmo
        if self._nocc is not None: cc_chk['_nocc'] = self._nocc

        lib.chkfile.save(self.chkfile, 'ccsd', cc_chk)

    def density_fit(self, auxbasis=None, with_df=None):
        from pyscf.cc import dfccsd
        mycc = dfccsd.RCCSD(self._scf, self.frozen, self.mo_coeff, self.mo_occ)
        if with_df is not None:
            mycc.with_df = with_df
        if mycc.with_df.auxbasis != auxbasis:
            import copy
            mycc.with_df = copy.copy(mycc.with_df)
            mycc.with_df.auxbasis = auxbasis
        return mycc

    def nuc_grad_method(self):
        from pyscf.grad import ccsd
        return ccsd.Gradients(self)

    def get_t1_diagnostic(self, t1=None):
        if t1 is None: t1 = self.t1
        return get_t1_diagnostic(t1)

    def get_d1_diagnostic(self, t1=None):
        if t1 is None: t1 = self.t1
        return get_d1_diagnostic(t1)

    def get_d2_diagnostic(self, t2=None):
        if t2 is None: t2 = self.t2
        return get_d2_diagnostic(t2)


CC = RCCSD = CCSD


class _ChemistsERIs:
    '''(pq|rs)'''
    def __init__(self, mol=None):
        self.mol = mol
        self.mo_coeff = None
        self.nocc = None
        self.fock = None
        self.e_hf = None

        self.oooo = None
        self.ovoo = None
        self.oovv = None
        self.ovvo = None
        self.ovov = None
        self.ovvv = None
        self.vvvv = None

    def _common_init_(self, mycc, mo_coeff=None):
        if mo_coeff is None:
            mo_coeff = mycc.mo_coeff
        self.mo_coeff = mo_coeff = _mo_without_core(mycc, mo_coeff)

# Note: Recomputed fock matrix and HF energy since SCF may not be fully converged.
        dm = mycc._scf.make_rdm1(mycc.mo_coeff, mycc.mo_occ)
        vhf = mycc._scf.get_veff(mycc.mol, dm)
        fockao = mycc._scf.get_fock(vhf=vhf, dm=dm)
        self.fock = reduce(numpy.dot, (mo_coeff.conj().T, fockao, mo_coeff))
        self.e_hf = mycc._scf.energy_tot(dm=dm, vhf=vhf)
        nocc = self.nocc = mycc.nocc
        self.mol = mycc.mol

        # Note self.mo_energy can be different to fock.diagonal().
        # self.mo_energy is used in the initial guess function (to generate
        # MP2 amplitudes) and CCSD update_amps preconditioner.
        # fock.diagonal() should only be used to compute the expectation value
        # of Slater determinants.
        mo_e = self.mo_energy = self.fock.diagonal().real
        try:
            gap = abs(mo_e[:nocc,None] - mo_e[None,nocc:]).min()
            if gap < 1e-5:
                logger.warn(mycc, 'HOMO-LUMO gap %s too small for CCSD.\n'
                            'CCSD may be difficult to converge. Increasing '
                            'CCSD Attribute level_shift may improve '
                            'convergence.', gap)
        except ValueError:  # gap.size == 0
            pass
        return self

    def get_ovvv(self, *slices):
        '''To access a subblock of ovvv tensor'''
        ovw = numpy.asarray(self.ovvv[slices])
        nocc, nvir, nvir_pair = ovw.shape
        ovvv = lib.unpack_tril(ovw.reshape(nocc*nvir,nvir_pair))
        nvir1 = ovvv.shape[2]
        return ovvv.reshape(nocc,nvir,nvir1,nvir1)

    def _contract_vvvv_t2(self, mycc, t2, vvvv_or_direct=False, out=None, verbose=None):
        if isinstance(vvvv_or_direct, numpy.ndarray):
            vvvv = vvvv_or_direct
        elif vvvv_or_direct:  # AO-direct contraction
            vvvv = None
        else:
            vvvv = self.vvvv
        return _contract_vvvv_t2(mycc, self.mol, vvvv, t2, out, verbose)

    def _contract_vvvv_oov(self, mycc, r2, out=None):
        raise NotImplementedError

    def _contract_vvvv_ovv(self, mycc, r2, out=None):
        raise NotImplementedError

def _make_eris_incore(mycc, mo_coeff=None):
    cput0 = (logger.process_clock(), logger.perf_counter())
    eris = _ChemistsERIs()
    eris._common_init_(mycc, mo_coeff)
    nocc = eris.nocc
    nmo = eris.fock.shape[0]
    nvir = nmo - nocc

    eri1 = ao2mo.incore.full(mycc._scf._eri, eris.mo_coeff)
    #:eri1 = ao2mo.restore(1, eri1, nmo)
    #:eris.oooo = eri1[:nocc,:nocc,:nocc,:nocc].copy()
    #:eris.ovoo = eri1[:nocc,nocc:,:nocc,:nocc].copy()
    #:eris.ovvo = eri1[:nocc,nocc:,nocc:,:nocc].copy()
    #:eris.ovov = eri1[:nocc,nocc:,:nocc,nocc:].copy()
    #:eris.oovv = eri1[:nocc,:nocc,nocc:,nocc:].copy()
    #:ovvv = eri1[:nocc,nocc:,nocc:,nocc:].copy()
    #:eris.ovvv = lib.pack_tril(ovvv.reshape(-1,nvir,nvir)).reshape(nocc,nvir,-1)
    #:eris.vvvv = ao2mo.restore(4, eri1[nocc:,nocc:,nocc:,nocc:], nvir)

    if eri1.ndim == 4:
        eri1 = ao2mo.restore(4, eri1, nmo)

    nvir_pair = nvir * (nvir+1) // 2
    eris.oooo = numpy.empty((nocc,nocc,nocc,nocc))
    eris.ovoo = numpy.empty((nocc,nvir,nocc,nocc))
    eris.ovvo = numpy.empty((nocc,nvir,nvir,nocc))
    eris.ovov = numpy.empty((nocc,nvir,nocc,nvir))
    eris.ovvv = numpy.empty((nocc,nvir,nvir_pair))
    eris.vvvv = numpy.empty((nvir_pair,nvir_pair))

    ij = 0
    outbuf = numpy.empty((nmo,nmo,nmo))
    oovv = numpy.empty((nocc,nocc,nvir,nvir))
    for i in range(nocc):
        buf = lib.unpack_tril(eri1[ij:ij+i+1], out=outbuf[:i+1])
        for j in range(i+1):
            eris.oooo[i,j] = eris.oooo[j,i] = buf[j,:nocc,:nocc]
            oovv[i,j] = oovv[j,i] = buf[j,nocc:,nocc:]
        ij += i + 1
    eris.oovv = oovv
    oovv = None

    ij1 = 0
    for i in range(nocc,nmo):
        buf = lib.unpack_tril(eri1[ij:ij+i+1], out=outbuf[:i+1])
        eris.ovoo[:,i-nocc] = buf[:nocc,:nocc,:nocc]
        eris.ovvo[:,i-nocc] = buf[:nocc,nocc:,:nocc]
        eris.ovov[:,i-nocc] = buf[:nocc,:nocc,nocc:]
        eris.ovvv[:,i-nocc] = lib.pack_tril(buf[:nocc,nocc:,nocc:])
        dij = i - nocc + 1
        lib.pack_tril(buf[nocc:i+1,nocc:,nocc:],
                      out=eris.vvvv[ij1:ij1+dij])
        ij += i + 1
        ij1 += dij
    logger.timer(mycc, 'CCSD integral transformation', *cput0)
    return eris

def _make_eris_outcore(mycc, mo_coeff=None):
    cput0 = (logger.process_clock(), logger.perf_counter())
    log = logger.Logger(mycc.stdout, mycc.verbose)
    eris = _ChemistsERIs()
    eris._common_init_(mycc, mo_coeff)

    mol = mycc.mol
    mo_coeff = numpy.asarray(eris.mo_coeff, order='F')
    nocc = eris.nocc
    nao, nmo = mo_coeff.shape
    nvir = nmo - nocc
    orbo = mo_coeff[:,:nocc]
    orbv = mo_coeff[:,nocc:]
    nvpair = nvir * (nvir+1) // 2
    eris.feri1 = lib.H5TmpFile()
    eris.oooo = eris.feri1.create_dataset('oooo', (nocc,nocc,nocc,nocc), 'f8')
    eris.oovv = eris.feri1.create_dataset('oovv', (nocc,nocc,nvir,nvir), 'f8', chunks=(nocc,nocc,1,nvir))
    eris.ovoo = eris.feri1.create_dataset('ovoo', (nocc,nvir,nocc,nocc), 'f8', chunks=(nocc,1,nocc,nocc))
    eris.ovvo = eris.feri1.create_dataset('ovvo', (nocc,nvir,nvir,nocc), 'f8', chunks=(nocc,1,nvir,nocc))
    eris.ovov = eris.feri1.create_dataset('ovov', (nocc,nvir,nocc,nvir), 'f8', chunks=(nocc,1,nocc,nvir))
    eris.ovvv = eris.feri1.create_dataset('ovvv', (nocc,nvir,nvpair), 'f8')

    def save_occ_frac(p0, p1, eri):
        eri = eri.reshape(p1-p0,nocc,nmo,nmo)
        eris.oooo[p0:p1] = eri[:,:,:nocc,:nocc]
        eris.oovv[p0:p1] = eri[:,:,nocc:,nocc:]

    def save_vir_frac(p0, p1, eri):
        eri = eri.reshape(p1-p0,nocc,nmo,nmo)
        eris.ovoo[:,p0:p1] = eri[:,:,:nocc,:nocc].transpose(1,0,2,3)
        eris.ovvo[:,p0:p1] = eri[:,:,nocc:,:nocc].transpose(1,0,2,3)
        eris.ovov[:,p0:p1] = eri[:,:,:nocc,nocc:].transpose(1,0,2,3)
        vvv = lib.pack_tril(eri[:,:,nocc:,nocc:].reshape((p1-p0)*nocc,nvir,nvir))
        eris.ovvv[:,p0:p1] = vvv.reshape(p1-p0,nocc,nvpair).transpose(1,0,2)

    cput1 = logger.process_clock(), logger.perf_counter()
    if not mycc.direct:
        max_memory = max(MEMORYMIN, mycc.max_memory-lib.current_memory()[0])
        eris.feri2 = lib.H5TmpFile()
        ao2mo.full(mol, orbv, eris.feri2, max_memory=max_memory, verbose=log)
        eris.vvvv = eris.feri2['eri_mo']
        cput1 = log.timer_debug1('transforming vvvv', *cput1)

    fswap = lib.H5TmpFile()
    max_memory = max(MEMORYMIN, mycc.max_memory-lib.current_memory()[0])
    int2e = mol._add_suffix('int2e')
    ao2mo.outcore.half_e1(mol, (mo_coeff,orbo), fswap, int2e,
                          's4', 1, max_memory, verbose=log)

    ao_loc = mol.ao_loc_nr()
    nao_pair = nao * (nao+1) // 2
    blksize = int(min(8e9,max_memory*.5e6)/8/(nao_pair+nmo**2)/nocc)
    blksize = min(nmo, max(BLKMIN, blksize))
    log.debug1('blksize %d', blksize)
    cput2 = cput1

    fload = ao2mo.outcore._load_from_h5g
    buf = numpy.empty((blksize*nocc,nao_pair))
    buf_prefetch = numpy.empty_like(buf)
    def load(buf_prefetch, p0, rowmax):
        if p0 < rowmax:
            p1 = min(rowmax, p0+blksize)
            fload(fswap['0'], p0*nocc, p1*nocc, buf_prefetch)

    outbuf = numpy.empty((blksize*nocc,nmo**2))
    with lib.call_in_background(load, sync=not mycc.async_io) as prefetch:
        prefetch(buf_prefetch, 0, nocc)
        for p0, p1 in lib.prange(0, nocc, blksize):
            buf, buf_prefetch = buf_prefetch, buf
            prefetch(buf_prefetch, p1, nocc)

            nrow = (p1 - p0) * nocc
            dat = ao2mo._ao2mo.nr_e2(buf[:nrow], mo_coeff, (0,nmo,0,nmo),
                                     's4', 's1', out=outbuf, ao_loc=ao_loc)
            save_occ_frac(p0, p1, dat)
        cput2 = log.timer_debug1('transforming oopp', *cput2)

        prefetch(buf_prefetch, nocc, nmo)
        for p0, p1 in lib.prange(0, nvir, blksize):
            buf, buf_prefetch = buf_prefetch, buf
            prefetch(buf_prefetch, nocc+p1, nmo)

            nrow = (p1 - p0) * nocc
            dat = ao2mo._ao2mo.nr_e2(buf[:nrow], mo_coeff, (0,nmo,0,nmo),
                                     's4', 's1', out=outbuf, ao_loc=ao_loc)
            save_vir_frac(p0, p1, dat)
            cput2 = log.timer_debug1('transforming ovpp [%d:%d]'%(p0,p1), *cput2)

    cput1 = log.timer_debug1('transforming oppp', *cput1)
    log.timer('CCSD integral transformation', *cput0)
    return eris

def _make_df_eris_outcore(mycc, mo_coeff=None):
    cput0 = (logger.process_clock(), logger.perf_counter())
    log = logger.Logger(mycc.stdout, mycc.verbose)
    eris = _ChemistsERIs()
    eris._common_init_(mycc, mo_coeff)

    mo_coeff = numpy.asarray(eris.mo_coeff, order='F')
    nocc = eris.nocc
    nao, nmo = mo_coeff.shape
    nvir = nmo - nocc
    nvir_pair = nvir*(nvir+1)//2

    naux = mycc._scf.with_df.get_naoaux()
    Loo = numpy.empty((naux,nocc,nocc))
    Lov = numpy.empty((naux,nocc,nvir))
    Lvo = numpy.empty((naux,nvir,nocc))
    Lvv = numpy.empty((naux,nvir_pair))
    ijslice = (0, nmo, 0, nmo)
    Lpq = None
    p1 = 0
    for eri1 in mycc._scf.with_df.loop():
        Lpq = _ao2mo.nr_e2(eri1, mo_coeff, ijslice, aosym='s2', out=Lpq).reshape(-1,nmo,nmo)
        p0, p1 = p1, p1 + Lpq.shape[0]
        Loo[p0:p1] = Lpq[:,:nocc,:nocc]
        Lov[p0:p1] = Lpq[:,:nocc,nocc:]
        Lvo[p0:p1] = Lpq[:,nocc:,:nocc]
        Lvv[p0:p1] = lib.pack_tril(Lpq[:,nocc:,nocc:].reshape(-1,nvir,nvir))
    Loo = Loo.reshape(naux,nocc*nocc)
    Lov = Lov.reshape(naux,nocc*nvir)
    Lvo = Lvo.reshape(naux,nocc*nvir)

    eris.feri1 = lib.H5TmpFile()
    eris.oooo = eris.feri1.create_dataset('oooo', (nocc,nocc,nocc,nocc), 'f8')
    eris.oovv = eris.feri1.create_dataset('oovv', (nocc,nocc,nvir,nvir), 'f8', chunks=(nocc,nocc,1,nvir))
    eris.ovoo = eris.feri1.create_dataset('ovoo', (nocc,nvir,nocc,nocc), 'f8', chunks=(nocc,1,nocc,nocc))
    eris.ovvo = eris.feri1.create_dataset('ovvo', (nocc,nvir,nvir,nocc), 'f8', chunks=(nocc,1,nvir,nocc))
    eris.ovov = eris.feri1.create_dataset('ovov', (nocc,nvir,nocc,nvir), 'f8', chunks=(nocc,1,nocc,nvir))
    eris.ovvv = eris.feri1.create_dataset('ovvv', (nocc,nvir,nvir_pair), 'f8')
    eris.vvvv = eris.feri1.create_dataset('vvvv', (nvir_pair,nvir_pair), 'f8')
    eris.oooo[:] = lib.ddot(Loo.T, Loo).reshape(nocc,nocc,nocc,nocc)
    eris.ovoo[:] = lib.ddot(Lov.T, Loo).reshape(nocc,nvir,nocc,nocc)
    eris.oovv[:] = lib.unpack_tril(lib.ddot(Loo.T, Lvv)).reshape(nocc,nocc,nvir,nvir)
    eris.ovvo[:] = lib.ddot(Lov.T, Lvo).reshape(nocc,nvir,nvir,nocc)
    eris.ovov[:] = lib.ddot(Lov.T, Lov).reshape(nocc,nvir,nocc,nvir)
    eris.ovvv[:] = lib.ddot(Lov.T, Lvv).reshape(nocc,nvir,nvir_pair)
    eris.vvvv[:] = lib.ddot(Lvv.T, Lvv)
    log.timer('CCSD integral transformation', *cput0)
    return eris

def _flops(nocc, nvir):
    '''Total float points'''
    return (nocc**3*nvir**2*2 + nocc**2*nvir**3*2 +     # Ftilde
            nocc**4*nvir*2 * 2 + nocc**4*nvir**2*2 +    # Wijkl
            nocc*nvir**4*2 * 2 +                        # Wabcd
            nocc**2*nvir**3*2 + nocc**3*nvir**2*2 +
            nocc**3*nvir**3*2 + nocc**3*nvir**3*2 +
            nocc**2*nvir**3*2 + nocc**3*nvir**2*2 +     # Wiabj
            nocc**2*nvir**3*2 + nocc**3*nvir**2*2 +     # t1
            nocc**3*nvir**2*2 * 2 + nocc**4*nvir**2*2 +
            nocc*(nocc+1)/2*nvir**4*2 +                 # vvvv
            nocc**2*nvir**3*2 * 2 + nocc**3*nvir**2*2 * 2 +     # t2
            nocc**3*nvir**3*2 +
            nocc**3*nvir**3*2 * 2 + nocc**3*nvir**2*2 * 4)      # Wiabj


if __name__ == '__main__':
    from pyscf import scf

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]

    mol.basis = {'H': 'cc-pvdz',
                 'O': 'cc-pvdz',}
    mol.build()
    rhf = scf.RHF(mol)
    rhf.scf() # -76.0267656731

    mf = rhf.density_fit(auxbasis='weigend')
    mf._eri = None
    mcc = CCSD(mf)
    eris = mcc.ao2mo()
    emp2, t1, t2 = mcc.init_amps(eris)
    print(abs(t2).sum() - 4.9318753386922278)
    print(emp2 - -0.20401737899811551)
    t1, t2 = update_amps(mcc, t1, t2, eris)
    print(abs(t1).sum() - 0.046961325647584914)
    print(abs(t2).sum() - 5.378260578551683   )

    mcc = CCSD(rhf)
    eris = mcc.ao2mo()
    emp2, t1, t2 = mcc.init_amps(eris)
    print(abs(t2).sum() - 4.9556571218177)
    print(emp2 - -0.2040199672883385)
    t1, t2 = update_amps(mcc, t1, t2, eris)
    print(abs(t1).sum()-0.0475038989126)
    print(abs(t2).sum()-5.401823846018721)
    print(energy(mcc, t1, t2, eris) - -0.208967840546667)
    t1, t2 = update_amps(mcc, t1, t2, eris)
    print(energy(mcc, t1, t2, eris) - -0.212173678670510)
    print(abs(t1).sum() - 0.05470123093500083)
    print(abs(t2).sum() - 5.5605208391876539)

    mcc.ccsd()
    print(mcc.ecc - -0.213343234198275)
    print(abs(mcc.t2).sum() - 5.63970304662375)

    mcc.max_memory = 1
    mcc.direct = True
    mcc.ccsd()
    print(mcc.ecc - -0.213343234198275)
    print(abs(mcc.t2).sum() - 5.63970304662375)

    e, v = mcc.ipccsd(nroots=3)
    print(e[0] - 0.43356041409195489)
    print(e[1] - 0.51876598058509493)
    print(e[2] - 0.6782879569941862 )

    e, v = mcc.eeccsd(nroots=4)
    print(e[0] - 0.2757159395886167)
    print(e[1] - 0.2757159395886167)
    print(e[2] - 0.2757159395886167)
    print(e[3] - 0.3005716731825082)


# pyscf/cc/ccsd_lambda.py

#!/usr/bin/env python
# Copyright 2014-2021 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: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Restricted CCSD implementation for real integrals.  Permutation symmetry for
the 4-index integrals (ij|kl) = (ij|lk) = (ji|kl) are assumed.

Note MO integrals are treated in chemist's notation
'''


from functools import reduce
import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf.cc import ccsd
from pyscf.cc import _ccsd

# t2,l2 as ijab
def kernel(mycc, eris=None, t1=None, t2=None, l1=None, l2=None,
           max_cycle=50, tol=1e-8, verbose=logger.INFO,
           fintermediates=None, fupdate=None):
    if eris is None: eris = mycc.ao2mo()
    cput0 = (logger.process_clock(), logger.perf_counter())
    log = logger.new_logger(mycc, verbose)

    if t1 is None: t1 = mycc.t1
    if t2 is None: t2 = mycc.t2
    if l1 is None: l1 = t1
    if l2 is None: l2 = t2
    if fintermediates is None:
        fintermediates = make_intermediates
    if fupdate is None:
        fupdate = update_lambda

    imds = fintermediates(mycc, t1, t2, eris)

    if isinstance(mycc.diis, lib.diis.DIIS):
        adiis = mycc.diis
    elif mycc.diis:
        adiis = lib.diis.DIIS(mycc, mycc.diis_file, incore=mycc.incore_complete)
        adiis.space = mycc.diis_space
    else:
        adiis = None
    cput0 = log.timer('CCSD lambda initialization', *cput0)

    conv = False
    for istep in range(max_cycle):
        l1new, l2new = fupdate(mycc, t1, t2, l1, l2, eris, imds)
        normt = numpy.linalg.norm(mycc.amplitudes_to_vector(l1new, l2new) -
                                  mycc.amplitudes_to_vector(l1, l2))
        l1, l2 = l1new, l2new
        l1new = l2new = None
        l1, l2 = mycc.run_diis(l1, l2, istep, normt, 0, adiis)
        log.info('cycle = %d  norm(lambda1,lambda2) = %.6g', istep+1, normt)
        cput0 = log.timer('CCSD iter', *cput0)
        if normt < tol:
            conv = True
            break
    return conv, l1, l2


# l2, t2 as ijab
def make_intermediates(mycc, t1, t2, eris):
    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocc, nvir = t1.shape
    foo = eris.fock[:nocc,:nocc]
    fov = eris.fock[:nocc,nocc:]
    #fvo = eris.fock[nocc:,:nocc]
    fvv = eris.fock[nocc:,nocc:]

    class _IMDS: pass
    imds = _IMDS()
    #TODO: mycc.incore_complete
    imds.ftmp = lib.H5TmpFile()
    imds.woooo = imds.ftmp.create_dataset('woooo', (nocc,nocc,nocc,nocc), 'f8')
    imds.wvooo = imds.ftmp.create_dataset('wvooo', (nvir,nocc,nocc,nocc), 'f8')
    imds.wVOov = imds.ftmp.create_dataset('wVOov', (nvir,nocc,nocc,nvir), 'f8')
    imds.wvOOv = imds.ftmp.create_dataset('wvOOv', (nvir,nocc,nocc,nvir), 'f8')
    imds.wvvov = imds.ftmp.create_dataset('wvvov', (nvir,nvir,nocc,nvir), 'f8')

    w1 = fvv - numpy.einsum('ja,jb->ba', fov, t1)
    w2 = foo + numpy.einsum('ib,jb->ij', fov, t1)
    w3 = numpy.einsum('kc,jkbc->bj', fov, t2) * 2 + fov.T
    w3 -= numpy.einsum('kc,kjbc->bj', fov, t2)
    w3 += lib.einsum('kc,kb,jc->bj', fov, t1, t1)
    w4 = fov.copy()

    unit = nocc*nvir**2*6
    max_memory = max(0, mycc.max_memory - lib.current_memory()[0])
    blksize = min(nvir, max(ccsd.BLKMIN, int((max_memory*.95e6/8-nocc**4-nvir*nocc**3)/unit)))
    log.debug1('ccsd lambda make_intermediates: block size = %d, nvir = %d in %d blocks',
               blksize, nvir, int((nvir+blksize-1)//blksize))

    fswap = lib.H5TmpFile()
    for istep, (p0, p1) in enumerate(lib.prange(0, nvir, blksize)):
        eris_ovvv = eris.get_ovvv(slice(None), slice(p0,p1))
        fswap['vvov/%d'%istep] = eris_ovvv.transpose(2,3,0,1)

    woooo = 0
    wvooo = numpy.zeros((nvir,nocc,nocc,nocc))
    for p0, p1 in lib.prange(0, nvir, blksize):
        eris_ovvv = eris.get_ovvv(slice(None), slice(p0,p1))
        eris_vvov = numpy.empty(((p1-p0),nvir,nocc,nvir))
        for istep, (q0, q1) in enumerate(lib.prange(0, nvir, blksize)):
            eris_vvov[:,:,:,q0:q1] = fswap['vvov/%d'%istep][p0:p1]

        w1 += numpy.einsum('jcba,jc->ba', eris_ovvv, t1[:,p0:p1]*2)
        w1[:,p0:p1] -= numpy.einsum('jabc,jc->ba', eris_ovvv, t1)
        theta = t2[:,:,:,p0:p1] * 2 - t2[:,:,:,p0:p1].transpose(1,0,2,3)
        w3 += lib.einsum('jkcd,kdcb->bj', theta, eris_ovvv)
        theta = None
        wVOov = lib.einsum('jbcd,kd->bjkc', eris_ovvv, t1)
        wvOOv = lib.einsum('cbjd,kd->cjkb', eris_vvov,-t1)
        g2vovv = eris_vvov.transpose(0,2,1,3) * 2 - eris_vvov.transpose(0,2,3,1)
        for i0, i1 in lib.prange(0, nocc, blksize):
            tau = t2[:,i0:i1] + numpy.einsum('ia,jb->ijab', t1, t1[i0:i1])
            wvooo[p0:p1,i0:i1] += lib.einsum('cibd,jkbd->ckij', g2vovv, tau)
        g2vovv = tau = None

        # Watch out memory usage here, due to the t2 transpose
        wvvov  = lib.einsum('jabd,jkcd->abkc', eris_ovvv, t2) * -1.5
        wvvov += eris_vvov.transpose(0,3,2,1) * 2
        wvvov -= eris_vvov

        g2vvov = eris_vvov * 2 - eris_ovvv.transpose(1,2,0,3)
        for i0, i1 in lib.prange(0, nocc, blksize):
            theta = t2[i0:i1] * 2 - t2[i0:i1].transpose(0,1,3,2)
            vackb = lib.einsum('acjd,kjbd->ackb', g2vvov, theta)
            wvvov[:,:,i0:i1] += vackb.transpose(0,3,2,1)
            wvvov[:,:,i0:i1] -= vackb * .5
        g2vvov = eris_ovvv = eris_vvov = theta = None

        eris_ovoo = _cp(eris.ovoo[:,p0:p1])
        w2 += numpy.einsum('kbij,kb->ij', eris_ovoo, t1[:,p0:p1]) * 2
        w2 -= numpy.einsum('ibkj,kb->ij', eris_ovoo, t1[:,p0:p1])
        theta = t2[:,:,p0:p1].transpose(1,0,2,3) * 2 - t2[:,:,p0:p1]
        w3 -= lib.einsum('lckj,klcb->bj', eris_ovoo, theta)

        tmp = lib.einsum('lc,jcik->ijkl', t1[:,p0:p1], eris_ovoo)
        woooo += tmp
        woooo += tmp.transpose(1,0,3,2)
        theta = tmp = None

        wvOOv += lib.einsum('lbjk,lc->bjkc', eris_ovoo, t1)
        wVOov -= lib.einsum('jbkl,lc->bjkc', eris_ovoo, t1)
        wvooo[p0:p1] += eris_ovoo.transpose(1,3,2,0) * 2
        wvooo[p0:p1] -= eris_ovoo.transpose(1,0,2,3)
        wvooo -= lib.einsum('klbc,iblj->ckij', t2[:,:,p0:p1], eris_ovoo*1.5)

        g2ovoo = eris_ovoo * 2 - eris_ovoo.transpose(2,1,0,3)
        theta = t2[:,:,:,p0:p1]*2 - t2[:,:,:,p0:p1].transpose(1,0,2,3)
        vcjik = lib.einsum('jlcb,lbki->cjki', theta, g2ovoo)
        wvooo += vcjik.transpose(0,3,2,1)
        wvooo -= vcjik*.5
        theta = g2ovoo = None

        eris_voov = _cp(eris.ovvo[:,p0:p1]).transpose(1,0,3,2)
        tau = t2[:,:,p0:p1] + numpy.einsum('ia,jb->ijab', t1[:,p0:p1], t1)
        woooo += lib.einsum('cijd,klcd->ijkl', eris_voov, tau)
        tau = None

        g2voov = eris_voov*2 - eris_voov.transpose(0,2,1,3)
        tmpw4 = numpy.einsum('ckld,ld->kc', g2voov, t1)
        w1 -= lib.einsum('ckja,kjcb->ba', g2voov, t2[:,:,p0:p1])
        w1[:,p0:p1] -= numpy.einsum('ja,jb->ba', tmpw4, t1)
        w2 += lib.einsum('jkbc,bikc->ij', t2[:,:,p0:p1], g2voov)
        w2 += numpy.einsum('ib,jb->ij', tmpw4, t1[:,p0:p1])
        w3 += reduce(numpy.dot, (t1.T, tmpw4, t1[:,p0:p1].T))
        w4[:,p0:p1] += tmpw4

        wvOOv += lib.einsum('bljd,kd,lc->bjkc', eris_voov, t1, t1)
        wVOov -= lib.einsum('bjld,kd,lc->bjkc', eris_voov, t1, t1)

        VOov  = lib.einsum('bjld,klcd->bjkc', g2voov, t2)
        VOov -= lib.einsum('bjld,kldc->bjkc', eris_voov, t2)
        VOov += eris_voov
        vOOv = lib.einsum('bljd,kldc->bjkc', eris_voov, t2)
        vOOv -= _cp(eris.oovv[:,:,p0:p1]).transpose(2,1,0,3)
        wVOov += VOov
        wvOOv += vOOv
        imds.wVOov[p0:p1] = wVOov
        imds.wvOOv[p0:p1] = wvOOv

        ov1 = vOOv*2 + VOov
        ov2 = VOov*2 + vOOv
        vOOv = VOov = None
        wvooo -= lib.einsum('jb,bikc->ckij', t1[:,p0:p1], ov1)
        wvooo += lib.einsum('kb,bijc->ckij', t1[:,p0:p1], ov2)
        w3 += numpy.einsum('ckjb,kc->bj', ov2, t1[:,p0:p1])

        wvvov += lib.einsum('ajkc,jb->abkc', ov1, t1)
        wvvov -= lib.einsum('ajkb,jc->abkc', ov2, t1)

        eris_ovoo = _cp(eris.ovoo[:,p0:p1])
        g2ovoo = eris_ovoo * 2 - eris_ovoo.transpose(2,1,0,3)
        tau = t2 + numpy.einsum('ia,jb->ijab', t1, t1)
        wvvov += lib.einsum('laki,klbc->abic', g2ovoo, tau)
        imds.wvvov[p0:p1] = wvvov
        wvvov = ov1 = ov2 = g2ovoo = None

    woooo += _cp(eris.oooo).transpose(0,2,1,3)
    imds.woooo[:] = woooo
    imds.wvooo[:] = wvooo
    woooo = wvooo = None

    w3 += numpy.einsum('bc,jc->bj', w1, t1)
    w3 -= numpy.einsum('kj,kb->bj', w2, t1)

    fswap = None

    imds.w1 = w1
    imds.w2 = w2
    imds.w3 = w3
    imds.w4 = w4
    imds.ftmp.flush()
    return imds


# update L1, L2
def update_lambda(mycc, t1, t2, l1, l2, eris=None, imds=None):
    if imds is None: imds = make_intermediates(mycc, t1, t2, eris)
    time0 = logger.process_clock(), logger.perf_counter()
    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocc, nvir = t1.shape
    fov = eris.fock[:nocc,nocc:]
    mo_e_o = eris.mo_energy[:nocc]
    mo_e_v = eris.mo_energy[nocc:] + mycc.level_shift

    theta = t2*2 - t2.transpose(0,1,3,2)
    mba = lib.einsum('klca,klcb->ba', l2, theta)
    mij = lib.einsum('ikcd,jkcd->ij', l2, theta)
    theta = None
    mba1 = numpy.einsum('jc,jb->bc', l1, t1) + mba
    mij1 = numpy.einsum('kb,jb->kj', l1, t1) + mij
    mia1 = t1 + numpy.einsum('kc,jkbc->jb', l1, t2) * 2
    mia1 -= numpy.einsum('kc,jkcb->jb', l1, t2)
    mia1 -= reduce(numpy.dot, (t1, l1.T, t1))
    mia1 -= numpy.einsum('bd,jd->jb', mba, t1)
    mia1 -= numpy.einsum('lj,lb->jb', mij, t1)

    l2new = mycc._add_vvvv(None, l2, eris, with_ovvv=False, t2sym='jiba')
    l1new  = numpy.einsum('ijab,jb->ia', l2new, t1) * 2
    l1new -= numpy.einsum('jiab,jb->ia', l2new, t1)
    l2new *= .5  # *.5 because of l2+l2.transpose(1,0,3,2) in the end
    tmp = None

    w1 = imds.w1 - numpy.diag(mo_e_v)
    w2 = imds.w2 - numpy.diag(mo_e_o)

    l1new += fov
    l1new += numpy.einsum('ib,ba->ia', l1, w1)
    l1new -= numpy.einsum('ja,ij->ia', l1, w2)
    l1new -= numpy.einsum('ik,ka->ia', mij, imds.w4)
    l1new -= numpy.einsum('ca,ic->ia', mba, imds.w4)
    l1new += numpy.einsum('ijab,bj->ia', l2, imds.w3) * 2
    l1new -= numpy.einsum('ijba,bj->ia', l2, imds.w3)

    l2new += numpy.einsum('ia,jb->ijab', l1, imds.w4)
    l2new += lib.einsum('jibc,ca->jiba', l2, w1)
    l2new -= lib.einsum('jk,kiba->jiba', w2, l2)

    eris_ovoo = _cp(eris.ovoo)
    l1new -= numpy.einsum('iajk,kj->ia', eris_ovoo, mij1) * 2
    l1new += numpy.einsum('jaik,kj->ia', eris_ovoo, mij1)
    l2new -= lib.einsum('jbki,ka->jiba', eris_ovoo, l1)
    eris_ovoo = None

    tau = _ccsd.make_tau(t2, t1, t1)
    l2tau = lib.einsum('ijcd,klcd->ijkl', l2, tau)
    tau = None
    l2t1 = lib.einsum('jidc,kc->ijkd', l2, t1)

    max_memory = max(0, mycc.max_memory - lib.current_memory()[0])
    unit = nocc*nvir**2*5
    blksize = min(nocc, max(ccsd.BLKMIN, int(max_memory*.95e6/8/unit)))
    log.debug1('block size = %d, nocc = %d is divided into %d blocks',
               blksize, nocc, int((nocc+blksize-1)/blksize))

    l1new -= numpy.einsum('jb,jiab->ia', l1, _cp(eris.oovv))
    for p0, p1 in lib.prange(0, nvir, blksize):
        eris_ovvv = eris.get_ovvv(slice(None), slice(p0,p1))
        l1new[:,p0:p1] += numpy.einsum('iabc,bc->ia', eris_ovvv, mba1) * 2
        l1new -= numpy.einsum('ibca,bc->ia', eris_ovvv, mba1[p0:p1])
        l2new[:,:,p0:p1] += lib.einsum('jbac,ic->jiba', eris_ovvv, l1)
        m4 = lib.einsum('ijkd,kadb->ijab', l2t1, eris_ovvv)
        l2new[:,:,p0:p1] -= m4
        l1new[:,p0:p1] -= numpy.einsum('ijab,jb->ia', m4, t1) * 2
        l1new -= numpy.einsum('ijab,ia->jb', m4, t1[:,p0:p1]) * 2
        l1new[:,p0:p1] += numpy.einsum('jiab,jb->ia', m4, t1)
        l1new += numpy.einsum('jiab,ia->jb', m4, t1[:,p0:p1])
        eris_ovvv = m4 = None

        eris_voov = _cp(eris.ovvo[:,p0:p1].transpose(1,0,3,2))
        l1new[:,p0:p1] += numpy.einsum('jb,aijb->ia', l1, eris_voov) * 2
        l2new[:,:,p0:p1] += eris_voov.transpose(1,2,0,3) * .5
        l2new[:,:,p0:p1] -= lib.einsum('bjic,ca->jiba', eris_voov, mba1)
        l2new[:,:,p0:p1] -= lib.einsum('bjka,ik->jiba', eris_voov, mij1)
        l1new[:,p0:p1] += numpy.einsum('aijb,jb->ia', eris_voov, mia1) * 2
        l1new -= numpy.einsum('bija,jb->ia', eris_voov, mia1[:,p0:p1])
        m4 = lib.einsum('ijkl,aklb->ijab', l2tau, eris_voov)
        l2new[:,:,p0:p1] += m4 * .5
        l1new[:,p0:p1] += numpy.einsum('ijab,jb->ia', m4, t1) * 2
        l1new -= numpy.einsum('ijba,jb->ia', m4, t1[:,p0:p1])

        saved_wvooo = _cp(imds.wvooo[p0:p1])
        l1new -= lib.einsum('ckij,jkca->ia', saved_wvooo, l2[:,:,p0:p1])
        saved_wvovv = _cp(imds.wvvov[p0:p1])
        # Watch out memory usage here, due to the l2 transpose
        l1new[:,p0:p1] += lib.einsum('abkc,kibc->ia', saved_wvovv, l2)
        saved_wvooo = saved_wvovv = None

        saved_wvOOv = _cp(imds.wvOOv[p0:p1])
        tmp_voov = _cp(imds.wVOov[p0:p1]) * 2
        tmp_voov += saved_wvOOv
        tmp = l2.transpose(0,2,1,3) - l2.transpose(0,3,1,2)*.5
        l2new[:,:,p0:p1] += lib.einsum('iakc,bjkc->jiba', tmp, tmp_voov)
        tmp = None

        tmp = lib.einsum('jkca,bikc->jiba', l2, saved_wvOOv)
        l2new[:,:,p0:p1] += tmp
        l2new[:,:,p0:p1] += tmp.transpose(1,0,2,3) * .5
        saved_wvOOv = tmp = None

    saved_woooo = _cp(imds.woooo)
    m3 = lib.einsum('ijkl,klab->ijab', saved_woooo, l2)
    l2new += m3 * .5
    l1new += numpy.einsum('ijab,jb->ia', m3, t1) * 2
    l1new -= numpy.einsum('ijba,jb->ia', m3, t1)
    saved_woooo = m3 = None
    #time1 = log.timer_debug1('lambda pass [%d:%d]'%(p0, p1), *time1)

    eia = lib.direct_sum('i-a->ia', mo_e_o, mo_e_v)
    l1new /= eia

#    l2new = l2new + l2new.transpose(1,0,3,2)
#    l2new /= lib.direct_sum('ia+jb->ijab', eia, eia)
#    l2new += l2
    for i in range(nocc):
        if i > 0:
            l2new[i,:i] += l2new[:i,i].transpose(0,2,1)
            l2new[i,:i] /= lib.direct_sum('a,jb->jab', eia[i], eia[:i])
            l2new[:i,i] = l2new[i,:i].transpose(0,2,1)
        l2new[i,i] = l2new[i,i] + l2new[i,i].T
        l2new[i,i] /= lib.direct_sum('a,b->ab', eia[i], eia[i])

    time0 = log.timer_debug1('update l1 l2', *time0)
    return l1new, l2new

def _cp(a):
    return numpy.array(a, copy=False, order='C')


if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf
    from pyscf import ao2mo

    mol = gto.M()
    mf = scf.RHF(mol)

    mcc = ccsd.CCSD(mf)

    numpy.random.seed(12)
    nocc = 5
    nmo = 12
    nvir = nmo - nocc
    eri0 = numpy.random.random((nmo,nmo,nmo,nmo))
    eri0 = ao2mo.restore(1, ao2mo.restore(8, eri0, nmo), nmo)
    fock0 = numpy.random.random((nmo,nmo))
    fock0 = fock0 + fock0.T + numpy.diag(range(nmo))*2
    t1 = numpy.random.random((nocc,nvir))
    t2 = numpy.random.random((nocc,nocc,nvir,nvir))
    t2 = t2 + t2.transpose(1,0,3,2)
    l1 = numpy.random.random((nocc,nvir))
    l2 = numpy.random.random((nocc,nocc,nvir,nvir))
    l2 = l2 + l2.transpose(1,0,3,2)

    eris = ccsd._ChemistsERIs()
    eris.oooo = eri0[:nocc,:nocc,:nocc,:nocc].copy()
    eris.ovoo = eri0[:nocc,nocc:,:nocc,:nocc].copy()
    eris.oovv = eri0[:nocc,:nocc,nocc:,nocc:].copy()
    eris.ovvo = eri0[:nocc,nocc:,nocc:,:nocc].copy()
    idx = numpy.tril_indices(nvir)
    eris.ovvv = eri0[:nocc,nocc:,nocc:,nocc:][:,:,idx[0],idx[1]].copy()
    eris.vvvv = ao2mo.restore(4,eri0[nocc:,nocc:,nocc:,nocc:],nvir)
    eris.fock = fock0
    eris.mo_energy = fock0.diagonal()

    imds = make_intermediates(mcc, t1, t2, eris)
    l1new, l2new = update_lambda(mcc, t1, t2, l1, l2, eris, imds)
    print(lib.finger(l1new) - -6699.5335665027187)
    print(lib.finger(l2new) - -514.7001243502192 )
    print(abs(l2new-l2new.transpose(1,0,3,2)).sum())

    mcc.max_memory = 0
    imds = make_intermediates(mcc, t1, t2, eris)
    l1new, l2new = update_lambda(mcc, t1, t2, l1, l2, eris, imds)
    print(lib.finger(l1new) - -6699.5335665027187)
    print(lib.finger(l2new) - -514.7001243502192 )
    print(abs(l2new-l2new.transpose(1,0,3,2)).sum())

    mol = gto.Mole()
    mol.verbose = 0
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]

    mol.basis = 'cc-pvdz'
    mol.build()
    rhf = scf.RHF(mol)
    rhf.conv_tol = 1e-16
    rhf.scf()

    mcc = ccsd.CCSD(rhf)
    mcc.conv_tol = 1e-12
    ecc, t1, t2 = mcc.kernel()

    nmo = rhf.mo_energy.size
    fock0 = numpy.diag(rhf.mo_energy)
    nocc = mol.nelectron // 2
    nvir = nmo - nocc

    eris = mcc.ao2mo()
    conv, l1, l2 = kernel(mcc, eris, t1, t2, tol=1e-8)
    print(numpy.linalg.norm(l1)-0.0132626841292)
    print(numpy.linalg.norm(l2)-0.212575609057)

    from pyscf.cc import ccsd_rdm
    dm1 = ccsd_rdm.make_rdm1(mcc, t1, t2, l1, l2)
    dm2 = ccsd_rdm.make_rdm2(mcc, t1, t2, l1, l2)
    h1 = reduce(numpy.dot, (rhf.mo_coeff.T, rhf.get_hcore(), rhf.mo_coeff))
    eri = ao2mo.full(rhf._eri, rhf.mo_coeff)
    eri = ao2mo.restore(1, eri, nmo).reshape((nmo,)*4)
    e1 = numpy.einsum('pq,pq', h1, dm1)
    e2 = numpy.einsum('pqrs,pqrs', eri, dm2) * .5
    print(e1+e2+mol.energy_nuc() - rhf.e_tot - ecc)


# pyscf/cc/ccsd_rdm.py

#!/usr/bin/env python
# Copyright 2014-2021 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: Qiming Sun <osirpt.sun@gmail.com>
#


import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf import ao2mo
from pyscf.cc import ccsd

#
# JCP 95, 2623 (1991); DOI:10.1063/1.460915
# JCP 95, 2639 (1991); DOI:10.1063/1.460916
#

def _gamma1_intermediates(mycc, t1, t2, l1, l2):
    nocc, nvir = t1.shape
    doo =-numpy.einsum('ja,ia->ij', t1, l1)
    dvv = numpy.einsum('ia,ib->ab', t1, l1)
    xtv = numpy.einsum('ie,me->im', t1, l1)
    dvo = t1.T - numpy.einsum('im,ma->ai', xtv, t1)
    theta = t2 * 2 - t2.transpose(0,1,3,2)
    doo -= lib.einsum('jkab,ikab->ij', theta, l2)
    dvv += lib.einsum('jica,jicb->ab', theta, l2)
    xt1  = lib.einsum('mnef,inef->mi', l2, theta)
    xt2  = lib.einsum('mnaf,mnef->ea', l2, theta)
    dvo += numpy.einsum('imae,me->ai', theta, l1)
    dvo -= numpy.einsum('mi,ma->ai', xt1, t1)
    dvo -= numpy.einsum('ie,ae->ai', t1, xt2)
    dov = l1
    return doo, dov, dvo, dvv

# gamma2 intermediates in Chemist's notation
def _gamma2_intermediates(mycc, t1, t2, l1, l2, compress_vvvv=False):
    f = lib.H5TmpFile()
    _gamma2_outcore(mycc, t1, t2, l1, l2, f, compress_vvvv)
    d2 = (f['dovov'][:], f['dvvvv'][:], f['doooo'][:], f['doovv'][:],
          f['dovvo'][:], None,          f['dovvv'][:], f['dooov'][:])
    return d2

def _gamma2_outcore(mycc, t1, t2, l1, l2, h5fobj, compress_vvvv=False):
    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocc, nvir = t1.shape
    nvir_pair = nvir * (nvir+1) //2
    dtype = numpy.result_type(t1, t2, l1, l2).char
    if compress_vvvv:
        dvvvv = h5fobj.create_dataset('dvvvv', (nvir_pair,nvir_pair), dtype)
    else:
        dvvvv = h5fobj.create_dataset('dvvvv', (nvir,nvir,nvir,nvir), dtype)
    dovvo = h5fobj.create_dataset('dovvo', (nocc,nvir,nvir,nocc), dtype,
                                  chunks=(nocc,1,nvir,nocc))
    fswap = lib.H5TmpFile()

    time1 = logger.process_clock(), logger.perf_counter()
    pvOOv = lib.einsum('ikca,jkcb->aijb', l2, t2)
    moo = numpy.einsum('dljd->jl', pvOOv) * 2
    mvv = numpy.einsum('blld->db', pvOOv) * 2
    gooov = lib.einsum('kc,cija->jkia', t1, pvOOv)
    fswap['mvOOv'] = pvOOv
    pvOOv = None

    pvoOV = -lib.einsum('ikca,jkbc->aijb', l2, t2)
    theta = t2 * 2 - t2.transpose(0,1,3,2)
    pvoOV += lib.einsum('ikac,jkbc->aijb', l2, theta)
    moo += numpy.einsum('dljd->jl', pvoOV)
    mvv += numpy.einsum('blld->db', pvoOV)
    gooov -= lib.einsum('jc,cika->jkia', t1, pvoOV)
    fswap['mvoOV'] = pvoOV
    pvoOV = None

    mia =(numpy.einsum('kc,ikac->ia', l1, t2) * 2 -
          numpy.einsum('kc,ikca->ia', l1, t2))
    mab = numpy.einsum('kc,kb->cb', l1, t1)
    mij = numpy.einsum('kc,jc->jk', l1, t1) + moo*.5

    tau = numpy.einsum('ia,jb->ijab', t1, t1)
    tau += t2
    goooo = lib.einsum('ijab,klab->ijkl', tau, l2)*.5
    h5fobj['doooo'] = (goooo.transpose(0,2,1,3)*2 -
                       goooo.transpose(0,3,1,2)).conj()

    gooov += numpy.einsum('ji,ka->jkia', -.5*moo, t1)
    gooov += lib.einsum('la,jkil->jkia', 2*t1, goooo)
    gooov -= lib.einsum('ib,jkba->jkia', l1, tau)
    gooov = gooov.conj()
    gooov -= lib.einsum('jkba,ib->jkia', l2, t1)
    h5fobj['dooov'] = gooov.transpose(0,2,1,3)*2 - gooov.transpose(1,2,0,3)
    tau = None
    time1 = log.timer_debug1('rdm intermediates pass1', *time1)

    goovv = numpy.einsum('ia,jb->ijab', mia.conj(), t1.conj())
    max_memory = max(0, mycc.max_memory - lib.current_memory()[0])
    unit = nocc**2*nvir*6
    blksize = min(nocc, nvir, max(ccsd.BLKMIN, int(max_memory*.95e6/8/unit)))
    doovv = h5fobj.create_dataset('doovv', (nocc,nocc,nvir,nvir), dtype,
                                  chunks=(nocc,nocc,1,nvir))

    log.debug1('rdm intermediates pass 2: block size = %d, nvir = %d in %d blocks',
               blksize, nvir, int((nvir+blksize-1)/blksize))
    for p0, p1 in lib.prange(0, nvir, blksize):
        tau = numpy.einsum('ia,jb->ijab', t1[:,p0:p1], t1)
        tau += t2[:,:,p0:p1]
        tmpoovv  = lib.einsum('ijkl,klab->ijab', goooo, tau)
        tmpoovv -= lib.einsum('jk,ikab->ijab', mij, tau)
        tmpoovv -= lib.einsum('cb,ijac->ijab', mab, t2[:,:,p0:p1])
        tmpoovv -= lib.einsum('bd,ijad->ijab', mvv*.5, tau)
        tmpoovv += .5 * tau
        tmpoovv = tmpoovv.conj()
        tmpoovv += .5 * l2[:,:,p0:p1]
        goovv[:,:,p0:p1] += tmpoovv

        pvOOv = fswap['mvOOv'][p0:p1]
        pvoOV = fswap['mvoOV'][p0:p1]
        gOvvO = lib.einsum('kiac,jc,kb->iabj', l2[:,:,p0:p1], t1, t1)
        gOvvO += numpy.einsum('aijb->iabj', pvOOv)
        govVO = numpy.einsum('ia,jb->iabj', l1[:,p0:p1], t1)
        govVO -= lib.einsum('ikac,jc,kb->iabj', l2[:,:,p0:p1], t1, t1)
        govVO += numpy.einsum('aijb->iabj', pvoOV)
        dovvo[:,p0:p1] = 2*govVO + gOvvO
        doovv[:,:,p0:p1] = (-2*gOvvO - govVO).transpose(3,0,1,2).conj()
        gOvvO = govVO = None

        tau -= t2[:,:,p0:p1] * .5
        for q0, q1 in lib.prange(0, nvir, blksize):
            goovv[:,:,q0:q1,:] += lib.einsum('dlib,jlda->ijab', pvOOv, tau[:,:,:,q0:q1]).conj()
            goovv[:,:,:,q0:q1] -= lib.einsum('dlia,jldb->ijab', pvoOV, tau[:,:,:,q0:q1]).conj()
            tmp = pvoOV[:,:,:,q0:q1] + pvOOv[:,:,:,q0:q1]*.5
            goovv[:,:,q0:q1,:] += lib.einsum('dlia,jlbd->ijab', tmp, t2[:,:,:,p0:p1]).conj()
        pvOOv = pvoOV = tau = None
        time1 = log.timer_debug1('rdm intermediates pass2 [%d:%d]'%(p0, p1), *time1)
    h5fobj['dovov'] = goovv.transpose(0,2,1,3) * 2 - goovv.transpose(1,2,0,3)
    goovv = goooo = None

    max_memory = max(0, mycc.max_memory - lib.current_memory()[0])
    unit = max(nocc**2*nvir*2+nocc*nvir**2*3,
               nvir**3*2+nocc*nvir**2*2+nocc**2*nvir*2)
    blksize = min(nvir, max(ccsd.BLKMIN, int(max_memory*.9e6/8/unit)))
    log.debug1('rdm intermediates pass 3: block size = %d, nvir = %d in %d blocks',
               blksize, nocc, int((nvir+blksize-1)/blksize))
    dovvv = h5fobj.create_dataset('dovvv', (nocc,nvir,nvir,nvir), dtype,
                                  chunks=(nocc,min(nocc,nvir),1,nvir))
    time1 = logger.process_clock(), logger.perf_counter()
    for istep, (p0, p1) in enumerate(lib.prange(0, nvir, blksize)):
        l2tmp = l2[:,:,p0:p1]
        gvvvv = lib.einsum('ijab,ijcd->abcd', l2tmp, t2)
        jabc = lib.einsum('ijab,ic->jabc', l2tmp, t1)
        gvvvv += lib.einsum('jabc,jd->abcd', jabc, t1)
        l2tmp = jabc = None

        if compress_vvvv:
            # symmetrize dvvvv because it does not affect the results of ccsd_grad
            # dvvvv = gvvvv.transpose(0,2,1,3)-gvvvv.transpose(0,3,1,2)*.5
            # dvvvv = (dvvvv+dvvvv.transpose(0,1,3,2)) * .5
            # dvvvv = (dvvvv+dvvvv.transpose(1,0,2,3)) * .5
            # now dvvvv == dvvvv.transpose(0,1,3,2) == dvvvv.transpose(1,0,3,2)
            tmp = numpy.empty((nvir,nvir,nvir))
            tmpvvvv = numpy.empty((p1-p0,nvir,nvir_pair))
            for i in range(p1-p0):
                vvv = gvvvv[i].conj().transpose(1,0,2)
                tmp[:] = vvv - vvv.transpose(2,1,0)*.5
                lib.pack_tril(tmp+tmp.transpose(0,2,1), out=tmpvvvv[i])
            # tril of (dvvvv[p0:p1,p0:p1]+dvvvv[p0:p1,p0:p1].T)
            for i in range(p0, p1):
                for j in range(p0, i):
                    tmpvvvv[i-p0,j] += tmpvvvv[j-p0,i]
                tmpvvvv[i-p0,i] *= 2
            for i in range(p1, nvir):
                off = i * (i+1) // 2
                dvvvv[off+p0:off+p1] = tmpvvvv[:,i]
            for i in range(p0, p1):
                off = i * (i+1) // 2
                if p0 > 0:
                    tmpvvvv[i-p0,:p0] += dvvvv[off:off+p0]
                dvvvv[off:off+i+1] = tmpvvvv[i-p0,:i+1] * .25
            tmp = tmpvvvv = None
        else:
            for i in range(p0, p1):
                vvv = gvvvv[i-p0].conj().transpose(1,0,2)
                dvvvv[i] = vvv - vvv.transpose(2,1,0)*.5

        gvovv = lib.einsum('adbc,id->aibc', gvvvv, -t1)
        gvvvv = None

        gvovv += lib.einsum('akic,kb->aibc', fswap['mvoOV'][p0:p1], t1)
        gvovv -= lib.einsum('akib,kc->aibc', fswap['mvOOv'][p0:p1], t1)

        gvovv += lib.einsum('ja,jibc->aibc', l1[:,p0:p1], t2)
        gvovv += lib.einsum('ja,jb,ic->aibc', l1[:,p0:p1], t1, t1)
        gvovv += numpy.einsum('ba,ic->aibc', mvv[:,p0:p1]*.5, t1)
        gvovv = gvovv.conj()
        gvovv += lib.einsum('ja,jibc->aibc', t1[:,p0:p1], l2)

        dovvv[:,:,p0:p1] = gvovv.transpose(1,3,0,2)*2 - gvovv.transpose(1,2,0,3)
        gvovv = None
        time1 = log.timer_debug1('rdm intermediates pass3 [%d:%d]'%(p0, p1), *time1)

    fswap = None
    dvvov = None
    return (h5fobj['dovov'], h5fobj['dvvvv'], h5fobj['doooo'], h5fobj['doovv'],
            h5fobj['dovvo'], dvvov          , h5fobj['dovvv'], h5fobj['dooov'])

def make_rdm1(mycc, t1, t2, l1, l2, ao_repr=False, with_frozen=True, with_mf=True):
    r'''
    Spin-traced one-particle density matrix in MO basis (the occupied-virtual
    blocks from the orbital response contribution are not included).

    dm1[p,q] = <q_alpha^\dagger p_alpha> + <q_beta^\dagger p_beta>

    The convention of 1-pdm is based on McWeeney's book, Eq (5.4.20).
    The contraction between 1-particle Hamiltonian and rdm1 is
    E = einsum('pq,qp', h1, rdm1)
    '''
    d1 = _gamma1_intermediates(mycc, t1, t2, l1, l2)
    return _make_rdm1(mycc, d1, with_frozen=with_frozen, ao_repr=ao_repr, with_mf=with_mf)

def make_rdm2(mycc, t1, t2, l1, l2, ao_repr=False, with_frozen=True, with_dm1=True):
    r'''
    Spin-traced two-particle density matrix in MO basis

    dm2[p,q,r,s] = \sum_{sigma,tau} <p_sigma^\dagger r_tau^\dagger s_tau q_sigma>

    Note the contraction between ERIs (in Chemist's notation) and rdm2 is
    E = einsum('pqrs,pqrs', eri, rdm2)
    '''
    d1 = _gamma1_intermediates(mycc, t1, t2, l1, l2)
    f = lib.H5TmpFile()
    d2 = _gamma2_outcore(mycc, t1, t2, l1, l2, f, False)
    return _make_rdm2(mycc, d1, d2, with_dm1=with_dm1, with_frozen=with_frozen,
                      ao_repr=ao_repr)

def _make_rdm1(mycc, d1, with_frozen=True, ao_repr=False, with_mf=True):
    r'''dm1[p,q] = <q_alpha^\dagger p_alpha> + <q_beta^\dagger p_beta>

    The convention of 1-pdm is based on McWeeney's book, Eq (5.4.20).
    The contraction between 1-particle Hamiltonian and rdm1 is
    E = einsum('pq,qp', h1, rdm1)
    '''
    doo, dov, dvo, dvv = d1
    nocc, nvir = dov.shape
    nmo = nocc + nvir
    dm1 = numpy.empty((nmo,nmo), dtype=doo.dtype)
    dm1[:nocc,:nocc] = doo + doo.conj().T
    dm1[:nocc,nocc:] = dov + dvo.conj().T
    dm1[nocc:,:nocc] = dm1[:nocc,nocc:].conj().T
    dm1[nocc:,nocc:] = dvv + dvv.conj().T
    if with_mf:
        dm1[numpy.diag_indices(nocc)] += 2

    if with_frozen and mycc.frozen is not None:
        nmo = mycc.mo_occ.size
        nocc = numpy.count_nonzero(mycc.mo_occ > 0)
        rdm1 = numpy.zeros((nmo,nmo), dtype=dm1.dtype)
        if with_mf:
            rdm1[numpy.diag_indices(nocc)] = 2
        moidx = numpy.where(mycc.get_frozen_mask())[0]
        rdm1[moidx[:,None],moidx] = dm1
        dm1 = rdm1

    if ao_repr:
        mo = mycc.mo_coeff
        dm1 = lib.einsum('pi,ij,qj->pq', mo, dm1, mo.conj())
    return dm1

# Note vvvv part of 2pdm have been symmetrized.  It does not correspond to
# vvvv part of CI 2pdm
def _make_rdm2(mycc, d1, d2, with_dm1=True, with_frozen=True, ao_repr=False):
    r'''
    dm2[p,q,r,s] = \sum_{sigma,tau} <p_sigma^\dagger r_tau^\dagger s_tau q_sigma>

    Note the contraction between ERIs (in Chemist's notation) and rdm2 is
    E = einsum('pqrs,pqrs', eri, rdm2)
    '''
    dovov, dvvvv, doooo, doovv, dovvo, dvvov, dovvv, dooov = d2
    nocc, nvir = dovov.shape[:2]
    nmo = nocc + nvir

    dm2 = numpy.empty((nmo,nmo,nmo,nmo), dtype=doovv.dtype)

    dovov = numpy.asarray(dovov)
    dm2[:nocc,nocc:,:nocc,nocc:] = dovov
    dm2[:nocc,nocc:,:nocc,nocc:]+= dovov.transpose(2,3,0,1)
    dm2[nocc:,:nocc,nocc:,:nocc] = dm2[:nocc,nocc:,:nocc,nocc:].transpose(1,0,3,2).conj()
    dovov = None

    doovv = numpy.asarray(doovv)
    dm2[:nocc,:nocc,nocc:,nocc:] = doovv
    dm2[:nocc,:nocc,nocc:,nocc:]+= doovv.transpose(1,0,3,2).conj()
    dm2[nocc:,nocc:,:nocc,:nocc] = dm2[:nocc,:nocc,nocc:,nocc:].transpose(2,3,0,1)
    doovv = None

    dovvo = numpy.asarray(dovvo)
    dm2[:nocc,nocc:,nocc:,:nocc] = dovvo
    dm2[:nocc,nocc:,nocc:,:nocc]+= dovvo.transpose(3,2,1,0).conj()
    dm2[nocc:,:nocc,:nocc,nocc:] = dm2[:nocc,nocc:,nocc:,:nocc].transpose(1,0,3,2).conj()
    dovvo = None

    if dvvvv.ndim == 2:
        # To handle the case of compressed vvvv, which is used in nuclear gradients
        dvvvv = ao2mo.restore(1, dvvvv, nvir)
        dm2[nocc:,nocc:,nocc:,nocc:] = dvvvv
        dm2[nocc:,nocc:,nocc:,nocc:]*= 4
    else:
        dvvvv = numpy.asarray(dvvvv)
        dm2[nocc:,nocc:,nocc:,nocc:] = dvvvv
        dm2[nocc:,nocc:,nocc:,nocc:]+= dvvvv.transpose(1,0,3,2).conj()
        dm2[nocc:,nocc:,nocc:,nocc:]*= 2
    dvvvv = None

    doooo = numpy.asarray(doooo)
    dm2[:nocc,:nocc,:nocc,:nocc] = doooo
    dm2[:nocc,:nocc,:nocc,:nocc]+= doooo.transpose(1,0,3,2).conj()
    dm2[:nocc,:nocc,:nocc,:nocc]*= 2
    doooo = None

    dovvv = numpy.asarray(dovvv)
    dm2[:nocc,nocc:,nocc:,nocc:] = dovvv
    dm2[nocc:,nocc:,:nocc,nocc:] = dovvv.transpose(2,3,0,1)
    dm2[nocc:,nocc:,nocc:,:nocc] = dovvv.transpose(3,2,1,0).conj()
    dm2[nocc:,:nocc,nocc:,nocc:] = dovvv.transpose(1,0,3,2).conj()
    dovvv = None

    dooov = numpy.asarray(dooov)
    dm2[:nocc,:nocc,:nocc,nocc:] = dooov
    dm2[:nocc,nocc:,:nocc,:nocc] = dooov.transpose(2,3,0,1)
    dm2[:nocc,:nocc,nocc:,:nocc] = dooov.transpose(1,0,3,2).conj()
    dm2[nocc:,:nocc,:nocc,:nocc] = dooov.transpose(3,2,1,0).conj()

    if with_frozen and mycc.frozen is not None:
        nmo, nmo0 = mycc.mo_occ.size, nmo
        nocc = numpy.count_nonzero(mycc.mo_occ > 0)
        rdm2 = numpy.zeros((nmo,nmo,nmo,nmo), dtype=dm2.dtype)
        moidx = numpy.where(mycc.get_frozen_mask())[0]
        idx = (moidx.reshape(-1,1) * nmo + moidx).ravel()
        lib.takebak_2d(rdm2.reshape(nmo**2,nmo**2),
                       dm2.reshape(nmo0**2,nmo0**2), idx, idx)
        dm2 = rdm2

    if with_dm1:
        dm1 = _make_rdm1(mycc, d1, with_frozen)
        dm1[numpy.diag_indices(nocc)] -= 2

        for i in range(nocc):
            dm2[i,i,:,:] += dm1 * 2
            dm2[:,:,i,i] += dm1 * 2
            dm2[:,i,i,:] -= dm1
            dm2[i,:,:,i] -= dm1.T

        for i in range(nocc):
            for j in range(nocc):
                dm2[i,i,j,j] += 4
                dm2[i,j,j,i] -= 2

    # dm2 was computed as dm2[p,q,r,s] = < p^\dagger r^\dagger s q > in the
    # above. Transposing it so that it be contracted with ERIs (in Chemist's
    # notation):
    #   E = einsum('pqrs,pqrs', eri, rdm2)
    dm2 = dm2.transpose(1,0,3,2)

    if ao_repr:
        dm2 = _rdm2_mo2ao(dm2.transpose(1,0,3,2), mycc.mo_coeff)
    return dm2


def _rdm2_mo2ao(dm2, mo):
    mo_C = mo.conj()
    return lib.einsum('ijkl,pi,qj,rk,sl->pqrs', dm2, mo, mo_C, mo, mo_C)


if __name__ == '__main__':
    from functools import reduce
    from pyscf import gto
    from pyscf import scf

    mol = gto.M()
    mf = scf.RHF(mol)
    mcc = ccsd.CCSD(mf)

    numpy.random.seed(2)
    nocc = 5
    nmo = 12
    nvir = nmo - nocc
    eri0 = numpy.random.random((nmo,nmo,nmo,nmo))
    eri0 = ao2mo.restore(1, ao2mo.restore(8, eri0, nmo), nmo)
    fock0 = numpy.random.random((nmo,nmo))
    fock0 = fock0 + fock0.T + numpy.diag(range(nmo))*2
    t1 = numpy.random.random((nocc,nvir))
    t2 = numpy.random.random((nocc,nocc,nvir,nvir))
    t2 = t2 + t2.transpose(1,0,3,2)
    l1 = numpy.random.random((nocc,nvir))
    l2 = numpy.random.random((nocc,nocc,nvir,nvir))
    l2 = l2 + l2.transpose(1,0,3,2)
    h1 = fock0 - (numpy.einsum('kkpq->pq', eri0[:nocc,:nocc])*2 -
                  numpy.einsum('pkkq->pq', eri0[:,:nocc,:nocc]))

    eris = lambda:None
    eris.oooo = eri0[:nocc,:nocc,:nocc,:nocc].copy()
    eris.ooov = eri0[:nocc,:nocc,:nocc,nocc:].copy()
    eris.ovoo = eri0[:nocc,nocc:,:nocc,:nocc].copy()
    eris.oovv = eri0[:nocc,:nocc,nocc:,nocc:].copy()
    eris.ovov = eri0[:nocc,nocc:,:nocc,nocc:].copy()
    eris.ovvo = eri0[:nocc,nocc:,nocc:,:nocc].copy()
    eris.ovvv = eri0[:nocc,nocc:,nocc:,nocc:].copy()
    eris.vvvv = eri0[nocc:,nocc:,nocc:,nocc:].copy()
    eris.fock = fock0

    doo, dov, dvo, dvv = _gamma1_intermediates(mcc, t1, t2, l1, l2)
    print((numpy.einsum('ij,ij', doo, fock0[:nocc,:nocc]))*2+20166.329861034799)
    print((numpy.einsum('ab,ab', dvv, fock0[nocc:,nocc:]))*2-58078.964019246778)
    print((numpy.einsum('ai,ia', dvo, fock0[:nocc,nocc:]))*2+74994.356886784764)
    print((numpy.einsum('ia,ai', dov, fock0[nocc:,:nocc]))*2-34.010188025702391)

    fdm2 = lib.H5TmpFile()
    dovov, dvvvv, doooo, doovv, dovvo, dvvov, dovvv, dooov = \
            _gamma2_outcore(mcc, t1, t2, l1, l2, fdm2, True)
    print('dovov', lib.finger(numpy.array(dovov)) - -14384.907042073517)
    print('dvvvv', lib.finger(numpy.array(dvvvv)) - -25.374007033024839)
    print('doooo', lib.finger(numpy.array(doooo)) -  60.114594698129963)
    print('doovv', lib.finger(numpy.array(doovv)) - -79.176348067958401)
    print('dovvo', lib.finger(numpy.array(dovvo)) -   9.864134457251815)
    print('dovvv', lib.finger(numpy.array(dovvv)) - -421.90333700061342)
    print('dooov', lib.finger(numpy.array(dooov)) - -592.66863759586136)
    fdm2 = None

    dovov, dvvvv, doooo, doovv, dovvo, dvvov, dovvv, dooov = \
            _gamma2_intermediates(mcc, t1, t2, l1, l2)
    print('dovov', lib.finger(numpy.array(dovov)) - -14384.907042073517)
    print('dvvvv', lib.finger(numpy.array(dvvvv)) -  45.872344902116758)
    print('doooo', lib.finger(numpy.array(doooo)) -  60.114594698129963)
    print('doovv', lib.finger(numpy.array(doovv)) - -79.176348067958401)
    print('dovvo', lib.finger(numpy.array(dovvo)) -   9.864134457251815)
    print('dovvv', lib.finger(numpy.array(dovvv)) - -421.90333700061342)
    print('dooov', lib.finger(numpy.array(dooov)) - -592.66863759586136)

    print('doooo',numpy.einsum('kilj,kilj', doooo, eris.oooo)*2-15939.9007625418)
    print('dvvvv',numpy.einsum('acbd,acbd', dvvvv, eris.vvvv)*2-37581.823919588 )
    print('dooov',numpy.einsum('jkia,jkia', dooov, eris.ooov)*2-128470.009687716)
    print('dovvv',numpy.einsum('icba,icba', dovvv, eris.ovvv)*2+166794.225195056)
    print('dovov',numpy.einsum('iajb,iajb', dovov, eris.ovov)*2+719279.812916893)
    print('dovvo',numpy.einsum('jbai,jbia', dovvo, eris.ovov)*2
                 +numpy.einsum('jiab,jiba', doovv, eris.oovv)*2+53634.0012286654)

    dm1 = make_rdm1(mcc, t1, t2, l1, l2)
    dm2 = make_rdm2(mcc, t1, t2, l1, l2)
    e2 = (numpy.einsum('ijkl,ijkl', doooo, eris.oooo)*2 +
          numpy.einsum('acbd,acbd', dvvvv, eris.vvvv)*2 +
          numpy.einsum('jkia,jkia', dooov, eris.ooov)*2 +
          numpy.einsum('icba,icba', dovvv, eris.ovvv)*2 +
          numpy.einsum('iajb,iajb', dovov, eris.ovov)*2 +
          numpy.einsum('jbai,jbia', dovvo, eris.ovov)*2 +
          numpy.einsum('ijab,ijab', doovv, eris.oovv)*2 +
          numpy.einsum('ij,ij', doo, fock0[:nocc,:nocc])*2 +
          numpy.einsum('ia,ia', dov, fock0[:nocc,nocc:])*2 +
          numpy.einsum('ai,ai', dvo, fock0[nocc:,:nocc])*2 +
          numpy.einsum('ab,ab', dvv, fock0[nocc:,nocc:])*2 +
          fock0[:nocc].trace()*2 -
          numpy.einsum('kkpq->pq', eri0[:nocc,:nocc,:nocc,:nocc]).trace()*2 +
          numpy.einsum('pkkq->pq', eri0[:nocc,:nocc,:nocc,:nocc]).trace())
    print(e2+794721.197459942)
    print(numpy.einsum('pqrs,pqrs', dm2, eri0)*.5 +
          numpy.einsum('pq,qp', dm1, h1) - e2)

    print(numpy.allclose(dm2, dm2.transpose(1,0,3,2)))
    print(numpy.allclose(dm2, dm2.transpose(2,3,0,1)))

    d1 = numpy.einsum('kkpq->qp', dm2) / 9
    print(numpy.allclose(d1, dm1))

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = '631g'
    mol.build()
    mf = scf.RHF(mol).run()

    mycc = ccsd.CCSD(mf)
    mycc.frozen = 2
    ecc, t1, t2 = mycc.kernel()
    l1, l2 = mycc.solve_lambda()
    dm1 = make_rdm1(mycc, t1, t2, l1, l2)
    dm2 = make_rdm2(mycc, t1, t2, l1, l2)
    nmo = mf.mo_coeff.shape[1]
    eri = ao2mo.kernel(mf._eri, mf.mo_coeff, compact=False).reshape([nmo]*4)
    hcore = mf.get_hcore()
    h1 = reduce(numpy.dot, (mf.mo_coeff.T, hcore, mf.mo_coeff))
    e1 = numpy.einsum('ij,ji', h1, dm1)
    e1+= numpy.einsum('ijkl,ijkl', eri, dm2) * .5
    e1+= mol.energy_nuc()
    print(e1 - mycc.e_tot)


# pyscf/cc/ccsd_rdm_slow.py

#!/usr/bin/env python
# Copyright 2014-2021 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: Qiming Sun <osirpt.sun@gmail.com>
#         Carlos Jimenez-Hoyos
#

import numpy
from pyscf.cc import ccsd_rdm

def _gamma1_intermediates(cc, t1, t2, l1, l2):
    nocc, nvir = t1.shape

    t1a  = t1
    t2ab = numpy.copy(t2)
    t2aa = numpy.copy(t2) \
         - t2.transpose(0,1,3,2)

    l1a  = l1
    l2ab = 2*numpy.copy(l2)
    l2aa = numpy.copy(l2) \
         - l2.transpose(0,1,3,2)

    doo  = numpy.zeros((nocc,nocc))
    doo += -2*numpy.einsum('ie,je->ij', t1a, l1a)
    doo +=   -numpy.einsum('imef,jmef->ij', t2ab, l2ab) \
             -numpy.einsum('imef,jmef->ij', t2aa, l2aa)

    dvv  = numpy.zeros((nvir,nvir))
    dvv +=  2*numpy.einsum('ma,mb->ab', l1a, t1a)
    dvv +=    numpy.einsum('mnae,mnbe->ab', l2ab, t2ab) \
         +    numpy.einsum('mnae,mnbe->ab', l2aa, t2aa)

    xt1  = numpy.einsum('mnef,inef->mi', l2aa, t2aa)
    xt1 += numpy.einsum('mnef,inef->mi', l2ab, t2ab)
    xt2  = numpy.einsum('mnaf,mnef->ae', t2aa, l2aa)
    xt2 += numpy.einsum('mnaf,mnef->ae', t2ab, l2ab)
    xtv  = numpy.einsum('ma,me->ae', t1a, l1a)

    dov  = numpy.zeros((nocc,nvir))
    dov +=  2*t1a
    dov +=  2*numpy.einsum('imae,me->ia', t2aa, l1a) \
         +  2*numpy.einsum('imae,me->ia', t2ab, l1a) \
         + -2*numpy.einsum('ie,ae->ia', t1a, xtv)
    dov +=   -numpy.einsum('mi,ma->ia', xt1, t1a) \
         +   -numpy.einsum('ie,ae->ia', t1a, xt2)

    dvo  = numpy.zeros((nvir,nocc))
    dvo += 2*l1a.transpose(1,0)

    return doo*.5, dov*.5, dvo*.5, dvv*.5

# gamma2 intermediates in Chemist's notation
def _gamma2_intermediates(cc, t1, t2, l1, l2):
    tau = t2 + numpy.einsum('ia,jb->ijab', t1, t1)
    tau2 = t2 + numpy.einsum('ia,jb->ijab', t1, t1*2)
    theta = t2*2 - t2.transpose(0,1,3,2)

    mOvOv = numpy.einsum('ikca,jkcb->jbia', l2, t2)
    mOVov = numpy.einsum('ikac,jkbc->jbia', l2, theta)
    mOVov -= numpy.einsum('ikca,jkbc->jbia', l2, t2)
    moo =(numpy.einsum('jdld->jl', mOvOv) * 2 +
          numpy.einsum('jdld->jl', mOVov))
    mvv =(numpy.einsum('lbld->bd', mOvOv) * 2 +
          numpy.einsum('lbld->bd', mOVov))

    gvvvv = numpy.einsum('ijab,ijcd->abcd', l2*.5, tau)

    goooo = numpy.einsum('ijab,klab->klij', l2, tau)*.5

    goovv = .5 * l2 + .5 * tau
    tmp = numpy.einsum('kc,ikac->ia', l1, theta)
    goovv += numpy.einsum('ia,jb->ijab', tmp, t1)
    tmp = numpy.einsum('kc,kb->cb', l1, t1)
    goovv -= numpy.einsum('cb,ijac->ijab', tmp, t2)
    tmp = numpy.einsum('kc,jc->kj', l1, t1)
    goovv -= numpy.einsum('kj,ikab->ijab', tmp, tau)
    goovv -= numpy.einsum('jl,ilab->ijab', moo*.5, tau)
    goovv -= numpy.einsum('bd,ijad->ijab', mvv*.5, tau)
    goovv += numpy.einsum('ibld,ljad->ijab', mOvOv, tau2) * .5
    goovv -= numpy.einsum('iald,ljbd->ijab', mOVov, tau2) * .5
    goovv += numpy.einsum('iald,ljdb->ijab', mOVov*2+mOvOv, t2) * .5
    goovv += numpy.einsum('ijkl,klab->ijab', goooo, tau)

    gooov = numpy.einsum('ib,kjab->jkia', -l1, tau)
    gooov += numpy.einsum('jkil,la->jkia', goooo, t1*2)
    gooov += numpy.einsum('ji,ka->jkia', moo*-.5, t1)
    gooov += numpy.einsum('jaic,kc->jkia', mOvOv, t1)
    gooov -= numpy.einsum('kaic,jc->jkia', mOVov, t1)
    gooov -= numpy.einsum('jkba,ib->jkia', l2, t1)

    govvv = numpy.einsum('ja,jibc->iacb', l1, tau)
    govvv -= numpy.einsum('adbc,id->iacb', gvvvv, t1*2)
    govvv += numpy.einsum('ba,ic->iacb', mvv, t1*.5)
    govvv -= numpy.einsum('ibka,kc->iacb', mOvOv, t1)
    govvv += numpy.einsum('icka,kb->iacb', mOVov, t1)
    govvv += numpy.einsum('jibc,ja->iacb', l2, t1)

    gOvVo = numpy.einsum('ia,jb->jabi', l1, t1) + mOVov.transpose(0,3,1,2)
    tmp = numpy.einsum('ikac,jc->jaik', l2, t1)
    gOvVo -= numpy.einsum('jaik,kb->jabi', tmp, t1)
    gOvvO = mOvOv.transpose(0,3,1,2) + numpy.einsum('jaki,kb->jabi', tmp, t1)

    doovv = goovv*2 - goovv.transpose(0,1,3,2)
    dvvvv = gvvvv*2 - gvvvv.transpose(0,1,3,2)
    doooo = goooo*2 - goooo.transpose(0,1,3,2)
    dovov = -2*gOvvO.transpose(0,1,3,2) - gOvVo.transpose(0,1,3,2)
    dovvo = gOvVo*2 + gOvvO
    dovvv = govvv*2 - govvv.transpose(0,1,3,2)
    dooov = gooov*2 - gooov.transpose(1,0,2,3)

    doovv, dovov = dovov.transpose(0,2,1,3), doovv.transpose(0,2,1,3)
    dvvvv = dvvvv.transpose(0,2,1,3)
    doooo = doooo.transpose(0,2,1,3)
    dovvo = dovvo.transpose(0,2,1,3)
    dovvv = dovvv.transpose(0,2,1,3)
    dooov = dooov.transpose(0,2,1,3)
    dvvov = None
    return (dovov, dvvvv, doooo, doovv, dovvo, dvvov, dovvv, dooov)

def make_rdm1(cc, t1, t2, l1, l2):
    d1 = _gamma1_intermediates(cc, t1, t2, l1, l2)
    return ccsd_rdm._make_rdm1(cc, d1, with_frozen=True)

def make_rdm2(cc, t1, t2, l1, l2, d1=None, d2=None):
    d1 = _gamma1_intermediates(cc, t1, t2, l1, l2)
    d2 = _gamma2_intermediates(cc, t1, t2, l1, l2)
    return ccsd_rdm._make_rdm2(cc, d1, d2, with_dm1=True, with_frozen=True)


if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf
    from pyscf.cc import ccsd
    from pyscf import ao2mo

    mol = gto.M()
    mf = scf.RHF(mol)

    mcc = ccsd.CCSD(mf)

    numpy.random.seed(2)
    nocc = 5
    nmo = 12
    nvir = nmo - nocc
    eri0 = numpy.random.random((nmo,nmo,nmo,nmo))
    eri0 = ao2mo.restore(1, ao2mo.restore(8, eri0, nmo), nmo)
    fock0 = numpy.random.random((nmo,nmo))
    fock0 = fock0 + fock0.T + numpy.diag(range(nmo))*2
    t1 = numpy.random.random((nocc,nvir))
    t2 = numpy.random.random((nocc,nocc,nvir,nvir))
    t2 = t2 + t2.transpose(1,0,3,2)
    l1 = numpy.random.random((nocc,nvir))
    l2 = numpy.random.random((nocc,nocc,nvir,nvir))
    l2 = l2 + l2.transpose(1,0,3,2)
    h1 = fock0 - (numpy.einsum('kkpq->pq', eri0[:nocc,:nocc])*2 -
                  numpy.einsum('pkkq->pq', eri0[:,:nocc,:nocc]))

    eris = lambda:None
    eris.oooo = eri0[:nocc,:nocc,:nocc,:nocc].copy()
    eris.ooov = eri0[:nocc,:nocc,:nocc,nocc:].copy()
    eris.ovoo = eri0[:nocc,nocc:,:nocc,:nocc].copy()
    eris.oovv = eri0[:nocc,:nocc,nocc:,nocc:].copy()
    eris.ovov = eri0[:nocc,nocc:,:nocc,nocc:].copy()
    eris.ovvo = eri0[:nocc,nocc:,nocc:,:nocc].copy()
    eris.ovvv = eri0[:nocc,nocc:,nocc:,nocc:].copy()
    eris.vvvv = eri0[nocc:,nocc:,nocc:,nocc:].copy()
    eris.fock = fock0

    doo, dov, dvo, dvv = _gamma1_intermediates(mcc, t1, t2, l1, l2)
    print((numpy.einsum('ij,ij', doo, fock0[:nocc,:nocc]))*2+20166.329861034799)
    print((numpy.einsum('ab,ab', dvv, fock0[nocc:,nocc:]))*2-58078.964019246778)
    print((numpy.einsum('ia,ia', dov, fock0[:nocc,nocc:]))*2+74994.356886784764)
    print((numpy.einsum('ai,ai', dvo, fock0[nocc:,:nocc]))*2-34.010188025702391)

    dovov, dvvvv, doooo, doovv, dovvo, dvvov, dovvv, dooov = \
            _gamma2_intermediates(mcc, t1, t2, l1, l2)

    print('doooo',numpy.einsum('ijkl,ijkl', doooo, eris.oooo)*2-15939.9007625418)
    print('dvvvv',numpy.einsum('acbd,acbd', dvvvv, eris.vvvv)*2-37581.823919588 )
    print('dooov',numpy.einsum('jkia,jkia', dooov, eris.ooov)*2-128470.009687716)
    print('dovvv',numpy.einsum('icab,icab', dovvv, eris.ovvv)*2+166794.225195056)
    print('doovv',numpy.einsum('iajb,iajb', dovov, eris.ovov)*2+719279.812916893)
    print('dovvo',numpy.einsum('jbai,jbia', dovvo, eris.ovov)*2
                 +numpy.einsum('ijab,ijab', doovv, eris.oovv)*2+53634.0012286654)

    dm1 = make_rdm1(mcc, t1, t2, l1, l2)
    dm2 = make_rdm2(mcc, t1, t2, l1, l2)
    e2 = (numpy.einsum('ijkl,ijkl', doooo, eris.oooo)*2 +
          numpy.einsum('acbd,acbd', dvvvv, eris.vvvv)*2 +
          numpy.einsum('jkia,jkia', dooov, eris.ooov)*2 +
          numpy.einsum('icab,icab', dovvv, eris.ovvv)*2 +
          numpy.einsum('iajb,iajb', dovov, eris.ovov)*2 +
          numpy.einsum('jbai,jbia', dovvo, eris.ovov)*2 +
          numpy.einsum('ijab,ijab', doovv, eris.oovv)*2 +
          numpy.einsum('ij,ij', doo, fock0[:nocc,:nocc])*2 +
          numpy.einsum('ia,ia', dov, fock0[:nocc,nocc:])*2 +
          numpy.einsum('ai,ai', dvo, fock0[nocc:,:nocc])*2 +
          numpy.einsum('ab,ab', dvv, fock0[nocc:,nocc:])*2 +
          fock0[:nocc].trace()*2 -
          numpy.einsum('kkpq->pq', eri0[:nocc,:nocc,:nocc,:nocc]).trace()*2 +
          numpy.einsum('pkkq->pq', eri0[:nocc,:nocc,:nocc,:nocc]).trace())
    print(e2+794721.197459942)
    print(numpy.einsum('pqrs,pqrs', dm2, eri0)*.5 +
          numpy.einsum('pq,qp', dm1, h1) - e2)

    print(numpy.allclose(dm2, dm2.transpose(1,0,3,2)))
    print(numpy.allclose(dm2, dm2.transpose(2,3,0,1)))

    d1 = numpy.einsum('kkpq->qp', dm2) / 9
    print(numpy.allclose(d1, dm1))


# pyscf/cc/ccsd_t.py

#!/usr/bin/env python
# Copyright 2014-2020 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: Qiming Sun <osirpt.sun@gmail.com>
#

'''
RHF-CCSD(T) for real integrals
'''


import ctypes
import numpy
from pyscf import lib
from pyscf import symm
from pyscf.lib import logger
from pyscf.cc import _ccsd

# t3 as ijkabc

# JCP 94, 442 (1991); DOI:10.1063/1.460359.  Error in Eq (1), should be [ia] >= [jb] >= [kc]
def kernel(mycc, eris, t1=None, t2=None, verbose=logger.NOTE):
    cpu1 = cpu0 = (logger.process_clock(), logger.perf_counter())
    log = logger.new_logger(mycc, verbose)
    if t1 is None: t1 = mycc.t1
    if t2 is None: t2 = mycc.t2

    nocc, nvir = t1.shape
    nmo = nocc + nvir

    dtype = numpy.result_type(t1, t2, eris.ovoo.dtype)
    if mycc.incore_complete:
        ftmp = None
        eris_vvop = numpy.zeros((nvir,nvir,nocc,nmo), dtype)
    else:
        ftmp = lib.H5TmpFile()
        eris_vvop = ftmp.create_dataset('vvop', (nvir,nvir,nocc,nmo), dtype)

    orbsym = _sort_eri(mycc, eris, nocc, nvir, eris_vvop, log)

    mo_energy, t1T, t2T, vooo, fvo, restore_t2_inplace = \
            _sort_t2_vooo_(mycc, orbsym, t1, t2, eris)
    cpu1 = log.timer_debug1('CCSD(T) sort_eri', *cpu1)

    cpu2 = list(cpu1)
    orbsym = numpy.hstack((numpy.sort(orbsym[:nocc]),numpy.sort(orbsym[nocc:])))
    o_ir_loc = numpy.append(0, numpy.cumsum(numpy.bincount(orbsym[:nocc], minlength=8)))
    v_ir_loc = numpy.append(0, numpy.cumsum(numpy.bincount(orbsym[nocc:], minlength=8)))
    o_sym = orbsym[:nocc]
    oo_sym = (o_sym[:,None] ^ o_sym).ravel()
    oo_ir_loc = numpy.append(0, numpy.cumsum(numpy.bincount(oo_sym, minlength=8)))
    nirrep = max(oo_sym) + 1

    orbsym   = orbsym.astype(numpy.int32)
    o_ir_loc = o_ir_loc.astype(numpy.int32)
    v_ir_loc = v_ir_loc.astype(numpy.int32)
    oo_ir_loc = oo_ir_loc.astype(numpy.int32)
    if dtype == numpy.complex128:
        drv = _ccsd.libcc.CCsd_t_zcontract
    else:
        drv = _ccsd.libcc.CCsd_t_contract
    et_sum = numpy.zeros(1, dtype=dtype)
    def contract(a0, a1, b0, b1, cache):
        cache_row_a, cache_col_a, cache_row_b, cache_col_b = cache
        drv(et_sum.ctypes.data_as(ctypes.c_void_p),
            mo_energy.ctypes.data_as(ctypes.c_void_p),
            t1T.ctypes.data_as(ctypes.c_void_p),
            t2T.ctypes.data_as(ctypes.c_void_p),
            vooo.ctypes.data_as(ctypes.c_void_p),
            fvo.ctypes.data_as(ctypes.c_void_p),
            ctypes.c_int(nocc), ctypes.c_int(nvir),
            ctypes.c_int(a0), ctypes.c_int(a1),
            ctypes.c_int(b0), ctypes.c_int(b1),
            ctypes.c_int(nirrep),
            o_ir_loc.ctypes.data_as(ctypes.c_void_p),
            v_ir_loc.ctypes.data_as(ctypes.c_void_p),
            oo_ir_loc.ctypes.data_as(ctypes.c_void_p),
            orbsym.ctypes.data_as(ctypes.c_void_p),
            cache_row_a.ctypes.data_as(ctypes.c_void_p),
            cache_col_a.ctypes.data_as(ctypes.c_void_p),
            cache_row_b.ctypes.data_as(ctypes.c_void_p),
            cache_col_b.ctypes.data_as(ctypes.c_void_p))
        cpu2[:] = log.timer_debug1('contract %d:%d,%d:%d'%(a0,a1,b0,b1), *cpu2)

    # The rest 20% memory for cache b
    mem_now = lib.current_memory()[0]
    max_memory = max(0, mycc.max_memory - mem_now)
    bufsize = (max_memory*.5e6/8-nocc**3*3*lib.num_threads())/(nocc*nmo)  #*.5 for async_io
    bufsize *= .5  #*.5 upper triangular part is loaded
    bufsize *= .8  #*.8 for [a0:a1]/[b0:b1] partition
    bufsize = max(8, bufsize)
    log.debug('max_memory %d MB (%d MB in use)', max_memory, mem_now)
    with lib.call_in_background(contract, sync=not mycc.async_io) as async_contract:
        for a0, a1 in reversed(list(lib.prange_tril(0, nvir, bufsize))):
            cache_row_a = numpy.asarray(eris_vvop[a0:a1,:a1], order='C')
            if a0 == 0:
                cache_col_a = cache_row_a
            else:
                cache_col_a = numpy.asarray(eris_vvop[:a0,a0:a1], order='C')
            async_contract(a0, a1, a0, a1, (cache_row_a,cache_col_a,
                                            cache_row_a,cache_col_a))

            for b0, b1 in lib.prange_tril(0, a0, bufsize/8):
                cache_row_b = numpy.asarray(eris_vvop[b0:b1,:b1], order='C')
                if b0 == 0:
                    cache_col_b = cache_row_b
                else:
                    cache_col_b = numpy.asarray(eris_vvop[:b0,b0:b1], order='C')
                async_contract(a0, a1, b0, b1, (cache_row_a,cache_col_a,
                                                cache_row_b,cache_col_b))

    t2 = restore_t2_inplace(t2T)
    et_sum *= 2
    if abs(et_sum[0].imag) > 1e-4:
        logger.warn(mycc, 'Non-zero imaginary part of CCSD(T) energy was found %s',
                    et_sum[0])
    et = et_sum[0].real
    log.timer('CCSD(T)', *cpu0)
    log.note('CCSD(T) correction = %.15g', et)
    return et

def _sort_eri(mycc, eris, nocc, nvir, vvop, log):
    cpu1 = (logger.process_clock(), logger.perf_counter())
    mol = mycc.mol
    nmo = nocc + nvir

    if mol.symmetry:
        orbsym = symm.addons.label_orb_symm(mol, mol.irrep_id, mol.symm_orb,
                                            eris.mo_coeff, check=False)
        orbsym = numpy.asarray(orbsym, dtype=numpy.int32) % 10
    else:
        orbsym = numpy.zeros(nmo, dtype=numpy.int32)

    o_sorted = _irrep_argsort(orbsym[:nocc])
    v_sorted = _irrep_argsort(orbsym[nocc:])
    vrank = numpy.argsort(v_sorted)

    max_memory = max(0, mycc.max_memory - lib.current_memory()[0])
    max_memory = min(8000, max_memory*.9)
    blksize = min(nvir, max(16, int(max_memory*1e6/8/(nvir*nocc*nmo))))
    log.debug1('_sort_eri max_memory %g  blksize %d', max_memory, blksize)
    dtype = vvop.dtype
    with lib.call_in_background(vvop.__setitem__, sync=not mycc.async_io) as save:
        bufopv = numpy.empty((nocc,nmo,nvir), dtype=dtype)
        buf1 = numpy.empty_like(bufopv)
        for j0, j1 in lib.prange(0, nvir, blksize):
            ovov = numpy.asarray(eris.ovov[:,j0:j1])
            #ovvv = numpy.asarray(eris.ovvv[:,j0:j1])
            ovvv = eris.get_ovvv(slice(None), slice(j0,j1))
            for j in range(j0,j1):
                oov = ovov[o_sorted,j-j0]
                ovv = ovvv[o_sorted,j-j0]
                #if ovv.ndim == 2:
                #    ovv = lib.unpack_tril(ovv, out=buf)
                bufopv[:,:nocc,:] = oov[:,o_sorted][:,:,v_sorted].conj()
                bufopv[:,nocc:,:] = ovv[:,v_sorted][:,:,v_sorted].conj()
                save(vrank[j], bufopv.transpose(2,0,1))
                bufopv, buf1 = buf1, bufopv
            cpu1 = log.timer_debug1('transpose %d:%d'%(j0,j1), *cpu1)

    return orbsym

def _sort_t2_vooo_(mycc, orbsym, t1, t2, eris):
    assert(t2.flags.c_contiguous)
    vooo = numpy.asarray(eris.ovoo).transpose(1,0,3,2).conj().copy()
    nocc, nvir = t1.shape
    if mycc.mol.symmetry:
        orbsym = numpy.asarray(orbsym, dtype=numpy.int32)
        o_sorted = _irrep_argsort(orbsym[:nocc])
        v_sorted = _irrep_argsort(orbsym[nocc:])
        mo_energy = eris.mo_energy
        mo_energy = numpy.hstack((mo_energy[:nocc][o_sorted],
                                  mo_energy[nocc:][v_sorted]))
        t1T = numpy.asarray(t1.T[v_sorted][:,o_sorted], order='C')
        fvo = eris.fock[nocc:,:nocc]
        fvo = numpy.asarray(fvo[v_sorted][:,o_sorted], order='C')

        o_sym = orbsym[o_sorted]
        oo_sym = (o_sym[:,None] ^ o_sym).ravel()
        oo_sorted = _irrep_argsort(oo_sym)
        #:vooo = eris.ovoo.transpose(1,0,2,3)
        #:vooo = vooo[v_sorted][:,o_sorted][:,:,o_sorted][:,:,:,o_sorted]
        #:vooo = vooo.reshape(nvir,-1,nocc)[:,oo_sorted]
        oo_idx = numpy.arange(nocc**2).reshape(nocc,nocc)[o_sorted][:,o_sorted]
        oo_idx = oo_idx.ravel()[oo_sorted]
        oo_idx = (oo_idx[:,None]*nocc+o_sorted).ravel()
        vooo = lib.take_2d(vooo.reshape(nvir,-1), v_sorted, oo_idx)
        vooo = vooo.reshape(nvir,nocc,nocc,nocc)

        #:t2T = t2.transpose(2,3,1,0)
        #:t2T = ref_t2T[v_sorted][:,v_sorted][:,:,o_sorted][:,:,:,o_sorted]
        #:t2T = ref_t2T.reshape(nvir,nvir,-1)[:,:,oo_sorted]
        t2T = lib.transpose(t2.reshape(nocc**2,-1))
        oo_idx = numpy.arange(nocc**2).reshape(nocc,nocc).T[o_sorted][:,o_sorted]
        oo_idx = oo_idx.ravel()[oo_sorted]
        vv_idx = (v_sorted[:,None]*nvir+v_sorted).ravel()
        t2T = lib.take_2d(t2T.reshape(nvir**2,nocc**2), vv_idx, oo_idx, out=t2)
        t2T = t2T.reshape(nvir,nvir,nocc,nocc)
        def restore_t2_inplace(t2T):
            tmp = numpy.zeros((nvir**2,nocc**2), dtype=t2T.dtype)
            lib.takebak_2d(tmp, t2T.reshape(nvir**2,nocc**2), vv_idx, oo_idx)
            t2 = lib.transpose(tmp.reshape(nvir**2,nocc**2), out=t2T)
            return t2.reshape(nocc,nocc,nvir,nvir)
    else:
        fvo = eris.fock[nocc:,:nocc].copy()
        t1T = t1.T.copy()
        t2T = lib.transpose(t2.reshape(nocc**2,nvir**2))
        t2T = lib.transpose(t2T.reshape(nvir**2,nocc,nocc), axes=(0,2,1), out=t2)
        mo_energy = numpy.asarray(eris.mo_energy, order='C')
        def restore_t2_inplace(t2T):
            tmp = lib.transpose(t2T.reshape(nvir**2,nocc,nocc), axes=(0,2,1))
            t2 = lib.transpose(tmp.reshape(nvir**2,nocc**2), out=t2T)
            return t2.reshape(nocc,nocc,nvir,nvir)
    t2T = t2T.reshape(nvir,nvir,nocc,nocc)
    return mo_energy, t1T, t2T, vooo, fvo, restore_t2_inplace

def _irrep_argsort(orbsym):
    return numpy.hstack([numpy.where(orbsym == i)[0] for i in range(8)])


if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf
    from pyscf import cc

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -.957 , .587)],
        [1 , (0.2,  .757 , .487)]]

    mol.basis = 'ccpvdz'
    mol.build()
    rhf = scf.RHF(mol)
    rhf.conv_tol = 1e-14
    rhf.scf()
    mcc = cc.CCSD(rhf)
    mcc.conv_tol = 1e-14
    mcc.ccsd()
    e3a = kernel(mcc, mcc.ao2mo())
    print(e3a - -0.0033300722704016289)

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -.757 , .587)],
        [1 , (0. ,  .757 , .587)]]
    mol.symmetry = True

    mol.basis = 'ccpvdz'
    mol.build()
    rhf = scf.RHF(mol)
    rhf.conv_tol = 1e-14
    rhf.scf()
    mcc = cc.CCSD(rhf)
    mcc.conv_tol = 1e-14
    mcc.ccsd()
    e3a = kernel(mcc, mcc.ao2mo())
    print(e3a - -0.003060022611584471)


# pyscf/cc/ccsd_t_lambda_slow.py

#!/usr/bin/env python
# Copyright 2014-2021 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: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Spin-free lambda equation of RHF-CCSD(T)

Ref:
JCP 147, 044104 (2017); DOI:10.1063/1.4994918
'''

import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf.cc import ccsd
from pyscf.cc import ccsd_lambda

# Note: not support fov != 0

def kernel(mycc, eris=None, t1=None, t2=None, l1=None, l2=None,
           max_cycle=50, tol=1e-8, verbose=logger.INFO):
    return ccsd_lambda.kernel(mycc, eris, t1, t2, l1, l2, max_cycle, tol,
                              verbose, make_intermediates, update_lambda)

def make_intermediates(mycc, t1, t2, eris):
    imds = ccsd_lambda.make_intermediates(mycc, t1, t2, eris)

    nocc, nvir = t1.shape
    eris_ovvv = numpy.asarray(eris.get_ovvv())
    eris_ovoo = numpy.asarray(eris.ovoo)
    eris_ovov = numpy.asarray(eris.ovov)

    mo_e = eris.mo_energy
    eia = lib.direct_sum('i-a->ia', mo_e[:nocc], mo_e[nocc:])
    d3 = lib.direct_sum('ia,jb,kc->ijkabc', eia, eia, eia)

    def p6(t):
        t1 = t + t.transpose(0,2,1,3,5,4)
        return t1 + t1.transpose(1,0,2,4,3,5) + t1.transpose(1,2,0,4,5,3)

    def r6(w):
        return (4 * w + w.transpose(0,1,2,4,5,3) + w.transpose(0,1,2,5,3,4)
                - 2 * w.transpose(0,1,2,5,4,3) - 2 * w.transpose(0,1,2,3,5,4)
                - 2 * w.transpose(0,1,2,4,3,5))

    w =(numpy.einsum('iafb,kjcf->ijkabc', eris_ovvv.conj(), t2) -
        numpy.einsum('iajm,mkbc->ijkabc', eris_ovoo.conj(), t2)) / d3
    v =(numpy.einsum('iajb,kc->ijkabc', eris_ovov.conj(), t1) +
        numpy.einsum('ck,ijab->ijkabc', eris.fock[nocc:,:nocc], t2)) / d3
    w = p6(w)
    v = p6(v)

    imds.l1_t = numpy.einsum('jbkc,ijkabc->ia', eris_ovov, r6(w)).conj() / eia * .5

    def as_r6(m):
        # When making derivative over t2, r6 should be called on the 6-index
        # tensor. It gives the equation for lambda2, but not corresponding to
        # the lambda equation used by RCCSD-lambda code.  A transformation was
        # applied in RCCSD-lambda equation  F(lambda)_{ijab} = 0:
        #       2/3 * # F(lambda)_{ijab} + 1/3 * F(lambda)_{jiab} = 0
        # Combining this transformation with r6 operation, leads to the
        # transformation code below
        return m * 2 - m.transpose(0,1,2,5,4,3) - m.transpose(0,1,2,3,5,4)

    m = as_r6(w * 2 + v * .5)
    joovv = numpy.einsum('kfbe,ijkaef->ijab', eris_ovvv, m.conj())
    joovv-= numpy.einsum('ncmj,imnabc->ijab', eris_ovoo, m.conj())
    joovv = joovv + joovv.transpose(1,0,3,2)
    rw = as_r6(w)
    joovv+= numpy.einsum('kc,ijkabc->ijab', eris.fock[:nocc,nocc:], rw.conj())
    imds.l2_t = joovv / lib.direct_sum('ia+jb->ijab', eia, eia)

    return imds

def update_lambda(mycc, t1, t2, l1, l2, eris=None, imds=None):
    if eris is None: eris = mycc.ao2mo()
    if imds is None: imds = make_intermediates(mycc, t1, t2, eris)
    l1, l2 = ccsd_lambda.update_lambda(mycc, t1, t2, l1, l2, eris, imds)
    l1 += imds.l1_t
    l2 += imds.l2_t
    return l1, l2


if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf

    mol = gto.Mole()
    mol.verbose = 0
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]

    mol.basis = 'cc-pvdz'
    mol.build()
    rhf = scf.RHF(mol)
    rhf.conv_tol = 1e-16
    rhf.scf()

    mcc = ccsd.CCSD(rhf)
    mcc.conv_tol = 1e-12
    ecc, t1, t2 = mcc.kernel()
    #l1, l2 = mcc.solve_lambda()
    #print(numpy.linalg.norm(l1)-0.0132626841292)
    #print(numpy.linalg.norm(l2)-0.212575609057)

    conv, l1, l2 = kernel(mcc, mcc.ao2mo(), t1, t2, tol=1e-8)
    print(numpy.linalg.norm(l1)-0.013575484203926739)
    print(numpy.linalg.norm(l2)-0.22029981372536928)



# pyscf/cc/ccsd_t_rdm_slow.py

#!/usr/bin/env python
# Copyright 2014-2021 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: Qiming Sun <osirpt.sun@gmail.com>
#

import numpy
from pyscf import lib
from pyscf.cc import ccsd_rdm

def _gamma1_intermediates(mycc, t1, t2, l1, l2, eris=None, for_grad=False):
    doo, dov, dvo, dvv = ccsd_rdm._gamma1_intermediates(mycc, t1, t2, l1, l2)

    if eris is None: eris = mycc.ao2mo()
    nocc, nvir = t1.shape
    eris_ovvv = numpy.asarray(eris.get_ovvv())
    eris_ovoo = numpy.asarray(eris.ovoo)
    eris_ovov = numpy.asarray(eris.ovov)

    mo_e = eris.mo_energy
    eia = lib.direct_sum('i-a->ia', mo_e[:nocc], mo_e[nocc:])
    d3 = lib.direct_sum('ia,jb,kc->ijkabc', eia, eia, eia)

    w = (numpy.einsum('iafb,kjcf->ijkabc', eris_ovvv.conj(), t2) -
         numpy.einsum('iajm,mkbc->ijkabc', eris_ovoo.conj(), t2)) / d3
    v = (numpy.einsum('iajb,kc->ijkabc', eris_ovov.conj(), t1) +
         numpy.einsum('ck,ijab->ijkabc', eris.fock[nocc:,:nocc], t2)) / d3
    w = p6(w)
    v = p6(v)
    wv = w + v * .5
    rw = r6(w)
    goo = numpy.einsum('iklabc,jklabc->ij', wv.conj(), rw)
    gvv = numpy.einsum('ijkacd,ijkbcd->ab', wv, rw.conj())

    if not for_grad:
        # t3 amplitudes in CCSD(T) is computed non-iteratively. The
        # off-diagonal blocks of fock matrix does not contribute to CCSD(T)
        # energy. To make Tr(H,D) consistent to the CCSD(T) total energy, the
        # density matrix off-diagonal parts are excluded.
        doo[numpy.diag_indices(nocc)] -= goo.diagonal() * .5
        dvv[numpy.diag_indices(nvir)] += gvv.diagonal() * .5

    else:
        # The off-diagonal blocks of fock matrix have small contributions to
        # analytical nuclear gradients.
        doo -= goo * .5
        dvv += gvv * .5

    dvo += numpy.einsum('ijab,ijkabc->ck', t2.conj(), rw) * .5
    return doo, dov, dvo, dvv

def _gamma2_intermediates(mycc, t1, t2, l1, l2, eris=None,
                          compress_vvvv=False):
    '''intermediates tensors for gamma2 are sorted in Chemist's notation
    '''
    dovov, dvvvv, doooo, doovv, dovvo, dvvov, dovvv, dooov = \
            ccsd_rdm._gamma2_intermediates(mycc, t1, t2, l1, l2)
    if eris is None: eris = mycc.ao2mo()

    nocc, nvir = t1.shape
    eris_ovvv = numpy.asarray(eris.get_ovvv())
    eris_ovoo = numpy.asarray(eris.ovoo)
    eris_ovov = numpy.asarray(eris.ovov)

    mo_e = eris.mo_energy
    eia = lib.direct_sum('i-a->ia', mo_e[:nocc], mo_e[nocc:])
    d3 = lib.direct_sum('ia,jb,kc->ijkabc', eia, eia, eia)

    w = (numpy.einsum('iafb,kjcf->ijkabc', eris_ovvv.conj(), t2) -
         numpy.einsum('iajm,mkbc->ijkabc', eris_ovoo.conj(), t2)) / d3
    v = (numpy.einsum('iajb,kc->ijkabc', eris_ovov.conj(), t1) +
         numpy.einsum('ck,ijab->ijkabc', eris.fock[nocc:,:nocc], t2)) / d3
    w = p6(w)
    v = p6(v)
    rw = r6(w)
    rwv = r6(w * 2 + v * .5)

    dovov += numpy.einsum('kc,ijkabc->iajb', t1, rw.conj()) * .5
    dooov -= numpy.einsum('mkbc,ijkabc->jmia', t2, rwv.conj())
    # Note "dovvv +=" also changes the value of dvvov
    dovvv += numpy.einsum('kjcf,ijkabc->iafb', t2, rwv.conj())
    dvvov = dovvv.transpose(2,3,0,1)

    if compress_vvvv:
        nvir = mycc.nmo - mycc.nocc
        idx = numpy.tril_indices(nvir)
        vidx = idx[0] * nvir + idx[1]
        dvvvv = dvvvv + dvvvv.transpose(1,0,2,3)
        dvvvv = dvvvv + dvvvv.transpose(0,1,3,2)
        dvvvv = lib.take_2d(dvvvv.reshape(nvir**2,nvir**2), vidx, vidx)
        dvvvv *= .25

    return dovov, dvvvv, doooo, doovv, dovvo, dvvov, dovvv, dooov

def _gamma2_outcore(mycc, t1, t2, l1, l2, eris, h5fobj, compress_vvvv=False):
    return _gamma2_intermediates(mycc, t1, t2, l1, l2, eris, compress_vvvv)

def make_rdm1(mycc, t1, t2, l1, l2, eris=None, ao_repr=False):
    d1 = _gamma1_intermediates(mycc, t1, t2, l1, l2, eris)
    return ccsd_rdm._make_rdm1(mycc, d1, True, ao_repr=ao_repr)

# rdm2 in Chemist's notation
def make_rdm2(mycc, t1, t2, l1, l2, eris=None):
    d1 = _gamma1_intermediates(mycc, t1, t2, l1, l2, eris)
    d2 = _gamma2_intermediates(mycc, t1, t2, l1, l2, eris)
    return ccsd_rdm._make_rdm2(mycc, d1, d2, True, True)


def p6(t):
    t1 = t + t.transpose(0,2,1,3,5,4)
    return t1 + t1.transpose(1,0,2,4,3,5) + t1.transpose(1,2,0,4,5,3)
def r6(w):
    return (4 * w + w.transpose(0,1,2,4,5,3) + w.transpose(0,1,2,5,3,4)
            - 2 * w.transpose(0,1,2,5,4,3) - 2 * w.transpose(0,1,2,3,5,4)
            - 2 * w.transpose(0,1,2,4,3,5))


if __name__ == '__main__':
    from functools import reduce
    from pyscf import gto
    from pyscf import scf
    from pyscf.cc import ccsd
    from pyscf.cc import ccsd_t_slow as ccsd_t
    from pyscf.cc import ccsd_t_lambda_slow as ccsd_t_lambda
    from pyscf import ao2mo

    mol = gto.Mole()
    #mol.verbose = 5
    #mol.output = 'out_h2o'
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -.957 , .587)],
        [1 , (0.2,  .757 , .487)]]

    #mol.basis = 'ccpvdz'
    mol.basis = '631g'
    mol.build()
    mf = scf.RHF(mol)
    mf.conv_tol = 1e-1
    mf.scf()
    mcc = ccsd.CCSD(mf)
    mcc.conv_tol = 1e-14
    ecc, t1, t2 = mcc.kernel()
    eris = mcc.ao2mo()
    e3ref = ccsd_t.kernel(mcc, eris, t1, t2)
    l1, l2 = ccsd_t_lambda.kernel(mcc, eris, t1, t2)[1:]
    print(ecc, e3ref)

    eri_mo = ao2mo.kernel(mf._eri, mf.mo_coeff, compact=False)
    nmo = mf.mo_coeff.shape[1]
    eri_mo = eri_mo.reshape(nmo,nmo,nmo,nmo)
    dm1 = make_rdm1(mcc, t1, t2, l1, l2, eris=eris)
    dm2 = make_rdm2(mcc, t1, t2, l1, l2, eris=eris)
    print(lib.finger(dm1) - 1.289951975176953 )
    print(lib.finger(dm2) - 6.6184784979411164)
    h1 = reduce(numpy.dot, (mf.mo_coeff.T, mf.get_hcore(), mf.mo_coeff))
    e3 =(numpy.einsum('ij,ji->', h1, dm1) +
         numpy.einsum('ijkl,ijkl->', eri_mo, dm2)*.5 + mf.mol.energy_nuc())
    print(e3ref, e3-(mf.e_tot+ecc))



# pyscf/cc/ccsd_t_slow.py

#!/usr/bin/env python
# Copyright 2014-2020 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: Qiming Sun <osirpt.sun@gmail.com>
#

import numpy
from pyscf import lib
from pyscf.lib import logger

'''
CCSD(T)
'''

# t3 as ijkabc

# JCP 94, 442 (1991); DOI:10.1063/1.460359.  Error in Eq (1), should be [ia] >= [jb] >= [kc]
def kernel(mycc, eris, t1=None, t2=None, verbose=logger.NOTE):
    if isinstance(verbose, logger.Logger):
        log = verbose
    else:
        log = logger.Logger(mycc.stdout, verbose)

    if t1 is None: t1 = mycc.t1
    if t2 is None: t2 = mycc.t2

    t1T = t1.T
    t2T = t2.transpose(2,3,0,1)

    nocc, nvir = t1.shape
    mo_e = eris.mo_energy
    e_occ, e_vir = mo_e[:nocc], mo_e[nocc:]
    eijk = lib.direct_sum('i,j,k->ijk', e_occ, e_occ, e_occ)

    eris_vvov = eris.get_ovvv().conj().transpose(1,3,0,2)
    eris_vooo = numpy.asarray(eris.ovoo).conj().transpose(1,0,3,2)
    eris_vvoo = numpy.asarray(eris.ovov).conj().transpose(1,3,0,2)
    fvo = eris.fock[nocc:,:nocc]
    def get_w(a, b, c):
        w = numpy.einsum('if,fkj->ijk', eris_vvov[a,b], t2T[c,:])
        w-= numpy.einsum('ijm,mk->ijk', eris_vooo[a,:], t2T[b,c])
        return w
    def get_v(a, b, c):
        v = numpy.einsum('ij,k->ijk', eris_vvoo[a,b], t1T[c])
        v+= numpy.einsum('ij,k->ijk', t2T[a,b], fvo[c])
        return v

    et = 0
    for a in range(nvir):
        for b in range(a+1):
            for c in range(b+1):
                d3 = eijk - e_vir[a] - e_vir[b] - e_vir[c]
                if a == c:  # a == b == c
                    d3 *= 6
                elif a == b or b == c:
                    d3 *= 2

                wabc = get_w(a, b, c)
                wacb = get_w(a, c, b)
                wbac = get_w(b, a, c)
                wbca = get_w(b, c, a)
                wcab = get_w(c, a, b)
                wcba = get_w(c, b, a)
                vabc = get_v(a, b, c)
                vacb = get_v(a, c, b)
                vbac = get_v(b, a, c)
                vbca = get_v(b, c, a)
                vcab = get_v(c, a, b)
                vcba = get_v(c, b, a)
                zabc = r3(wabc + .5 * vabc) / d3
                zacb = r3(wacb + .5 * vacb) / d3
                zbac = r3(wbac + .5 * vbac) / d3
                zbca = r3(wbca + .5 * vbca) / d3
                zcab = r3(wcab + .5 * vcab) / d3
                zcba = r3(wcba + .5 * vcba) / d3

                et+= numpy.einsum('ijk,ijk', wabc, zabc.conj())
                et+= numpy.einsum('ikj,ijk', wacb, zabc.conj())
                et+= numpy.einsum('jik,ijk', wbac, zabc.conj())
                et+= numpy.einsum('jki,ijk', wbca, zabc.conj())
                et+= numpy.einsum('kij,ijk', wcab, zabc.conj())
                et+= numpy.einsum('kji,ijk', wcba, zabc.conj())

                et+= numpy.einsum('ijk,ijk', wacb, zacb.conj())
                et+= numpy.einsum('ikj,ijk', wabc, zacb.conj())
                et+= numpy.einsum('jik,ijk', wcab, zacb.conj())
                et+= numpy.einsum('jki,ijk', wcba, zacb.conj())
                et+= numpy.einsum('kij,ijk', wbac, zacb.conj())
                et+= numpy.einsum('kji,ijk', wbca, zacb.conj())

                et+= numpy.einsum('ijk,ijk', wbac, zbac.conj())
                et+= numpy.einsum('ikj,ijk', wbca, zbac.conj())
                et+= numpy.einsum('jik,ijk', wabc, zbac.conj())
                et+= numpy.einsum('jki,ijk', wacb, zbac.conj())
                et+= numpy.einsum('kij,ijk', wcba, zbac.conj())
                et+= numpy.einsum('kji,ijk', wcab, zbac.conj())

                et+= numpy.einsum('ijk,ijk', wbca, zbca.conj())
                et+= numpy.einsum('ikj,ijk', wbac, zbca.conj())
                et+= numpy.einsum('jik,ijk', wcba, zbca.conj())
                et+= numpy.einsum('jki,ijk', wcab, zbca.conj())
                et+= numpy.einsum('kij,ijk', wabc, zbca.conj())
                et+= numpy.einsum('kji,ijk', wacb, zbca.conj())

                et+= numpy.einsum('ijk,ijk', wcab, zcab.conj())
                et+= numpy.einsum('ikj,ijk', wcba, zcab.conj())
                et+= numpy.einsum('jik,ijk', wacb, zcab.conj())
                et+= numpy.einsum('jki,ijk', wabc, zcab.conj())
                et+= numpy.einsum('kij,ijk', wbca, zcab.conj())
                et+= numpy.einsum('kji,ijk', wbac, zcab.conj())

                et+= numpy.einsum('ijk,ijk', wcba, zcba.conj())
                et+= numpy.einsum('ikj,ijk', wcab, zcba.conj())
                et+= numpy.einsum('jik,ijk', wbca, zcba.conj())
                et+= numpy.einsum('jki,ijk', wbac, zcba.conj())
                et+= numpy.einsum('kij,ijk', wacb, zcba.conj())
                et+= numpy.einsum('kji,ijk', wabc, zcba.conj())
    et *= 2
    log.info('CCSD(T) correction = %.15g', et)
    return et

def r3(w):
    return (4 * w + w.transpose(1,2,0) + w.transpose(2,0,1)
            - 2 * w.transpose(2,1,0) - 2 * w.transpose(0,2,1)
            - 2 * w.transpose(1,0,2))


if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf
    from pyscf import cc

    mol = gto.M()
    numpy.random.seed(12)
    nocc, nvir = 5, 12
    eris = cc.ccsd._ChemistsERIs()
    eris.ovvv = numpy.random.random((nocc,nvir,nvir*(nvir+1)//2)) * .1
    eris.ovoo = numpy.random.random((nocc,nvir,nocc,nocc)) * .1
    eris.ovov = numpy.random.random((nocc,nvir,nocc,nvir)) * .1
    t1 = numpy.random.random((nocc,nvir)) * .1
    t2 = numpy.random.random((nocc,nocc,nvir,nvir)) * .1
    t2 = t2 + t2.transpose(1,0,3,2)
    mf = scf.RHF(mol)
    mcc = cc.CCSD(mf)
    f = numpy.random.random((nocc+nvir,nocc+nvir)) * .1
    eris.fock = f+f.T + numpy.diag(numpy.arange(nocc+nvir))
    print(kernel(mcc, eris, t1, t2) - -8.0038781018306828)

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -.957 , .587)],
        [1 , (0.2,  .757 , .487)]]

    mol.basis = 'ccpvdz'
    mol.build()
    rhf = scf.RHF(mol)
    rhf.conv_tol = 1e-14
    rhf.scf()
    mcc = cc.CCSD(rhf)
    mcc.conv_tol = 1e-12
    mcc.ccsd()

    e3a = kernel(mcc, mcc.ao2mo())
    print(e3a - -0.0033300722698513989)


# pyscf/cc/dfccsd.py

#!/usr/bin/env python
# Copyright 2014-2020 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.


import ctypes
import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf.ao2mo import _ao2mo
from pyscf import df
from pyscf.cc import ccsd
from pyscf.cc import _ccsd
from pyscf import __config__

MEMORYMIN = getattr(__config__, 'cc_ccsd_memorymin', 2000)

class RCCSD(ccsd.CCSD):
    def __init__(self, mf, frozen=None, mo_coeff=None, mo_occ=None):
        ccsd.CCSD.__init__(self, mf, frozen, mo_coeff, mo_occ)
        if getattr(mf, 'with_df', None):
            self.with_df = mf.with_df
        else:
            self.with_df = df.DF(mf.mol)
            self.with_df.auxbasis = df.make_auxbasis(mf.mol, mp2fit=True)
        self._keys.update(['with_df'])

    def reset(self, mol=None):
        self.with_df.reset(mol)
        return ccsd.CCSD.reset(self, mol)

    def ao2mo(self, mo_coeff=None):
        return _make_df_eris(self, mo_coeff)

    def _add_vvvv(self, t1, t2, eris, out=None, with_ovvv=False, t2sym=None):
        assert(not self.direct)
        return ccsd.CCSD._add_vvvv(self, t1, t2, eris, out, with_ovvv, t2sym)


def _contract_vvvv_t2(mycc, mol, vvL, t2, out=None, verbose=None):
    '''Ht2 = numpy.einsum('ijcd,acdb->ijab', t2, vvvv)

    Args:
        vvvv : None or integral object
            if vvvv is None, contract t2 to AO-integrals using AO-direct algorithm
    '''
    _dgemm = lib.numpy_helper._dgemm
    time0 = logger.process_clock(), logger.perf_counter()
    log = logger.new_logger(mol, verbose)

    naux = vvL.shape[-1]
    nvira, nvirb = t2.shape[-2:]
    x2 = t2.reshape(-1,nvira,nvirb)
    nocc2 = x2.shape[0]
    nvir2 = nvira * nvirb
    Ht2 = numpy.ndarray(x2.shape, buffer=out)
    Ht2[:] = 0

    max_memory = max(MEMORYMIN, mycc.max_memory - lib.current_memory()[0])
    def contract_blk_(eri, i0, i1, j0, j1):
        ic = i1 - i0
        jc = j1 - j0
        #:Ht2[:,j0:j1] += numpy.einsum('xef,efab->xab', x2[:,i0:i1], eri)
        _dgemm('N', 'N', nocc2, jc*nvirb, ic*nvirb,
               x2.reshape(-1,nvir2), eri.reshape(-1,jc*nvirb),
               Ht2.reshape(-1,nvir2), 1, 1, i0*nvirb, 0, j0*nvirb)

        if i0 > j0:
            #:Ht2[:,i0:i1] += numpy.einsum('xef,abef->xab', x2[:,j0:j1], eri)
            _dgemm('N', 'T', nocc2, ic*nvirb, jc*nvirb,
                   x2.reshape(-1,nvir2), eri.reshape(-1,jc*nvirb),
                   Ht2.reshape(-1,nvir2), 1, 1, j0*nvirb, 0, i0*nvirb)

#TODO: check if vvL can be entirely loaded into memory
    nvir_pair = nvirb * (nvirb+1) // 2
    dmax = numpy.sqrt(max_memory*.7e6/8/nvirb**2/2)
    dmax = int(min((nvira+3)//4, max(ccsd.BLKMIN, dmax)))
    vvblk = (max_memory*1e6/8 - dmax**2*(nvirb**2*1.5+naux))/naux
    vvblk = int(min((nvira+3)//4, max(ccsd.BLKMIN, vvblk/naux)))
    eribuf = numpy.empty((dmax,dmax,nvir_pair))
    loadbuf = numpy.empty((dmax,dmax,nvirb,nvirb))
    tril2sq = lib.square_mat_in_trilu_indices(nvira)

    for i0, i1 in lib.prange(0, nvira, dmax):
        off0 = i0*(i0+1)//2
        off1 = i1*(i1+1)//2
        vvL0 = _cp(vvL[off0:off1])
        for j0, j1 in lib.prange(0, i1, dmax):
            ijL = vvL0[tril2sq[i0:i1,j0:j1] - off0].reshape(-1,naux)
            eri = numpy.ndarray(((i1-i0)*(j1-j0),nvir_pair), buffer=eribuf)
            for p0, p1 in lib.prange(0, nvir_pair, vvblk):
                vvL1 = _cp(vvL[p0:p1])
                eri[:,p0:p1] = lib.ddot(ijL, vvL1.T)
                vvL1 = None

            tmp = numpy.ndarray((i1-i0,nvirb,j1-j0,nvirb), buffer=loadbuf)
            _ccsd.libcc.CCload_eri(tmp.ctypes.data_as(ctypes.c_void_p),
                                   eri.ctypes.data_as(ctypes.c_void_p),
                                   (ctypes.c_int*4)(i0, i1, j0, j1),
                                   ctypes.c_int(nvirb))
            contract_blk_(tmp, i0, i1, j0, j1)
            time0 = log.timer_debug1('vvvv [%d:%d,%d:%d]'%(i0,i1,j0,j1), *time0)
    return Ht2.reshape(t2.shape)


class _ChemistsERIs(ccsd._ChemistsERIs):
    def _contract_vvvv_t2(self, mycc, t2, direct=False, out=None, verbose=None):
        assert(not direct)
        return _contract_vvvv_t2(mycc, self.mol, self.vvL, t2, out, verbose)

def _make_df_eris(cc, mo_coeff=None):
    eris = _ChemistsERIs()
    eris._common_init_(cc, mo_coeff)
    nocc = eris.nocc
    nmo = eris.fock.shape[0]
    nvir = nmo - nocc
    nvir_pair = nvir*(nvir+1)//2
    with_df = cc.with_df
    naux = eris.naux = with_df.get_naoaux()

    eris.feri = lib.H5TmpFile()
    eris.oooo = eris.feri.create_dataset('oooo', (nocc,nocc,nocc,nocc), 'f8')
    eris.ovoo = eris.feri.create_dataset('ovoo', (nocc,nvir,nocc,nocc), 'f8', chunks=(nocc,1,nocc,nocc))
    eris.ovov = eris.feri.create_dataset('ovov', (nocc,nvir,nocc,nvir), 'f8', chunks=(nocc,1,nocc,nvir))
    eris.ovvo = eris.feri.create_dataset('ovvo', (nocc,nvir,nvir,nocc), 'f8', chunks=(nocc,1,nvir,nocc))
    eris.oovv = eris.feri.create_dataset('oovv', (nocc,nocc,nvir,nvir), 'f8', chunks=(nocc,nocc,1,nvir))
    # nrow ~ 4e9/8/blockdim to ensure hdf5 chunk < 4GB
    chunks = (min(nvir_pair,int(4e8/with_df.blockdim)), min(naux,with_df.blockdim))
    eris.vvL = eris.feri.create_dataset('vvL', (nvir_pair,naux), 'f8', chunks=chunks)

    Loo = numpy.empty((naux,nocc,nocc))
    Lov = numpy.empty((naux,nocc,nvir))
    mo = numpy.asarray(eris.mo_coeff, order='F')
    ijslice = (0, nmo, 0, nmo)
    p1 = 0
    Lpq = None
    for k, eri1 in enumerate(with_df.loop()):
        Lpq = _ao2mo.nr_e2(eri1, mo, ijslice, aosym='s2', mosym='s1', out=Lpq)
        p0, p1 = p1, p1 + Lpq.shape[0]
        Lpq = Lpq.reshape(p1-p0,nmo,nmo)
        Loo[p0:p1] = Lpq[:,:nocc,:nocc]
        Lov[p0:p1] = Lpq[:,:nocc,nocc:]
        Lvv = lib.pack_tril(Lpq[:,nocc:,nocc:])
        eris.vvL[:,p0:p1] = Lvv.T
    Lpq = Lvv = None
    Loo = Loo.reshape(naux,nocc**2)
    #Lvo = Lov.transpose(0,2,1).reshape(naux,nvir*nocc)
    Lov = Lov.reshape(naux,nocc*nvir)
    eris.oooo[:] = lib.ddot(Loo.T, Loo).reshape(nocc,nocc,nocc,nocc)
    eris.ovoo[:] = lib.ddot(Lov.T, Loo).reshape(nocc,nvir,nocc,nocc)
    ovov = lib.ddot(Lov.T, Lov).reshape(nocc,nvir,nocc,nvir)
    eris.ovov[:] = ovov
    eris.ovvo[:] = ovov.transpose(0,1,3,2)
    ovov = None

    mem_now = lib.current_memory()[0]
    max_memory = max(0, cc.max_memory - mem_now)
    blksize = max(ccsd.BLKMIN, int((max_memory*.9e6/8-nocc**2*nvir_pair)/(nocc**2+naux)))
    oovv_tril = numpy.empty((nocc*nocc,nvir_pair))
    for p0, p1 in lib.prange(0, nvir_pair, blksize):
        oovv_tril[:,p0:p1] = lib.ddot(Loo.T, _cp(eris.vvL[p0:p1]).T)
    eris.oovv[:] = lib.unpack_tril(oovv_tril).reshape(nocc,nocc,nvir,nvir)
    oovv_tril = Loo = None

    Lov = Lov.reshape(naux,nocc,nvir)
    vblk = max(nocc, int((max_memory*.15e6/8)/(nocc*nvir_pair)))
    vvblk = int(min(nvir_pair, 4e8/nocc, max(4, (max_memory*.8e6/8)/(vblk*nocc+naux))))
    eris.ovvv = eris.feri.create_dataset('ovvv', (nocc,nvir,nvir_pair), 'f8',
                                         chunks=(nocc,1,vvblk))
    for q0, q1 in lib.prange(0, nvir_pair, vvblk):
        vvL = _cp(eris.vvL[q0:q1])
        for p0, p1 in lib.prange(0, nvir, vblk):
            tmpLov = _cp(Lov[:,:,p0:p1]).reshape(naux,-1)
            eris.ovvv[:,p0:p1,q0:q1] = lib.ddot(tmpLov.T, vvL.T).reshape(nocc,p1-p0,q1-q0)
        vvL = None
    return eris

def _cp(a):
    return numpy.array(a, copy=False, order='C')

if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]

    mol.basis = 'cc-pvdz'
    mol.build()
    mf = scf.RHF(mol).density_fit('weigend').run()
    mycc = RCCSD(mf).run()
    print(mycc.e_corr - -0.21337100025961622)

    print("IP energies... (right eigenvector)")
    part = None
    e,v = mycc.ipccsd(nroots=3, partition=part)
    print(e)
    print(e[0] - 0.43364287418576897)
    print(e[1] - 0.5188001071775572 )
    print(e[2] - 0.67851590275796392)

    print("IP energies... (left eigenvector)")
    e,lv = mycc.ipccsd(nroots=3,left=True,partition=part)
    print(e)
    print(e[0] - 0.43364286531878882)
    print(e[1] - 0.51879999865136994)
    print(e[2] - 0.67851587320495355)

    print("EA energies... (right eigenvector)")
    e,v = mycc.eaccsd(nroots=3, partition=part)
    print(e)
    print(e[0] - 0.16730125785810035)
    print(e[1] - 0.23999823045518162)
    print(e[2] - 0.50960183439619933)

    print("EA energies... (left eigenvector)")
    e,lv = mycc.eaccsd(nroots=3, left=True, partition=part)
    print(e)
    print(e[0] - 0.16730137808538076)
    print(e[1] - 0.23999845448276602)
    print(e[2] - 0.50960182130968001)

    e, v = mycc.eeccsd(nroots=4)
    print(e[0] - 0.27575637238275519)
    print(e[1] - 0.27575637238275519)
    print(e[2] - 0.27575637238275519)
    print(e[3] - 0.30068967373840394)


# pyscf/cc/eom_gccsd.py

#!/usr/bin/env python
# Copyright 2014-2021 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.


import numpy as np

from pyscf import lib
from pyscf.lib import logger
from pyscf.cc import ccsd
from pyscf.cc import eom_rccsd
from pyscf.cc import gccsd
from pyscf.cc import gintermediates as imd


########################################
# EOM-IP-CCSD
########################################

def vector_to_amplitudes_ip(vector, nmo, nocc):
    nvir = nmo - nocc
    r1 = vector[:nocc].copy()
    r2 = np.zeros((nocc,nocc,nvir), dtype=vector.dtype)
    idx, idy = np.tril_indices(nocc, -1)
    r2[idx,idy] = vector[nocc:].reshape(nocc*(nocc-1)//2,nvir)
    r2[idy,idx] =-vector[nocc:].reshape(nocc*(nocc-1)//2,nvir)
    return r1, r2

def amplitudes_to_vector_ip(r1, r2):
    nocc = r1.size
    return np.hstack((r1, r2[np.tril_indices(nocc, -1)].ravel()))

def ipccsd_matvec(eom, vector, imds=None, diag=None):
    # Ref: Tu, Wang, and Li, J. Chem. Phys. 136, 174102 (2012) Eqs.(8)-(9)
    if imds is None: imds = eom.make_imds()
    nocc = eom.nocc
    nmo = eom.nmo
    r1, r2 = vector_to_amplitudes_ip(vector, nmo, nocc)

    # Eq. (8)
    Hr1 = -np.einsum('mi,m->i', imds.Foo, r1)
    Hr1 += np.einsum('me,mie->i', imds.Fov, r2)
    Hr1 += -0.5*np.einsum('nmie,mne->i', imds.Wooov, r2)
    # Eq. (9)
    Hr2 =  lib.einsum('ae,ije->ija', imds.Fvv, r2)
    tmp1 = lib.einsum('mi,mja->ija', imds.Foo, r2)
    Hr2 -= tmp1 - tmp1.transpose(1,0,2)
    Hr2 -= np.einsum('maji,m->ija', imds.Wovoo, r1)
    Hr2 += 0.5*lib.einsum('mnij,mna->ija', imds.Woooo, r2)
    tmp2 = lib.einsum('maei,mje->ija', imds.Wovvo, r2)
    Hr2 += tmp2 - tmp2.transpose(1,0,2)
    Hr2 += 0.5*lib.einsum('mnef,mnf,ijae->ija', imds.Woovv, r2, imds.t2)

    vector = amplitudes_to_vector_ip(Hr1, Hr2)
    return vector

def lipccsd_matvec(eom, vector, imds=None, diag=None):
    '''IP-CCSD left eigenvector equation.

    For description of args, see ipccsd_matvec.'''
    if imds is None:
        imds = eom.make_imds()
    nocc = eom.nocc
    nmo = eom.nmo
    r1, r2 = vector_to_amplitudes_ip(vector, nmo, nocc)

    Hr1 = -lib.einsum('mi,i->m', imds.Foo, r1)
    Hr1 += -0.5 * lib.einsum('maji,ija->m', imds.Wovoo, r2)

    Hr2 = lib.einsum('me,i->mie', imds.Fov, r1)
    Hr2 -= lib.einsum('ie,m->mie', imds.Fov, r1)
    Hr2 += -lib.einsum('nmie,i->mne', imds.Wooov, r1)
    Hr2 += lib.einsum('ae,ija->ije', imds.Fvv, r2)
    tmp1 = lib.einsum('mi,ija->mja', imds.Foo, r2)
    Hr2 += (-tmp1 + tmp1.transpose(1, 0, 2))
    Hr2 += 0.5 * lib.einsum('mnij,ija->mna', imds.Woooo, r2)
    tmp2 = lib.einsum('maei,ija->mje', imds.Wovvo, r2)
    Hr2 += (tmp2 - tmp2.transpose(1, 0, 2))
    Hr2 += 0.5 * lib.einsum('mnef,ija,ijae->mnf', imds.Woovv, r2, imds.t2)

    vector = amplitudes_to_vector_ip(Hr1, Hr2)
    return vector

def ipccsd_diag(eom, imds=None):
    if imds is None: imds = eom.make_imds()
    t1, t2 = imds.t1, imds.t2
    nocc, nvir = t1.shape

    Hr1 = -np.diag(imds.Foo)
    Hr2 = np.zeros((nocc,nocc,nvir), dtype=t1.dtype)
    for i in range(nocc):
        for j in range(nocc):
            for a in range(nvir):
                Hr2[i,j,a] += imds.Fvv[a,a]
                Hr2[i,j,a] += -imds.Foo[i,i]
                Hr2[i,j,a] += -imds.Foo[j,j]
                Hr2[i,j,a] += 0.5*(imds.Woooo[i,j,i,j]-imds.Woooo[j,i,i,j])
                Hr2[i,j,a] += imds.Wovvo[i,a,a,i]
                Hr2[i,j,a] += imds.Wovvo[j,a,a,j]
                Hr2[i,j,a] += 0.5*(np.dot(imds.Woovv[i,j,:,a], t2[i,j,a,:]) -
                                   np.dot(imds.Woovv[j,i,:,a], t2[i,j,a,:]))

    vector = amplitudes_to_vector_ip(Hr1, Hr2)
    return vector

def ipccsd_star_contract(eom, ipccsd_evals, ipccsd_evecs, lipccsd_evecs, imds=None):
    """
    Returns:
        e_star (list of float):
            The IP-CCSD* energy.

    Notes:
        The user should check to make sure the right and left eigenvalues
        before running the perturbative correction.

        The 2hp right amplitudes are assumed to be of the form s^{a }_{ij}, i.e.
        the (ia) indices are coupled while the left are assumed to be of the form
        s^{ b}_{ij}, i.e. the (jb) indices are coupled.

    Reference:
        Saeh, Stanton "...energy surfaces of radicals" JCP 111, 8275 (1999); DOI:10.1063/1.480171
    """
    assert eom.partition is None
    if imds is None:
        imds = eom.make_imds()
    t1, t2 = imds.t1, imds.t2
    eris = imds.eris
    assert (isinstance(eris, gccsd._PhysicistsERIs))
    fock = eris.fock
    nocc, nvir = t1.shape
    nmo = nocc + nvir

    #fov = fock[:nocc, nocc:]
    foo = fock[:nocc, :nocc].diagonal()
    fvv = fock[nocc:, nocc:].diagonal()

    oovv = _cp(eris.oovv)
    ovvv = _cp(eris.ovvv)
    ovov = _cp(eris.ovov)
    #ovvo = -_cp(eris.ovov).transpose(0,1,3,2)
    ooov = _cp(eris.ooov)
    vooo = _cp(ooov).conj().transpose(3,2,1,0)
    vvvo = _cp(ovvv).conj().transpose(3,2,1,0)
    oooo = _cp(eris.oooo)

    # Create denominator
    eijk = foo[:, None, None] + foo[None, :, None] + foo[None, None, :]
    eab = fvv[:, None] + fvv[None, :]
    eijkab = eijk[:, :, :, None, None] - eab[None, None, None, :, :]

    # Permutation operators
    def pijk(tmp):
        '''P(ijk)'''
        return tmp + tmp.transpose(1,2,0,3,4) + tmp.transpose(2,0,1,3,4)

    def pab(tmp):
        '''P(ab)'''
        return tmp - tmp.transpose(0,1,2,4,3)

    def pij(tmp):
        '''P(ij)'''
        return tmp - tmp.transpose(1,0,2,3,4)

    ipccsd_evecs = np.array(ipccsd_evecs)
    lipccsd_evecs = np.array(lipccsd_evecs)
    e_star = []
    ipccsd_evecs, lipccsd_evecs = [np.atleast_2d(x) for x in [ipccsd_evecs, lipccsd_evecs]]
    ipccsd_evals = np.atleast_1d(ipccsd_evals)
    for ip_eval, ip_evec, ip_levec in zip(ipccsd_evals, ipccsd_evecs, lipccsd_evecs):
        # Enforcing <L|R> = 1
        l1, l2 = vector_to_amplitudes_ip(ip_levec, nmo, nocc)
        r1, r2 = vector_to_amplitudes_ip(ip_evec, nmo, nocc)
        ldotr = np.dot(l1, r1) + 0.5 * np.dot(l2.ravel(), r2.ravel())

        logger.info(eom, 'Left-right amplitude overlap : %14.8e', ldotr)
        if abs(ldotr) < 1e-7:
            logger.warn(eom, 'Small %s left-right amplitude overlap. Results '
                             'may be inaccurate.', ldotr)

        l1 /= ldotr
        l2 /= ldotr

        # Denominator + eigenvalue(IP-CCSD)
        denom = eijkab + ip_eval
        denom = 1. / denom

        tmp = lib.einsum('ijab,k->ijkab', oovv, l1)
        lijkab = pijk(tmp)
        tmp = -lib.einsum('jima,mkb->ijkab', ooov, l2)
        tmp = pijk(tmp)
        lijkab += pab(tmp)
        tmp = lib.einsum('ieab,jke->ijkab', ovvv, l2)
        lijkab += pijk(tmp)

        tmp = lib.einsum('mbke,m->bke', ovov, r1)
        tmp = lib.einsum('bke,ijae->ijkab', tmp, t2)
        tmp = pijk(tmp)
        rijkab = -pab(tmp)
        tmp = lib.einsum('mnjk,n->mjk', oooo, r1)
        tmp = lib.einsum('mjk,imab->ijkab', tmp, t2)
        rijkab += pijk(tmp)
        tmp = lib.einsum('amij,mkb->ijkab', vooo, r2)
        tmp = pijk(tmp)
        rijkab -= pab(tmp)
        tmp = lib.einsum('baei,jke->ijkab', vvvo, r2)
        rijkab += pijk(tmp)

        deltaE = (1. / 12) * lib.einsum('ijkab,ijkab,ijkab', lijkab, rijkab, denom)
        deltaE = deltaE.real
        logger.info(eom, "Exc. energy, delta energy = %16.12f, %16.12f",
                    ip_eval + deltaE, deltaE)
        e_star.append(ip_eval + deltaE)
    return e_star

class EOMIP(eom_rccsd.EOMIP):
    matvec = ipccsd_matvec
    l_matvec = lipccsd_matvec
    get_diag = ipccsd_diag
    ccsd_star_contract = ipccsd_star_contract

    def vector_to_amplitudes(self, vector, nmo=None, nocc=None):
        if nmo is None: nmo = self.nmo
        if nocc is None: nocc = self.nocc
        return vector_to_amplitudes_ip(vector, nmo, nocc)

    def amplitudes_to_vector(self, r1, r2):
        return amplitudes_to_vector_ip(r1, r2)

    def vector_size(self):
        nocc = self.nocc
        nvir = self.nmo - nocc
        return nocc + nocc*(nocc-1)//2*nvir

    def make_imds(self, eris=None):
        imds = _IMDS(self._cc, eris)
        imds.make_ip()
        return imds

class EOMIP_Ta(EOMIP):
    '''Class for EOM IPCCSD(T)*(a) method by Matthews and Stanton.'''
    def make_imds(self, eris=None):
        imds = _IMDS(self._cc, eris=eris)
        imds.make_t3p2_ip(self._cc)
        return imds

########################################
# EOM-EA-CCSD
########################################

def vector_to_amplitudes_ea(vector, nmo, nocc):
    nvir = nmo - nocc
    r1 = vector[:nvir].copy()
    r2 = np.zeros((nocc,nvir,nvir), vector.dtype)
    idx, idy = np.tril_indices(nvir, -1)
    r2[:,idx,idy] = vector[nvir:].reshape(nocc,-1)
    r2[:,idy,idx] =-vector[nvir:].reshape(nocc,-1)
    return r1, r2

def amplitudes_to_vector_ea(r1, r2):
    nvir = r1.size
    idx, idy = np.tril_indices(nvir, -1)
    return np.hstack((r1, r2[:,idx,idy].ravel()))

def eaccsd_matvec(eom, vector, imds=None, diag=None):
    # Ref: Nooijen and Bartlett, J. Chem. Phys. 102, 3629 (1994) Eqs.(30)-(31)
    if imds is None: imds = eom.make_imds()
    nocc = eom.nocc
    nmo = eom.nmo
    r1, r2 = vector_to_amplitudes_ea(vector, nmo, nocc)

    # Eq. (30)
    Hr1  = np.einsum('ac,c->a', imds.Fvv, r1)
    Hr1 += np.einsum('ld,lad->a', imds.Fov, r2)
    Hr1 += 0.5*np.einsum('alcd,lcd->a', imds.Wvovv, r2)
    # Eq. (31)
    Hr2 = np.einsum('abcj,c->jab', imds.Wvvvo, r1)
    tmp1 = lib.einsum('ac,jcb->jab', imds.Fvv, r2)
    Hr2 += tmp1 - tmp1.transpose(0,2,1)
    Hr2 -= lib.einsum('lj,lab->jab', imds.Foo, r2)
    tmp2 = lib.einsum('lbdj,lad->jab', imds.Wovvo, r2)
    Hr2 += tmp2 - tmp2.transpose(0,2,1)
    Hr2 += 0.5*lib.einsum('abcd,jcd->jab', imds.Wvvvv, r2)
    Hr2 -= 0.5*lib.einsum('klcd,lcd,kjab->jab', imds.Woovv, r2, imds.t2)

    vector = amplitudes_to_vector_ea(Hr1, Hr2)
    return vector

def leaccsd_matvec(eom, vector, imds=None, diag=None):
    '''EA-CCSD left eigenvector equation.

    For description of args, see eaccsd_matvec.'''
    # Ref: Nooijen and Bartlett, J. Chem. Phys. 102, 3629 (1994) Eqs.(32)-(33)
    if imds is None:
        imds = eom.make_imds()
    nocc = eom.nocc
    nmo = eom.nmo
    r1, r2 = vector_to_amplitudes_ea(vector, nmo, nocc)

    # Eq. (32)
    Hr1 = lib.einsum('ac,a->c',imds.Fvv,r1)
    Hr1 += 0.5*lib.einsum('abcj,jab->c',imds.Wvvvo,r2)
    # Eq. (33)
    Hr2 = lib.einsum('alcd,a->lcd',imds.Wvovv,r1)
    Hr2 += lib.einsum('ld,a->lad',imds.Fov,r1)
    Hr2 -= lib.einsum('la,d->lad',imds.Fov,r1)
    tmp1 = lib.einsum('ac,jab->jcb',imds.Fvv,r2)
    Hr2 += (tmp1 - tmp1.transpose(0,2,1))
    Hr2 += -lib.einsum('lj,jab->lab',imds.Foo,r2)
    tmp2 = lib.einsum('lbdj,jab->lad',imds.Wovvo,r2)
    Hr2 += (tmp2 - tmp2.transpose(0,2,1))
    Hr2 += 0.5*lib.einsum('abcd,jab->jcd',imds.Wvvvv,r2)
    Hr2 += -0.5*lib.einsum('klcd,jab,kjab->lcd',imds.Woovv,r2,imds.t2)

    vector = amplitudes_to_vector_ea(Hr1,Hr2)
    return vector

def eaccsd_diag(eom, imds=None):
    if imds is None: imds = eom.make_imds()
    t1, t2 = imds.t1, imds.t2
    nocc, nvir = t1.shape

    Hr1 = np.diag(imds.Fvv)
    Hr2 = np.zeros((nocc,nvir,nvir),dtype=t1.dtype)
    for a in range(nvir):
        _Wvvvva = np.array(imds.Wvvvv[a])
        for b in range(a):
            for j in range(nocc):
                Hr2[j,a,b] += imds.Fvv[a,a]
                Hr2[j,a,b] += imds.Fvv[b,b]
                Hr2[j,a,b] += -imds.Foo[j,j]
                Hr2[j,a,b] += imds.Wovvo[j,b,b,j]
                Hr2[j,a,b] += imds.Wovvo[j,a,a,j]
                Hr2[j,a,b] += 0.5*(_Wvvvva[b,a,b]-_Wvvvva[b,b,a])
                Hr2[j,a,b] -= 0.5*(np.dot(imds.Woovv[:,j,a,b], t2[:,j,a,b]) -
                                   np.dot(imds.Woovv[:,j,b,a], t2[:,j,a,b]))

    vector = amplitudes_to_vector_ea(Hr1, Hr2)
    return vector

def eaccsd_star_contract(eom, eaccsd_evals, eaccsd_evecs, leaccsd_evecs, imds=None):
    """
    Returns:
        e_star (list of float):
            The EA-CCSD* energy.

    Notes:
        See `ipccsd_star_contract` for description of arguments.

    Reference:
        Saeh, Stanton "...energy surfaces of radicals" JCP 111, 8275 (1999); DOI:10.1063/1.480171
    """
    assert eom.partition is None
    if imds is None:
        imds = eom.make_imds()
    t1, t2 = imds.t1, imds.t2
    eris = imds.eris
    assert isinstance(eris, gccsd._PhysicistsERIs)
    fock = eris.fock
    nocc, nvir = t1.shape
    nmo = nocc + nvir

    #fov = fock[:nocc, nocc:].diagonal()
    foo = fock[:nocc, :nocc].diagonal()
    fvv = fock[nocc:, nocc:].diagonal()

    vvvv = _cp(eris.vvvv)
    oovv = _cp(eris.oovv)
    ovvv = _cp(eris.ovvv)
    ovov = _cp(eris.ovov)
    #ovvo = -_cp(eris.ovov).transpose(0,1,3,2)
    ooov = _cp(eris.ooov)
    vooo = _cp(ooov).conj().transpose(3,2,1,0)
    vvvo = _cp(ovvv).conj().transpose(3,2,1,0)

    # Create denominator
    eabc = fvv[:, None, None] + fvv[None, :, None] + fvv[None, None, :]
    eij = foo[:, None] + foo[None, :]
    eijabc = eij[:, :, None, None, None] - eabc[None, None, :, :, :]

    # Permutation operators
    def pabc(tmp):
        '''P(abc)'''
        return tmp + tmp.transpose(0,1,3,4,2) + tmp.transpose(0,1,4,2,3)

    def pij(tmp):
        '''P(ij)'''
        return tmp - tmp.transpose(1,0,2,3,4)

    def pab(tmp):
        '''P(ab)'''
        return tmp - tmp.transpose(0,1,3,2,4)

    eaccsd_evecs = np.array(eaccsd_evecs)
    leaccsd_evecs = np.array(leaccsd_evecs)
    e_star = []
    eaccsd_evecs, leaccsd_evecs = [np.atleast_2d(x) for x in [eaccsd_evecs, leaccsd_evecs]]
    eaccsd_evals = np.atleast_1d(eaccsd_evals)
    for ea_eval, ea_evec, ea_levec in zip(eaccsd_evals, eaccsd_evecs, leaccsd_evecs):
        # Enforcing <L|R> = 1
        l1, l2 = vector_to_amplitudes_ea(ea_levec, nmo, nocc)
        r1, r2 = vector_to_amplitudes_ea(ea_evec, nmo, nocc)
        ldotr = np.dot(l1, r1) + 0.5 * np.dot(l2.ravel(), r2.ravel())

        logger.info(eom, 'Left-right amplitude overlap : %14.8e', ldotr)
        if abs(ldotr) < 1e-7:
            logger.warn(eom, 'Small %s left-right amplitude overlap. Results '
                             'may be inaccurate.', ldotr)

        l1 /= ldotr
        l2 /= ldotr

        # Denominator + eigenvalue(EA-CCSD)
        denom = eijabc + ea_eval
        denom = 1. / denom

        tmp = lib.einsum('c,ijab->ijabc', l1, oovv)
        lijabc = -pabc(tmp)
        tmp = lib.einsum('jima,mbc->ijabc', ooov, l2)
        lijabc += -pabc(tmp)
        tmp = lib.einsum('ieab,jce->ijabc', ovvv, l2)
        tmp = pabc(tmp)
        lijabc += -pij(tmp)

        tmp = lib.einsum('bcef,f->bce', vvvv, r1)
        tmp = lib.einsum('bce,ijae->ijabc', tmp, t2)
        rijabc = -pabc(tmp)
        tmp = lib.einsum('mcje,e->mcj', ovov, r1)
        tmp = lib.einsum('mcj,imab->ijabc', tmp, t2)
        tmp = pabc(tmp)
        rijabc += pij(tmp)
        tmp = lib.einsum('amij,mcb->ijabc', vooo, r2)
        rijabc += pabc(tmp)
        tmp = lib.einsum('baei,jce->ijabc', vvvo, r2)
        tmp = pabc(tmp)
        rijabc -= pij(tmp)

        deltaE = (1. / 12) * lib.einsum('ijabc,ijabc,ijabc', lijabc, rijabc, denom)
        deltaE = deltaE.real
        logger.info(eom, "Exc. energy, delta energy = %16.12f, %16.12f",
                    ea_eval + deltaE, deltaE)
        e_star.append(ea_eval + deltaE)

    return e_star


class EOMEA(eom_rccsd.EOMEA):
    matvec = eaccsd_matvec
    l_matvec = leaccsd_matvec
    get_diag = eaccsd_diag
    ccsd_star_contract = eaccsd_star_contract

    def vector_to_amplitudes(self, vector, nmo=None, nocc=None):
        if nmo is None: nmo = self.nmo
        if nocc is None: nocc = self.nocc
        return vector_to_amplitudes_ea(vector, nmo, nocc)

    def amplitudes_to_vector(self, r1, r2):
        return amplitudes_to_vector_ea(r1, r2)

    def vector_size(self):
        nocc = self.nocc
        nvir = self.nmo - nocc
        return nvir + nocc*nvir*(nvir-1)//2

    def make_imds(self, eris=None):
        imds = _IMDS(self._cc, eris)
        imds.make_ea()
        return imds


class EOMEA_Ta(EOMEA):
    '''Class for EOM EACCSD(T)*(a) method by Matthews and Stanton.'''
    def make_imds(self, eris=None):
        imds = _IMDS(self._cc, eris=eris)
        imds.make_t3p2_ea(self._cc)
        return imds

########################################
# EOM-EE-CCSD
########################################

vector_to_amplitudes_ee = ccsd.vector_to_amplitudes_s4
amplitudes_to_vector_ee = ccsd.amplitudes_to_vector_s4

def eeccsd_matvec(eom, vector, imds=None, diag=None):
    # Ref: Wang, Tu, and Wang, J. Chem. Theory Comput. 10, 5567 (2014) Eqs.(9)-(10)
    # Note: Last line in Eq. (10) is superfluous.
    # See, e.g. Gwaltney, Nooijen, and Barlett, Chem. Phys. Lett. 248, 189 (1996)
    if imds is None: imds = eom.make_imds()
    nocc = eom.nocc
    nmo = eom.nmo
    r1, r2 = vector_to_amplitudes_ee(vector, nmo, nocc)

    # Eq. (9)
    Hr1  = lib.einsum('ae,ie->ia', imds.Fvv, r1)
    Hr1 -= lib.einsum('mi,ma->ia', imds.Foo, r1)
    Hr1 += lib.einsum('me,imae->ia', imds.Fov, r2)
    Hr1 += lib.einsum('maei,me->ia', imds.Wovvo, r1)
    Hr1 -= 0.5*lib.einsum('mnie,mnae->ia', imds.Wooov, r2)
    Hr1 += 0.5*lib.einsum('amef,imef->ia', imds.Wvovv, r2)
    # Eq. (10)
    tmpab = lib.einsum('be,ijae->ijab', imds.Fvv, r2)
    tmpab -= 0.5*lib.einsum('mnef,ijae,mnbf->ijab', imds.Woovv, imds.t2, r2)
    tmpab -= lib.einsum('mbij,ma->ijab', imds.Wovoo, r1)
    tmpab -= lib.einsum('amef,ijfb,me->ijab', imds.Wvovv, imds.t2, r1)
    tmpij  = lib.einsum('mj,imab->ijab', -imds.Foo, r2)
    tmpij -= 0.5*lib.einsum('mnef,imab,jnef->ijab', imds.Woovv, imds.t2, r2)
    tmpij += lib.einsum('abej,ie->ijab', imds.Wvvvo, r1)
    tmpij += lib.einsum('mnie,njab,me->ijab', imds.Wooov, imds.t2, r1)

    tmpabij = lib.einsum('mbej,imae->ijab', imds.Wovvo, r2)
    tmpabij = tmpabij - tmpabij.transpose(1,0,2,3)
    tmpabij = tmpabij - tmpabij.transpose(0,1,3,2)
    Hr2 = tmpabij

    Hr2 += tmpab - tmpab.transpose(0,1,3,2)
    Hr2 += tmpij - tmpij.transpose(1,0,2,3)
    Hr2 += 0.5*lib.einsum('mnij,mnab->ijab', imds.Woooo, r2)
    Hr2 += 0.5*lib.einsum('abef,ijef->ijab', imds.Wvvvv, r2)

    vector = amplitudes_to_vector_ee(Hr1, Hr2)
    return vector

def eeccsd_diag(eom, imds=None):
    if imds is None: imds = eom.make_imds()
    t1, t2 = imds.t1, imds.t2
    nocc, nvir = t1.shape

    Hr1 = np.zeros((nocc,nvir), dtype=t1.dtype)
    Hr2 = np.zeros((nocc,nocc,nvir,nvir), dtype=t1.dtype)
    for i in range(nocc):
        for a in range(nvir):
            Hr1[i,a] = imds.Fvv[a,a] - imds.Foo[i,i] + imds.Wovvo[i,a,a,i]
    for a in range(nvir):
        tmp = 0.5*(np.einsum('ijeb,ijbe->ijb', imds.Woovv, t2) -
                   np.einsum('jieb,ijbe->ijb', imds.Woovv, t2))
        Hr2[:,:,:,a] += imds.Fvv[a,a] + tmp
        Hr2[:,:,a,:] += imds.Fvv[a,a] + tmp
        _Wvvvva = np.array(imds.Wvvvv[a])
        for b in range(a):
            Hr2[:,:,a,b] += 0.5*(_Wvvvva[b,a,b]-_Wvvvva[b,b,a])
        for i in range(nocc):
            tmp = imds.Wovvo[i,a,a,i]
            Hr2[:,i,:,a] += tmp
            Hr2[i,:,:,a] += tmp
            Hr2[:,i,a,:] += tmp
            Hr2[i,:,a,:] += tmp
    for i in range(nocc):
        tmp = 0.5*(np.einsum('kjab,jkab->jab', imds.Woovv, t2) -
                   np.einsum('kjba,jkab->jab', imds.Woovv, t2))
        Hr2[:,i,:,:] += -imds.Foo[i,i] + tmp
        Hr2[i,:,:,:] += -imds.Foo[i,i] + tmp
        for j in range(i):
            Hr2[i,j,:,:] += 0.5*(imds.Woooo[i,j,i,j]-imds.Woooo[j,i,i,j])

    vector = amplitudes_to_vector_ee(Hr1, Hr2)
    return vector


def eeccsd(eom, nroots=1, koopmans=False, guess=None, eris=None, imds=None):
    '''Calculate N-electron neutral excitations via EOM-EE-CCSD.

    Kwargs:
        nroots : int
            Number of roots (eigenvalues) requested
        koopmans : bool
            Calculate Koopmans'-like (1p1h) excitations only, targeting via
            overlap.
        guess : list of ndarray
            List of guess vectors to use for targeting via overlap.
    '''
    return eom_rccsd.eomee_ccsd_singlet(eom, nroots, koopmans, guess, eris, imds)


class EOMEE(eom_rccsd.EOMEE):

    kernel = eeccsd
    eeccsd = eeccsd
    matvec = eeccsd_matvec
    get_diag = eeccsd_diag

    def gen_matvec(self, imds=None, **kwargs):
        if imds is None: imds = self.make_imds()
        diag = self.get_diag(imds)
        matvec = lambda xs: [self.matvec(x, imds) for x in xs]
        return matvec, diag

    def vector_to_amplitudes(self, vector, nmo=None, nocc=None):
        if nmo is None: nmo = self.nmo
        if nocc is None: nocc = self.nocc
        return vector_to_amplitudes_ee(vector, nmo, nocc)

    def amplitudes_to_vector(self, r1, r2):
        return amplitudes_to_vector_ee(r1, r2)

    def vector_size(self):
        '''size of the vector based on spin-orbital basis'''
        nocc = self.nocc
        nvir = self.nmo - nocc
        return nocc*nvir + nocc*(nocc-1)//2*nvir*(nvir-1)//2

    def make_imds(self, eris=None):
        imds = _IMDS(self._cc, eris)
        imds.make_ee()
        return imds

gccsd.GCCSD.EOMIP    = lib.class_as_method(EOMIP)
gccsd.GCCSD.EOMIP_Ta = lib.class_as_method(EOMIP_Ta)
gccsd.GCCSD.EOMEA    = lib.class_as_method(EOMEA)
gccsd.GCCSD.EOMEA_Ta = lib.class_as_method(EOMEA_Ta)
gccsd.GCCSD.EOMEE    = lib.class_as_method(EOMEE)

class _IMDS:
    # Exactly the same as RCCSD IMDS except
    # -- rintermediates --> gintermediates
    # -- Loo, Lvv, cc_Fov --> Foo, Fvv, Fov
    # -- One less 2-virtual intermediate
    def __init__(self, cc, eris=None):
        self.verbose = cc.verbose
        self.stdout = cc.stdout
        self.t1 = cc.t1
        self.t2 = cc.t2
        if eris is None:
            eris = cc.ao2mo()
        self.eris = eris
        self._made_shared = False
        self.made_ip_imds = False
        self.made_ea_imds = False
        self.made_ee_imds = False

    def _make_shared(self):
        cput0 = (logger.process_clock(), logger.perf_counter())

        t1, t2, eris = self.t1, self.t2, self.eris
        self.Foo = imd.Foo(t1, t2, eris)
        self.Fvv = imd.Fvv(t1, t2, eris)
        self.Fov = imd.Fov(t1, t2, eris)

        # 2 virtuals
        self.Wovvo = imd.Wovvo(t1, t2, eris)
        self.Woovv = eris.oovv

        self._made_shared = True
        logger.timer_debug1(self, 'EOM-CCSD shared intermediates', *cput0)
        return self

    def make_ip(self):
        if not self._made_shared:
            self._make_shared()

        cput0 = (logger.process_clock(), logger.perf_counter())

        t1, t2, eris = self.t1, self.t2, self.eris

        # 0 or 1 virtuals
        self.Woooo = imd.Woooo(t1, t2, eris)
        self.Wooov = imd.Wooov(t1, t2, eris)
        self.Wovoo = imd.Wovoo(t1, t2, eris)

        self.made_ip_imds = True
        logger.timer_debug1(self, 'EOM-CCSD IP intermediates', *cput0)
        return self

    def make_t3p2_ip(self, cc):
        cput0 = (logger.process_clock(), logger.perf_counter())

        t1, t2, eris = cc.t1, cc.t2, self.eris
        delta_E_corr, pt1, pt2, Wovoo, Wvvvo = \
            imd.get_t3p2_imds_slow(cc, t1, t2, eris)
        self.t1 = pt1
        self.t2 = pt2

        self._made_shared = False  # Force update
        self.make_ip()  # Make after t1/t2 updated
        self.Wovoo = self.Wovoo + Wovoo

        self.made_ip_imds = True
        logger.timer_debug1(self, 'EOM-CCSD(T)a IP intermediates', *cput0)
        return self


    def make_ea(self):
        if not self._made_shared:
            self._make_shared()

        cput0 = (logger.process_clock(), logger.perf_counter())

        t1, t2, eris = self.t1, self.t2, self.eris

        # 3 or 4 virtuals
        self.Wvovv = imd.Wvovv(t1, t2, eris)
        self.Wvvvv = imd.Wvvvv(t1, t2, eris)
        self.Wvvvo = imd.Wvvvo(t1, t2, eris,self.Wvvvv)

        self.made_ea_imds = True
        logger.timer_debug1(self, 'EOM-CCSD EA intermediates', *cput0)
        return self

    def make_t3p2_ea(self, cc):
        cput0 = (logger.process_clock(), logger.perf_counter())

        t1, t2, eris = cc.t1, cc.t2, self.eris
        delta_E_corr, pt1, pt2, Wovoo, Wvvvo = \
            imd.get_t3p2_imds_slow(cc, t1, t2, eris)
        self.t1 = pt1
        self.t2 = pt2

        self._made_shared = False  # Force update
        self.make_ea()  # Make after t1/t2 updated
        self.Wvvvo = self.Wvvvo + Wvvvo

        self.made_ea_imds = True
        logger.timer_debug1(self, 'EOM-CCSD(T)a EA intermediates', *cput0)
        return self


    def make_ee(self):
        if not self._made_shared:
            self._make_shared()

        cput0 = (logger.process_clock(), logger.perf_counter())

        t1, t2, eris = self.t1, self.t2, self.eris

        if not self.made_ip_imds:
            # 0 or 1 virtuals
            self.Woooo = imd.Woooo(t1, t2, eris)
            self.Wooov = imd.Wooov(t1, t2, eris)
            self.Wovoo = imd.Wovoo(t1, t2, eris)
        if not self.made_ea_imds:
            # 3 or 4 virtuals
            self.Wvovv = imd.Wvovv(t1, t2, eris)
            self.Wvvvv = imd.Wvvvv(t1, t2, eris)
            self.Wvvvo = imd.Wvvvo(t1, t2, eris,self.Wvvvv)

        self.made_ee_imds = True
        logger.timer(self, 'EOM-CCSD EE intermediates', *cput0)
        return self

def _cp(a):
    return np.array(a, copy=False, order='C')

if __name__ == '__main__':
    from pyscf import scf
    from pyscf import gto
    from pyscf.cc import gccsd
    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = 'cc-pvdz'
    mol.spin = 0
    mol.build()
    mf = scf.UHF(mol).run()
    mf = scf.addons.convert_to_ghf(mf)

    mycc = gccsd.GCCSD(mf)
    ecc, t1, t2 = mycc.kernel()
    print(ecc - -0.2133432712431435)
    e,v = mycc.ipccsd(nroots=8)
    print(e[0] - 0.4335604332073799)
    print(e[2] - 0.5187659896045407)
    print(e[4] - 0.6782876002229172)

    #mycc.verbose = 5
    e,v = mycc.eaccsd(nroots=8)
    print(e[0] - 0.16737886338859731)
    print(e[2] - 0.24027613852009164)
    print(e[4] - 0.51006797826488071)

    e,v = mycc.eeccsd(nroots=4)
    print(e[0] - 0.2757159395886167)
    print(e[1] - 0.2757159395886167)
    print(e[2] - 0.2757159395886167)
    print(e[3] - 0.3005716731825082)


# pyscf/cc/eom_rccsd.py

#!/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: Qiming Sun <osirpt.sun@gmail.com>
#         James D. McClain
#         Timothy Berkelbach <tim.berkelbach@gmail.com>
#


import numpy as np

from pyscf import lib
from pyscf import ao2mo
from pyscf.lib import logger
from pyscf.cc import ccsd
from pyscf.cc import rintermediates as imd
from pyscf import __config__


def kernel(eom, nroots=1, koopmans=False, guess=None, left=False,
           eris=None, imds=None, **kwargs):
    cput0 = (logger.process_clock(), logger.perf_counter())
    log = logger.Logger(eom.stdout, eom.verbose)
    if eom.verbose >= logger.WARN:
        eom.check_sanity()
    eom.dump_flags()

    if imds is None:
        imds = eom.make_imds(eris)

    matvec, diag = eom.gen_matvec(imds, left=left, **kwargs)

    size = eom.vector_size()
    nroots = min(nroots, size)
    if guess is not None:
        user_guess = True
        for g in guess:
            assert g.size == size
    else:
        user_guess = False
        guess = eom.get_init_guess(nroots, koopmans, diag)

    def precond(r, e0, x0):
        return r/(e0-diag+1e-12)

    # GHF or customized RHF/UHF may be of complex type
    real_system = (eom._cc._scf.mo_coeff[0].dtype == np.double)

    eig = lib.davidson_nosym1
    if user_guess or koopmans:
        assert len(guess) == nroots
        def eig_close_to_init_guess(w, v, nroots, envs):
            x0 = lib.linalg_helper._gen_x0(envs['v'], envs['xs'])
            s = np.dot(np.asarray(guess).conj(), np.asarray(x0).T)
            snorm = np.einsum('pi,pi->i', s.conj(), s)
            idx = np.argsort(-snorm)[:nroots]
            return lib.linalg_helper._eigs_cmplx2real(w, v, idx, real_system)
        conv, es, vs = eig(matvec, guess, precond, pick=eig_close_to_init_guess,
                           tol=eom.conv_tol, max_cycle=eom.max_cycle,
                           max_space=eom.max_space, nroots=nroots, verbose=log)
    else:
        def pickeig(w, v, nroots, envs):
            real_idx = np.where(abs(w.imag) < 1e-3)[0]
            return lib.linalg_helper._eigs_cmplx2real(w, v, real_idx, real_system)
        conv, es, vs = eig(matvec, guess, precond, pick=pickeig,
                           tol=eom.conv_tol, max_cycle=eom.max_cycle,
                           max_space=eom.max_space, nroots=nroots, verbose=log)

    if eom.verbose >= logger.INFO:
        for n, en, vn, convn in zip(range(nroots), es, vs, conv):
            r1, r2 = eom.vector_to_amplitudes(vn)
            if isinstance(r1, np.ndarray):
                qp_weight = np.linalg.norm(r1)**2
            else: # for EOM-UCCSD
                r1 = np.hstack([x.ravel() for x in r1])
                qp_weight = np.linalg.norm(r1)**2
            logger.info(eom, 'EOM-CCSD root %d E = %.16g  qpwt = %.6g  conv = %s',
                        n, en, qp_weight, convn)
        log.timer('EOM-CCSD', *cput0)
    if nroots == 1:
        return conv[0], es[0].real, vs[0]
    else:
        return conv, es.real, vs


class EOM(lib.StreamObject):
    def __init__(self, cc):
        self.mol = cc.mol
        self._cc = cc
        self.verbose = cc.verbose
        self.stdout = cc.stdout
        self.max_memory = cc.max_memory

        self.max_space = getattr(__config__, 'eom_rccsd_EOM_max_space', 20)
        self.max_cycle = getattr(__config__, 'eom_rccsd_EOM_max_cycle', cc.max_cycle)
        self.conv_tol = getattr(__config__, 'eom_rccsd_EOM_conv_tol', cc.conv_tol)
        self.partition = getattr(__config__, 'eom_rccsd_EOM_partition', None)

##################################################
# don't modify the following attributes, they are not input options
        self.e = None
        self.v = None
        self.nocc = cc.nocc
        self.nmo = cc.nmo
        self._keys = set(self.__dict__.keys())

    def dump_flags(self, verbose=None):
        logger.info(self, '')
        logger.info(self, '******** %s ********', self.__class__)
        logger.info(self, 'max_space = %d', self.max_space)
        logger.info(self, 'max_cycle = %d', self.max_cycle)
        logger.info(self, 'conv_tol = %s', self.conv_tol)
        logger.info(self, 'partition = %s', self.partition)
        #logger.info(self, 'nocc = %d', self.nocc)
        #logger.info(self, 'nmo = %d', self.nmo)
        logger.info(self, 'max_memory %d MB (current use %d MB)',
                    self.max_memory, lib.current_memory()[0])
        return self

    def reset(self, mol=None):
        self._cc.reset(mol)
        return self


def _sort_left_right_eigensystem(eom, right_converged, right_evals, right_evecs,
                                 left_converged, left_evals, left_evecs, tol=1e-6):
    '''Ensures the left and right eigenvectors correspond to the same eigenvalue.

    Note:
        Useful for perturbative methods that need both eigenstates.  Right now, just
        simply checks for equality between left and right eigenvalues, but can be
        extended to make sure the overlap between states is sufficiently large.

    Kwargs:
        eom : :class:`EOM`
            Class holding EOM results.
        right_converged : array-like of bool
            Whether the right eigenstates converged.
        right_evals : array-like
            Eigenvalues of right eigenstates.
        right_evecs : array-like of ndarray
            Eigenvectors of right eigenstates.
        left_converged : array-like of bool
            Whether the left eigenstates converged.
        left_evals : array-like
            Eigenvalues of left eigenstates.
        left_evecs : array-like of ndarray
            Eigenvectors of left eigenstates.
        tol : float
            Tolerance for determining whether a left and right eigenvalue
            should be considered equal.
    '''
    log = logger.Logger(eom.stdout, eom.verbose)

    right_evecs, left_evecs = [np.atleast_2d(x) for x in [right_evecs, left_evecs]]
    right_evals, left_evals = [np.atleast_1d(x) for x in [right_evals, left_evals]]
    right_converged, left_converged = [np.atleast_1d(x) for x in [right_converged, left_converged]]

    srt_right_idx = []
    srt_left_idx = []
    left_idx = [idx for idx in range(len(left_evals)) if left_converged[idx]]
    right_idx = [idx for idx in range(len(right_evals)) if right_converged[idx]]
    if len(right_idx) != len(left_idx):
        log.warn('Number of converged left and right eigenvalues are not equal.\n'
                 '    No. Left = %3d, No. Right = %3d.' %
                 (len(left_idx), len(right_idx)))

    for ir_idx, ir in enumerate(right_idx):
        found = False
        for il_idx, il in enumerate(left_idx):
            if abs(right_evals[ir] - left_evals[il]) < tol:
                found = True
                srt_right_idx.append(ir)
                srt_left_idx.append(il)
                break
        if found:
            left_idx.pop(il_idx)
        else:
            log.warn('No converged left eigenvalue corresponding to right eigenvalue '
                     '%.6g (right idx=%3d).\nWill not perform perturbation on this state.'
                     % (right_evals[ir], ir))

    log.info('Resulting left/right eigenstates:')
    log.info('Left Eigen (idx) <-> Right Eigen (idx)')
    for il, ir in zip(srt_left_idx, srt_right_idx):
        log.info('%10.6g (%3d)      %10.6g (%3d)',
                 left_evals[il], il, right_evals[ir], ir)
    return (right_evals[srt_right_idx], right_evecs[srt_right_idx], left_evecs[srt_left_idx])


def perturbed_ccsd_kernel(eom, nroots=1, koopmans=False, right_guess=None,
                          left_guess=None, eris=None, imds=None):
    '''Wrapper for running perturbative excited-states that require both left
    and right amplitudes.'''
    if imds is None:
        imds = eom.make_imds(eris=eris)

    # Right eigenvectors
    r_converged, r_e, r_v = \
               kernel(eom, nroots, koopmans=koopmans, guess=right_guess, left=False,
                      eris=eris, imds=imds)
    # Left eigenvectors
    l_converged, l_e, l_v = \
               kernel(eom, nroots, koopmans=koopmans, guess=right_guess, left=True,
                      eris=eris, imds=imds)

    e, r_v, l_v = _sort_left_right_eigensystem(eom, r_converged, r_e, r_v, l_converged, l_e, l_v)
    e_star = eom.ccsd_star_contract(e, r_v, l_v, imds=imds)
    return e_star


########################################
# EOM-IP-CCSD
########################################

def ipccsd(eom, nroots=1, left=False, koopmans=False, guess=None,
           partition=None, eris=None, imds=None):
    '''Calculate (N-1)-electron charged excitations via IP-EOM-CCSD.

    Kwargs:
        nroots : int
            Number of roots (eigenvalues) requested
        partition : bool or str
            Use a matrix-partitioning for the doubles-doubles block.
            Can be None, 'mp' (Moller-Plesset, i.e. orbital energies on the diagonal),
            or 'full' (full diagonal elements).
        koopmans : bool
            Calculate Koopmans'-like (quasiparticle) excitations only, targeting via
            overlap.
        guess : list of ndarray
            List of guess vectors to use for targeting via overlap.
    '''
    if partition is not None:
        eom.partition = partition.lower()
        assert eom.partition in ['mp','full']
    eom.converged, eom.e, eom.v \
            = kernel(eom, nroots, koopmans, guess, left, eris=eris, imds=imds)
    return eom.e, eom.v

def ipccsd_star(eom, nroots=1, koopmans=False, right_guess=None,
                left_guess=None, eris=None, imds=None):
    """Calculates CCSD* perturbative correction.

    Simply calls the relevant `kernel()` function and `perturb_star` of the
    `eom` class.

    Returns:
        e_t_a_star (list of float):
            The IP-CCSD* energy.
    """
    return perturbed_ccsd_kernel(eom, nroots=nroots, koopmans=koopmans,
                                 right_guess=right_guess, left_guess=left_guess, eris=eris,
                                 imds=imds)

def vector_to_amplitudes_ip(vector, nmo, nocc):
    nvir = nmo - nocc
    r1 = vector[:nocc].copy()
    r2 = vector[nocc:].copy().reshape(nocc,nocc,nvir)
    return r1, r2

def amplitudes_to_vector_ip(r1, r2):
    vector = np.hstack((r1, r2.ravel()))
    return vector

def ipccsd_matvec(eom, vector, imds=None, diag=None):
    # Ref: Nooijen and Snijders, J. Chem. Phys. 102, 1681 (1995) Eqs.(8)-(9)
    if imds is None: imds = eom.make_imds()
    nocc = eom.nocc
    nmo = eom.nmo
    r1, r2 = vector_to_amplitudes_ip(vector, nmo, nocc)

    # 1h-1h block
    Hr1 = -np.einsum('ki,k->i', imds.Loo, r1)
    #1h-2h1p block
    Hr1 += 2*np.einsum('ld,ild->i', imds.Fov, r2)
    Hr1 +=  -np.einsum('kd,kid->i', imds.Fov, r2)
    Hr1 += -2*np.einsum('klid,kld->i', imds.Wooov, r2)
    Hr1 +=    np.einsum('lkid,kld->i', imds.Wooov, r2)

    # 2h1p-1h block
    Hr2 = -np.einsum('kbij,k->ijb', imds.Wovoo, r1)
    # 2h1p-2h1p block
    if eom.partition == 'mp':
        fock = imds.eris.fock
        foo = fock[:nocc,:nocc]
        fvv = fock[nocc:,nocc:]
        Hr2 += lib.einsum('bd,ijd->ijb', fvv, r2)
        Hr2 += -lib.einsum('ki,kjb->ijb', foo, r2)
        Hr2 += -lib.einsum('lj,ilb->ijb', foo, r2)
    elif eom.partition == 'full':
        diag_matrix2 = vector_to_amplitudes_ip(diag, nmo, nocc)[1]
        Hr2 += diag_matrix2 * r2
    else:
        Hr2 += lib.einsum('bd,ijd->ijb', imds.Lvv, r2)
        Hr2 += -lib.einsum('ki,kjb->ijb', imds.Loo, r2)
        Hr2 += -lib.einsum('lj,ilb->ijb', imds.Loo, r2)
        Hr2 +=  lib.einsum('klij,klb->ijb', imds.Woooo, r2)
        Hr2 += 2*lib.einsum('lbdj,ild->ijb', imds.Wovvo, r2)
        Hr2 +=  -lib.einsum('kbdj,kid->ijb', imds.Wovvo, r2)
        Hr2 +=  -lib.einsum('lbjd,ild->ijb', imds.Wovov, r2) #typo in Ref
        Hr2 +=  -lib.einsum('kbid,kjd->ijb', imds.Wovov, r2)
        tmp = 2*np.einsum('lkdc,kld->c', imds.Woovv, r2)
        tmp += -np.einsum('kldc,kld->c', imds.Woovv, r2)
        Hr2 += -np.einsum('c,ijcb->ijb', tmp, imds.t2)

    vector = amplitudes_to_vector_ip(Hr1, Hr2)
    return vector

def lipccsd_matvec(eom, vector, imds=None, diag=None):
    '''For left eigenvector'''
    # Note this is not the same left EA equations used by Nooijen and Bartlett.
    # Small changes were made so that the same type L2 basis was used for both the
    # left EA and left IP equations.  You will note more similarity for these
    # equations to the left IP equations than for the left EA equations by Nooijen.
    if imds is None: imds = eom.make_imds()
    nocc = eom.nocc
    nmo = eom.nmo
    r1, r2 = vector_to_amplitudes_ip(vector, nmo, nocc)

    # 1h-1h block
    Hr1 = -np.einsum('ki,i->k', imds.Loo, r1)
    #1h-2h1p block
    Hr1 += -np.einsum('kbij,ijb->k', imds.Wovoo, r2)

    # 2h1p-1h block
    Hr2 = -np.einsum('kd,l->kld', imds.Fov, r1)
    Hr2 += 2.*np.einsum('ld,k->kld', imds.Fov, r1)
    Hr2 += -np.einsum('klid,i->kld', 2.*imds.Wooov-imds.Wooov.transpose(1,0,2,3), r1)
    # 2h1p-2h1p block
    if eom.partition == 'mp':
        fock = imds.eris.fock
        foo = fock[:nocc,:nocc]
        fvv = fock[nocc:,nocc:]
        Hr2 += lib.einsum('bd,klb->kld', fvv, r2)
        Hr2 += -lib.einsum('ki,ild->kld', foo, r2)
        Hr2 += -lib.einsum('lj,kjd->kld', foo, r2)
    elif eom.partition == 'full':
        diag_matrix2 = vector_to_amplitudes_ip(diag, nmo, nocc)[1]
        Hr2 += diag_matrix2 * r2
    else:
        Hr2 += lib.einsum('bd,klb->kld', imds.Lvv, r2)
        Hr2 += -lib.einsum('ki,ild->kld', imds.Loo, r2)
        Hr2 += -lib.einsum('lj,kjd->kld', imds.Loo, r2)
        Hr2 += lib.einsum('lbdj,kjb->kld', 2.*imds.Wovvo-imds.Wovov.transpose(0,1,3,2), r2)
        Hr2 += -lib.einsum('kbdj,ljb->kld', imds.Wovvo, r2)
        Hr2 += lib.einsum('klij,ijd->kld', imds.Woooo, r2)
        Hr2 += -lib.einsum('kbid,ilb->kld', imds.Wovov, r2)
        tmp = np.einsum('ijcb,ijb->c', imds.t2, r2)
        Hr2 += -np.einsum('lkdc,c->kld', 2.*imds.Woovv-imds.Woovv.transpose(1,0,2,3), tmp)

    vector = amplitudes_to_vector_ip(Hr1, Hr2)
    return vector

def ipccsd_diag(eom, imds=None):
    if imds is None: imds = eom.make_imds()
    t1, t2 = imds.t1, imds.t2
    dtype = np.result_type(t1, t2)
    nocc, nvir = t1.shape
    fock = imds.eris.fock
    foo = fock[:nocc,:nocc]
    fvv = fock[nocc:,nocc:]

    Hr1 = -np.diag(imds.Loo)
    Hr2 = np.zeros((nocc,nocc,nvir), dtype)
    for i in range(nocc):
        for j in range(nocc):
            for b in range(nvir):
                if eom.partition == 'mp':
                    Hr2[i,j,b] += fvv[b,b]
                    Hr2[i,j,b] += -foo[i,i]
                    Hr2[i,j,b] += -foo[j,j]
                else:
                    Hr2[i,j,b] += imds.Lvv[b,b]
                    Hr2[i,j,b] += -imds.Loo[i,i]
                    Hr2[i,j,b] += -imds.Loo[j,j]
                    Hr2[i,j,b] +=  imds.Woooo[i,j,i,j]
                    Hr2[i,j,b] +=2*imds.Wovvo[j,b,b,j]
                    Hr2[i,j,b] += -imds.Wovvo[i,b,b,i]*(i==j)
                    Hr2[i,j,b] += -imds.Wovov[j,b,j,b]
                    Hr2[i,j,b] += -imds.Wovov[i,b,i,b]
                    Hr2[i,j,b] += -2*np.dot(imds.Woovv[j,i,b,:], t2[i,j,:,b])
                    Hr2[i,j,b] += np.dot(imds.Woovv[i,j,b,:], t2[i,j,:,b])

    vector = amplitudes_to_vector_ip(Hr1, Hr2)
    return vector

def ipccsd_star_contract(eom, ipccsd_evals, ipccsd_evecs, lipccsd_evecs, imds=None):
    from pyscf.cc.ccsd_t import _sort_eri, _sort_t2_vooo_
    cpu1 = (logger.process_clock(), logger.perf_counter())
    log = logger.Logger(eom.stdout, eom.verbose)
    if imds is None:
        imds = eom.make_imds()
    t1, t2 = imds.t1, imds.t2
    eris = imds.eris

    nocc, nvir = t1.shape
    nmo = nocc + nvir

    dtype = np.result_type(t1, t2, eris.ovoo.dtype)
    if eom._cc.incore_complete:
        ftmp = None
        eris_vvop = np.zeros((nvir,nvir,nocc,nmo), dtype)
    else:
        ftmp = lib.H5TmpFile()
        eris_vvop = ftmp.create_dataset('vvop', (nvir,nvir,nocc,nmo), dtype)

    orbsym = _sort_eri(eom._cc, eris, nocc, nvir, eris_vvop, log)
    mo_energy, t1T, t2T, vooo, fvo, restore_t2_inplace = \
            _sort_t2_vooo_(eom._cc, orbsym, t1, t2, eris)

    cpu1 = log.timer_debug1('CCSD(T) sort_eri', *cpu1)

    mem_now = lib.current_memory()[0]
    max_memory = max(0, eom.max_memory - mem_now)
    blksize = min(nvir, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nocc**3*6))))

    mo_e_occ = np.asarray(mo_energy[:nocc])
    mo_e_vir = np.asarray(mo_energy[nocc:])

    def contract_l2p(l1, l2, a0, a1, b0, b1, cache_vvop, out=None):
        '''Create perturbed l2.'''
        if out is None:
            out = np.zeros((nocc,)*3 + (a1-a0,b1-b0), dtype=dtype)
        out += 0.5*np.einsum('abij,k->ijkab', cache_vvop[:,:,:,:nocc].conj(), l1)
        out += lib.einsum('abie,jke->ijkab', cache_vvop[:,:,:,nocc:].conj(), l2)
        out += -lib.einsum('bjkm,ima->ijkab', vooo[b0:b1], l2[:,:,a0:a1])
        out += -lib.einsum('bjim,mka->ijkab', vooo[b0:b1], l2[:,:,a0:a1])
        return out

    def contract_pl2p(l1, l2, a0, a1, b0, b1, cache_vvop_a, cache_vvop_b):
        '''Create P(ia|jb) of perturbed l2.'''
        out = contract_l2p(l1, l2, a0, a1, b0, b1, cache_vvop_a)
        out += contract_l2p(l1, l2, b0, b1, a0, a1, cache_vvop_b).transpose(1,0,2,4,3)  # P(ia|jb)
        return out

    def contract_r2p(r1, r2, a0, a1, b0, b1, cache_vvop, out=None):
        '''Create perturbed r2.'''
        if out is None:
            out = np.zeros((nocc,)*3 + (a1-a0,b1-b0), dtype=dtype)
        tmp = np.einsum('mkbe,m->bke', eris.oovv[:,:,b0:b1,:], r1)
        out += -lib.einsum('bke,aeji->ijkab', tmp, t2T[a0:a1])
        tmp = np.einsum('mebj,m->bej', eris.ovvo[:,:,b0:b1,:], r1)
        out += -lib.einsum('bej,aeki->ijkab', tmp, t2T[a0:a1])
        tmp = np.einsum('mjnk,n->mjk', eris.oooo, r1)
        out += lib.einsum('mjk,abmi->ijkab', tmp, t2T[a0:a1,b0:b1])
        out += lib.einsum('abie,kje->ijkab', cache_vvop[:,:,:,nocc:], r2)
        out += -lib.einsum('bjkm,mia->ijkab', vooo[b0:b1].conj(), r2[:,:,a0:a1])
        out += -lib.einsum('bjim,kma->ijkab', vooo[b0:b1].conj(), r2[:,:,a0:a1])
        return out

    def contract_pr2p(r1, r2, a0, a1, b0, b1, cache_vvop_a, cache_vvop_b):
        '''Create P(ia|jb) of perturbed r2.'''
        out = contract_r2p(r1, r2, a0, a1, b0, b1, cache_vvop_a)
        out += contract_r2p(r1, r2, b0, b1, a0, a1, cache_vvop_b).transpose(1,0,2,4,3)  # P(ia|jb)
        return out

    ipccsd_evecs  = np.array(ipccsd_evecs)
    lipccsd_evecs = np.array(lipccsd_evecs)
    e = []
    ipccsd_evecs, lipccsd_evecs = [np.atleast_2d(x) for x in [ipccsd_evecs, lipccsd_evecs]]
    ipccsd_evals = np.atleast_1d(ipccsd_evals)
    for eval_, evec_, levec_ in zip(ipccsd_evals, ipccsd_evecs, lipccsd_evecs):
        l1, l2 = eom.vector_to_amplitudes(levec_)
        r1, r2 = eom.vector_to_amplitudes(evec_)
        ldotr = np.dot(l1, r1) + np.dot(l2.ravel(), r2.ravel())
        l1 /= ldotr
        l2 /= ldotr
        l2 = 1./3*(l2 + 2.*l2.transpose(1,0,2))

        deltaE = 0.0
        eijk = (mo_e_occ[:,None,None,None,None] +
                mo_e_occ[None,:,None,None,None] +
                mo_e_occ[None,None,:,None,None] + eval_)
        for a0, a1 in lib.prange_tril(0, nvir, blksize):
            b0, b1 = 0, a1
            eijkab = (eijk - mo_e_vir[a0:a1][None,None,None,:,None] -
                      mo_e_vir[b0:b1][None,None,None,None,:])
            eijkab = 1./eijkab
            vvov_a = eris_vvop[a0:a1,b0:b1,:,:]
            vvov_b = eris_vvop[b0:b1,a0:a1,:,:]
            lijkab = contract_pl2p(l1, l2, a0, a1, b0, b1, vvov_a, vvov_b)
            rijkab = contract_pr2p(r1, r2, a0, a1, b0, b1, vvov_a, vvov_b)

            lijkab = 4.*lijkab \
                   - 2.*lijkab.transpose(1,0,2,3,4) \
                   - 2.*lijkab.transpose(2,1,0,3,4) \
                   - 2.*lijkab.transpose(0,2,1,3,4) \
                   + 1.*lijkab.transpose(1,2,0,3,4) \
                   + 1.*lijkab.transpose(2,0,1,3,4)

            # Symmetry factors (1 for a == b, 2 for a < b)
            fac = 2*np.ones_like(rijkab, dtype=int)
            triu_idx = np.triu_indices(a1-a0, a0+1, m=b1-b0)
            fac[:,:,:,triu_idx[0],triu_idx[1]] = 0
            fac[:,:,:,np.arange(a1-a0),np.arange(a0,b1)] = 1
            eijkab *= fac

            deltaE += np.einsum('ijkab,ijkab,ijkab', lijkab, rijkab, eijkab)
        deltaE = 0.5*deltaE.real
        logger.info(eom, "ipccsd energy, star energy, delta energy = %16.12f, %16.12f, %16.12f",
                    eval_, eval_+deltaE, deltaE)
        e.append(eval_+deltaE)
    t2 = restore_t2_inplace(t2T)
    return e

class EOMIP(EOM):
    def get_init_guess(self, nroots=1, koopmans=True, diag=None):
        size = self.vector_size()
        dtype = getattr(diag, 'dtype', np.double)
        nroots = min(nroots, size)
        guess = []
        if koopmans:
            for n in range(nroots):
                g = np.zeros(int(size), dtype)
                g[self.nocc-n-1] = 1.0
                guess.append(g)
        else:
            idx = diag.argsort()[:nroots]
            for i in idx:
                g = np.zeros(int(size), dtype)
                g[i] = 1.0
                guess.append(g)
        return guess

    kernel = ipccsd
    ipccsd = ipccsd
    ipccsd_star = ipccsd_star

    matvec = ipccsd_matvec
    l_matvec = lipccsd_matvec
    get_diag = ipccsd_diag
    ccsd_star_contract = ipccsd_star_contract

    def ipccsd_star_contract(self, ipccsd_evals, ipccsd_evecs, lipccsd_evecs, imds=None):
        return self.ccsd_star_contract(ipccsd_evals, ipccsd_evecs, lipccsd_evecs, imds=imds)

    def gen_matvec(self, imds=None, left=False, **kwargs):
        if imds is None: imds = self.make_imds()
        diag = self.get_diag(imds)
        if left:
            matvec = lambda xs: [self.l_matvec(x, imds, diag) for x in xs]
        else:
            matvec = lambda xs: [self.matvec(x, imds, diag) for x in xs]
        return matvec, diag

    def vector_to_amplitudes(self, vector, nmo=None, nocc=None):
        if nmo is None: nmo = self.nmo
        if nocc is None: nocc = self.nocc
        return vector_to_amplitudes_ip(vector, nmo, nocc)

    def amplitudes_to_vector(self, r1, r2):
        return amplitudes_to_vector_ip(r1, r2)

    def vector_size(self):
        nocc = self.nocc
        nvir = self.nmo - nocc
        return nocc + nocc*nocc*nvir

    def make_imds(self, eris=None):
        imds = _IMDS(self._cc, eris=eris)
        imds.make_ip(self.partition)
        return imds

    @property
    def eip(self):
        return self.e


class EOMIP_Ta(EOMIP):
    '''Class for EOM IPCCSD(T)*(a) method by Matthews and Stanton.'''
    def make_imds(self, eris=None):
        imds = _IMDS(self._cc, eris=eris)
        imds.make_t3p2_ip(self._cc, self.partition)
        return imds

########################################
# EOM-EA-CCSD
########################################

def eaccsd(eom, nroots=1, left=False, koopmans=False, guess=None,
           partition=None, eris=None, imds=None):
    '''Calculate (N+1)-electron charged excitations via EA-EOM-CCSD.

    Args:
        See also ipccd()
    '''
    return ipccsd(eom, nroots, left, koopmans, guess, partition, eris, imds)

def eaccsd_star(eom, nroots=1, koopmans=False, right_guess=None,
                left_guess=None, eris=None, imds=None, **kwargs):
    """Calculates CCSD* perturbative correction.

    Args:
        See also ipccd_star()
    """
    return perturbed_ccsd_kernel(eom, nroots=nroots, koopmans=koopmans,
                                 right_guess=right_guess, left_guess=left_guess, eris=eris,
                                 imds=imds)

def vector_to_amplitudes_ea(vector, nmo, nocc):
    nvir = nmo - nocc
    r1 = vector[:nvir].copy()
    r2 = vector[nvir:].copy().reshape(nocc,nvir,nvir)
    return r1, r2

def amplitudes_to_vector_ea(r1, r2):
    vector = np.hstack((r1, r2.ravel()))
    return vector

def eaccsd_matvec(eom, vector, imds=None, diag=None):
    # Ref: Nooijen and Bartlett, J. Chem. Phys. 102, 3629 (1995) Eqs.(30)-(31)
    if imds is None: imds = eom.make_imds()
    nocc = eom.nocc
    nmo = eom.nmo
    nvir = nmo - nocc
    r1, r2 = vector_to_amplitudes_ea(vector, nmo, nocc)

    # Eq. (37)
    # 1p-1p block
    Hr1 =  np.einsum('ac,c->a', imds.Lvv, r1)
    # 1p-2p1h block
    Hr1 += np.einsum('ld,lad->a', 2.*imds.Fov, r2)
    Hr1 += np.einsum('ld,lda->a',   -imds.Fov, r2)
    Hr1 += np.einsum('alcd,lcd->a', 2.*imds.Wvovv-imds.Wvovv.transpose(0,1,3,2), r2)
    # Eq. (38)
    # 2p1h-1p block
    Hr2 = np.einsum('abcj,c->jab', imds.Wvvvo, r1)
    # 2p1h-2p1h block
    if eom.partition == 'mp':
        fock = imds.eris.fock
        foo = fock[:nocc,:nocc]
        fvv = fock[nocc:,nocc:]
        Hr2 +=  lib.einsum('ac,jcb->jab', fvv, r2)
        Hr2 +=  lib.einsum('bd,jad->jab', fvv, r2)
        Hr2 += -lib.einsum('lj,lab->jab', foo, r2)
    elif eom.partition == 'full':
        diag_matrix2 = vector_to_amplitudes_ea(diag, nmo, nocc)[1]
        Hr2 += diag_matrix2 * r2
    else:
        Hr2 +=  lib.einsum('ac,jcb->jab', imds.Lvv, r2)
        Hr2 +=  lib.einsum('bd,jad->jab', imds.Lvv, r2)
        Hr2 += -lib.einsum('lj,lab->jab', imds.Loo, r2)
        Hr2 += lib.einsum('lbdj,lad->jab', 2.*imds.Wovvo-imds.Wovov.transpose(0,1,3,2), r2)
        Hr2 += -lib.einsum('lajc,lcb->jab', imds.Wovov, r2)
        Hr2 += -lib.einsum('lbcj,lca->jab', imds.Wovvo, r2)
        for a in range(nvir):
            Hr2[:,a,:] += lib.einsum('bcd,jcd->jb', imds.Wvvvv[a], r2)
        tmp = np.einsum('klcd,lcd->k', 2.*imds.Woovv-imds.Woovv.transpose(0,1,3,2), r2)
        Hr2 += -np.einsum('k,kjab->jab', tmp, imds.t2)

    vector = amplitudes_to_vector_ea(Hr1,Hr2)
    return vector

def leaccsd_matvec(eom, vector, imds=None, diag=None):
    # Note this is not the same left EA equations used by Nooijen and Bartlett.
    # Small changes were made so that the same type L2 basis was used for both the
    # left EA and left IP equations.  You will note more similarity for these
    # equations to the left IP equations than for the left EA equations by Nooijen.
    if imds is None: imds = eom.make_imds()
    nocc = eom.nocc
    nmo = eom.nmo
    nvir = nmo - nocc
    r1, r2 = vector_to_amplitudes_ea(vector, nmo, nocc)

    # Eq. (30)
    # 1p-1p block
    Hr1 = np.einsum('ac,a->c', imds.Lvv, r1)
    # 1p-2p1h block
    Hr1 += np.einsum('abcj,jab->c', imds.Wvvvo, r2)
    # Eq. (31)
    # 2p1h-1p block
    Hr2 = 2.*np.einsum('c,ld->lcd', r1, imds.Fov)
    Hr2 +=  -np.einsum('d,lc->lcd', r1, imds.Fov)
    Hr2 += np.einsum('a,alcd->lcd', r1, 2.*imds.Wvovv-imds.Wvovv.transpose(0,1,3,2))
    # 2p1h-2p1h block
    if eom.partition == 'mp':
        fock = imds.eris.fock
        foo = fock[:nocc,:nocc]
        fvv = fock[nocc:,nocc:]
        Hr2 += lib.einsum('lad,ac->lcd', r2, fvv)
        Hr2 += lib.einsum('lcb,bd->lcd', r2, fvv)
        Hr2 += -lib.einsum('jcd,lj->lcd', r2, foo)
    elif eom.partition == 'full':
        diag_matrix2 = vector_to_amplitudes_ea(diag, nmo, nocc)[1]
        Hr2 += diag_matrix2 * r2
    else:
        Hr2 += lib.einsum('lad,ac->lcd', r2, imds.Lvv)
        Hr2 += lib.einsum('lcb,bd->lcd', r2, imds.Lvv)
        Hr2 += -lib.einsum('jcd,lj->lcd', r2, imds.Loo)
        Hr2 += lib.einsum('jcb,lbdj->lcd', r2, 2.*imds.Wovvo-imds.Wovov.transpose(0,1,3,2))
        Hr2 += -lib.einsum('lajc,jab->lcb', imds.Wovov, r2)
        Hr2 += -lib.einsum('lbcj,jab->lca', imds.Wovvo, r2)
        for a in range(nvir):
            Hr2 += lib.einsum('lb,bcd->lcd', r2[:,a,:], imds.Wvvvv[a])
        tmp = np.einsum('ijcb,ibc->j', imds.t2, r2)
        Hr2 += -np.einsum('kjfe,j->kef', 2.*imds.Woovv-imds.Woovv.transpose(0,1,3,2),tmp)

    vector = amplitudes_to_vector_ea(Hr1,Hr2)
    return vector

def eaccsd_diag(eom, imds=None):
    if imds is None: imds = eom.make_imds()
    t1, t2 = imds.t1, imds.t2
    dtype = np.result_type(t1, t2)
    nocc, nvir = t1.shape

    fock = imds.eris.fock
    foo = fock[:nocc,:nocc]
    fvv = fock[nocc:,nocc:]

    Hr1 = np.diag(imds.Lvv)
    Hr2 = np.zeros((nocc,nvir,nvir), dtype)
    for a in range(nvir):
        if eom.partition != 'mp':
            _Wvvvva = np.array(imds.Wvvvv[a])
        for b in range(nvir):
            for j in range(nocc):
                if eom.partition == 'mp':
                    Hr2[j,a,b] += fvv[a,a]
                    Hr2[j,a,b] += fvv[b,b]
                    Hr2[j,a,b] += -foo[j,j]
                else:
                    Hr2[j,a,b] += imds.Lvv[a,a]
                    Hr2[j,a,b] += imds.Lvv[b,b]
                    Hr2[j,a,b] += -imds.Loo[j,j]
                    Hr2[j,a,b] += 2*imds.Wovvo[j,b,b,j]
                    Hr2[j,a,b] += -imds.Wovov[j,b,j,b]
                    Hr2[j,a,b] += -imds.Wovov[j,a,j,a]
                    Hr2[j,a,b] += -imds.Wovvo[j,b,b,j]*(a==b)
                    Hr2[j,a,b] += _Wvvvva[b,a,b]
                    Hr2[j,a,b] += -2*np.dot(imds.Woovv[:,j,a,b], t2[:,j,a,b])
                    Hr2[j,a,b] += np.dot(imds.Woovv[:,j,b,a], t2[:,j,a,b])

    vector = amplitudes_to_vector_ea(Hr1,Hr2)
    return vector

def eaccsd_star_contract(eom, eaccsd_evals, eaccsd_evecs, leaccsd_evecs, imds=None):
    cpu1 = (logger.process_clock(), logger.perf_counter())
    log = logger.Logger(eom.stdout, eom.verbose)
    if imds is None:
        imds = eom.make_imds()
    t1, t2 = imds.t1, imds.t2
    eris = imds.eris

    nocc, nvir = t1.shape
    dtype = np.result_type(t1, t2, eris.ovoo.dtype)
    # Notice we do not use `sort_eri` as compared to the eaccsd_star.
    # The sort_eri does not produce eri's that are read-in quickly for the current contraction
    # scheme.  Here, we have that the block loop is over occupied indices whereas in the
    # sort_eri it is done over virtual indices (due to the permutation over occupied indices
    # in ipccsd_star versus virtual indices in eaccsd_star).
    cpu1 = log.timer_debug1('CCSD(T) sort_eri', *cpu1)  # Left if new sort_eri implemented

    mem_now = lib.current_memory()[0]
    max_memory = max(0, eom.max_memory - mem_now)
    blksize = min(nocc, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvir**3*6))))

    mo_energy = np.asarray(eris.mo_energy)
    mo_e_occ = np.asarray(mo_energy[:nocc])
    mo_e_vir = np.asarray(mo_energy[nocc:])

    def contract_l2p(l1, l2, i0, i1, j0, j1, cache_ovvv_i, cache_ovvv_j, out=None):
        '''Create perturbed l2.'''
        if out is None:
            out = np.zeros((i1-i0,j1-j0) + (nvir,)*3, dtype=dtype)
        out += -0.5*np.einsum('iajb,c->ijabc', eris.ovov[i0:i1,:,j0:j1], l1)
        out += lib.einsum('iajm,mbc->ijabc', eris.ovoo[i0:i1,:,j0:j1], l2)
        out -= lib.einsum('iaeb,jec->ijabc', cache_ovvv_i, l2[j0:j1])
        out -= lib.einsum('jbec,iae->ijabc', cache_ovvv_j, l2[i0:i1])
        return out

    def contract_pl2p(l1, l2, i0, i1, j0, j1, cache_ovvv_i, cache_ovvv_j):
        '''Create P(ia|jb) of perturbed l2.'''
        out = contract_l2p(l1, l2, i0, i1, j0, j1, cache_ovvv_i, cache_ovvv_j)
        tmp = contract_l2p(l1, l2, j0, j1, i0, i1, cache_ovvv_j, cache_ovvv_i)
        tmp = tmp.transpose(1,0,3,2,4)  # P(ia|jb)
        out = out + tmp
        return out

    def _get_vvvv(eris):
        if eris.vvvv is None and getattr(eris, 'vvL', None) is not None:  # DF eris
            vvL = np.asarray(eris.vvL)
            nvir = int(np.sqrt(eris.vvL.shape[0]*2))
            return ao2mo.restore(1, lib.dot(vvL, vvL.T), nvir)
        elif eris.vvvv.ndim == 2:
            nvir = int(np.sqrt(eris.vvvv.shape[0]*2))
            return ao2mo.restore(1, np.asarray(eris.vvvv), nvir)
        else:
            return eris.vvvv

    def contract_r2p(r1, r2, i0, i1, j0, j1, cache_ovvv_i, cache_ovvv_j, out=None):
        '''Create perturbed r2.'''
        if out is None:
            out = np.zeros((i1-i0,j1-j0) + (nvir,)*3, dtype=dtype)
        tmp = np.einsum('becf,f->bce', _get_vvvv(eris), r1)
        out += -lib.einsum('bce,ijae->ijabc', tmp, t2[i0:i1,j0:j1])
        tmp = np.einsum('mjce,e->mcj', eris.oovv[:,j0:j1], r1)
        out += lib.einsum('mcj,imab->ijabc', tmp, t2[i0:i1])
        tmp = np.einsum('jbem,e->mbj', eris.ovvo[j0:j1], r1)
        out += lib.einsum('mbj,imac->ijabc', tmp, t2[i0:i1])
        out += lib.einsum('iajm,mbc->ijabc', eris.ovoo[i0:i1,:,j0:j1].conj(), r2)
        out += -lib.einsum('iaeb,jec->ijabc', cache_ovvv_i.conj(), r2[j0:j1])
        out += -lib.einsum('jbec,iae->ijabc', cache_ovvv_j.conj(), r2[i0:i1])
        return out

    def contract_pr2p(r1, r2, i0, i1, j0, j1, cache_ovvv_i, cache_ovvv_j):
        '''Create P(ia|jb) of perturbed r2.'''
        out = contract_r2p(r1, r2, i0, i1, j0, j1, cache_ovvv_i, cache_ovvv_j)
        tmp = contract_r2p(r1, r2, j0, j1, i0, i1, cache_ovvv_j, cache_ovvv_i)
        tmp = tmp.transpose(1,0,3,2,4)  # P(ia|jb)
        out = out + tmp
        return out

    eaccsd_evecs  = np.array(eaccsd_evecs)
    leaccsd_evecs = np.array(leaccsd_evecs)
    e = []
    eaccsd_evecs, leaccsd_evecs = [np.atleast_2d(x) for x in [eaccsd_evecs, leaccsd_evecs]]
    eaccsd_evals = np.atleast_1d(eaccsd_evals)
    for eval_, evec_, levec_ in zip(eaccsd_evals, eaccsd_evecs, leaccsd_evecs):
        l1, l2 = eom.vector_to_amplitudes(levec_)
        r1, r2 = eom.vector_to_amplitudes(evec_)
        ldotr = np.dot(l1, r1) + np.dot(l2.ravel(),r2.ravel())
        l1 /= ldotr
        l2 /= ldotr
        l2 = 1./3*(1.*l2 + 2.*l2.transpose(0,2,1))
        r2 = r2.transpose(0,2,1)

        deltaE = 0.0
        eabc = (mo_e_vir[None,None,:,None,None] +
                mo_e_vir[None,None,None,:,None] +
                mo_e_vir[None,None,None,None,:] - eval_)

        for i0, i1 in lib.prange_tril(0, nocc, blksize):
            j0, j1 = 0, i1
            eijabc = (mo_e_occ[i0:i1][:,None,None,None,None] +
                      mo_e_occ[j0:j1][None,:,None,None,None] - eabc)
            eijabc = 1./eijabc
            cache_ovvv_i = eris.get_ovvv(slice(i0,i1))
            cache_ovvv_j = eris.get_ovvv(slice(j0,j1))
            lijabc = contract_pl2p(l1, l2, i0, i1, j0, j1, cache_ovvv_i, cache_ovvv_j)
            rijabc = contract_pr2p(r1, r2, i0, i1, j0, j1, cache_ovvv_i, cache_ovvv_j)

            lijabc =  4.*lijabc \
                    - 2.*lijabc.transpose(0,1,3,2,4) \
                    - 2.*lijabc.transpose(0,1,4,3,2) \
                    - 2.*lijabc.transpose(0,1,2,4,3) \
                    + 1.*lijabc.transpose(0,1,3,4,2) \
                    + 1.*lijabc.transpose(0,1,4,2,3)

            # Symmetry factors (1 for a == b, 2 for a < b)
            fac = 2*np.ones_like(rijabc, dtype=int)
            triu_idx = np.triu_indices(i1-i0,i0+1,m=j1-j0)
            fac[triu_idx[0],triu_idx[1],:,:,:] = 0
            fac[np.arange(i1-i0),np.arange(i0,j1)] = 1
            eijabc *= fac

            deltaE += np.einsum('ijabc,ijabc,ijabc',lijabc,rijabc,eijabc)
        deltaE = 0.5*deltaE.real
        logger.info(eom, "eaccsd energy, star energy, delta energy = %16.12f, %16.12f, %16.12f",
                    eval_, eval_+deltaE, deltaE)
        e.append(eval_+deltaE)
    return e


class EOMEA(EOM):
    def get_init_guess(self, nroots=1, koopmans=True, diag=None):
        size = self.vector_size()
        dtype = getattr(diag, 'dtype', np.double)
        nroots = min(nroots, size)
        guess = []
        if koopmans:
            for n in range(nroots):
                g = np.zeros(size, dtype)
                g[n] = 1.0
                guess.append(g)
        else:
            idx = diag.argsort()[:nroots]
            for i in idx:
                g = np.zeros(size, dtype)
                g[i] = 1.0
                guess.append(g)
        return guess

    kernel = eaccsd
    eaccsd = eaccsd
    eaccsd_star = eaccsd_star

    matvec = eaccsd_matvec
    l_matvec = leaccsd_matvec
    get_diag = eaccsd_diag
    ccsd_star_contract = eaccsd_star_contract

    def eaccsd_star_contract(self, eaccsd_evals, eaccsd_evecs, leaccsd_evecs, imds=None):
        return self.ccsd_star_contract(eaccsd_evals, eaccsd_evecs, leaccsd_evecs, imds=imds)

    def gen_matvec(self, imds=None, left=False, **kwargs):
        if imds is None: imds = self.make_imds()
        diag = self.get_diag(imds)
        if left:
            matvec = lambda xs: [self.l_matvec(x, imds, diag) for x in xs]
        else:
            matvec = lambda xs: [self.matvec(x, imds, diag) for x in xs]
        return matvec, diag

    def vector_to_amplitudes(self, vector, nmo=None, nocc=None):
        if nmo is None: nmo = self.nmo
        if nocc is None: nocc = self.nocc
        return vector_to_amplitudes_ea(vector, nmo, nocc)

    def amplitudes_to_vector(self, r1, r2):
        return amplitudes_to_vector_ea(r1, r2)

    def vector_size(self):
        nocc = self.nocc
        nvir = self.nmo - nocc
        return nvir + nocc*nvir*nvir

    def make_imds(self, eris=None):
        imds = _IMDS(self._cc, eris=eris)
        imds.make_ea(self.partition)
        return imds

    @property
    def eea(self):
        return self.e


class EOMEA_Ta(EOMEA):
    '''Class for EOM EACCSD(T)*(a) method by Matthews and Stanton.'''
    def make_imds(self, eris=None):
        imds = _IMDS(self._cc, eris=eris)
        imds.make_t3p2_ea(self._cc, self.partition)
        return imds

########################################
# EOM-EE-CCSD
########################################

#TODO: double spin-flip EOM-EE

def eeccsd(eom, nroots=1, koopmans=False, guess=None, eris=None, imds=None):
    '''Calculate N-electron neutral excitations via EOM-EE-CCSD.

    Kwargs:
        nroots : int
            Number of roots (eigenvalues) requested
        koopmans : bool
            Calculate Koopmans'-like (1p1h) excitations only, targeting via
            overlap.
        guess : list of ndarray
            List of guess vectors to use for targeting via overlap.
    '''
    if eris is None: eris = eom._cc.ao2mo()
    if imds is None: imds = eom.make_imds(eris)

    spinvec_size = eom.vector_size()
    nroots = min(nroots, spinvec_size)

    diag_eeS, diag_eeT, diag_sf = eom.get_diag(imds)
    guess_eeS = []
    guess_eeT = []
    guess_sf = []
    if guess:
        for g in guess:
            if g is None: # beta->alpha spin-flip excitation
                pass
            elif g.size == diag_eeS.size:
                guess_eeS.append(g)
            elif g.size == diag_eeT.size:
                guess_eeT.append(g)
            else:
                guess_sf.append(g)
        nroots_eeS = len(guess_eeS)
        nroots_eeT = len(guess_eeT)
        nroots_sf = len(guess_sf)
        if len(guess) != nroots:
            logger.warn(eom, 'Number of states in initial guess %d does not '
                        'equal to nroots %d.', len(guess), nroots)
    else:
        deeS = np.sort(diag_eeS)[:nroots]
        deeT = np.sort(diag_eeT)[:nroots]
        dsf = np.sort(diag_sf)[:nroots]
        dmax = np.sort(np.hstack([deeS,deeT,dsf,dsf]))[nroots-1]
        nroots_eeS = np.count_nonzero(deeS <= dmax)
        nroots_eeT = np.count_nonzero(deeT <= dmax)
        nroots_sf = np.count_nonzero(dsf <= dmax)
        guess_eeS = guess_eeT = guess_sf = None

    def eomee_sub(cls, nroots, guess, diag):
        ee_sub = cls(eom._cc)
        ee_sub.__dict__.update(eom.__dict__)
        e, v = ee_sub.kernel(nroots, koopmans, guess, eris, imds, diag=diag)
        if nroots == 1:
            e, v = [e], [v]
            ee_sub.converged = [ee_sub.converged]
        return list(ee_sub.converged), list(e), list(v)

    e0 = e1 = e2 = []
    v0 = v1 = v2 = []
    conv0 = conv1 = conv2 = []
    if nroots_eeS > 0:
        conv0, e0, v0 = eomee_sub(EOMEESinglet, nroots_eeS, guess_eeS, diag_eeS)
    if nroots_eeT > 0:
        conv2, e2, v2 = eomee_sub(EOMEETriplet, nroots_eeT, guess_eeT, diag_eeT)
    if nroots_sf > 0:
        conv1, e1, v1 = eomee_sub(EOMEESpinFlip, nroots_sf, guess_sf, diag_sf)
        # The associated solutions of beta->alpha excitations
        e1 = e1 + e1
        conv1 = conv1 + conv1
        v1 = v1 + [None] * len(v1)
# beta->alpha spin-flip excitations, the coefficients are (-r1, (-r2[0], r2[1]))
# as below.  The EOMEESpinFlip class only handles alpha->beta excitations.
# Setting beta->alpha to None to bypass the vectors in initial guess
        #for i in range(nroots_sf):
        #    r1, r2 = vector_to_amplitudes_eomsf(v1[i], eom.nmo, eom.nocc)
        #    v1.append(amplitudes_to_vector_eomsf(-r1, (-r2[0], r2[1])))

    e = np.hstack([e0,e2,e1])
    idx = e.argsort()
    e = e[idx]
    conv = conv0 + conv2 + conv1
    conv = [conv[x] for x in idx]
    v = v0 + v2 + v1
    v = [v[x] for x in idx]

    if nroots == 1:
        conv = conv[0]
        e = e[0]
        v = v[0]
    eom.converged = conv
    eom.e = e
    eom.v = v
    return eom.e, eom.v


def eomee_ccsd_singlet(eom, nroots=1, koopmans=False, guess=None,
                       eris=None, imds=None, diag=None):
    '''EOM-EE-CCSD singlet
    '''
    eom.converged, eom.e, eom.v \
            = kernel(eom, nroots, koopmans, guess, eris=eris, imds=imds, diag=diag)
    return eom.e, eom.v

def eomee_ccsd_triplet(eom, nroots=1, koopmans=False, guess=None,
                       eris=None, imds=None, diag=None):
    '''EOM-EE-CCSD triplet
    '''
    return eomee_ccsd_singlet(eom, nroots, koopmans, guess, eris, imds, diag)

def eomsf_ccsd(eom, nroots=1, koopmans=False, guess=None,
               eris=None, imds=None, diag=None):
    '''Spin flip EOM-EE-CCSD
    '''
    return eomee_ccsd_singlet(eom, nroots, koopmans, guess, eris, imds, diag)

vector_to_amplitudes_ee = vector_to_amplitudes_singlet = ccsd.vector_to_amplitudes
amplitudes_to_vector_ee = amplitudes_to_vector_singlet = ccsd.amplitudes_to_vector

def amplitudes_to_vector_eomsf(t1, t2, out=None):
    nocc, nvir = t1.shape
    t2baaa, t2aaba = t2
    otril = np.tril_indices(nocc, k=-1)
    vtril = np.tril_indices(nvir, k=-1)
    baaa = np.take(t2baaa.reshape(nocc*nocc,nvir*nvir),
                   vtril[0]*nvir+vtril[1], axis=1)
    vector = np.hstack((t1.ravel(), baaa.ravel(), t2aaba[otril].ravel()))
    return vector

def vector_to_amplitudes_eomsf(vector, nmo, nocc):
    nvir = nmo - nocc
    t1 = vector[:nocc*nvir].reshape(nocc,nvir).copy()
    pvec = vector[t1.size:]

    nbaaa = nocc*nocc*nvir*(nvir-1)//2
    naaba = nocc*(nocc-1)//2*nvir*nvir
    t2baaa = np.zeros((nocc*nocc,nvir*nvir), dtype=vector.dtype)
    t2aaba = np.zeros((nocc*nocc,nvir*nvir), dtype=vector.dtype)
    otril = np.tril_indices(nocc, k=-1)
    vtril = np.tril_indices(nvir, k=-1)

    v = pvec[:nbaaa].reshape(nocc*nocc,nvir*(nvir-1)//2)
    t2baaa[:,vtril[0]*nvir+vtril[1]] = v
    t2baaa[:,vtril[1]*nvir+vtril[0]] = -v

    v = pvec[nbaaa:nbaaa+naaba].reshape(-1,nvir*nvir)
    t2aaba[otril[0]*nocc+otril[1]] = v
    t2aaba[otril[1]*nocc+otril[0]] = -v

    t2baaa = t2baaa.reshape(nocc,nocc,nvir,nvir)
    t2aaba = t2aaba.reshape(nocc,nocc,nvir,nvir)
    return t1, (t2baaa, t2aaba)

def amplitudes_to_vector_triplet(t1, t2, out=None):
    t2aa, t2ab = t2
    dtype = np.result_type(t1, t2aa, t2ab)
    nocc, nvir = t1.shape
    nov = nocc * nvir
    size1 = nov + nocc*(nocc-1)//2*nvir*(nvir-1)//2
    size = size1 + nov*(nov+1)//2
    vector = np.ndarray(size, dtype, buffer=out)
    ccsd.amplitudes_to_vector_s4(t1, t2[0], out=vector)
    t2ab = t2[1].transpose(0,2,1,3).reshape(nov,nov)
    lib.pack_tril(t2ab, out=vector[size1:])
    return vector

def vector_to_amplitudes_triplet(vector, nmo, nocc):
    nvir = nmo - nocc
    nov = nocc * nvir
    size1 = nov + nocc*(nocc-1)//2*nvir*(nvir-1)//2
    size = size1 + nov*(nov+1)//2
    t1, t2aa = ccsd.vector_to_amplitudes_s4(vector[:size1], nmo, nocc)
    t2ab = lib.unpack_tril(vector[size1:size], filltriu=2)
    t2ab = t2ab.reshape(nocc,nvir,nocc,nvir).transpose(0,2,1,3).copy()
    return t1, (t2aa, t2ab)

def eeccsd_matvec_singlet(eom, vector, imds=None):
    if imds is None: imds = eom.make_imds()
    nocc = eom.nocc
    nmo = eom.nmo
    nvir = nmo - nocc

    r1, r2 = vector_to_amplitudes_singlet(vector, nmo, nocc)
    t1, t2, eris = imds.t1, imds.t2, imds.eris
    nocc, nvir = t1.shape

    Hr1  = lib.einsum('ae,ie->ia', imds.Fvv, r1)
    Hr1 -= lib.einsum('mi,ma->ia', imds.Foo, r1)
    Hr1 += np.einsum('me,imae->ia',imds.Fov, r2) * 2
    Hr1 -= np.einsum('me,imea->ia',imds.Fov, r2)

    #:eris_vvvv = ao2mo.restore(1,np.asarray(eris.vvvv), t1.shape[1])
    #:Hr2 += lib.einsum('ijef,aebf->ijab', tau2, eris_vvvv) * .5
    tau2 = _make_tau(r2, r1, t1, fac=2)
    Hr2 = eom._cc._add_vvvv(None, tau2, eris, with_ovvv=False, t2sym='jiba')

    woOoO = np.asarray(imds.woOoO)
    Hr2 += lib.einsum('mnij,mnab->ijab', woOoO, r2)
    Hr2 *= .5
    woOoO = None

    Hr2 += lib.einsum('be,ijae->ijab', imds.Fvv   , r2)
    Hr2 -= lib.einsum('mj,imab->ijab', imds.Foo   , r2)

    mem_now = lib.current_memory()[0]
    max_memory = max(0, eom.max_memory - mem_now - Hr2.size*8e-6)
    blksize = min(nocc, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvir**3*3))))
    for p0,p1 in lib.prange(0, nocc, blksize):
        ovvv = eris.get_ovvv(slice(p0,p1))  # ovvv = eris.ovvv[p0:p1]
        theta = r2[p0:p1] * 2 - r2[p0:p1].transpose(0,1,3,2)
        Hr1 += lib.einsum('mfae,mife->ia', ovvv, theta)
        theta = None
        tmp = lib.einsum('meaf,ijef->maij', ovvv, tau2)
        Hr2 -= lib.einsum('ma,mbij->ijab', t1[p0:p1], tmp)
        tmp  = lib.einsum('meaf,me->af', ovvv, r1[p0:p1]) * 2
        tmp -= lib.einsum('mfae,me->af', ovvv, r1[p0:p1])
        Hr2 += lib.einsum('af,ijfb->ijab', tmp, t2)
        ovvv = tmp = None
    tau2 = None
    Hr2 -= lib.einsum('mbij,ma->ijab', imds.woVoO, r1)

    blksize = min(nvir, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nocc*nvir**2*2))))
    for p0, p1 in lib.prange(0, nvir, nocc):
        Hr2 += lib.einsum('ejab,ie->ijab', np.asarray(imds.wvOvV[p0:p1]), r1[:,p0:p1])

    woVVo = np.asarray(imds.woVVo)
    tmp = lib.einsum('mbej,imea->jiab', woVVo, r2)
    Hr2 += tmp
    tmp *= .5
    Hr2 += tmp.transpose(0,1,3,2)
    tmp = None

    woVvO = woVVo * .5
    woVVo = None
    woVvO += np.asarray(imds.woVvO)
    theta = r2*2 - r2.transpose(0,1,3,2)
    Hr1 += np.einsum('maei,me->ia', woVvO, r1) * 2
    Hr2 += lib.einsum('mbej,imae->ijab', woVvO, theta)
    woVvO = None

    woOoV = np.asarray(imds.woOoV)
    Hr1-= lib.einsum('mnie,mnae->ia', woOoV, theta)
    tmp = lib.einsum('nmie,me->ni', woOoV, r1) * 2
    tmp-= lib.einsum('mnie,me->ni', woOoV, r1)
    Hr2 -= lib.einsum('ni,njab->ijab', tmp, t2)
    tmp = woOoV = None

    eris_ovov = np.asarray(eris.ovov)
    tmp  = np.einsum('mfne,mf->en', eris_ovov, r1) * 2
    tmp -= np.einsum('menf,mf->en', eris_ovov, r1)
    tmp  = np.einsum('en,nb->eb', tmp, t1)
    tmp += lib.einsum('menf,mnbf->eb', eris_ovov, theta)
    Hr2 -= lib.einsum('eb,ijea->jiab', tmp, t2)
    tmp = None

    tmp = lib.einsum('nemf,imef->ni', eris_ovov, theta)
    Hr1 -= lib.einsum('na,ni->ia', t1, tmp)
    Hr2 -= lib.einsum('mj,miab->ijba', tmp, t2)
    tmp = theta = None

    tau2 = _make_tau(r2, r1, t1, fac=2)
    tmp = lib.einsum('menf,ijef->mnij', eris_ovov, tau2)
    tau2 = None

    tau = _make_tau(t2, t1, t1)
    tau *= .5
    Hr2 += lib.einsum('mnij,mnab->ijab', tmp, tau)
    tau = tmp = eris_ovov = None

    Hr2 = Hr2 + Hr2.transpose(1,0,3,2)
    vector = amplitudes_to_vector_ee(Hr1, Hr2)
    return vector

def eeccsd_matvec_triplet(eom, vector, imds=None):
    if imds is None: imds = eom.make_imds()
    nocc = eom.nocc
    nmo = eom.nmo
    nvir = nmo - nocc

    r1, r2 = vector_to_amplitudes_triplet(vector, nmo, nocc)
    r2aa, r2ab = r2
    t1, t2, eris = imds.t1, imds.t2, imds.eris
    nocc, nvir = t1.shape

    Hr1  = lib.einsum('ae,ie->ia', imds.Fvv, r1)
    Hr1 -= lib.einsum('mi,ma->ia', imds.Foo, r1)
    Hr1 += np.einsum('me,imae->ia',imds.Fov, r2aa)
    Hr1 += np.einsum('ME,iMaE->ia',imds.Fov, r2ab)

    tau2ab = np.einsum('ia,jb->ijab', r1, t1)
    tau2ab-= np.einsum('ia,jb->ijab', t1, r1)
    tau2ab+= r2ab
    tau2aa = np.einsum('ia,jb->ijab', r1, t1)
    tau2aa-= np.einsum('ia,jb->jiab', r1, t1)
    tau2aa = tau2aa - tau2aa.transpose(0,1,3,2)
    tau2aa+= r2aa

    #:eris_vvvv = ao2mo.restore(1,np.asarray(eris.vvvv), t1.shape[1])
    #:Hr2aa += lib.einsum('ijef,aebf->ijab', tau2aa, eris_vvvv) * .25
    #:Hr2ab += lib.einsum('ijef,aebf->ijab', tau2ab, eris_vvvv) * .5
    Hr2aa = eom._cc._add_vvvv(None, tau2aa, eris, with_ovvv=False, t2sym='jiba')
    Hr2ab = eom._cc._add_vvvv(None, tau2ab, eris, with_ovvv=False, t2sym='-jiba')

    woOoO = np.asarray(imds.woOoO)
    Hr2aa += lib.einsum('mnij,mnab->ijab', woOoO, r2aa)
    Hr2ab += lib.einsum('mNiJ,mNaB->iJaB', woOoO, r2ab)
    Hr2aa *= .25
    Hr2ab *= .5
    woOoO = None

    Hr2aa += lib.einsum('be,ijae->ijab', imds.Fvv*.5, r2aa)
    Hr2aa -= lib.einsum('mj,imab->ijab', imds.Foo*.5, r2aa)
    Hr2ab += lib.einsum('BE,iJaE->iJaB', imds.Fvv, r2ab)
    Hr2ab -= lib.einsum('MJ,iMaB->iJaB', imds.Foo, r2ab)

    mem_now = lib.current_memory()[0]
    max_memory = max(0, eom.max_memory - mem_now - Hr2aa.size*8e-6)
    blksize = min(nocc, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvir**3*3))))
    tmp1 = np.zeros((nvir,nvir), dtype=r1.dtype)
    for p0,p1 in lib.prange(0, nocc, blksize):
        ovvv = eris.get_ovvv(slice(p0,p1))  # ovvv = eris.ovvv[p0:p1]
        theta = r2aa[:,p0:p1] + r2ab[:,p0:p1]
        Hr1 += lib.einsum('mfae,imef->ia', ovvv, theta)
        theta = None
        tmpaa = lib.einsum('meaf,ijef->maij', ovvv, tau2aa)
        tmpab = lib.einsum('meAF,iJeF->mAiJ', ovvv, tau2ab)
        tmp1 += lib.einsum('mfae,me->af', ovvv, r1[p0:p1])
        Hr2aa+= lib.einsum('mb,maij->ijab', t1[p0:p1]*.5, tmpaa)
        Hr2ab-= lib.einsum('mb,mAiJ->iJbA', t1[p0:p1], tmpab)
        ovvv = tmpaa = tmpab = None
    tau2aa = tau2ab = None

    woVVo = np.asarray(imds.woVVo)
    Hr1 += np.einsum('maei,me->ia', woVVo, r1)
    Hr2aa += lib.einsum('mbej,imae->ijba', woVVo, r2ab)
    Hr2ab += lib.einsum('MBEJ,iMEa->iJaB', woVVo, r2aa)
    Hr2ab += lib.einsum('MbeJ,iMeA->iJbA', woVVo, r2ab)

    woVVo = woVVo + np.asarray(imds.woVvO)
    theta = r2aa + r2ab
    tmp = lib.einsum('mbej,imae->ijab', woVVo, theta)
    woVVo = None

    woOoV = np.asarray(imds.woOoV)
    Hr1 -= lib.einsum('mnie,mnae->ia', woOoV, theta)
    tmpa = lib.einsum('mnie,me->ni', woOoV, r1)
    tmp += lib.einsum('ni,njab->ijab', tmpa, t2)
    tmp -= lib.einsum('af,ijfb->ijab', tmp1, t2)
    tmp -= lib.einsum('mbij,ma->ijab', imds.woVoO, r1)

    blksize = min(nvir, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nocc*nvir**2*2))))
    for p0,p1 in lib.prange(0, nvir, blksize):
        tmp += lib.einsum('ejab,ie->ijab', np.asarray(imds.wvOvV[p0:p1]), r1[:,p0:p1])

    Hr2aa += tmp
    Hr2ab += tmp
    tmp = woOoV = None

    eris_ovov = np.asarray(eris.ovov)
    tmpa = -lib.einsum('menf,imfe->ni', eris_ovov, theta)
    Hr1 += lib.einsum('na,ni->ia', t1, tmpa)
    tmp  = lib.einsum('mj,imab->ijab', tmpa, t2)
    tmp1 = np.einsum('menf,mf->en', eris_ovov, r1)
    tmpa = lib.einsum('en,nb->eb', tmp1, t1)
    tmpa-= lib.einsum('menf,mnbf->eb', eris_ovov, theta)
    tmp += lib.einsum('eb,ijae->ijab', tmpa, t2)
    Hr2aa += tmp
    Hr2ab -= tmp
    tmp = theta = tmp1 = tmpa = None

    tau2aa = np.einsum('ia,jb->ijab', r1, t1)
    tau2aa-= np.einsum('ia,jb->jiab', r1, t1)
    tau2aa = tau2aa - tau2aa.transpose(0,1,3,2)
    tau2aa+= r2aa
    tmpaa = lib.einsum('menf,ijef->mnij', eris_ovov, tau2aa)
    tau2aa = None
    tmpaa *= .25
    tau = _make_tau(t2, t1, t1)
    Hr2aa += lib.einsum('mnij,mnab->ijab', tmpaa, tau)
    tmpaa = tau = None

    tau2ab = np.einsum('ia,jb->ijab', r1, t1)
    tau2ab-= np.einsum('ia,jb->ijab', t1, r1)
    tau2ab+= r2ab
    tmpab = lib.einsum('meNF,iJeF->mNiJ', eris_ovov, tau2ab)
    tau2ab = None
    tmpab *= .5
    tau = _make_tau(t2, t1, t1)
    Hr2ab += lib.einsum('mNiJ,mNaB->iJaB', tmpab, tau)
    tmpab = tau = None
    eris_ovov = None

    Hr2aa = Hr2aa - Hr2aa.transpose(0,1,3,2)
    Hr2aa = Hr2aa - Hr2aa.transpose(1,0,2,3)
    Hr2ab = Hr2ab - Hr2ab.transpose(1,0,3,2)
    vector = amplitudes_to_vector_triplet(Hr1, (Hr2aa,Hr2ab))
    return vector

def eeccsd_matvec_sf(eom, vector, imds=None):
    '''Spin flip EOM-CCSD'''
    if imds is None: imds = eom.make_imds()
    nocc = eom.nocc
    nmo = eom.nmo
    nvir = nmo - nocc

    t1, t2, eris = imds.t1, imds.t2, imds.eris
    r1, r2 = vector_to_amplitudes_eomsf(vector, nmo, nocc)
    r2baaa, r2aaba = r2
    nocc, nvir = t1.shape

    Hr1  = np.einsum('ae,ie->ia', imds.Fvv, r1)
    Hr1 -= np.einsum('mi,ma->ia', imds.Foo, r1)
    Hr1 += np.einsum('me,imae->ia', imds.Fov, r2baaa)
    Hr1 += np.einsum('me,imae->ia', imds.Fov, r2aaba)

    tau2baaa = np.einsum('ia,jb->ijab', r1, t1)
    tau2baaa += r2baaa * .5
    tau2baaa = tau2baaa - tau2baaa.transpose(0,1,3,2)
    tau2aaba = np.einsum('ia,jb->ijab', r1, t1)
    tau2aaba += r2aaba * .5
    tau2aaba = tau2aaba - tau2aaba.transpose(1,0,2,3)

    #:eris_vvvv = ao2mo.restore(1,np.asarray(eris.vvvv), t1.shape[1])
    #:Hr2baaa += .5*lib.einsum('ijef,aebf->ijab', tau2baaa, eris_vvvv)
    #:Hr2aaba += .5*lib.einsum('ijef,aebf->ijab', tau2aaba, eris_vvvv)
    Hr2aaba = eom._cc._add_vvvv(None, tau2aaba, eris, with_ovvv=False, t2sym='-jiab')
    Hr2baaa = eom._cc._add_vvvv(None, tau2baaa, eris, with_ovvv=False, t2sym=False)

    woOoO = np.asarray(imds.woOoO)
    Hr2baaa += lib.einsum('mnij,mnab->ijab', woOoO, r2baaa)
    Hr2aaba += lib.einsum('mnij,mnab->ijab', woOoO, r2aaba)
    Hr2aaba *= .5
    Hr2baaa *= .5
    woOoO = None

    Hr2baaa -= lib.einsum('mj,imab->ijab', imds.Foo*.5, r2baaa)
    Hr2aaba -= lib.einsum('mj,imab->ijab', imds.Foo*.5, r2aaba)
    Hr2baaa -= lib.einsum('mj,miab->jiab', imds.Foo*.5, r2baaa)
    Hr2aaba -= lib.einsum('mj,miab->jiab', imds.Foo*.5, r2aaba)
    Hr2baaa += lib.einsum('be,ijae->ijab', imds.Fvv*.5, r2baaa)
    Hr2aaba += lib.einsum('be,ijae->ijab', imds.Fvv*.5, r2aaba)
    Hr2baaa += lib.einsum('be,ijea->ijba', imds.Fvv*.5, r2baaa)
    Hr2aaba += lib.einsum('be,ijea->ijba', imds.Fvv*.5, r2aaba)

    mem_now = lib.current_memory()[0]
    max_memory = max(0, eom.max_memory - mem_now - Hr2aaba.size*8e-6)
    blksize = min(nocc, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvir**3*3))))
    for p0,p1 in lib.prange(0, nocc, blksize):
        ovvv = eris.get_ovvv(slice(p0,p1))  # ovvv = eris.ovvv[p0:p1]
        theta = r2baaa[:,p0:p1] + r2aaba[:,p0:p1]
        Hr1 += lib.einsum('mfae,imef->ia', ovvv, theta)
        theta = None

        tmp1aaba = lib.einsum('meaf,ijef->maij', ovvv, tau2baaa)
        Hr2baaa -= lib.einsum('mb,maij->ijba', t1[p0:p1]*.5, tmp1aaba)
        tmp1aaba = tmp1baaa = tmp1abaa = tmp2aaba = None

        tmp2aaba = lib.einsum('meaf,ijfe->maij', ovvv, tau2baaa)
        Hr2baaa -= lib.einsum('mb,maij->ijab', t1[p0:p1]*.5, tmp2aaba)
        tmp1aaba = tmp1baaa = tmp1abaa = tmp2aaba = None

        tmp1baaa = lib.einsum('meaf,ijef->maij', ovvv, tau2aaba)
        Hr2aaba -= lib.einsum('mb,maij->ijba', t1[p0:p1]*.5, tmp1baaa)
        tmp1aaba = tmp1baaa = tmp1abaa = tmp2aaba = None

        tmp1abaa = lib.einsum('meaf,ijfe->maij', ovvv, tau2aaba)
        Hr2aaba -= lib.einsum('mb,maij->ijab', t1[p0:p1]*.5, tmp1abaa)
        tmp1aaba = tmp1baaa = tmp1abaa = tmp2aaba = None

        tmp = lib.einsum('mfae,me->af', ovvv, r1[p0:p1])
        tmp = lib.einsum('af,jibf->ijab', tmp, t2)
        Hr2baaa -= tmp
        Hr2aaba -= tmp
        tmp = ovvv = None
    tau2aaba = tau2baaa = None

    tmp = lib.einsum('mbij,ma->ijab', imds.woVoO, r1)
    Hr2baaa -= tmp
    Hr2aaba -= tmp
    tmp = None

    woOoV = np.asarray(imds.woOoV)
    Hr1 -= lib.einsum('mnie,mnae->ia', woOoV, r2aaba)
    Hr1 -= lib.einsum('mnie,mnae->ia', woOoV, r2baaa)
    tmp = lib.einsum('mnie,me->ni', woOoV, r1)
    tmp = lib.einsum('ni,njab->ijab', tmp, t2)
    Hr2baaa += tmp
    Hr2aaba += tmp
    tmp = woOoV = None

    for p0,p1 in lib.prange(0, nvir, nocc):
        tmp = lib.einsum('ejab,ie->ijab', np.asarray(imds.wvOvV[p0:p1]), r1[:,p0:p1])
        Hr2baaa += tmp
        Hr2aaba += tmp
        tmp = None

    woVVo = np.asarray(imds.woVVo)
    Hr1 += np.einsum('maei,me->ia', woVVo, r1)
    Hr2baaa += lib.einsum('mbej,miea->jiba', woVVo, r2baaa)
    Hr2aaba += lib.einsum('mbej,miea->jiba', woVVo, r2aaba)
    woVVo = None
    woVvO = np.asarray(imds.woVvO)
    Hr2baaa += lib.einsum('mbej,imae->ijab', woVvO, r2aaba)
    Hr2aaba += lib.einsum('mbej,imae->ijab', woVvO, r2baaa)
    woVvO = woVvO + np.asarray(imds.woVVo)
    Hr2baaa += lib.einsum('mbej,imae->ijab', woVvO, r2baaa)
    Hr2aaba += lib.einsum('mbej,imae->ijab', woVvO, r2aaba)
    woVvO = None

    eris_ovov = np.asarray(eris.ovov)
    theta = r2aaba + r2baaa
    tmp = lib.einsum('nfme,imfe->ni', eris_ovov, theta)
    Hr1 -= np.einsum('na,ni->ia', t1, tmp)
    Hr2baaa -= lib.einsum('mj,imba->jiab', tmp, t2)
    Hr2aaba -= lib.einsum('mj,imba->jiab', tmp, t2)

    tmp = np.einsum('menf,mf->en', eris_ovov, r1)
    tmp = np.einsum('en,nb->eb', tmp, t1)
    tmp-= lib.einsum('menf,mnbf->eb', eris_ovov, theta)
    Hr2baaa += lib.einsum('ea,ijbe->jiab', tmp, t2)
    Hr2aaba += lib.einsum('ea,ijbe->jiab', tmp, t2)
    theta = tmp = None

    tau2baaa = np.einsum('ia,jb->ijab', r1, t1)
    tau2baaa += r2baaa * .5
    tau2baaa = tau2baaa - tau2baaa.transpose(0,1,3,2)
    tau = _make_tau(t2, t1, t1)
    tmp1aaba = lib.einsum('menf,ijef->mnij', eris_ovov, tau2baaa)
    tau2baaa = None
    Hr2baaa += .5*lib.einsum('mnij,mnab->ijab', tmp1aaba, tau)
    tau = tmp1aaba = None

    tau2aaba = np.einsum('ia,jb->ijab', r1, t1)
    tau2aaba += r2aaba * .5
    tau2aaba = tau2aaba - tau2aaba.transpose(1,0,2,3)
    tau = _make_tau(t2, t1, t1)
    tmp1baaa = lib.einsum('menf,ijef->mnij', eris_ovov, tau2aaba)
    tau2aaba = None
    Hr2aaba += .5*lib.einsum('mnij,mnab->ijab', tmp1baaa, tau)
    tau = tmp1baaa = None
    eris_ovov = None

    Hr2baaa = Hr2baaa - Hr2baaa.transpose(0,1,3,2)
    Hr2aaba = Hr2aaba - Hr2aaba.transpose(1,0,2,3)
    vector = amplitudes_to_vector_eomsf(Hr1, (Hr2baaa,Hr2aaba))
    return vector

def eeccsd_diag(eom, imds=None):
    if imds is None: imds = eom.make_imds()
    eris = imds.eris
    t1, t2 = imds.t1, imds.t2
    dtype = np.result_type(t1, t2)
    tau = _make_tau(t2, t1, t1)
    nocc, nvir = t1.shape

    Fo = imds.Foo.diagonal()
    Fv = imds.Fvv.diagonal()
    Wovab = np.einsum('iaai->ia', imds.woVVo)
    Wovaa = Wovab + np.einsum('iaai->ia', imds.woVvO)

    eia = lib.direct_sum('-i+a->ia', Fo, Fv)
    Hr1aa = eia + Wovaa
    Hr1ab = eia + Wovab

    eris_ovov = np.asarray(eris.ovov)
    Wvvab = np.einsum('mnab,manb->ab', tau, eris_ovov)
    Wvvaa = .5*Wvvab - .5*np.einsum('mnba,manb->ab', tau, eris_ovov)
    ijb = np.einsum('iejb,ijeb->ijb', eris_ovov, t2)
    Hr2ab = lib.direct_sum('iJB+a->iJaB',-ijb, Fv)
    jab = np.einsum('kajb,kjab->jab', eris_ovov, t2)
    Hr2ab+= lib.direct_sum('-i-jab->ijab', Fo, jab)

    jib = np.einsum('iejb,ijbe->jib', eris_ovov, t2)
    jib = jib + jib.transpose(1,0,2)
    jib-= ijb + ijb.transpose(1,0,2)
    jba = np.einsum('kajb,jkab->jba', eris_ovov, t2)
    jba = jba + jba.transpose(0,2,1)
    jba-= jab + jab.transpose(0,2,1)
    Hr2aa = lib.direct_sum('jib+a->jiba', jib, Fv)
    Hr2aa+= lib.direct_sum('-i+jba->ijba', Fo, jba)
    eris_ovov = None

    Hr2baaa = lib.direct_sum('ijb+a->ijba',-ijb, Fv)
    Hr2baaa += Wovaa.reshape(1,nocc,1,nvir)
    Hr2baaa += Wovab.reshape(nocc,1,1,nvir)
    Hr2baaa = Hr2baaa + Hr2baaa.transpose(0,1,3,2)
    Hr2baaa+= lib.direct_sum('-i+jab->ijab', Fo, jba)
    Hr2baaa-= Fo.reshape(1,-1,1,1)
    Hr2aaba = lib.direct_sum('-i-jab->ijab', Fo, jab)
    Hr2aaba += Wovaa.reshape(1,nocc,1,nvir)
    Hr2aaba += Wovab.reshape(1,nocc,nvir,1)
    Hr2aaba = Hr2aaba + Hr2aaba.transpose(1,0,2,3)
    Hr2aaba+= lib.direct_sum('ijb+a->ijab', jib, Fv)
    Hr2aaba+= Fv.reshape(1,1,1,-1)
    Hr2ab += Wovaa.reshape(1,nocc,1,nvir)
    Hr2ab += Wovab.reshape(nocc,1,1,nvir)
    Hr2ab = Hr2ab + Hr2ab.transpose(1,0,3,2)
    Hr2aa += Wovaa.reshape(1,nocc,1,nvir) * 2
    Hr2aa = Hr2aa + Hr2aa.transpose(0,1,3,2)
    Hr2aa = Hr2aa + Hr2aa.transpose(1,0,2,3)
    Hr2aa *= .5

    Wooab = np.einsum('ijij->ij', imds.woOoO)
    Wooaa = Wooab - np.einsum('ijji->ij', imds.woOoO)
    Hr2aa += Wooaa.reshape(nocc,nocc,1,1)
    Hr2ab += Wooab.reshape(nocc,nocc,1,1)
    Hr2baaa += Wooab.reshape(nocc,nocc,1,1)
    Hr2aaba += Wooaa.reshape(nocc,nocc,1,1)

    #:eris_ovvv = lib.unpack_tril(np.asarray(eris.ovvv).reshape(nocc*nvir,nvir**2)).reshape(nocc,nvir,nvir,nvir)
    #:tmp = np.einsum('mb,mbaa->ab', t1, eris_ovvv)
    #:Wvvaa += np.einsum('mb,maab->ab', t1, eris_ovvv)
    mem_now = lib.current_memory()[0]
    max_memory = max(0, eom.max_memory - mem_now)
    blksize = min(nocc, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvir**3*3))))
    tmp = np.zeros((nvir,nvir), dtype=dtype)
    for p0,p1 in lib.prange(0, nocc, blksize):
        ovvv = eris.get_ovvv(slice(p0,p1))  # ovvv = eris.ovvv[p0:p1]
        tmp += np.einsum('mb,mbaa->ab', t1[p0:p1], ovvv)
        Wvvaa += np.einsum('mb,maab->ab', t1[p0:p1], ovvv)
        ovvv = None
    Wvvaa -= tmp
    Wvvab -= tmp
    Wvvab -= tmp.T
    Wvvaa = Wvvaa + Wvvaa.T
    if eris.vvvv is None: # AO-direct CCSD, vvvv is not generated.
        pass
    elif eris.vvvv.ndim == 4:
        eris_vvvv = ao2mo.restore(1,np.asarray(eris.vvvv), t1.shape[1])
        tmp = np.einsum('aabb->ab', eris_vvvv)
        Wvvaa += tmp
        Wvvaa -= np.einsum('abba->ab', eris_vvvv)
        Wvvab += tmp
    else:
        for i in range(nvir):
            i0 = i*(i+1)//2
            vvv = lib.unpack_tril(np.asarray(eris.vvvv[i0:i0+i+1]))
            tmp = np.einsum('bb->b', vvv[i])
            Wvvaa[i] += tmp
            Wvvab[i] += tmp
            tmp = np.einsum('bb->b', vvv[:,:i+1,i])
            Wvvaa[i,:i+1] -= tmp
            Wvvaa[:i  ,i] -= tmp[:i]
            vvv = None

    Hr2aa += Wvvaa.reshape(1,1,nvir,nvir)
    Hr2ab += Wvvab.reshape(1,1,nvir,nvir)
    Hr2baaa += Wvvaa.reshape(1,1,nvir,nvir)
    Hr2aaba += Wvvab.reshape(1,1,nvir,nvir)

    vec_eeS = amplitudes_to_vector_singlet(Hr1aa, Hr2ab)
    vec_eeT = amplitudes_to_vector_triplet(Hr1aa, (Hr2aa,Hr2ab))
    vec_sf = amplitudes_to_vector_eomsf(Hr1ab, (Hr2baaa,Hr2aaba))
    return vec_eeS, vec_eeT, vec_sf


class EOMEE(EOM):
    def get_init_guess(self, nroots=1, koopmans=True, diag=None):
        if koopmans:
            nocc = self.nocc
            nvir = self.nmo - nocc
            idx = diag[:nocc*nvir].argsort()
        else:
            idx = diag.argsort()

        size = self.vector_size()
        dtype = getattr(diag, 'dtype', np.double)
        nroots = min(nroots, size)
        guess = []
        for i in idx[:nroots]:
            g = np.zeros(size, dtype)
            g[i] = 1.0
            guess.append(g)
        return guess

    kernel = eeccsd
    eeccsd = eeccsd
    get_diag = eeccsd_diag

    def vector_size(self):
        '''size of the vector based on spin-orbital basis'''
        nocc = self.nocc
        nvir = self.nmo - nocc
        return nocc*nvir + nocc*nocc*nvir*nvir

    def make_imds(self, eris=None):
        imds = _IMDS(self._cc, eris=eris)
        imds.make_ee()
        return imds

    @property
    def eee(self):
        return self.e


class EOMEESinglet(EOMEE):
    kernel = eomee_ccsd_singlet
    eomee_ccsd_singlet = eomee_ccsd_singlet
    matvec = eeccsd_matvec_singlet

    def gen_matvec(self, imds=None, diag=None, **kwargs):
        if imds is None: imds = self.make_imds()
        if diag is None: diag = self.get_diag(imds)[0]
        matvec = lambda xs: [self.matvec(x, imds) for x in xs]
        return matvec, diag

    def vector_to_amplitudes(self, vector, nmo=None, nocc=None):
        if nmo is None: nmo = self.nmo
        if nocc is None: nocc = self.nocc
        return vector_to_amplitudes_singlet(vector, nmo, nocc)

    def amplitudes_to_vector(self, r1, r2):
        return amplitudes_to_vector_singlet(r1, r2)

    def vector_size(self):
        nocc = self.nocc
        nvir = self.nmo - nocc
        nov = nocc * nvir
        return nov + nov*(nov+1)//2


class EOMEETriplet(EOMEE):
    kernel = eomee_ccsd_triplet
    eomee_ccsd_triplet = eomee_ccsd_triplet
    matvec = eeccsd_matvec_triplet

    def gen_matvec(self, imds=None, diag=None, **kwargs):
        if imds is None: imds = self.make_imds()
        if diag is None: diag = self.get_diag(imds)[1]
        matvec = lambda xs: [self.matvec(x, imds) for x in xs]
        return matvec, diag

    def vector_to_amplitudes(self, vector, nmo=None, nocc=None):
        if nmo is None: nmo = self.nmo
        if nocc is None: nocc = self.nocc
        return vector_to_amplitudes_triplet(vector, nmo, nocc)

    def amplitudes_to_vector(self, r1, r2):
        return amplitudes_to_vector_triplet(r1, r2)

    def vector_size(self):
        nocc = self.nocc
        nvir = self.nmo - nocc
        nov = nocc * nvir
        return nov + nocc*(nocc-1)//2*nvir*(nvir-1)//2 + nov*(nov+1)//2


class EOMEESpinFlip(EOMEE):
    kernel = eomsf_ccsd
    eomsf_ccsd = eomsf_ccsd
    matvec = eeccsd_matvec_sf

    def gen_matvec(self, imds=None, diag=None, **kwargs):
        if imds is None: imds = self.make_imds()
        if diag is None: diag = self.get_diag(imds)[2]
        matvec = lambda xs: [self.matvec(x, imds) for x in xs]
        return matvec, diag

    def vector_to_amplitudes(self, vector, nmo=None, nocc=None):
        if nmo is None: nmo = self.nmo
        if nocc is None: nocc = self.nocc
        return vector_to_amplitudes_eomsf(vector, nmo, nocc)

    def amplitudes_to_vector(self, r1, r2):
        return amplitudes_to_vector_eomsf(r1, r2)

    def vector_size(self):
        nocc = self.nocc
        nvir = self.nmo - nocc
        nbaaa = nocc*nocc*nvir*(nvir-1)//2
        naaba = nocc*(nocc-1)//2*nvir*nvir
        return nocc*nvir + nbaaa + naaba

#TODO: Check whether EOM methods works with rccsd.RCCSD when orbitals are complex
ccsd.CCSD.EOMIP         = lib.class_as_method(EOMIP)
ccsd.CCSD.EOMIP_Ta      = lib.class_as_method(EOMIP_Ta)
ccsd.CCSD.EOMEA         = lib.class_as_method(EOMEA)
ccsd.CCSD.EOMEA_Ta      = lib.class_as_method(EOMEA_Ta)
ccsd.CCSD.EOMEE         = lib.class_as_method(EOMEE)
ccsd.CCSD.EOMEESinglet  = lib.class_as_method(EOMEESinglet)
ccsd.CCSD.EOMEETriplet  = lib.class_as_method(EOMEETriplet)
ccsd.CCSD.EOMEESpinFlip = lib.class_as_method(EOMEESpinFlip)


class _IMDS:
    def __init__(self, cc, eris=None):
        self.verbose = cc.verbose
        self.stdout = cc.stdout
        self.max_memory = cc.max_memory
        self.t1 = cc.t1
        self.t2 = cc.t2
        if eris is None:
            eris = cc.ao2mo()
        self.eris = eris
        self._made_shared_2e = False

    def _make_shared_1e(self):
        cput0 = (logger.process_clock(), logger.perf_counter())

        t1, t2, eris = self.t1, self.t2, self.eris
        self.Loo = imd.Loo(t1, t2, eris)
        self.Lvv = imd.Lvv(t1, t2, eris)
        self.Fov = imd.cc_Fov(t1, t2, eris)

        logger.timer_debug1(self, 'EOM-CCSD shared one-electron '
                            'intermediates', *cput0)
        return self

    def _make_shared_2e(self):
        cput0 = (logger.process_clock(), logger.perf_counter())
        log = logger.Logger(self.stdout, self.verbose)

        t1, t2, eris = self.t1, self.t2, self.eris
        # 2 virtuals
        self.Wovov = imd.Wovov(t1, t2, eris)
        self.Wovvo = imd.Wovvo(t1, t2, eris)
        self.Woovv = np.asarray(eris.ovov).transpose(0,2,1,3)

        self._made_shared_2e = True
        log.timer_debug1('EOM-CCSD shared two-electron intermediates', *cput0)
        return self

    def make_ip(self, ip_partition=None):
        self._make_shared_1e()
        if not self._made_shared_2e and ip_partition != 'mp':
            self._make_shared_2e()

        cput0 = (logger.process_clock(), logger.perf_counter())
        log = logger.Logger(self.stdout, self.verbose)

        t1, t2, eris = self.t1, self.t2, self.eris

        # 0 or 1 virtuals
        if ip_partition != 'mp':
            self.Woooo = imd.Woooo(t1, t2, eris)
        self.Wooov = imd.Wooov(t1, t2, eris)
        self.Wovoo = imd.Wovoo(t1, t2, eris)
        log.timer_debug1('EOM-CCSD IP intermediates', *cput0)
        return self

    def make_t3p2_ip(self, cc, ip_partition=None):
        assert(ip_partition is None)
        cput0 = (logger.process_clock(), logger.perf_counter())

        t1, t2, eris = cc.t1, cc.t2, self.eris
        delta_E_corr, pt1, pt2, Wovoo, Wvvvo = \
            imd.get_t3p2_imds_slow(cc, t1, t2, eris)
        self.t1 = pt1
        self.t2 = pt2

        self._made_shared_2e = False  # Force update
        self.make_ip()  # Make after t1/t2 updated
        self.Wovoo = self.Wovoo + Wovoo

        logger.timer_debug1(self, 'EOM-CCSD(T)a IP intermediates', *cput0)
        return self


    def make_ea(self, ea_partition=None):
        self._make_shared_1e()
        if not self._made_shared_2e and ea_partition != 'mp':
            self._make_shared_2e()

        cput0 = (logger.process_clock(), logger.perf_counter())
        log = logger.Logger(self.stdout, self.verbose)

        t1, t2, eris = self.t1, self.t2, self.eris

        # 3 or 4 virtuals
        self.Wvovv = imd.Wvovv(t1, t2, eris)
        if ea_partition == 'mp':
            self.Wvvvo = imd.Wvvvo(t1, t2, eris)
        else:
            self.Wvvvv = imd.Wvvvv(t1, t2, eris)
            self.Wvvvo = imd.Wvvvo(t1, t2, eris, self.Wvvvv)
        log.timer_debug1('EOM-CCSD EA intermediates', *cput0)
        return self

    def make_t3p2_ea(self, cc, ea_partition=None):
        assert(ea_partition is None)
        cput0 = (logger.process_clock(), logger.perf_counter())

        t1, t2, eris = cc.t1, cc.t2, self.eris
        delta_E_corr, pt1, pt2, Wovoo, Wvvvo = \
            imd.get_t3p2_imds_slow(cc, t1, t2, eris)
        self.t1 = pt1
        self.t2 = pt2

        self._made_shared_2e = False  # Force update
        self.make_ea()  # Make after t1/t2 updated
        self.Wvvvo = self.Wvvvo + Wvvvo

        logger.timer_debug1(self, 'EOM-CCSD(T)a EA intermediates', *cput0)
        return self


    def make_ee(self):
        cput0 = (logger.process_clock(), logger.perf_counter())
        log = logger.Logger(self.stdout, self.verbose)

        t1, t2, eris = self.t1, self.t2, self.eris
        dtype = np.result_type(t1, t2)
        if np.iscomplexobj(t2):
            raise NotImplementedError('Complex integrals are not supported in EOM-EE-CCSD')

        nocc, nvir = t1.shape

        fswap = lib.H5TmpFile()
        self.saved = lib.H5TmpFile()
        self.wvOvV = self.saved.create_dataset('wvOvV', (nvir,nocc,nvir,nvir), dtype.char)
        self.woVvO = self.saved.create_dataset('woVvO', (nocc,nvir,nvir,nocc), dtype.char)
        self.woVVo = self.saved.create_dataset('woVVo', (nocc,nvir,nvir,nocc), dtype.char)
        self.woOoV = self.saved.create_dataset('woOoV', (nocc,nocc,nocc,nvir), dtype.char)

        foo = eris.fock[:nocc,:nocc]
        fov = eris.fock[:nocc,nocc:]
        fvv = eris.fock[nocc:,nocc:]

        self.Fov = np.zeros((nocc,nvir), dtype=dtype)
        self.Foo = np.zeros((nocc,nocc), dtype=dtype)
        self.Fvv = np.zeros((nvir,nvir), dtype=dtype)

        #:eris_ovvv = lib.unpack_tril(np.asarray(eris.ovvv).reshape(nocc*nvir,nvir**2)).reshape(nocc,nvir,nvir,nvir)
        #:self.Fvv  = np.einsum('mf,mfae->ae', t1, eris_ovvv) * 2
        #:self.Fvv -= np.einsum('mf,meaf->ae', t1, eris_ovvv)
        #:self.woVvO = lib.einsum('jf,mebf->mbej', t1, eris_ovvv)
        #:self.woVVo = lib.einsum('jf,mfbe->mbej',-t1, eris_ovvv)
        #:tau = _make_tau(t2, t1, t1)
        #:self.woVoO  = 0.5 * lib.einsum('mebf,ijef->mbij', eris_ovvv, tau)
        #:self.woVoO += 0.5 * lib.einsum('mfbe,ijfe->mbij', eris_ovvv, tau)
        eris_ovoo = np.asarray(eris.ovoo)
        woVoO = np.empty((nocc,nvir,nocc,nocc), dtype=dtype)
        tau = _make_tau(t2, t1, t1)
        theta = t2*2 - t2.transpose(0,1,3,2)

        mem_now = lib.current_memory()[0]
        max_memory = max(0, self.max_memory - mem_now)
        blksize = min(nocc, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvir**3*3))))
        for seg, (p0,p1) in enumerate(lib.prange(0, nocc, blksize)):
            ovvv = eris.get_ovvv(slice(p0,p1))  # ovvv = eris.ovvv[p0:p1]
            # transform integrals (ia|bc) -> (ac|ib)
            fswap['ebmf/%d'%seg] = np.einsum('mebf->ebmf', ovvv)

            self.Fvv += np.einsum('mf,mfae->ae', t1[p0:p1], ovvv) * 2
            self.Fvv -= np.einsum('mf,meaf->ae', t1[p0:p1], ovvv)
            woVoO[p0:p1] = lib.einsum('mebf,ijef->mbij', ovvv, tau)
            woVvO = lib.einsum('jf,mebf->mbej', t1, ovvv)
            woVVo = lib.einsum('jf,mfbe->mbej',-t1, ovvv)
            ovvv = None

            eris_ovov = np.asarray(eris.ovov[p0:p1])
            woOoV = lib.einsum('if,mfne->mnie', t1, eris_ovov)
            woOoV+= eris_ovoo[:,:,p0:p1].transpose(2,0,3,1)
            self.woOoV[p0:p1] = woOoV
            woOoV = None

            tmp = lib.einsum('njbf,mfne->mbej', t2, eris_ovov)
            woVvO -= tmp * .5
            woVVo += tmp

            ovoo = lib.einsum('menf,jf->menj', eris_ovov, t1)
            woVvO -= lib.einsum('nb,menj->mbej', t1, ovoo)
            ovoo = lib.einsum('mfne,jf->menj', eris_ovov, t1)
            woVVo += lib.einsum('nb,menj->mbej', t1, ovoo)
            ovoo = None

            ovov = eris_ovov * 2 - eris_ovov.transpose(0,3,2,1)
            woVvO += lib.einsum('njfb,menf->mbej', theta, ovov) * .5

            self.Fov[p0:p1] = np.einsum('nf,menf->me', t1, ovov)
            tilab = np.einsum('ia,jb->ijab', t1[p0:p1], t1) * .5
            tilab += t2[p0:p1]
            self.Foo += lib.einsum('mief,menf->ni', tilab, ovov)
            self.Fvv -= lib.einsum('mnaf,menf->ae', tilab, ovov)
            eris_ovov = ovov = tilab = None

            woVvO -= lib.einsum('nb,menj->mbej', t1, eris_ovoo[p0:p1,:,:])
            woVVo += lib.einsum('nb,nemj->mbej', t1, eris_ovoo[:,:,p0:p1])

            woVvO += np.asarray(eris.ovvo[p0:p1]).transpose(0,2,1,3)
            woVVo -= np.asarray(eris.oovv[p0:p1]).transpose(0,2,3,1)

            self.woVvO[p0:p1] = woVvO
            self.woVVo[p0:p1] = woVVo

        self.Foo += foo + 0.5*np.einsum('me,ie->mi', self.Fov+fov, t1)
        self.Fvv += fvv - 0.5*np.einsum('me,ma->ae', self.Fov+fov, t1)

        # 0 or 1 virtuals
        woOoO = lib.einsum('je,nemi->mnij', t1, eris_ovoo)
        woOoO = woOoO + woOoO.transpose(1,0,3,2)
        woOoO += np.asarray(eris.oooo).transpose(0,2,1,3)

        tmp = lib.einsum('meni,jneb->mbji', eris_ovoo, t2)
        woVoO -= tmp.transpose(0,1,3,2) * .5
        woVoO -= tmp
        tmp = None
        ovoo = eris_ovoo*2 - eris_ovoo.transpose(2,1,0,3)
        woVoO += lib.einsum('nemi,njeb->mbij', ovoo, theta) * .5
        self.Foo += np.einsum('ne,nemi->mi', t1, ovoo)
        ovoo = None

        eris_ovov = np.asarray(eris.ovov)
        woOoO += lib.einsum('ijef,menf->mnij', tau, eris_ovov)
        self.woOoO = self.saved['woOoO'] = woOoO
        woVoO -= lib.einsum('nb,mnij->mbij', t1, woOoO)
        woOoO = None

        tmpoovv = lib.einsum('njbf,nemf->ejmb', t2, eris_ovov)
        ovov = eris_ovov*2 - eris_ovov.transpose(0,3,2,1)
        eris_ovov = None

        tmpovvo = lib.einsum('nifb,menf->eimb', theta, ovov)
        ovov = None

        tmpovvo *= -.5
        tmpovvo += tmpoovv * .5
        woVoO -= lib.einsum('ie,ejmb->mbij', t1, tmpovvo)
        woVoO -= lib.einsum('ie,ejmb->mbji', t1, tmpoovv)
        woVoO += eris_ovoo.transpose(3,1,2,0)

        # 3 or 4 virtuals
        eris_ovvo = np.asarray(eris.ovvo)
        tmpovvo -= eris_ovvo.transpose(1,3,0,2)
        fswap['ovvo'] = tmpovvo
        tmpovvo = None

        eris_oovv = np.asarray(eris.oovv)
        tmpoovv -= eris_oovv.transpose(3,1,0,2)
        fswap['oovv'] = tmpoovv
        tmpoovv = None

        woVoO += lib.einsum('mebj,ie->mbij', eris_ovvo, t1)
        woVoO += lib.einsum('mjbe,ie->mbji', eris_oovv, t1)
        woVoO += lib.einsum('me,ijeb->mbij', self.Fov, t2)
        self.woVoO = self.saved['woVoO'] = woVoO
        woVoO = eris_ovvo = eris_oovv = None

        #:theta = t2*2 - t2.transpose(0,1,3,2)
        #:eris_ovvv = lib.unpack_tril(np.asarray(eris.ovvv).reshape(nocc*nvir,nvir**2)).reshape(nocc,nvir,nvir,nvir)
        #:ovvv = eris_ovvv*2 - eris_ovvv.transpose(0,3,2,1)
        #:tmpab = lib.einsum('mebf,miaf->eiab', eris_ovvv, t2)
        #:tmpab = tmpab + tmpab.transpose(0,1,3,2) * .5
        #:tmpab-= lib.einsum('mfbe,mifa->eiba', ovvv, theta) * .5
        #:self.wvOvV += eris_ovvv.transpose(2,0,3,1).conj()
        #:self.wvOvV -= tmpab
        nsegs = len(fswap['ebmf'])
        def load_ebmf(slice):
            dat = [fswap['ebmf/%d'%i][slice] for i in range(nsegs)]
            return np.concatenate(dat, axis=2)

        mem_now = lib.current_memory()[0]
        max_memory = max(0, self.max_memory - mem_now)
        blksize = min(nocc, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nocc*nvir**2*4))))
        for p0, p1 in lib.prange(0, nvir, blksize):
            #:wvOvV  = lib.einsum('mebf,miaf->eiab', ovvv, t2)
            #:wvOvV += lib.einsum('mfbe,miaf->eiba', ovvv, t2)
            #:wvOvV -= lib.einsum('mfbe,mifa->eiba', ovvv, t2)*2
            #:wvOvV += lib.einsum('mebf,mifa->eiba', ovvv, t2)

            ebmf = load_ebmf(slice(p0, p1))
            wvOvV = lib.einsum('ebmf,miaf->eiab', ebmf, t2)
            wvOvV = -.5 * wvOvV.transpose(0,1,3,2) - wvOvV

            # Using the permutation symmetry (em|fb) = (em|bf)
            efmb = load_ebmf((slice(None), slice(p0, p1)))
            wvOvV += np.einsum('ebmf->bmfe', efmb.conj())

            # tmp = (mf|be) - (me|bf)*.5
            tmp = -.5 * ebmf
            tmp += efmb.transpose(1,0,2,3)
            ebmf = None
            wvOvV += lib.einsum('efmb,mifa->eiba', tmp, theta)
            tmp = None

            wvOvV += lib.einsum('meni,mnab->eiab', eris_ovoo[:,p0:p1], tau)
            wvOvV -= lib.einsum('me,miab->eiab', self.Fov[:,p0:p1], t2)
            wvOvV += lib.einsum('ma,eimb->eiab', t1, fswap['ovvo'][p0:p1])
            wvOvV += lib.einsum('ma,eimb->eiba', t1, fswap['oovv'][p0:p1])

            self.wvOvV[p0:p1] = wvOvV

        self.made_ee_imds = True
        log.timer('EOM-CCSD EE intermediates', *cput0)
        return self

def _make_tau(t2, t1, r1, fac=1, out=None):
    tau = np.einsum('ia,jb->ijab', t1, r1)
    tau = tau + tau.transpose(1,0,3,2)
    tau *= fac * .5
    tau += t2
    return tau

def _cp(a):
    return np.array(a, copy=False, order='C')


if __name__ == '__main__':
    from pyscf import scf
    from pyscf import gto
    from pyscf.cc import rccsd

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = 'cc-pvdz'
    mol.verbose = 0
    mol.spin = 0
    mol.build()
    mf = scf.RHF(mol).run(conv_tol=1e-14)

    mycc = rccsd.RCCSD(mf)
    ecc, t1, t2 = mycc.kernel()
    print(ecc - -0.21334326214236796)

    myeom = EOMIP(mycc)
    print("IP energies... (right eigenvector)")
    e,v = ipccsd(myeom, nroots=3)
    print(e[0] - 0.43356041409195489)
    print(e[1] - 0.51876598058509493)
    print(e[2] - 0.6782879569941862 )

    print("IP energies... (left eigenvector)")
    le,lv = ipccsd(myeom, nroots=3,left=True)
    print(le[0] - 0.43356040428879794)
    print(le[1] - 0.51876597800180335)
    print(le[2] - 0.67828755013874864)

    e = myeom.ipccsd_star_contract(e, v, lv)
    print(e[0] - 0.43793202073189047)
    print(e[1] - 0.52287073446559729)
    print(e[2] - 0.67994597948852287)

    myeom = EOMEA(mycc)
    print("EA energies... (right eigenvector)")
    e,v = eaccsd(myeom, nroots=3)
    print(e[0] - 0.16737886282063008)
    print(e[1] - 0.24027622989542635)
    print(e[2] - 0.51006796667905585)

    print("EA energies... (left eigenvector)")
    le,lv = eaccsd(myeom, nroots=3, left=True)
    print(le[0] - 0.16737896537079733)
    print(le[1] - 0.24027634198123343)
    print(le[2] - 0.51006809015066612)

    e = myeom.eaccsd_star_contract(e,v,lv)
    print(e[0] - 0.16656250953550664)
    print(e[1] - 0.23944144521387614)
    print(e[2] - 0.41399436888830721)

    myeom = EOMEESpinFlip(mycc)
    np.random.seed(1)
    v = np.random.random(myeom.vector_size())
    r1, r2 = vector_to_amplitudes_eomsf(v, myeom.nmo, myeom.nocc)
    print(lib.finger(r1)    - 0.017703197938757409)
    print(lib.finger(r2[0]) --21.605764517401415)
    print(lib.finger(r2[1]) - 6.5857056438834842)
    print(abs(amplitudes_to_vector_eomsf(r1, r2) - v).max())

    myeom = EOMEE(mycc)
    e,v = myeom.eeccsd(nroots=1)
    print(e - 0.2757159395886167)

    e,v = myeom.eeccsd(nroots=4)
    print(e[0] - 0.2757159395886167)
    print(e[1] - 0.2757159395886167)
    print(e[2] - 0.2757159395886167)
    print(e[3] - 0.3005716731825082)

    e,v = myeom.eeccsd(nroots=4, koopmans=True)
    print(e[0] - 0.2757159395886167)
    print(e[1] - 0.2757159395886167)
    print(e[2] - 0.2757159395886167)
    print(e[3] - 0.3005716731825082)

    e,v = myeom.eeccsd(nroots=4, guess=v[:4])
    print(e[0] - 0.2757159395886167)
    print(e[1] - 0.2757159395886167)
    print(e[2] - 0.2757159395886167)
    print(e[3] - 0.3005716731825082)


    mycc = ccsd.CCSD(mf)
    ecc, t1, t2 = mycc.kernel()
    print(ecc - -0.21334326214236796)

    myeom = EOMIP(mycc)
    print("IP energies... (right eigenvector)")
    e,v = ipccsd(myeom, nroots=3)
    print(e[0] - 0.43356041409195489)
    print(e[1] - 0.51876598058509493)
    print(e[2] - 0.6782879569941862 )

    print("IP energies... (left eigenvector)")
    le,lv = ipccsd(myeom, nroots=3,left=True)
    print(le[0] - 0.43356040428879794)
    print(le[1] - 0.51876597800180335)
    print(le[2] - 0.67828755013874864)

    e = myeom.ipccsd_star_contract(e, v, lv)
    print(e[0] - 0.43793202073189047)
    print(e[1] - 0.52287073446559729)
    print(e[2] - 0.67994597948852287)

    myeom = EOMEA(mycc)
    print("EA energies... (right eigenvector)")
    e,v = eaccsd(myeom, nroots=3)
    print(e[0] - 0.16737886282063008)
    print(e[1] - 0.24027622989542635)
    print(e[2] - 0.51006796667905585)

    print("EA energies... (left eigenvector)")
    le,lv = eaccsd(myeom, nroots=3, left=True)
    print(le[0] - 0.16737896537079733)
    print(le[1] - 0.24027634198123343)
    print(le[2] - 0.51006809015066612)

    e = myeom.eaccsd_star_contract(e,v,lv)
    print(e[0] - 0.16656250953550664)
    print(e[1] - 0.23944144521387614)
    print(e[2] - 0.41399436888830721)

    myeom = EOMEE(mycc)
    e,v = myeom.eeccsd(nroots=1)
    print(e - 0.2757159395886167)

    e,v = myeom.eeccsd(nroots=4)
    print(e[0] - 0.2757159395886167)
    print(e[1] - 0.2757159395886167)
    print(e[2] - 0.2757159395886167)
    print(e[3] - 0.3005716731825082)

    e,v = myeom.eeccsd(nroots=4, koopmans=True)
    print(e[0] - 0.2757159395886167)
    print(e[1] - 0.2757159395886167)
    print(e[2] - 0.2757159395886167)
    print(e[3] - 0.3005716731825082)

    e,v = myeom.eeccsd(nroots=4, guess=v[:4])
    print(e[0] - 0.2757159395886167)
    print(e[1] - 0.2757159395886167)
    print(e[2] - 0.2757159395886167)
    print(e[3] - 0.3005716731825082)


# pyscf/cc/eom_uccsd.py

#!/usr/bin/env python
# Copyright 2014-2021 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: Qiming Sun <osirpt.sun@gmail.com>
#         James D. McClain
#         Jason Yu
#         Shining Sun
#         Mario Motta
#         Chong Sun
#


import numpy as np

from pyscf import lib
from pyscf.lib import logger
from pyscf import ao2mo
from pyscf.cc import ccsd
from pyscf.cc import uccsd
from pyscf.cc import eom_rccsd
from pyscf.cc import uintermediates

########################################
# EOM-IP-CCSD
########################################

def vector_to_amplitudes_ip(vector, nmo, nocc):
    '''For spin orbitals'''
    nocca, noccb = nocc
    nmoa, nmob = nmo
    nvira, nvirb = nmoa-nocca, nmob-noccb

    sizes = (nocca, noccb, nocca*(nocca-1)//2*nvira, noccb*nocca*nvira,
             nocca*noccb*nvirb, noccb*(noccb-1)//2*nvirb)
    sections = np.cumsum(sizes[:-1])
    r1a, r1b, r2a, r2baa, r2abb, r2b = np.split(vector, sections)
    r2a = r2a.reshape(nocca*(nocca-1)//2,nvira)
    r2b = r2b.reshape(noccb*(noccb-1)//2,nvirb)
    r2baa = r2baa.reshape(noccb,nocca,nvira).copy()
    r2abb = r2abb.reshape(nocca,noccb,nvirb).copy()

    idxa = np.tril_indices(nocca, -1)
    idxb = np.tril_indices(noccb, -1)
    r2aaa = np.zeros((nocca,nocca,nvira), vector.dtype)
    r2bbb = np.zeros((noccb,noccb,nvirb), vector.dtype)
    r2aaa[idxa[0],idxa[1]] = r2a
    r2aaa[idxa[1],idxa[0]] =-r2a
    r2bbb[idxb[0],idxb[1]] = r2b
    r2bbb[idxb[1],idxb[0]] =-r2b

    r1 = (r1a.copy(), r1b.copy())
    r2 = (r2aaa, r2baa, r2abb, r2bbb)
    return r1, r2

def amplitudes_to_vector_ip(r1, r2):
    '''For spin orbitals'''
    r1a, r1b = r1
    r2aaa, r2baa, r2abb, r2bbb = r2
    nocca, noccb, nvirb = r2abb.shape
    idxa = np.tril_indices(nocca, -1)
    idxb = np.tril_indices(noccb, -1)
    return np.hstack((r1a, r1b,
                      r2aaa[idxa].ravel(), r2baa.ravel(),
                      r2abb.ravel(), r2bbb[idxb].ravel()))

def spatial2spin_ip(r1, r2, orbspin=None):
    '''Convert R1/R2 of spatial orbital representation to R1/R2 of
    spin-orbital representation
    '''
    r1a, r1b = r1
    r2aaa, r2baa, r2abb, r2bbb = r2
    nocc_a, nvir_a = r2aaa.shape[1:]
    nocc_b, nvir_b = r2bbb.shape[1:]

    if orbspin is None:
        orbspin = np.zeros((nocc_a+nvir_a)*2, dtype=int)
        orbspin[1::2] = 1

    nocc = nocc_a + nocc_b
    nvir = nvir_a + nvir_b
    idxoa = np.where(orbspin[:nocc] == 0)[0]
    idxob = np.where(orbspin[:nocc] == 1)[0]
    idxva = np.where(orbspin[nocc:] == 0)[0]
    idxvb = np.where(orbspin[nocc:] == 1)[0]

    r1 = np.zeros((nocc), dtype=r1a.dtype)
    r1[idxoa] = r1a
    r1[idxob] = r1b

    r2 = np.zeros((nocc**2, nvir), dtype=r2aaa.dtype)
    idxoaa = idxoa[:,None] * nocc + idxoa
    idxoab = idxoa[:,None] * nocc + idxob
    idxoba = idxob[:,None] * nocc + idxoa
    idxobb = idxob[:,None] * nocc + idxob
    # idxvaa = idxva[:,None] * nvir + idxva
    # idxvab = idxva[:,None] * nvir + idxvb
    # idxvba = idxvb[:,None] * nvir + idxva
    # idxvbb = idxvb[:,None] * nvir + idxvb
    r2aaa = r2aaa.reshape(nocc_a*nocc_a, nvir_a)
    r2baa = r2baa.reshape(nocc_b*nocc_a, nvir_a)
    r2abb = r2abb.reshape(nocc_a*nocc_b, nvir_b)
    r2bbb = r2bbb.reshape(nocc_b*nocc_b, nvir_b)
    lib.takebak_2d(r2, r2aaa, idxoaa.ravel(), idxva.ravel())
    lib.takebak_2d(r2, r2baa, idxoba.ravel(), idxva.ravel())
    lib.takebak_2d(r2, r2abb, idxoab.ravel(), idxvb.ravel())
    lib.takebak_2d(r2, r2bbb, idxobb.ravel(), idxvb.ravel())
    r2aba = -r2baa
    r2bab = -r2abb
    lib.takebak_2d(r2, r2aba, idxoab.T.ravel(), idxva.ravel())
    lib.takebak_2d(r2, r2bab, idxoba.T.ravel(), idxvb.ravel())
    return r1, r2.reshape(nocc, nocc, nvir)

def spin2spatial_ip(r1, r2, orbspin):
    nocc, nvir = r2.shape[1:]

    idxoa = np.where(orbspin[:nocc] == 0)[0]
    idxob = np.where(orbspin[:nocc] == 1)[0]
    idxva = np.where(orbspin[nocc:] == 0)[0]
    idxvb = np.where(orbspin[nocc:] == 1)[0]
    nocc_a = len(idxoa)
    nocc_b = len(idxob)
    nvir_a = len(idxva)
    nvir_b = len(idxvb)

    r1a = r1[idxoa]
    r1b = r1[idxob]

    idxoaa = idxoa[:,None] * nocc + idxoa
    idxoab = idxoa[:,None] * nocc + idxob
    idxoba = idxob[:,None] * nocc + idxoa
    idxobb = idxob[:,None] * nocc + idxob
    #idxvaa = idxva[:,None] * nvir + idxva
    #idxvab = idxva[:,None] * nvir + idxvb
    #idxvba = idxvb[:,None] * nvir + idxva
    #idxvbb = idxvb[:,None] * nvir + idxvb

    r2 = r2.reshape(nocc**2, nvir)
    r2aaa = lib.take_2d(r2, idxoaa.ravel(), idxva.ravel())
    r2baa = lib.take_2d(r2, idxoba.ravel(), idxva.ravel())
    r2abb = lib.take_2d(r2, idxoab.ravel(), idxvb.ravel())
    r2bbb = lib.take_2d(r2, idxobb.ravel(), idxvb.ravel())

    r2aaa = r2aaa.reshape(nocc_a, nocc_a, nvir_a)
    r2baa = r2baa.reshape(nocc_b, nocc_a, nvir_a)
    r2abb = r2abb.reshape(nocc_a, nocc_b, nvir_b)
    r2bbb = r2bbb.reshape(nocc_b, nocc_b, nvir_b)
    return [r1a, r1b], [r2aaa, r2baa, r2abb, r2bbb]

def ipccsd_matvec(eom, vector, imds=None, diag=None):
    '''For spin orbitals
    R2 operators of the form s_{ij}^{ b}, i.e. indices jb are coupled.'''
    # Ref: Tu, Wang, and Li, J. Chem. Phys. 136, 174102 (2012) Eqs.(8)-(9)
    if imds is None: imds = eom.make_imds()
    t1, t2 = imds.t1, imds.t2
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nocca, noccb, nvira, nvirb = t2ab.shape
    nmoa, nmob = nocca+nvira, noccb+nvirb
    r1, r2 = vector_to_amplitudes_ip(vector, (nmoa,nmob), (nocca,noccb))
    r1a, r1b = r1
    r2aaa, r2baa, r2abb, r2bbb = r2

    #Foo, Fov, and Wooov
    Hr1a  = np.einsum('me,mie->i', imds.Fov, r2aaa)
    Hr1a -= np.einsum('ME,iME->i', imds.FOV, r2abb)
    Hr1b  = np.einsum('ME,MIE->I', imds.FOV, r2bbb)
    Hr1b -= np.einsum('me,Ime->I', imds.Fov, r2baa)

    Hr1a += -np.einsum('mi,m->i', imds.Foo, r1a)
    Hr1b += -np.einsum('MI,M->I', imds.FOO, r1b)

    Hr1a += -0.5*np.einsum('nime,mne->i', imds.Wooov, r2aaa)
    Hr1b +=      np.einsum('NIme,Nme->I', imds.WOOov, r2baa)
    Hr1b += -0.5*np.einsum('NIME,MNE->I', imds.WOOOV, r2bbb)
    Hr1a +=      np.einsum('niME,nME->i', imds.WooOV, r2abb)

    # Fvv term
    Hr2aaa = lib.einsum('be,ije->ijb', imds.Fvv, r2aaa)
    Hr2abb = lib.einsum('BE,iJE->iJB', imds.FVV, r2abb)
    Hr2bbb = lib.einsum('BE,IJE->IJB', imds.FVV, r2bbb)
    Hr2baa = lib.einsum('be,Ije->Ijb', imds.Fvv, r2baa)

    # Foo term
    tmpa = lib.einsum('mi,mjb->ijb', imds.Foo, r2aaa)
    Hr2aaa -= tmpa - tmpa.transpose((1,0,2))
    Hr2abb -= lib.einsum('mi,mJB->iJB', imds.Foo, r2abb)
    Hr2abb -= lib.einsum('MJ,iMB->iJB', imds.FOO, r2abb)
    Hr2baa -= lib.einsum('MI,Mjb->Ijb', imds.FOO, r2baa)
    Hr2baa -= lib.einsum('mj,Imb->Ijb', imds.Foo, r2baa)
    tmpb = lib.einsum('MI,MJB->IJB', imds.FOO, r2bbb)
    Hr2bbb -= tmpb - tmpb.transpose((1,0,2))

    # Wovoo term
    Hr2aaa -= np.einsum('mjbi,m->ijb', imds.Woovo, r1a)
    Hr2abb += np.einsum('miBJ,m->iJB', imds.WooVO, r1a)
    Hr2baa += np.einsum('MIbj,M->Ijb', imds.WOOvo, r1b)
    Hr2bbb -= np.einsum('MJBI,M->IJB', imds.WOOVO, r1b)

    # Woooo term
    Hr2aaa += .5 * lib.einsum('minj,mnb->ijb', imds.Woooo, r2aaa)
    Hr2abb +=      lib.einsum('miNJ,mNB->iJB', imds.WooOO, r2abb)
    Hr2bbb += .5 * lib.einsum('MINJ,MNB->IJB', imds.WOOOO, r2bbb)
    Hr2baa +=      lib.einsum('njMI,Mnb->Ijb', imds.WooOO, r2baa)

    # Wovvo terms
    tmp = lib.einsum('mebj,ime->ijb', imds.Wovvo, r2aaa)
    tmp += lib.einsum('MEbj,iME->ijb', imds.WOVvo, r2abb)
    Hr2aaa += tmp - tmp.transpose(1, 0, 2)

    WooVV = -imds.WoVVo.transpose(0,3,2,1)
    WOOvv = -imds.WOvvO.transpose(0,3,2,1)
    Hr2abb += lib.einsum('MEBJ,iME->iJB', imds.WOVVO, r2abb)
    Hr2abb += lib.einsum('meBJ,ime->iJB', imds.WovVO, r2aaa)
    Hr2abb += -lib.einsum('miBE,mJE->iJB', WooVV, r2abb)

    Hr2baa += lib.einsum('meaj,Ime->Ija', imds.Wovvo, r2baa)
    Hr2baa += lib.einsum('MEaj,IME->Ija', imds.WOVvo, r2bbb)
    Hr2baa += -lib.einsum('MIab,Mjb->Ija', WOOvv, r2baa)

    tmp = lib.einsum('MEBJ,IME->IJB', imds.WOVVO, r2bbb)
    tmp += lib.einsum('meBJ,Ime->IJB', imds.WovVO, r2baa)
    Hr2bbb += tmp - tmp.transpose(1, 0, 2)

    # T2 term
    Hr2aaa -= 0.5 * lib.einsum('menf,mnf,jibe->ijb', imds.Wovov, r2aaa, t2aa)
    Hr2aaa -= lib.einsum('meNF,mNF,jibe->ijb', imds.WovOV, r2abb, t2aa)

    Hr2abb -= 0.5 * lib.einsum('menf,mnf,iJeB->iJB', imds.Wovov, r2aaa, t2ab)
    Hr2abb -= lib.einsum('meNF,mNF,iJeB->iJB', imds.WovOV, r2abb, t2ab)

    Hr2baa -= 0.5 * lib.einsum('MENF,MNF,jIbE->Ijb', imds.WOVOV, r2bbb, t2ab)
    Hr2baa -= lib.einsum('nfME,Mnf,jIbE->Ijb', imds.WovOV, r2baa, t2ab)

    Hr2bbb -= 0.5 * lib.einsum('MENF,MNF,JIBE->IJB', imds.WOVOV, r2bbb, t2bb)
    Hr2bbb -= lib.einsum('nfME,Mnf,JIBE->IJB', imds.WovOV, r2baa, t2bb)

    vector = amplitudes_to_vector_ip([Hr1a, Hr1b], [Hr2aaa, Hr2baa, Hr2abb, Hr2bbb])
    return vector

def ipccsd_diag(eom, imds=None):
    if imds is None: imds = eom.make_imds()
    t1, t2 = imds.t1, imds.t2
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2

    nocc_a, nvir_a = t1a.shape
    nocc_b, nvir_b = t1b.shape

    Hr1a = -np.diag(imds.Foo)
    Hr1b = -np.diag(imds.FOO)

    Fvv_diag = np.diag(imds.Fvv)
    Foo_diag = np.diag(imds.Foo)
    FOO_diag = np.diag(imds.FOO)
    FVV_diag = np.diag(imds.FVV)

    Woooo_slice = np.einsum('iijj->ij',imds.Woooo)
    Wovvo_slice = np.einsum('iaai->ia',imds.Wovvo)
    WooOO_slice = np.einsum('jjii->ij',imds.WooOO)
    WOvvO_slice = np.einsum('iaai->ia',imds.WOvvO)
    WooOO_slice_T = np.einsum('iijj->ij',imds.WooOO)
    WoVVo_slice = np.einsum('iaai->ia',imds.WoVVo)
    WOVVO_slice = np.einsum('jaaj->ja',imds.WOVVO)
    WOOOO_slice = np.einsum('iijj->ij',imds.WOOOO)

    Wovov_t2_dot = np.einsum('jaib,jiab->ija',imds.Wovov,t2aa)
    WovOV_t2_dot = np.einsum('ibja,ijba->ija',imds.WovOV,t2ab)
    WovOV_t2_dot_T = np.einsum('jaib,jiab->ija',imds.WovOV,t2ab)
    WOVOV_t2_dot = np.einsum('jaib,jiab->ija',imds.WOVOV,t2bb)

    Hr2aaa = Fvv_diag[None,None,:] - Foo_diag[:,None,None] - Foo_diag[None,:,None] \
             + Woooo_slice[:,:,None] + Wovvo_slice[:,None,:] + Wovvo_slice[None,:,:] \
             - Wovov_t2_dot

    Hr2baa = Fvv_diag[None,None,:] - FOO_diag[:,None,None] - Foo_diag[None,:,None] \
             + WooOO_slice[:,:,None] + WOvvO_slice[:,None,:] + Wovvo_slice[None,:,:] \
             - WovOV_t2_dot_T

    Hr2abb = FVV_diag[None,None,:] - Foo_diag[:,None,None] - FOO_diag[None,:,None] \
             + WooOO_slice_T[:,:,None] + WoVVo_slice[:,None,:] + WOVVO_slice[None,:,:] \
             - WovOV_t2_dot

    Hr2bbb = FVV_diag[None,None,:] - FOO_diag[:,None,None] - FOO_diag[None,:,None] \
             + WOOOO_slice[:,:,None] + WOVVO_slice[:,None,:] + WOVVO_slice[None,:,:] \
             - WOVOV_t2_dot

    vector = amplitudes_to_vector_ip([Hr1a, Hr1b], [Hr2aaa, Hr2baa, Hr2abb, Hr2bbb])
    return vector


class EOMIP(eom_rccsd.EOMIP):
    matvec = ipccsd_matvec
    l_matvec = None
    get_diag = ipccsd_diag
    ipccsd_star = None
    ccsd_star_contract = None

    def __init__(self, cc):
        eom_rccsd.EOMIP.__init__(self, cc)
        self.nocc = cc.get_nocc()
        self.nmo = cc.get_nmo()

    def get_init_guess(self, nroots=1, koopmans=True, diag=None):
        if koopmans:
            nocca, noccb = self.nocc
            idx = diag[:nocca+noccb].argsort()
        else:
            idx = diag.argsort()

        size = self.vector_size()
        dtype = getattr(diag, 'dtype', np.double)
        nroots = min(nroots, size)
        guess = []
        for i in idx[:nroots]:
            g = np.zeros(size, dtype)
            g[i] = 1.0
            guess.append(g)
        return guess

    def vector_to_amplitudes(self, vector, nmo=None, nocc=None):
        if nmo is None: nmo = self.nmo
        if nocc is None: nocc = self.nocc
        return vector_to_amplitudes_ip(vector, nmo, nocc)

    def amplitudes_to_vector(self, r1, r2):
        return amplitudes_to_vector_ip(r1, r2)

    def vector_size(self):
        '''size of the vector based on spin-orbital basis'''
        nocca, noccb = self.nocc
        nmoa, nmob = self.nmo
        nvira, nvirb = nmoa-nocca, nmob-noccb
        return (nocca + noccb
                + nocca*(nocca-1)//2*nvira + noccb*nocca*nvira
                + nocca*noccb*nvirb + noccb*(noccb-1)//2*nvirb)

    def make_imds(self, eris=None):
        imds = _IMDS(self._cc, eris)
        imds.make_ip()
        return imds

########################################
# EOM-EA-CCSD
########################################

def vector_to_amplitudes_ea(vector, nmo, nocc):
    nocca, noccb = nocc
    nmoa, nmob = nmo
    nvira, nvirb = nmoa-nocca, nmob-noccb

    sizes = (nvira, nvirb, nocca*nvira*(nvira-1)//2, nocca*nvirb*nvira,
             noccb*nvira*nvirb, noccb*nvirb*(nvirb-1)//2)
    sections = np.cumsum(sizes[:-1])
    r1a, r1b, r2a, r2aba, r2bab, r2b = np.split(vector, sections)
    r2a = r2a.reshape(nocca,nvira*(nvira-1)//2)
    r2b = r2b.reshape(noccb,nvirb*(nvirb-1)//2)
    r2aba = r2aba.reshape(nocca,nvirb,nvira).copy()
    r2bab = r2bab.reshape(noccb,nvira,nvirb).copy()

    idxa = np.tril_indices(nvira, -1)
    idxb = np.tril_indices(nvirb, -1)
    r2aaa = np.zeros((nocca,nvira,nvira), vector.dtype)
    r2bbb = np.zeros((noccb,nvirb,nvirb), vector.dtype)
    r2aaa[:,idxa[0],idxa[1]] = r2a
    r2aaa[:,idxa[1],idxa[0]] =-r2a
    r2bbb[:,idxb[0],idxb[1]] = r2b
    r2bbb[:,idxb[1],idxb[0]] =-r2b

    r1 = (r1a.copy(), r1b.copy())
    r2 = (r2aaa, r2aba, r2bab, r2bbb)
    return r1, r2

def amplitudes_to_vector_ea(r1, r2):
    r1a, r1b = r1
    r2aaa, r2aba, r2bab, r2bbb = r2
    nocca, nvirb, nvira = r2aba.shape
    idxa = np.tril_indices(nvira, -1)
    idxb = np.tril_indices(nvirb, -1)
    return np.hstack((r1a, r1b,
                      r2aaa[:,idxa[0],idxa[1]].ravel(),
                      r2aba.ravel(), r2bab.ravel(),
                      r2bbb[:,idxb[0],idxb[1]].ravel()))

def spatial2spin_ea(r1, r2, orbspin=None):
    '''Convert R1/R2 of spatial orbital representation to R1/R2 of
    spin-orbital representation
    '''
    r1a, r1b = r1
    r2aaa, r2aba, r2bab, r2bbb = r2
    nocc_a, nvir_a = r2aaa.shape[:2]
    nocc_b, nvir_b = r2bbb.shape[:2]

    if orbspin is None:
        orbspin = np.zeros((nocc_a+nvir_a)*2, dtype=int)
        orbspin[1::2] = 1

    nocc = nocc_a + nocc_b
    nvir = nvir_a + nvir_b
    idxoa = np.where(orbspin[:nocc] == 0)[0]
    idxob = np.where(orbspin[:nocc] == 1)[0]
    idxva = np.where(orbspin[nocc:] == 0)[0]
    idxvb = np.where(orbspin[nocc:] == 1)[0]

    r1 = np.zeros((nvir), dtype=r1a.dtype)
    r1[idxva] = r1a
    r1[idxvb] = r1b

    r2 = np.zeros((nocc, nvir**2), dtype=r2aaa.dtype)
    #idxoaa = idxoa[:,None] * nocc + idxoa
    #idxoab = idxoa[:,None] * nocc + idxob
    #idxoba = idxob[:,None] * nocc + idxoa
    #idxobb = idxob[:,None] * nocc + idxob
    idxvaa = idxva[:,None] * nvir + idxva
    idxvab = idxva[:,None] * nvir + idxvb
    idxvba = idxvb[:,None] * nvir + idxva
    idxvbb = idxvb[:,None] * nvir + idxvb

    r2aaa = r2aaa.reshape(nocc_a, nvir_a*nvir_a)
    r2aba = r2aba.reshape(nocc_a, nvir_b*nvir_a)
    r2bab = r2bab.reshape(nocc_b, nvir_a*nvir_b)
    r2bbb = r2bbb.reshape(nocc_b, nvir_b*nvir_b)

    lib.takebak_2d(r2, r2aaa, idxoa.ravel(), idxvaa.ravel())
    lib.takebak_2d(r2, r2aba, idxoa.ravel(), idxvba.ravel())
    lib.takebak_2d(r2, r2bab, idxob.ravel(), idxvab.ravel())
    lib.takebak_2d(r2, r2bbb, idxob.ravel(), idxvbb.ravel())
    r2aab = -r2aba
    r2bba = -r2bab
    lib.takebak_2d(r2, r2bba, idxob.ravel(), idxvba.T.ravel())
    lib.takebak_2d(r2, r2aab, idxoa.ravel(), idxvab.T.ravel())
    r2 = r2.reshape(nocc, nvir, nvir)
    return r1, r2

def spin2spatial_ea(r1, r2, orbspin):
    nocc, nvir = r2.shape[:2]

    idxoa = np.where(orbspin[:nocc] == 0)[0]
    idxob = np.where(orbspin[:nocc] == 1)[0]
    idxva = np.where(orbspin[nocc:] == 0)[0]
    idxvb = np.where(orbspin[nocc:] == 1)[0]
    nocc_a = len(idxoa)
    nocc_b = len(idxob)
    nvir_a = len(idxva)
    nvir_b = len(idxvb)

    r1a = r1[idxva]
    r1b = r1[idxvb]

    #idxoaa = idxoa[:,None] * nocc + idxoa
    #idxoab = idxoa[:,None] * nocc + idxob
    #idxoba = idxob[:,None] * nocc + idxoa
    #idxobb = idxob[:,None] * nocc + idxob
    idxvaa = idxva[:,None] * nvir + idxva
    idxvab = idxva[:,None] * nvir + idxvb
    idxvba = idxvb[:,None] * nvir + idxva
    idxvbb = idxvb[:,None] * nvir + idxvb

    r2 = r2.reshape(nocc, nvir**2)
    r2aaa = lib.take_2d(r2, idxoa.ravel(), idxvaa.ravel())
    r2aba = lib.take_2d(r2, idxoa.ravel(), idxvba.ravel())
    r2bab = lib.take_2d(r2, idxob.ravel(), idxvab.ravel())
    r2bbb = lib.take_2d(r2, idxob.ravel(), idxvbb.ravel())

    r2aaa = r2aaa.reshape(nocc_a, nvir_a, nvir_a)
    r2aba = r2aba.reshape(nocc_a, nvir_b, nvir_a)
    r2bab = r2bab.reshape(nocc_b, nvir_a, nvir_b)
    r2bbb = r2bbb.reshape(nocc_b, nvir_b, nvir_b)
    return [r1a, r1b], [r2aaa, r2aba, r2bab, r2bbb]

def eaccsd_matvec(eom, vector, imds=None, diag=None):
    '''For spin orbitals.

    R2 operators of the form s_{ j}^{ab}, i.e. indices jb are coupled.'''
    # Ref: Nooijen and Bartlett, J. Chem. Phys. 102, 3629 (1994) Eqs.(30)-(31)
    if imds is None: imds = eom.make_imds()
    t1, t2, eris = imds.t1, imds.t2, imds.eris
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nocca, noccb, nvira, nvirb = t2ab.shape
    nmoa, nmob = nocca+nvira, noccb+nvirb
    r1, r2 = vector_to_amplitudes_ea(vector, (nmoa,nmob), (nocca,noccb))
    r1a, r1b = r1
    r2aaa, r2aba, r2bab, r2bbb = r2

    # Fov terms
    Hr1a  = np.einsum('ld,lad->a', imds.Fov, r2aaa)
    Hr1a += np.einsum('LD,LaD->a', imds.FOV, r2bab)
    Hr1b  = np.einsum('ld,lAd->A', imds.Fov, r2aba)
    Hr1b += np.einsum('LD,LAD->A', imds.FOV, r2bbb)

    # Fvv terms
    Hr1a += np.einsum('ac,c->a', imds.Fvv, r1a)
    Hr1b += np.einsum('AC,C->A', imds.FVV, r1b)

    # Wvovv
    Hr1a += 0.5*lib.einsum('acld,lcd->a', imds.Wvvov, r2aaa)
    Hr1a +=     lib.einsum('acLD,LcD->a', imds.WvvOV, r2bab)
    Hr1b += 0.5*lib.einsum('ACLD,LCD->A', imds.WVVOV, r2bbb)
    Hr1b +=     lib.einsum('ACld,lCd->A', imds.WVVov, r2aba)

    #** Wvvvv term
    #:Hr2aaa = lib.einsum('acbd,jcd->jab', eris_vvvv, r2aaa)
    #:Hr2aba = lib.einsum('bdac,jcd->jab', eris_vvVV, r2aba)
    #:Hr2bab = lib.einsum('acbd,jcd->jab', eris_vvVV, r2bab)
    #:Hr2bbb = lib.einsum('acbd,jcd->jab', eris_VVVV, r2bbb)
    u2 = (r2aaa + np.einsum('c,jd->jcd', r1a, t1a) - np.einsum('d,jc->jcd', r1a, t1a),
          r2aba + np.einsum('c,jd->jcd', r1b, t1a),
          r2bab + np.einsum('c,jd->jcd', r1a, t1b),
          r2bbb + np.einsum('c,jd->jcd', r1b, t1b) - np.einsum('d,jc->jcd', r1b, t1b))
    Hr2aaa, Hr2aba, Hr2bab, Hr2bbb = _add_vvvv_ea(eom._cc, u2, eris)
    u2 = None

    tauaa, tauab, taubb = uccsd.make_tau(t2, t1, t1)
    eris_ovov = np.asarray(eris.ovov)
    eris_OVOV = np.asarray(eris.OVOV)
    eris_ovOV = np.asarray(eris.ovOV)
    tmpaaa = lib.einsum('menf,jef->mnj', eris_ovov, r2aaa) * .5
    Hr2aaa += lib.einsum('mnj,mnab->jab', tmpaaa, tauaa)
    tmpaaa = tauaa = None

    tmpbbb = lib.einsum('menf,jef->mnj', eris_OVOV, r2bbb) * .5
    Hr2bbb += lib.einsum('mnj,mnab->jab', tmpbbb, taubb)
    tmpbbb = taubb = None

    tmpabb = lib.einsum('menf,jef->mnj', eris_ovOV, r2bab)
    Hr2bab += lib.einsum('mnj,mnab->jab', tmpabb, tauab)
    tmpaba = lib.einsum('nfme,jef->nmj', eris_ovOV, r2aba)
    Hr2aba += lib.einsum('nmj,nmba->jab', tmpaba, tauab)
    tmpaba = tauab = None
    eris_ovov = eris_OVOV = eris_ovOV = None

    eris_ovvv = imds.eris.get_ovvv(slice(None))
    tmpaaa = lib.einsum('mebf,jef->mjb', eris_ovvv, r2aaa)
    tmpaaa = lib.einsum('mjb,ma->jab', tmpaaa, t1a)
    Hr2aaa-= tmpaaa - tmpaaa.transpose(0,2,1)
    tmpaaa = eris_ovvv = None

    eris_OVVV = imds.eris.get_OVVV(slice(None))
    tmpbbb = lib.einsum('mebf,jef->mjb', eris_OVVV, r2bbb)
    tmpbbb = lib.einsum('mjb,ma->jab', tmpbbb, t1b)
    Hr2bbb-= tmpbbb - tmpbbb.transpose(0,2,1)
    tmpbbb = eris_OVVV = None

    eris_ovVV = imds.eris.get_ovVV(slice(None))
    eris_OVvv = imds.eris.get_OVvv(slice(None))
    tmpaab = lib.einsum('meBF,jFe->mjB', eris_ovVV, r2aba)
    Hr2aba-= lib.einsum('mjB,ma->jBa', tmpaab, t1a)
    tmpabb = lib.einsum('meBF,JeF->mJB', eris_ovVV, r2bab)
    Hr2bab-= lib.einsum('mJB,ma->JaB', tmpabb, t1a)
    tmpaab = tmpabb = eris_ovVV = None

    tmpbaa = lib.einsum('MEbf,jEf->Mjb', eris_OVvv, r2aba)
    Hr2aba-= lib.einsum('Mjb,MA->jAb', tmpbaa, t1b)
    tmpbba = lib.einsum('MEbf,JfE->MJb', eris_OVvv, r2bab)
    Hr2bab-= lib.einsum('MJb,MA->JbA', tmpbba, t1b)
    tmpbaa = tmpbba = eris_OVvv = None
    #** Wvvvv term end

    # Wvvvo
    Hr2aaa += np.einsum('acbj,c->jab', imds.Wvvvo, r1a)
    Hr2bbb += np.einsum('ACBJ,C->JAB', imds.WVVVO, r1b)
    Hr2bab += np.einsum('acBJ,c->JaB', imds.WvvVO, r1a)
    Hr2aba += np.einsum('ACbj,C->jAb', imds.WVVvo, r1b)

    # Wovvo
    tmp2aa = lib.einsum('ldbj,lad->jab', imds.Wovvo, r2aaa)
    tmp2aa += lib.einsum('ldbj,lad->jab', imds.WOVvo, r2bab)
    Hr2aaa += tmp2aa - tmp2aa.transpose(0,2,1)

    Hr2bab += lib.einsum('ldbj,lad->jab', imds.WovVO, r2aaa)
    Hr2bab += lib.einsum('ldbj,lad->jab', imds.WOVVO, r2bab)
    Hr2bab += lib.einsum('ldaj,ldb->jab', imds.WOvvO, r2bab)

    Hr2aba += lib.einsum('ldbj,lad->jab', imds.WOVvo, r2bbb)
    Hr2aba += lib.einsum('ldbj,lad->jab', imds.Wovvo, r2aba)
    Hr2aba += lib.einsum('ldaj,ldb->jab', imds.WoVVo, r2aba)

    tmp2bb = lib.einsum('ldbj,lad->jab', imds.WOVVO, r2bbb)
    tmp2bb += lib.einsum('ldbj,lad->jab', imds.WovVO, r2aba)
    Hr2bbb += tmp2bb - tmp2bb.transpose(0,2,1)

    #Fvv Term
    tmpa = lib.einsum('ac,jcb->jab', imds.Fvv, r2aaa)
    Hr2aaa += tmpa - tmpa.transpose((0,2,1))
    Hr2aba += lib.einsum('AC,jCb->jAb', imds.FVV, r2aba)
    Hr2bab += lib.einsum('ac,JcB->JaB', imds.Fvv, r2bab)
    Hr2aba += lib.einsum('bc, jAc -> jAb', imds.Fvv, r2aba)
    Hr2bab += lib.einsum('BC, JaC -> JaB', imds.FVV, r2bab)
    tmpb = lib.einsum('AC,JCB->JAB', imds.FVV, r2bbb)
    Hr2bbb += tmpb - tmpb.transpose((0,2,1))

    #Foo Term
    Hr2aaa -= lib.einsum('lj,lab->jab', imds.Foo, r2aaa)
    Hr2bbb -= lib.einsum('LJ,LAB->JAB', imds.FOO, r2bbb)
    Hr2bab -= lib.einsum('LJ,LaB->JaB', imds.FOO, r2bab)
    Hr2aba -= lib.einsum('lj,lAb->jAb', imds.Foo, r2aba)

    # Woovv term
    Hr2aaa -= 0.5 * lib.einsum('kcld,lcd,kjab->jab', imds.Wovov, r2aaa, t2aa)
    Hr2bab -= 0.5 * lib.einsum('kcld,lcd,kJaB->JaB', imds.Wovov, r2aaa, t2ab)

    Hr2aba -= lib.einsum('ldKC,lCd,jKbA->jAb', imds.WovOV, r2aba, t2ab)
    Hr2aaa -= lib.einsum('kcLD,LcD,kjab->jab', imds.WovOV, r2bab, t2aa)

    Hr2aba -= 0.5 * lib.einsum('KCLD,LCD,jKbA->jAb', imds.WOVOV, r2bbb, t2ab)
    Hr2bbb -= 0.5 * lib.einsum('KCLD,LCD,KJAB->JAB', imds.WOVOV, r2bbb, t2bb)

    Hr2bbb -= lib.einsum('ldKC,lCd,KJAB->JAB', imds.WovOV, r2aba, t2bb)
    Hr2bab -= lib.einsum('kcLD,LcD,kJaB->JaB', imds.WovOV, r2bab, t2ab)

    vector = amplitudes_to_vector_ea([Hr1a, Hr1b], [Hr2aaa, Hr2aba, Hr2bab, Hr2bbb])
    return vector

def _add_vvvv_ea(mycc, r2, eris):
    time0 = logger.process_clock(), logger.perf_counter()
    log = logger.Logger(mycc.stdout, mycc.verbose)
    r2aaa, r2aba, r2bab, r2bbb = r2
    nocca, noccb = mycc.nocc

    if mycc.direct:
        if getattr(eris, 'mo_coeff', None) is not None:
            mo_a, mo_b = eris.mo_coeff
        else:
            moidxa, moidxb = mycc.get_frozen_mask()
            mo_a = mycc.mo_coeff[0][:,moidxa]
            mo_b = mycc.mo_coeff[1][:,moidxb]

        r2aaa = lib.einsum('xab,pa->xpb', r2aaa, mo_a[:,nocca:])
        r2aaa = lib.einsum('xab,pb->xap', r2aaa, mo_a[:,nocca:])
        r2aba = lib.einsum('xab,pa->xpb', r2aba, mo_b[:,noccb:])
        r2aba = lib.einsum('xab,pb->xap', r2aba, mo_a[:,nocca:])
        r2bab = lib.einsum('xab,pa->xpb', r2bab, mo_a[:,nocca:])
        r2bab = lib.einsum('xab,pb->xap', r2bab, mo_b[:,noccb:])
        r2bbb = lib.einsum('xab,pa->xpb', r2bbb, mo_b[:,noccb:])
        r2bbb = lib.einsum('xab,pb->xap', r2bbb, mo_b[:,noccb:])

        r2 = np.vstack((r2aaa, r2aba, r2bab, r2bbb))
        r2aaa = r2aba = r2bab = r2bbb = None
        time0 = log.timer_debug1('vvvv-tau', *time0)

        buf = ccsd._contract_vvvv_t2(mycc, mycc.mol, None, r2, verbose=log)
        sections = np.cumsum([nocca,nocca,noccb])
        Hr2aaa, Hr2aba, Hr2bab, Hr2bbb = np.split(buf, sections)
        buf = None

        Hr2aaa = lib.einsum('xpb,pa->xab', Hr2aaa, mo_a[:,nocca:])
        Hr2aaa = lib.einsum('xap,pb->xab', Hr2aaa, mo_a[:,nocca:])
        Hr2aba = lib.einsum('xpb,pa->xab', Hr2aba, mo_b[:,noccb:])
        Hr2aba = lib.einsum('xap,pb->xab', Hr2aba, mo_a[:,nocca:])
        Hr2bab = lib.einsum('xpb,pa->xab', Hr2bab, mo_a[:,nocca:])
        Hr2bab = lib.einsum('xap,pb->xab', Hr2bab, mo_b[:,noccb:])
        Hr2bbb = lib.einsum('xpb,pa->xab', Hr2bbb, mo_b[:,noccb:])
        Hr2bbb = lib.einsum('xap,pb->xab', Hr2bbb, mo_b[:,noccb:])

    elif r2aaa.dtype == np.double:
        r2aab = np.asarray(r2aba.transpose(0,2,1), order='C')
        Hr2aab = eris._contract_vvVV_t2(mycc, r2aab, mycc.direct, None)
        Hr2aba = np.asarray(Hr2aab.transpose(0,2,1), order='C')
        r2aab = Hr2aab = None
        Hr2bab = eris._contract_vvVV_t2(mycc, r2bab, mycc.direct, None)
        Hr2aaa = eris._contract_vvvv_t2(mycc, r2aaa, mycc.direct, None)
        Hr2bbb = eris._contract_VVVV_t2(mycc, r2bbb, mycc.direct, None)

    else:
        noccb, nvira, nvirb = r2bab.shape
        eris_vvvv = ao2mo.restore(1, np.asarray(eris.vvvv), nvira)
        Hr2aaa = lib.einsum('acbd,jcd->jab', eris_vvvv, r2aaa)
        eris_vvvv = None

        eris_VVVV = ao2mo.restore(1, np.asarray(eris.VVVV), nvirb)
        Hr2bbb = lib.einsum('acbd,jcd->jab', eris_VVVV, r2bbb)
        eris_VVVV = None

        sqa = lib.square_mat_in_trilu_indices(nvira)
        sqb = lib.square_mat_in_trilu_indices(nvirb)
        eris_vvVV = np.asarray(eris.vvVV)[:,sqb][sqa]
        Hr2aba = lib.einsum('bdac,jcd->jab', eris_vvVV, r2aba)
        Hr2bab = lib.einsum('acbd,jcd->jab', eris_vvVV, r2bab)
        eris_vvVV = None

    return Hr2aaa, Hr2aba, Hr2bab, Hr2bbb

def eaccsd_diag(eom, imds=None):
    if imds is None: imds = eom.make_imds()
    eris = imds.eris
    t1, t2 = imds.t1, imds.t2
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    t2ba = t2ab.transpose(1,0,3,2)

    nocca, nvira = t1a.shape
    noccb, nvirb = t1b.shape

    Hr1a = np.diag(imds.Fvv)
    Hr1b = np.diag(imds.FVV)

    #-------------- intermediates

    Fvv_diag = np.diag(imds.Fvv)
    Foo_diag = np.diag(imds.Foo)
    FOO_diag = np.diag(imds.FOO)
    FVV_diag = np.diag(imds.FVV)

    Wovvo_slice = np.einsum('jbbj->jb',imds.Wovvo)
    Wovov_t2_dot = np.einsum('iajb,ijab->jab',imds.Wovov,t2aa)
    WoVVo_slice  = np.einsum('jaaj->ja',imds.WoVVo)
    WovOV_t2_dot = np.einsum('jbia,ijab->jab',imds.WovOV,t2ba)
    WOVVO_slice = np.einsum('jaaj->ja',imds.WOVVO)
    WOvvO_slice = np.einsum('jbbj->jb',imds.WOvvO)
    WovOV_t2_dot_T = np.einsum('ibja,ijba->jab',imds.WovOV,t2ab)
    WOVOV_t2_dot = np.einsum('iajb,ijab->jab',imds.WOVOV,t2bb)

    #-------------- contraction

    Hr2aaa = Fvv_diag[None,:,None]+Fvv_diag[None,None,:]-Foo_diag[:,None,None]+ \
             Wovvo_slice[:,None,:]+Wovvo_slice[:,:,None]-Wovov_t2_dot

    Hr2aba = FVV_diag[None,:,None]+Fvv_diag[None,None,:]-Foo_diag[:,None,None]+ \
             Wovvo_slice[:,None,:]+WoVVo_slice[:,:,None]-WovOV_t2_dot

    Hr2bab = -FOO_diag[:,None,None]+FVV_diag[None,:,None]+Fvv_diag[None,None,:]+ \
             WOVVO_slice[:,:,None]+WOvvO_slice[:,None,:]-WovOV_t2_dot_T
    Hr2bab = Hr2bab.transpose(0,2,1)

    Hr2bbb = -FOO_diag[:,None,None]+FVV_diag[None,:,None]+FVV_diag[None,None,:]+ \
             WOVVO_slice[:,:,None]+WOVVO_slice[:,None,:]-WOVOV_t2_dot

#    if imds.Wvvvv is not None:
#        Wvvvv_slice = np.einsum('aabb->ab',imds.Wvvvv)
#        Hr2aaa += 0.5 * Wvvvv_slice[None,:,:]
#        WVVvv_slice = np.einsum('aabb->ba',imds.WvvVV)
#        Hr2aba += WVVvv_slice[None,:,:]
#        WvvVV_slice = np.einsum('aabb->ab',imds.WvvVV)
#        Hr2bab += WvvVV_slice[None,:,:]
#        WVVVV_slice = np.einsum('aabb->ab',imds.WVVVV)
#        Hr2bbb += 0.5 * WVVVV_slice[None,:,:]

# TODO: test Wvvvv contribution
    # See also the code for Wvvvv contribution in function eeccsd_diag
    tauaa, tauab, taubb = uccsd.make_tau(t2, t1, t1)
    eris_ovov = np.asarray(eris.ovov)
    eris_OVOV = np.asarray(eris.OVOV)
    eris_ovOV = np.asarray(eris.ovOV)
    Wvvaa = .5*np.einsum('mnab,manb->ab', tauaa, eris_ovov)
    Wvvbb = .5*np.einsum('mnab,manb->ab', taubb, eris_OVOV)
    Wvvab =    np.einsum('mNaB,maNB->aB', tauab, eris_ovOV)
    eris_ovov = eris_OVOV = eris_ovOV = None

    mem_now = lib.current_memory()[0]
    max_memory = max(0, eom.max_memory - mem_now)
    blksize = min(nocca, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvira**3*3))))
    for p0,p1 in lib.prange(0, nocca, blksize):
        ovvv = eris.get_ovvv(slice(p0,p1))  # ovvv = eris.ovvv[p0:p1]
        Wvvaa += np.einsum('mb,maab->ab', t1a[p0:p1], ovvv)
        Wvvaa -= np.einsum('mb,mbaa->ab', t1a[p0:p1], ovvv)
        ovvv = None
    blksize = min(noccb, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvirb**3*3))))
    for p0, p1 in lib.prange(0, noccb, blksize):
        OVVV = eris.get_OVVV(slice(p0,p1))  # OVVV = eris.OVVV[p0:p1]
        Wvvbb += np.einsum('mb,maab->ab', t1b[p0:p1], OVVV)
        Wvvbb -= np.einsum('mb,mbaa->ab', t1b[p0:p1], OVVV)
        OVVV = None
    blksize = min(nocca, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvira*nvirb**2*3))))
    for p0,p1 in lib.prange(0, nocca, blksize):
        ovVV = eris.get_ovVV(slice(p0,p1))  # ovVV = eris.ovVV[p0:p1]
        Wvvab -= np.einsum('mb,mbaa->ba', t1a[p0:p1], ovVV)
        ovVV = None
    blksize = min(noccb, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvirb*nvira**2*3))))
    for p0, p1 in lib.prange(0, noccb, blksize):
        OVvv = eris.get_OVvv(slice(p0,p1))  # OVvv = eris.OVvv[p0:p1]
        Wvvab -= np.einsum('mb,mbaa->ab', t1b[p0:p1], OVvv)
        OVvv = None
    Wvvaa = Wvvaa + Wvvaa.T
    Wvvbb = Wvvbb + Wvvbb.T
    if eris.vvvv is not None:
        for i in range(nvira):
            i0 = i*(i+1)//2
            vvv = lib.unpack_tril(np.asarray(eris.vvvv[i0:i0+i+1]))
            tmp = np.einsum('bb->b', vvv[i])
            Wvvaa[i] += tmp
            tmp = np.einsum('bb->b', vvv[:,:i+1,i])
            Wvvaa[i,:i+1] -= tmp
            Wvvaa[:i  ,i] -= tmp[:i]
            vvv = lib.unpack_tril(np.asarray(eris.vvVV[i0:i0+i+1]))
            Wvvab[i] += np.einsum('bb->b', vvv[i])
            vvv = None
        for i in range(nvirb):
            i0 = i*(i+1)//2
            vvv = lib.unpack_tril(np.asarray(eris.VVVV[i0:i0+i+1]))
            tmp = np.einsum('bb->b', vvv[i])
            Wvvbb[i] += tmp
            tmp = np.einsum('bb->b', vvv[:,:i+1,i])
            Wvvbb[i,:i+1] -= tmp
            Wvvbb[:i  ,i] -= tmp[:i]
            vvv = None
    Wvvba = Wvvab.T

    Hr2aaa += Wvvaa[None,:,:]
    Hr2aba += Wvvba[None,:,:]
    Hr2bab += Wvvab[None,:,:]
    Hr2bbb += Wvvbb[None,:,:]
    # Wvvvv contribution end

    vector = amplitudes_to_vector_ea((Hr1a,Hr1b), (Hr2aaa,Hr2aba,Hr2bab,Hr2bbb))
    return vector


class EOMEA(eom_rccsd.EOMEA):
    matvec = eaccsd_matvec
    l_matvec = None
    get_diag = eaccsd_diag
    eaccsd_star = None
    ccsd_star_contract = None

    def __init__(self, cc):
        eom_rccsd.EOMEA.__init__(self, cc)
        self.nocc = cc.get_nocc()
        self.nmo = cc.get_nmo()

    def get_init_guess(self, nroots=1, koopmans=True, diag=None):
        if koopmans:
            nocca, noccb = self.nocc
            nmoa, nmob = self.nmo
            nvira, nvirb = nmoa-nocca, nmob-noccb
            idx = diag[:nvira+nvirb].argsort()
        else:
            idx = diag.argsort()

        size = self.vector_size()
        dtype = getattr(diag, 'dtype', np.double)
        nroots = min(nroots, size)
        guess = []
        for i in idx[:nroots]:
            g = np.zeros(size, dtype)
            g[i] = 1.0
            guess.append(g)
        return guess

    def vector_to_amplitudes(self, vector, nmo=None, nocc=None):
        if nmo is None: nmo = self.nmo
        if nocc is None: nocc = self.nocc
        return vector_to_amplitudes_ea(vector, nmo, nocc)

    def amplitudes_to_vector(self, r1, r2):
        return amplitudes_to_vector_ea(r1, r2)

    def vector_size(self):
        '''size of the vector based on spin-orbital basis'''
        nocca, noccb = self.nocc
        nmoa, nmob = self.nmo
        nvira, nvirb = nmoa-nocca, nmob-noccb
        return (nvira + nvirb
                + nocca*nvira*(nvira-1)//2 + nocca*nvirb*nvira
                + noccb*nvira*nvirb + noccb*nvirb*(nvirb-1)//2)

    def make_imds(self, eris=None):
        imds = _IMDS(self._cc, eris=eris)
        imds.make_ea()
        return imds

########################################
# EOM-EE-CCSD
########################################

def eeccsd(eom, nroots=1, koopmans=False, guess=None, eris=None, imds=None):
    '''Calculate N-electron neutral excitations via EOM-EE-CCSD.

    Kwargs:
        nroots : int
            Number of roots (eigenvalues) requested
        koopmans : bool
            Calculate Koopmans'-like (1p1h) excitations only, targeting via
            overlap.
        guess : list of ndarray
            List of guess vectors to use for targeting via overlap.
    '''
    if eris is None: eris = eom._cc.ao2mo()
    if imds is None: imds = eom.make_imds(eris)

    spinvec_size = eom.vector_size()
    nroots = min(nroots, spinvec_size)

    diag_ee, diag_sf = eom.get_diag(imds)
    guess_ee = []
    guess_sf = []
    if guess and guess[0].size == spinvec_size:
        raise NotImplementedError
        #TODO: initial guess from GCCSD EOM amplitudes
        #from pyscf.cc import addons
        #from pyscf.cc import eom_gccsd
        #orbspin = scf.addons.get_ghf_orbspin(eris.mo_coeff)
        #nmo = np.sum(eom.nmo)
        #nocc = np.sum(eom.nocc)
        #for g in guess:
        #    r1, r2 = eom_gccsd.vector_to_amplitudes_ee(g, nmo, nocc)
        #    r1aa = r1[orbspin==0][:,orbspin==0]
        #    r1ab = r1[orbspin==0][:,orbspin==1]
        #    if abs(r1aa).max() > 1e-7:
        #        r1 = addons.spin2spatial(r1, orbspin)
        #        r2 = addons.spin2spatial(r2, orbspin)
        #        guess_ee.append(eom.amplitudes_to_vector(r1, r2))
        #    else:
        #        r1 = spin2spatial_eomsf(r1, orbspin)
        #        r2 = spin2spatial_eomsf(r2, orbspin)
        #        guess_sf.append(amplitudes_to_vector_eomsf(r1, r2))
        #    r1 = r2 = r1aa = r1ab = g = None
        #nroots_ee = len(guess_ee)
        #nroots_sf = len(guess_sf)
    elif guess:
        for g in guess:
            if g.size == diag_ee.size:
                guess_ee.append(g)
            else:
                guess_sf.append(g)
        nroots_ee = len(guess_ee)
        nroots_sf = len(guess_sf)
    else:
        dee = np.sort(diag_ee)[:nroots]
        dsf = np.sort(diag_sf)[:nroots]
        dmax = np.sort(np.hstack([dee,dsf]))[nroots-1]
        nroots_ee = np.count_nonzero(dee <= dmax)
        nroots_sf = np.count_nonzero(dsf <= dmax)
        guess_ee = guess_sf = None

    def eomee_sub(cls, nroots, guess, diag):
        ee_sub = cls(eom._cc)
        ee_sub.__dict__.update(eom.__dict__)
        e, v = ee_sub.kernel(nroots, koopmans, guess, eris, imds, diag=diag)
        if nroots == 1:
            e, v = [e], [v]
            ee_sub.converged = [ee_sub.converged]
        return list(ee_sub.converged), list(e), list(v)

    e0 = e1 = []
    v0 = v1 = []
    conv0 = conv1 = []
    if nroots_ee > 0:
        conv0, e0, v0 = eomee_sub(EOMEESpinKeep, nroots_ee, guess_ee, diag_ee)
    if nroots_sf > 0:
        conv1, e1, v1 = eomee_sub(EOMEESpinFlip, nroots_sf, guess_sf, diag_sf)

    e = np.hstack([e0,e1])
    idx = e.argsort()
    e = e[idx]
    conv = conv0 + conv1
    conv = [conv[x] for x in idx]
    v = v0 + v1
    v = [v[x] for x in idx]

    if nroots == 1:
        conv = conv[0]
        e = e[0]
        v = v[0]
    eom.converged = conv
    eom.e = e
    eom.v = v
    return eom.e, eom.v

def eomee_ccsd(eom, nroots=1, koopmans=False, guess=None,
               eris=None, imds=None, diag=None):
    if eris is None: eris = eom._cc.ao2mo()
    if imds is None: imds = eom.make_imds(eris)
    eom.converged, eom.e, eom.v \
            = eom_rccsd.kernel(eom, nroots, koopmans, guess, imds=imds, diag=diag)
    return eom.e, eom.v

def eomsf_ccsd(eom, nroots=1, koopmans=False, guess=None,
               eris=None, imds=None, diag=None):
    '''Spin flip EOM-EE-CCSD
    '''
    return eomee_ccsd(eom, nroots, koopmans, guess, eris, imds, diag)

amplitudes_to_vector_ee = uccsd.amplitudes_to_vector
vector_to_amplitudes_ee = uccsd.vector_to_amplitudes

def amplitudes_to_vector_eomsf(t1, t2, out=None):
    t1ab, t1ba = t1
    t2baaa, t2aaba, t2abbb, t2bbab = t2
    nocca, nvirb = t1ab.shape
    noccb, nvira = t1ba.shape

    otrila = np.tril_indices(nocca, k=-1)
    otrilb = np.tril_indices(noccb, k=-1)
    vtrila = np.tril_indices(nvira, k=-1)
    vtrilb = np.tril_indices(nvirb, k=-1)
    baaa = np.take(t2baaa.reshape(noccb*nocca,nvira*nvira),
                   vtrila[0]*nvira+vtrila[1], axis=1)
    abbb = np.take(t2abbb.reshape(nocca*noccb,nvirb*nvirb),
                   vtrilb[0]*nvirb+vtrilb[1], axis=1)
    vector = np.hstack((t1ab.ravel(), t1ba.ravel(),
                        baaa.ravel(), t2aaba[otrila].ravel(),
                        abbb.ravel(), t2bbab[otrilb].ravel()))
    return vector

def vector_to_amplitudes_eomsf(vector, nmo, nocc):
    nocca, noccb = nocc
    nmoa, nmob = nmo
    nvira, nvirb = nmoa-nocca, nmob-noccb

    nbaaa = noccb*nocca*nvira*(nvira-1)//2
    naaba = nocca*(nocca-1)//2*nvirb*nvira
    nabbb = nocca*noccb*nvirb*(nvirb-1)//2
    nbbab = noccb*(noccb-1)//2*nvira*nvirb
    sizes = (nocca*nvirb, noccb*nvira, nbaaa, naaba, nabbb, nbbab)
    sections = np.cumsum(sizes[:-1])
    t1ab, t1ba, vbaaa, vaaba, vabbb, vbbab = np.split(vector, sections)

    t1ab = t1ab.reshape(nocca,nvirb).copy()
    t1ba = t1ba.reshape(noccb,nvira).copy()

    t2baaa = np.zeros((noccb*nocca,nvira*nvira), dtype=vector.dtype)
    t2aaba = np.zeros((nocca*nocca,nvirb*nvira), dtype=vector.dtype)
    t2abbb = np.zeros((nocca*noccb,nvirb*nvirb), dtype=vector.dtype)
    t2bbab = np.zeros((noccb*noccb,nvira*nvirb), dtype=vector.dtype)
    otrila = np.tril_indices(nocca, k=-1)
    otrilb = np.tril_indices(noccb, k=-1)
    vtrila = np.tril_indices(nvira, k=-1)
    vtrilb = np.tril_indices(nvirb, k=-1)
    oidxab = np.arange(nocca*noccb, dtype=np.int32)
    vidxab = np.arange(nvira*nvirb, dtype=np.int32)

    vbaaa = vbaaa.reshape(noccb*nocca,-1)
    lib.takebak_2d(t2baaa, vbaaa, oidxab, vtrila[0]*nvira+vtrila[1])
    lib.takebak_2d(t2baaa,-vbaaa, oidxab, vtrila[1]*nvira+vtrila[0])
    vaaba = vaaba.reshape(-1,nvirb*nvira)
    lib.takebak_2d(t2aaba, vaaba, otrila[0]*nocca+otrila[1], vidxab)
    lib.takebak_2d(t2aaba,-vaaba, otrila[1]*nocca+otrila[0], vidxab)
    vabbb = vabbb.reshape(nocca*noccb,-1)
    lib.takebak_2d(t2abbb, vabbb, oidxab, vtrilb[0]*nvirb+vtrilb[1])
    lib.takebak_2d(t2abbb,-vabbb, oidxab, vtrilb[1]*nvirb+vtrilb[0])
    vbbab = vbbab.reshape(-1,nvira*nvirb)
    lib.takebak_2d(t2bbab, vbbab, otrilb[0]*noccb+otrilb[1], vidxab)
    lib.takebak_2d(t2bbab,-vbbab, otrilb[1]*noccb+otrilb[0], vidxab)
    t2baaa = t2baaa.reshape(noccb,nocca,nvira,nvira)
    t2aaba = t2aaba.reshape(nocca,nocca,nvirb,nvira)
    t2abbb = t2abbb.reshape(nocca,noccb,nvirb,nvirb)
    t2bbab = t2bbab.reshape(noccb,noccb,nvira,nvirb)
    return (t1ab,t1ba), (t2baaa, t2aaba, t2abbb, t2bbab)

def spatial2spin_eomsf(rx, orbspin):
    '''Convert EOM spatial R1,R2 to spin-orbital R1,R2'''
    if len(rx) == 2:  # r1
        r1ab, r1ba = rx
        nocca, nvirb = r1ab.shape
        noccb, nvira = r1ba.shape
    else:
        r2baaa,r2aaba,r2abbb,r2bbab = rx
        noccb, nocca, nvira = r2baaa.shape[:3]
        nvirb = r2aaba.shape[2]

    nocc = nocca + noccb
    nvir = nvira + nvirb
    idxoa = np.where(orbspin[:nocc] == 0)[0]
    idxob = np.where(orbspin[:nocc] == 1)[0]
    idxva = np.where(orbspin[nocc:] == 0)[0]
    idxvb = np.where(orbspin[nocc:] == 1)[0]

    if len(rx) == 2:  # r1
        r1 = np.zeros((nocc,nvir), dtype=r1ab.dtype)
        lib.takebak_2d(r1, r1ab, idxoa, idxvb)
        lib.takebak_2d(r1, r1ba, idxob, idxva)
        return r1

    else:
        r2 = np.zeros((nocc**2,nvir**2), dtype=r2aaba.dtype)
        idxoaa = idxoa[:,None] * nocc + idxoa
        idxoab = idxoa[:,None] * nocc + idxob
        idxoba = idxob[:,None] * nocc + idxoa
        idxobb = idxob[:,None] * nocc + idxob
        idxvaa = idxva[:,None] * nvir + idxva
        idxvab = idxva[:,None] * nvir + idxvb
        idxvba = idxvb[:,None] * nvir + idxva
        idxvbb = idxvb[:,None] * nvir + idxvb
        r2baaa = r2baaa.reshape(noccb*nocca,nvira*nvira)
        r2aaba = r2aaba.reshape(nocca*nocca,nvirb*nvira)
        r2abbb = r2abbb.reshape(nocca*noccb,nvirb*nvirb)
        r2bbab = r2bbab.reshape(noccb*noccb,nvira*nvirb)
        lib.takebak_2d(r2, r2baaa, idxoba.ravel(), idxvaa.ravel())
        lib.takebak_2d(r2, r2aaba, idxoaa.ravel(), idxvba.ravel())
        lib.takebak_2d(r2, r2abbb, idxoab.ravel(), idxvbb.ravel())
        lib.takebak_2d(r2, r2bbab, idxobb.ravel(), idxvab.ravel())
        lib.takebak_2d(r2, r2baaa, idxoab.T.ravel(), idxvaa.T.ravel())
        lib.takebak_2d(r2, r2aaba, idxoaa.T.ravel(), idxvab.T.ravel())
        lib.takebak_2d(r2, r2abbb, idxoba.T.ravel(), idxvbb.T.ravel())
        lib.takebak_2d(r2, r2bbab, idxobb.T.ravel(), idxvba.T.ravel())
        return r2.reshape(nocc,nocc,nvir,nvir)

def spin2spatial_eomsf(rx, orbspin):
    '''Convert EOM spin-orbital R1,R2 to spatial R1,R2'''
    if rx.ndim == 2:  # r1
        nocc, nvir = rx.shape
    else:
        nocc, nvir = rx.shape[1:3]

    idxoa = np.where(orbspin[:nocc] == 0)[0]
    idxob = np.where(orbspin[:nocc] == 1)[0]
    idxva = np.where(orbspin[nocc:] == 0)[0]
    idxvb = np.where(orbspin[nocc:] == 1)[0]
    nocca = len(idxoa)
    noccb = len(idxob)
    nvira = len(idxva)
    nvirb = len(idxvb)

    if rx.ndim == 2:
        r1ab = lib.take_2d(rx, idxoa, idxvb)
        r1ba = lib.take_2d(rx, idxob, idxva)
        return r1ab, r1ba
    else:
        idxoaa = idxoa[:,None] * nocc + idxoa
        idxoab = idxoa[:,None] * nocc + idxob
        idxoba = idxob[:,None] * nocc + idxoa
        idxobb = idxob[:,None] * nocc + idxob
        idxvaa = idxva[:,None] * nvir + idxva
        idxvab = idxva[:,None] * nvir + idxvb
        idxvba = idxvb[:,None] * nvir + idxva
        idxvbb = idxvb[:,None] * nvir + idxvb
        r2 = rx.reshape(nocc**2,nvir**2)
        r2baaa = lib.take_2d(r2, idxoba.ravel(), idxvaa.ravel())
        r2aaba = lib.take_2d(r2, idxoaa.ravel(), idxvba.ravel())
        r2abbb = lib.take_2d(r2, idxoab.ravel(), idxvbb.ravel())
        r2bbab = lib.take_2d(r2, idxobb.ravel(), idxvab.ravel())
        r2baaa = r2baaa.reshape(noccb,nocca,nvira,nvira)
        r2aaba = r2aaba.reshape(nocca,nocca,nvirb,nvira)
        r2abbb = r2abbb.reshape(nocca,noccb,nvirb,nvirb)
        r2bbab = r2bbab.reshape(noccb,noccb,nvira,nvirb)
        return r2baaa,r2aaba,r2abbb,r2bbab

# Ref: Wang, Tu, and Wang, J. Chem. Theory Comput. 10, 5567 (2014) Eqs.(9)-(10)
# Note: Last line in Eq. (10) is superfluous.
# See, e.g. Gwaltney, Nooijen, and Barlett, Chem. Phys. Lett. 248, 189 (1996)
def eomee_ccsd_matvec(eom, vector, imds=None):
    if imds is None: imds = eom.make_imds()

    t1, t2, eris = imds.t1, imds.t2, imds.eris
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nocca, noccb, nvira, nvirb = t2ab.shape
    nmoa, nmob = nocca+nvira, noccb+nvirb
    r1, r2 = vector_to_amplitudes_ee(vector, (nmoa,nmob), (nocca,noccb))
    r1a, r1b = r1
    r2aa, r2ab, r2bb = r2

    #:Hr2aa += lib.einsum('ijef,aebf->ijab', tau2aa, eris.vvvv) * .5
    #:Hr2bb += lib.einsum('ijef,aebf->ijab', tau2bb, eris.VVVV) * .5
    #:Hr2ab += lib.einsum('iJeF,aeBF->iJaB', tau2ab, eris.vvVV)
    tau2aa, tau2ab, tau2bb = uccsd.make_tau(r2, r1, t1, 2)
    Hr2aa, Hr2ab, Hr2bb = eom._cc._add_vvvv(None, (tau2aa,tau2ab,tau2bb), eris)
    Hr2aa *= .5
    Hr2bb *= .5
    tau2aa = tau2ab = tau2bb = None

    Hr1a  = lib.einsum('ae,ie->ia', imds.Fvva, r1a)
    Hr1a -= lib.einsum('mi,ma->ia', imds.Fooa, r1a)
    Hr1a += np.einsum('me,imae->ia',imds.Fova, r2aa)
    Hr1a += np.einsum('ME,iMaE->ia',imds.Fovb, r2ab)
    Hr1b  = lib.einsum('ae,ie->ia', imds.Fvvb, r1b)
    Hr1b -= lib.einsum('mi,ma->ia', imds.Foob, r1b)
    Hr1b += np.einsum('me,imae->ia',imds.Fovb, r2bb)
    Hr1b += np.einsum('me,mIeA->IA',imds.Fova, r2ab)

    Hr2aa += lib.einsum('minj,mnab->ijab', imds.woooo, r2aa) * .25
    Hr2bb += lib.einsum('minj,mnab->ijab', imds.wOOOO, r2bb) * .25
    Hr2ab += lib.einsum('miNJ,mNaB->iJaB', imds.wooOO, r2ab)
    Hr2aa += lib.einsum('be,ijae->ijab', imds.Fvva, r2aa)
    Hr2bb += lib.einsum('be,ijae->ijab', imds.Fvvb, r2bb)
    Hr2ab += lib.einsum('BE,iJaE->iJaB', imds.Fvvb, r2ab)
    Hr2ab += lib.einsum('be,iJeA->iJbA', imds.Fvva, r2ab)
    Hr2aa -= lib.einsum('mj,imab->ijab', imds.Fooa, r2aa)
    Hr2bb -= lib.einsum('mj,imab->ijab', imds.Foob, r2bb)
    Hr2ab -= lib.einsum('MJ,iMaB->iJaB', imds.Foob, r2ab)
    Hr2ab -= lib.einsum('mj,mIaB->jIaB', imds.Fooa, r2ab)

    #:tau2aa, tau2ab, tau2bb = uccsd.make_tau(r2, r1, t1, 2)
    #:eris_ovvv = lib.unpack_tril(np.asarray(eris.ovvv).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvira,nvira)
    #:eris_ovVV = lib.unpack_tril(np.asarray(eris.ovVV).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvirb,nvirb)
    #:eris_OVvv = lib.unpack_tril(np.asarray(eris.OVvv).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvira,nvira)
    #:eris_OVVV = lib.unpack_tril(np.asarray(eris.OVVV).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvirb,nvirb)
    #:Hr1a += lib.einsum('mfae,imef->ia', eris_ovvv, r2aa)
    #:tmpaa = lib.einsum('meaf,ijef->maij', eris_ovvv, tau2aa)
    #:Hr2aa+= lib.einsum('mb,maij->ijab', t1a, tmpaa)
    #:tmpa = lib.einsum('mfae,me->af', eris_ovvv, r1a)
    #:tmpa-= lib.einsum('meaf,me->af', eris_ovvv, r1a)

    #:Hr1b += lib.einsum('mfae,imef->ia', eris_OVVV, r2bb)
    #:tmpbb = lib.einsum('meaf,ijef->maij', eris_OVVV, tau2bb)
    #:Hr2bb+= lib.einsum('mb,maij->ijab', t1b, tmpbb)
    #:tmpb = lib.einsum('mfae,me->af', eris_OVVV, r1b)
    #:tmpb-= lib.einsum('meaf,me->af', eris_OVVV, r1b)

    #:Hr1b += lib.einsum('mfAE,mIfE->IA', eris_ovVV, r2ab)
    #:tmpab = lib.einsum('meAF,iJeF->mAiJ', eris_ovVV, tau2ab)
    #:Hr2ab-= lib.einsum('mb,mAiJ->iJbA', t1a, tmpab)
    #:tmpb-= lib.einsum('meAF,me->AF', eris_ovVV, r1a)

    #:Hr1a += lib.einsum('MFae,iMeF->ia', eris_OVvv, r2ab)
    #:tmpba =-lib.einsum('MEaf,iJfE->MaiJ', eris_OVvv, tau2ab)
    #:Hr2ab+= lib.einsum('MB,MaiJ->iJaB', t1b, tmpba)
    #:tmpa-= lib.einsum('MEaf,ME->af', eris_OVvv, r1b)
    tau2aa = uccsd.make_tau_aa(r2aa, r1a, t1a, 2)
    mem_now = lib.current_memory()[0]
    max_memory = max(0, eom.max_memory - mem_now)
    tmpa = np.zeros((nvira,nvira))
    tmpb = np.zeros((nvirb,nvirb))
    blksize = min(nocca, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvira**3*3))))
    for p0, p1 in lib.prange(0, nocca, blksize):
        ovvv = eris.get_ovvv(slice(p0,p1))  # ovvv = eris.ovvv[p0:p1]
        Hr1a += lib.einsum('mfae,imef->ia', ovvv, r2aa[:,p0:p1])
        tmpaa = lib.einsum('meaf,ijef->maij', ovvv, tau2aa)
        Hr2aa+= lib.einsum('mb,maij->ijab', t1a[p0:p1], tmpaa)
        tmpa+= lib.einsum('mfae,me->af', ovvv, r1a[p0:p1])
        tmpa-= lib.einsum('meaf,me->af', ovvv, r1a[p0:p1])
        ovvv = tmpaa = None
    tau2aa = None

    tau2bb = uccsd.make_tau_aa(r2bb, r1b, t1b, 2)
    blksize = min(noccb, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvirb**3*3))))
    for p0, p1 in lib.prange(0, noccb, blksize):
        OVVV = eris.get_OVVV(slice(p0,p1))  # OVVV = eris.OVVV[p0:p1]
        Hr1b += lib.einsum('mfae,imef->ia', OVVV, r2bb[:,p0:p1])
        tmpbb = lib.einsum('meaf,ijef->maij', OVVV, tau2bb)
        Hr2bb+= lib.einsum('mb,maij->ijab', t1b[p0:p1], tmpbb)
        tmpb+= lib.einsum('mfae,me->af', OVVV, r1b[p0:p1])
        tmpb-= lib.einsum('meaf,me->af', OVVV, r1b[p0:p1])
        OVVV = tmpbb = None
    tau2bb = None

    tau2ab = uccsd.make_tau_ab(r2ab, r1 , t1 , 2)
    blksize = min(nocca, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvira*nvirb**2*3))))
    for p0, p1 in lib.prange(0, nocca, blksize):
        ovVV = eris.get_ovVV(slice(p0,p1))  # ovVV = eris.ovVV[p0:p1]
        Hr1b += lib.einsum('mfAE,mIfE->IA', ovVV, r2ab[p0:p1])
        tmpab = lib.einsum('meAF,iJeF->mAiJ', ovVV, tau2ab)
        Hr2ab-= lib.einsum('mb,mAiJ->iJbA', t1a[p0:p1], tmpab)
        tmpb-= lib.einsum('meAF,me->AF', ovVV, r1a[p0:p1])
        ovVV = tmpab = None

    blksize = min(noccb, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvirb*nvira**2*3))))
    for p0, p1 in lib.prange(0, noccb, blksize):
        OVvv = eris.get_OVvv(slice(p0,p1))  # OVvv = eris.OVvv[p0:p1]
        Hr1a += lib.einsum('MFae,iMeF->ia', OVvv, r2ab[:,p0:p1])
        tmpba = lib.einsum('MEaf,iJfE->MaiJ', OVvv, tau2ab)
        Hr2ab-= lib.einsum('MB,MaiJ->iJaB', t1b[p0:p1], tmpba)
        tmpa-= lib.einsum('MEaf,ME->af', OVvv, r1b[p0:p1])
        OVvv = tmpba = None
    tau2ab = None

    Hr2aa-= lib.einsum('af,ijfb->ijab', tmpa, t2aa)
    Hr2bb-= lib.einsum('af,ijfb->ijab', tmpb, t2bb)
    Hr2ab-= lib.einsum('af,iJfB->iJaB', tmpa, t2ab)
    Hr2ab-= lib.einsum('AF,iJbF->iJbA', tmpb, t2ab)

    eris_ovov = np.asarray(eris.ovov)
    eris_OVOV = np.asarray(eris.OVOV)
    eris_ovOV = np.asarray(eris.ovOV)
    tau2aa = uccsd.make_tau_aa(r2aa, r1a, t1a, 2)
    tauaa = uccsd.make_tau_aa(t2aa, t1a, t1a)
    tmpaa = lib.einsum('menf,ijef->mnij', eris_ovov, tau2aa)
    Hr2aa += lib.einsum('mnij,mnab->ijab', tmpaa, tauaa) * 0.25
    tmpaa = tau2aa = tauaa = None

    tau2bb = uccsd.make_tau_aa(r2bb, r1b, t1b, 2)
    taubb = uccsd.make_tau_aa(t2bb, t1b, t1b)
    tmpbb = lib.einsum('menf,ijef->mnij', eris_OVOV, tau2bb)
    Hr2bb += lib.einsum('mnij,mnab->ijab', tmpbb, taubb) * 0.25
    tmpbb = tau2bb = taubb = None

    tau2ab = uccsd.make_tau_ab(r2ab, r1 , t1 , 2)
    tauab = uccsd.make_tau_ab(t2ab, t1 , t1)
    tmpab = lib.einsum('meNF,iJeF->mNiJ', eris_ovOV, tau2ab)
    Hr2ab += lib.einsum('mNiJ,mNaB->iJaB', tmpab, tauab)
    tmpab = tau2ab = tauab = None

    tmpa = lib.einsum('menf,imef->ni', eris_ovov, r2aa)
    tmpa-= lib.einsum('neMF,iMeF->ni', eris_ovOV, r2ab)
    tmpb = lib.einsum('menf,imef->ni', eris_OVOV, r2bb)
    tmpb-= lib.einsum('mfNE,mIfE->NI', eris_ovOV, r2ab)
    Hr1a += lib.einsum('na,ni->ia', t1a, tmpa)
    Hr1b += lib.einsum('na,ni->ia', t1b, tmpb)
    Hr2aa+= lib.einsum('mj,imab->ijab', tmpa, t2aa)
    Hr2bb+= lib.einsum('mj,imab->ijab', tmpb, t2bb)
    Hr2ab+= lib.einsum('MJ,iMaB->iJaB', tmpb, t2ab)
    Hr2ab+= lib.einsum('mj,mIaB->jIaB', tmpa, t2ab)

    tmp1a = np.einsum('menf,mf->en', eris_ovov, r1a)
    tmp1a-= np.einsum('mfne,mf->en', eris_ovov, r1a)
    tmp1a-= np.einsum('neMF,MF->en', eris_ovOV, r1b)
    tmp1b = np.einsum('menf,mf->en', eris_OVOV, r1b)
    tmp1b-= np.einsum('mfne,mf->en', eris_OVOV, r1b)
    tmp1b-= np.einsum('mfNE,mf->EN', eris_ovOV, r1a)
    tmpa = np.einsum('en,nb->eb', tmp1a, t1a)
    tmpa+= lib.einsum('menf,mnfb->eb', eris_ovov, r2aa)
    tmpa-= lib.einsum('meNF,mNbF->eb', eris_ovOV, r2ab)
    tmpb = np.einsum('en,nb->eb', tmp1b, t1b)
    tmpb+= lib.einsum('menf,mnfb->eb', eris_OVOV, r2bb)
    tmpb-= lib.einsum('nfME,nMfB->EB', eris_ovOV, r2ab)
    Hr2aa+= lib.einsum('eb,ijae->ijab', tmpa, t2aa)
    Hr2bb+= lib.einsum('eb,ijae->ijab', tmpb, t2bb)
    Hr2ab+= lib.einsum('EB,iJaE->iJaB', tmpb, t2ab)
    Hr2ab+= lib.einsum('eb,iJeA->iJbA', tmpa, t2ab)
    eris_ovOV = eris_OVOV = None

    Hr2aa-= lib.einsum('mbij,ma->ijab', imds.wovoo, r1a)
    Hr2bb-= lib.einsum('mbij,ma->ijab', imds.wOVOO, r1b)
    Hr2ab-= lib.einsum('mBiJ,ma->iJaB', imds.woVoO, r1a)
    Hr2ab-= lib.einsum('MbJi,MA->iJbA', imds.wOvOo, r1b)

    Hr1a-= 0.5*lib.einsum('mine,mnae->ia', imds.wooov, r2aa)
    Hr1a-=     lib.einsum('miNE,mNaE->ia', imds.wooOV, r2ab)
    Hr1b-= 0.5*lib.einsum('mine,mnae->ia', imds.wOOOV, r2bb)
    Hr1b-=     lib.einsum('MIne,nMeA->IA', imds.wOOov, r2ab)
    tmpa = lib.einsum('mine,me->ni', imds.wooov, r1a)
    tmpa-= lib.einsum('niME,ME->ni', imds.wooOV, r1b)
    tmpb = lib.einsum('mine,me->ni', imds.wOOOV, r1b)
    tmpb-= lib.einsum('NIme,me->NI', imds.wOOov, r1a)
    Hr2aa+= lib.einsum('ni,njab->ijab', tmpa, t2aa)
    Hr2bb+= lib.einsum('ni,njab->ijab', tmpb, t2bb)
    Hr2ab+= lib.einsum('ni,nJaB->iJaB', tmpa, t2ab)
    Hr2ab+= lib.einsum('NI,jNaB->jIaB', tmpb, t2ab)
    for p0, p1 in lib.prange(0, nvira, nocca):
        Hr2aa+= lib.einsum('ejab,ie->ijab', imds.wvovv[p0:p1], r1a[:,p0:p1])
        Hr2ab+= lib.einsum('eJaB,ie->iJaB', imds.wvOvV[p0:p1], r1a[:,p0:p1])
    for p0, p1 in lib.prange(0, nvirb, noccb):
        Hr2bb+= lib.einsum('ejab,ie->ijab', imds.wVOVV[p0:p1], r1b[:,p0:p1])
        Hr2ab+= lib.einsum('EjBa,IE->jIaB', imds.wVoVv[p0:p1], r1b[:,p0:p1])

    Hr1a += np.einsum('maei,me->ia',imds.wovvo,r1a)
    Hr1a += np.einsum('MaEi,ME->ia',imds.wOvVo,r1b)
    Hr1b += np.einsum('maei,me->ia',imds.wOVVO,r1b)
    Hr1b += np.einsum('mAeI,me->IA',imds.woVvO,r1a)
    Hr2aa+= lib.einsum('mbej,imae->ijab', imds.wovvo, r2aa) * 2
    Hr2aa+= lib.einsum('MbEj,iMaE->ijab', imds.wOvVo, r2ab) * 2
    Hr2bb+= lib.einsum('mbej,imae->ijab', imds.wOVVO, r2bb) * 2
    Hr2bb+= lib.einsum('mBeJ,mIeA->IJAB', imds.woVvO, r2ab) * 2
    Hr2ab+= lib.einsum('mBeJ,imae->iJaB', imds.woVvO, r2aa)
    Hr2ab+= lib.einsum('MBEJ,iMaE->iJaB', imds.wOVVO, r2ab)
    Hr2ab+= lib.einsum('mBEj,mIaE->jIaB', imds.woVVo, r2ab)
    Hr2ab+= lib.einsum('mbej,mIeA->jIbA', imds.wovvo, r2ab)
    Hr2ab+= lib.einsum('MbEj,IMAE->jIbA', imds.wOvVo, r2bb)
    Hr2ab+= lib.einsum('MbeJ,iMeA->iJbA', imds.wOvvO, r2ab)

    Hr2aa *= .5
    Hr2bb *= .5
    Hr2aa = Hr2aa - Hr2aa.transpose(0,1,3,2)
    Hr2aa = Hr2aa - Hr2aa.transpose(1,0,2,3)
    Hr2bb = Hr2bb - Hr2bb.transpose(0,1,3,2)
    Hr2bb = Hr2bb - Hr2bb.transpose(1,0,2,3)

    vector = amplitudes_to_vector_ee((Hr1a,Hr1b), (Hr2aa,Hr2ab,Hr2bb))
    return vector

def eomsf_ccsd_matvec(eom, vector, imds=None):
    '''Spin flip EOM-CCSD'''
    if imds is None: imds = eom.make_imds()

    t1, t2, eris = imds.t1, imds.t2, imds.eris
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nocca, noccb, nvira, nvirb = t2ab.shape
    nmoa, nmob = nocca+nvira, noccb+nvirb
    r1, r2 = vector_to_amplitudes_eomsf(vector, (nmoa,nmob), (nocca,noccb))
    r1ab, r1ba = r1
    r2baaa, r2aaba, r2abbb, r2bbab = r2

    Hr1ab  = np.einsum('ae,ie->ia', imds.Fvvb, r1ab)
    Hr1ab -= np.einsum('mi,ma->ia', imds.Fooa, r1ab)
    Hr1ab += np.einsum('me,imae->ia', imds.Fovb, r2abbb)
    Hr1ab += np.einsum('me,imae->ia', imds.Fova, r2aaba)
    Hr1ba  = np.einsum('ae,ie->ia', imds.Fvva, r1ba)
    Hr1ba -= np.einsum('mi,ma->ia', imds.Foob, r1ba)
    Hr1ba += np.einsum('me,imae->ia', imds.Fova, r2baaa)
    Hr1ba += np.einsum('me,imae->ia', imds.Fovb, r2bbab)
    Hr2baaa = .5 *lib.einsum('njMI,Mnab->Ijab', imds.wooOO, r2baaa)
    Hr2aaba = .25*lib.einsum('minj,mnAb->ijAb', imds.woooo, r2aaba)
    Hr2abbb = .5 *lib.einsum('miNJ,mNAB->iJAB', imds.wooOO, r2abbb)
    Hr2bbab = .25*lib.einsum('MINJ,MNaB->IJaB', imds.wOOOO, r2bbab)
    Hr2baaa += lib.einsum('be,Ijae->Ijab', imds.Fvva   , r2baaa)
    Hr2baaa -= lib.einsum('mj,imab->ijab', imds.Fooa*.5, r2baaa)
    Hr2baaa -= lib.einsum('MJ,Miab->Jiab', imds.Foob*.5, r2baaa)
    Hr2bbab -= lib.einsum('mj,imab->ijab', imds.Foob   , r2bbab)
    Hr2bbab += lib.einsum('BE,IJaE->IJaB', imds.Fvvb*.5, r2bbab)
    Hr2bbab += lib.einsum('be,IJeA->IJbA', imds.Fvva*.5, r2bbab)
    Hr2aaba -= lib.einsum('mj,imab->ijab', imds.Fooa   , r2aaba)
    Hr2aaba += lib.einsum('be,ijAe->ijAb', imds.Fvva*.5, r2aaba)
    Hr2aaba += lib.einsum('BE,ijEa->ijBa', imds.Fvvb*.5, r2aaba)
    Hr2abbb += lib.einsum('BE,iJAE->iJAB', imds.Fvvb   , r2abbb)
    Hr2abbb -= lib.einsum('mj,imab->ijab', imds.Foob*.5, r2abbb)
    Hr2abbb -= lib.einsum('mj,mIAB->jIAB', imds.Fooa*.5, r2abbb)

    tau2baaa = np.einsum('ia,jb->ijab', r1ba, t1a)
    tau2baaa = tau2baaa - tau2baaa.transpose(0,1,3,2)
    tau2abbb = np.einsum('ia,jb->ijab', r1ab, t1b)
    tau2abbb = tau2abbb - tau2abbb.transpose(0,1,3,2)
    tau2aaba = np.einsum('ia,jb->ijab', r1ab, t1a)
    tau2aaba = tau2aaba - tau2aaba.transpose(1,0,2,3)
    tau2bbab = np.einsum('ia,jb->ijab', r1ba, t1b)
    tau2bbab = tau2bbab - tau2bbab.transpose(1,0,2,3)
    tau2baaa += r2baaa
    tau2bbab += r2bbab
    tau2abbb += r2abbb
    tau2aaba += r2aaba
    #:eris_ovvv = lib.unpack_tril(np.asarray(eris.ovvv).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvira,nvira)
    #:Hr1ba += lib.einsum('mfae,Imef->Ia', eris_ovvv, r2baaa)
    #:tmp1aaba = lib.einsum('meaf,Ijef->maIj', eris_ovvv, tau2baaa)
    #:Hr2baaa += lib.einsum('mb,maIj->Ijab', t1a   , tmp1aaba)
    mem_now = lib.current_memory()[0]
    max_memory = max(0, eom.max_memory - mem_now)
    blksize = min(nocca, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvira**3*3))))
    for p0,p1 in lib.prange(0, nocca, blksize):
        ovvv = eris.get_ovvv(slice(p0,p1))  # ovvv = eris.ovvv[p0:p1]
        Hr1ba += lib.einsum('mfae,Imef->Ia', ovvv, r2baaa[:,p0:p1])
        tmp1aaba = lib.einsum('meaf,Ijef->maIj', ovvv, tau2baaa)
        Hr2baaa += lib.einsum('mb,maIj->Ijab', t1a[p0:p1], tmp1aaba)
        ovvv = tmp1aaba = None

    #:eris_OVVV = lib.unpack_tril(np.asarray(eris.OVVV).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvirb,nvirb)
    #:Hr1ab += lib.einsum('MFAE,iMEF->iA', eris_OVVV, r2abbb)
    #:tmp1bbab = lib.einsum('MEAF,iJEF->MAiJ', eris_OVVV, tau2abbb)
    #:Hr2abbb += lib.einsum('MB,MAiJ->iJAB', t1b   , tmp1bbab)
    blksize = min(noccb, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvirb**3*3))))
    for p0, p1 in lib.prange(0, noccb, blksize):
        OVVV = eris.get_OVVV(slice(p0,p1))  # OVVV = eris.OVVV[p0:p1]
        Hr1ab += lib.einsum('MFAE,iMEF->iA', OVVV, r2abbb[:,p0:p1])
        tmp1bbab = lib.einsum('MEAF,iJEF->MAiJ', OVVV, tau2abbb)
        Hr2abbb += lib.einsum('MB,MAiJ->iJAB', t1b[p0:p1], tmp1bbab)
        OVVV = tmp1bbab = None

    #:eris_ovVV = lib.unpack_tril(np.asarray(eris.ovVV).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvirb,nvirb)
    #:Hr1ab += lib.einsum('mfAE,imEf->iA', eris_ovVV, r2aaba)
    #:tmp1abaa = lib.einsum('meAF,ijFe->mAij', eris_ovVV, tau2aaba)
    #:tmp1abbb = lib.einsum('meAF,IJeF->mAIJ', eris_ovVV, tau2bbab)
    #:tmp1ba = lib.einsum('mfAE,mE->Af', eris_ovVV, r1ab)
    #:Hr2bbab -= lib.einsum('mb,mAIJ->IJbA', t1a*.5, tmp1abbb)
    #:Hr2aaba -= lib.einsum('mb,mAij->ijAb', t1a*.5, tmp1abaa)
    tmp1ba = np.zeros((nvirb,nvira))
    blksize = min(nocca, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvira*nvirb**2*3))))
    for p0,p1 in lib.prange(0, nocca, blksize):
        ovVV = eris.get_ovVV(slice(p0,p1))  # ovVV = eris.ovVV[p0:p1]
        Hr1ab += lib.einsum('mfAE,imEf->iA', ovVV, r2aaba[:,p0:p1])
        tmp1abaa = lib.einsum('meAF,ijFe->mAij', ovVV, tau2aaba)
        tmp1abbb = lib.einsum('meAF,IJeF->mAIJ', ovVV, tau2bbab)
        tmp1ba += lib.einsum('mfAE,mE->Af', ovVV, r1ab[p0:p1])
        Hr2bbab -= lib.einsum('mb,mAIJ->IJbA', t1a[p0:p1]*.5, tmp1abbb)
        Hr2aaba -= lib.einsum('mb,mAij->ijAb', t1a[p0:p1]*.5, tmp1abaa)

    #:eris_OVvv = lib.unpack_tril(np.asarray(eris.OVvv).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvira,nvira)
    #:Hr1ba += lib.einsum('MFae,IMeF->Ia', eris_OVvv, r2bbab)
    #:tmp1baaa = lib.einsum('MEaf,ijEf->Maij', eris_OVvv, tau2aaba)
    #:tmp1babb = lib.einsum('MEaf,IJfE->MaIJ', eris_OVvv, tau2bbab)
    #:tmp1ab = lib.einsum('MFae,Me->aF', eris_OVvv, r1ba)
    #:Hr2aaba -= lib.einsum('MB,Maij->ijBa', t1b*.5, tmp1baaa)
    #:Hr2bbab -= lib.einsum('MB,MaIJ->IJaB', t1b*.5, tmp1babb)
    tmp1ab = np.zeros((nvira,nvirb))
    blksize = min(noccb, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvirb*nvira**2*3))))
    for p0, p1 in lib.prange(0, noccb, blksize):
        OVvv = eris.get_OVvv(slice(p0,p1))  # OVvv = eris.OVvv[p0:p1]
        Hr1ba += lib.einsum('MFae,IMeF->Ia', OVvv, r2bbab[:,p0:p1])
        tmp1baaa = lib.einsum('MEaf,ijEf->Maij', OVvv, tau2aaba)
        tmp1babb = lib.einsum('MEaf,IJfE->MaIJ', OVvv, tau2bbab)
        tmp1ab+= lib.einsum('MFae,Me->aF', OVvv, r1ba[p0:p1])
        Hr2aaba -= lib.einsum('MB,Maij->ijBa', t1b[p0:p1]*.5, tmp1baaa)
        Hr2bbab -= lib.einsum('MB,MaIJ->IJaB', t1b[p0:p1]*.5, tmp1babb)

    Hr2baaa += lib.einsum('aF,jIbF->Ijba', tmp1ab   , t2ab)
    Hr2bbab -= lib.einsum('aF,IJFB->IJaB', tmp1ab*.5, t2bb)
    Hr2abbb += lib.einsum('Af,iJfB->iJBA', tmp1ba   , t2ab)
    Hr2aaba -= lib.einsum('Af,ijfb->ijAb', tmp1ba*.5, t2aa)
    Hr2baaa -= lib.einsum('MbIj,Ma->Ijab', imds.wOvOo, r1ba   )
    Hr2bbab -= lib.einsum('MBIJ,Ma->IJaB', imds.wOVOO, r1ba*.5)
    Hr2abbb -= lib.einsum('mBiJ,mA->iJAB', imds.woVoO, r1ab   )
    Hr2aaba -= lib.einsum('mbij,mA->ijAb', imds.wovoo, r1ab*.5)

    Hr1ab -= 0.5*lib.einsum('mine,mnAe->iA', imds.wooov, r2aaba)
    Hr1ab -=     lib.einsum('miNE,mNAE->iA', imds.wooOV, r2abbb)
    Hr1ba -= 0.5*lib.einsum('MINE,MNaE->Ia', imds.wOOOV, r2bbab)
    Hr1ba -=     lib.einsum('MIne,Mnae->Ia', imds.wOOov, r2baaa)
    tmp1ab = lib.einsum('MIne,Me->nI', imds.wOOov, r1ba)
    tmp1ba = lib.einsum('miNE,mE->Ni', imds.wooOV, r1ab)
    Hr2baaa += lib.einsum('nI,njab->Ijab', tmp1ab*.5, t2aa)
    Hr2bbab += lib.einsum('nI,nJaB->IJaB', tmp1ab   , t2ab)
    Hr2abbb += lib.einsum('Ni,NJAB->iJAB', tmp1ba*.5, t2bb)
    Hr2aaba += lib.einsum('Ni,jNbA->ijAb', tmp1ba   , t2ab)
    for p0, p1 in lib.prange(0, nvira, nocca):
        Hr2baaa += lib.einsum('ejab,Ie->Ijab', imds.wvovv[p0:p1], r1ba[:,p0:p1]*.5)
        Hr2bbab += lib.einsum('eJaB,Ie->IJaB', imds.wvOvV[p0:p1], r1ba[:,p0:p1]   )
    for p0, p1 in lib.prange(0, nvirb, noccb):
        Hr2abbb += lib.einsum('EJAB,iE->iJAB', imds.wVOVV[p0:p1], r1ab[:,p0:p1]*.5)
        Hr2aaba += lib.einsum('EjAb,iE->ijAb', imds.wVoVv[p0:p1], r1ab[:,p0:p1]   )

    Hr1ab += np.einsum('mAEi,mE->iA', imds.woVVo, r1ab)
    Hr1ba += np.einsum('MaeI,Me->Ia', imds.wOvvO, r1ba)
    Hr2baaa += lib.einsum('mbej,Imae->Ijab', imds.wovvo, r2baaa)
    Hr2baaa += lib.einsum('MbeJ,Miae->Jiab', imds.wOvvO, r2baaa)
    Hr2baaa += lib.einsum('MbEj,IMaE->Ijab', imds.wOvVo, r2bbab)
    Hr2bbab += lib.einsum('MBEJ,IMaE->IJaB', imds.wOVVO, r2bbab)
    Hr2bbab += lib.einsum('MbeJ,IMeA->IJbA', imds.wOvvO, r2bbab)
    Hr2bbab += lib.einsum('mBeJ,Imae->IJaB', imds.woVvO, r2baaa)
    Hr2aaba += lib.einsum('mbej,imAe->ijAb', imds.wovvo, r2aaba)
    Hr2aaba += lib.einsum('mBEj,imEa->ijBa', imds.woVVo, r2aaba)
    Hr2aaba += lib.einsum('MbEj,iMAE->ijAb', imds.wOvVo, r2abbb)
    Hr2abbb += lib.einsum('MBEJ,iMAE->iJAB', imds.wOVVO, r2abbb)
    Hr2abbb += lib.einsum('mBEj,mIAE->jIAB', imds.woVVo, r2abbb)
    Hr2abbb += lib.einsum('mBeJ,imAe->iJAB', imds.woVvO, r2aaba)

    eris_ovov = np.asarray(eris.ovov)
    eris_OVOV = np.asarray(eris.OVOV)
    eris_ovOV = np.asarray(eris.ovOV)
    tauaa, tauab, taubb = uccsd.make_tau(t2, t1, t1)
    tmp1baaa = lib.einsum('nfME,ijEf->Mnij', eris_ovOV, tau2aaba)
    tmp1aaba = lib.einsum('menf,Ijef->mnIj', eris_ovov, tau2baaa)
    tmp1abbb = lib.einsum('meNF,IJeF->mNIJ', eris_ovOV, tau2bbab)
    tmp1bbab = lib.einsum('MENF,iJEF->MNiJ', eris_OVOV, tau2abbb)
    Hr2baaa += 0.5*.5*lib.einsum('mnIj,mnab->Ijab', tmp1aaba, tauaa)
    Hr2bbab +=     .5*lib.einsum('nMIJ,nMaB->IJaB', tmp1abbb, tauab)
    Hr2aaba +=     .5*lib.einsum('Nmij,mNbA->ijAb', tmp1baaa, tauab)
    Hr2abbb += 0.5*.5*lib.einsum('MNiJ,MNAB->iJAB', tmp1bbab, taubb)
    tauaa = tauab = taubb = None

    tmpab  = lib.einsum('menf,Imef->nI', eris_ovov, r2baaa)
    tmpab -= lib.einsum('nfME,IMfE->nI', eris_ovOV, r2bbab)
    tmpba  = lib.einsum('MENF,iMEF->Ni', eris_OVOV, r2abbb)
    tmpba -= lib.einsum('meNF,imFe->Ni', eris_ovOV, r2aaba)
    Hr1ab += np.einsum('NA,Ni->iA', t1b, tmpba)
    Hr1ba += np.einsum('na,nI->Ia', t1a, tmpab)
    Hr2baaa -= lib.einsum('mJ,imab->Jiab', tmpab*.5, t2aa)
    Hr2bbab -= lib.einsum('mJ,mIaB->IJaB', tmpab*.5, t2ab) * 2
    Hr2aaba -= lib.einsum('Mj,iMbA->ijAb', tmpba*.5, t2ab) * 2
    Hr2abbb -= lib.einsum('Mj,IMAB->jIAB', tmpba*.5, t2bb)

    tmp1ab = np.einsum('meNF,mF->eN', eris_ovOV, r1ab)
    tmp1ba = np.einsum('nfME,Mf->En', eris_ovOV, r1ba)
    tmpab = np.einsum('eN,NB->eB', tmp1ab, t1b)
    tmpba = np.einsum('En,nb->Eb', tmp1ba, t1a)
    tmpab -= lib.einsum('menf,mnBf->eB', eris_ovov, r2aaba)
    tmpab += lib.einsum('meNF,mNFB->eB', eris_ovOV, r2abbb)
    tmpba -= lib.einsum('MENF,MNbF->Eb', eris_OVOV, r2bbab)
    tmpba += lib.einsum('nfME,Mnfb->Eb', eris_ovOV, r2baaa)
    Hr2baaa -= lib.einsum('Eb,jIaE->Ijab', tmpba*.5, t2ab) * 2
    Hr2bbab -= lib.einsum('Eb,IJAE->IJbA', tmpba*.5, t2bb)
    Hr2aaba -= lib.einsum('eB,ijae->ijBa', tmpab*.5, t2aa)
    Hr2abbb -= lib.einsum('eB,iJeA->iJAB', tmpab*.5, t2ab) * 2
    eris_ovov = eris_OVOV = eris_ovOV = None

    #:Hr2baaa += .5*lib.einsum('Ijef,aebf->Ijab', tau2baaa, eris.vvvv)
    #:Hr2abbb += .5*lib.einsum('iJEF,AEBF->iJAB', tau2abbb, eris.VVVV)
    #:Hr2bbab += .5*lib.einsum('IJeF,aeBF->IJaB', tau2bbab, eris.vvVV)
    #:Hr2aaba += .5*lib.einsum('ijEf,bfAE->ijAb', tau2aaba, eris.vvVV)
    fakeri = uccsd._ChemistsERIs()
    fakeri.mol = eris.mol

    if eom._cc.direct:
        orbva = eris.mo_coeff[0][:,nocca:]
        orbvb = eris.mo_coeff[1][:,noccb:]
        tau2baaa = lib.einsum('ijab,pa,qb->ijpq', tau2baaa, .5*orbva, orbva)
        tmp = eris._contract_vvvv_t2(eom._cc, tau2baaa, True)
        Hr2baaa += lib.einsum('ijpq,pa,qb->ijab', tmp, orbva.conj(), orbva.conj())
        tmp = None

        tau2abbb = lib.einsum('ijab,pa,qb->ijpq', tau2abbb, .5*orbvb, orbvb)
        tmp = eris._contract_VVVV_t2(eom._cc, tau2abbb, True)
        Hr2abbb += lib.einsum('ijpq,pa,qb->ijab', tmp, orbvb.conj(), orbvb.conj())
        tmp = None
    else:
        tau2baaa *= .5
        Hr2baaa += eris._contract_vvvv_t2(eom._cc, tau2baaa, False)
        tau2abbb *= .5
        Hr2abbb += eris._contract_VVVV_t2(eom._cc, tau2abbb, False)

    tau2bbab *= .5
    Hr2bbab += eom._cc._add_vvVV(None, tau2bbab, eris)
    tau2aaba = tau2aaba.transpose(0,1,3,2)*.5
    Hr2aaba += eom._cc._add_vvVV(None, tau2aaba, eris).transpose(0,1,3,2)

    Hr2baaa = Hr2baaa - Hr2baaa.transpose(0,1,3,2)
    Hr2bbab = Hr2bbab - Hr2bbab.transpose(1,0,2,3)
    Hr2abbb = Hr2abbb - Hr2abbb.transpose(0,1,3,2)
    Hr2aaba = Hr2aaba - Hr2aaba.transpose(1,0,2,3)
    vector = amplitudes_to_vector_eomsf((Hr1ab, Hr1ba), (Hr2baaa,Hr2aaba,Hr2abbb,Hr2bbab))
    return vector

def eeccsd_diag(eom, imds=None):
    if imds is None: imds = eom.make_imds()
    eris = imds.eris
    t1, t2 = imds.t1, imds.t2
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    tauaa, tauab, taubb = uccsd.make_tau(t2, t1, t1)
    nocca, noccb, nvira, nvirb = t2ab.shape

    Foa = imds.Fooa.diagonal()
    Fob = imds.Foob.diagonal()
    Fva = imds.Fvva.diagonal()
    Fvb = imds.Fvvb.diagonal()
    Wovaa = np.einsum('iaai->ia', imds.wovvo)
    Wovbb = np.einsum('iaai->ia', imds.wOVVO)
    Wovab = np.einsum('iaai->ia', imds.woVVo)
    Wovba = np.einsum('iaai->ia', imds.wOvvO)

    Hr1aa = lib.direct_sum('-i+a->ia', Foa, Fva)
    Hr1bb = lib.direct_sum('-i+a->ia', Fob, Fvb)
    Hr1ab = lib.direct_sum('-i+a->ia', Foa, Fvb)
    Hr1ba = lib.direct_sum('-i+a->ia', Fob, Fva)
    Hr1aa += Wovaa
    Hr1bb += Wovbb
    Hr1ab += Wovab
    Hr1ba += Wovba

    eris_ovov = np.asarray(eris.ovov)
    eris_OVOV = np.asarray(eris.OVOV)
    eris_ovOV = np.asarray(eris.ovOV)
    ovov = eris_ovov - eris_ovov.transpose(0,3,2,1)
    OVOV = eris_OVOV - eris_OVOV.transpose(0,3,2,1)
    Wvvaa = .5*np.einsum('mnab,manb->ab', tauaa, eris_ovov)
    Wvvbb = .5*np.einsum('mnab,manb->ab', taubb, eris_OVOV)
    Wvvab =    np.einsum('mNaB,maNB->aB', tauab, eris_ovOV)
    ijb = np.einsum('iejb,ijbe->ijb',      ovov, t2aa)
    IJB = np.einsum('iejb,ijbe->ijb',      OVOV, t2bb)
    iJB =-np.einsum('ieJB,iJeB->iJB', eris_ovOV, t2ab)
    Ijb =-np.einsum('jbIE,jIbE->Ijb', eris_ovOV, t2ab)
    iJb =-np.einsum('ibJE,iJbE->iJb', eris_ovOV, t2ab)
    jab = np.einsum('kajb,jkab->jab',      ovov, t2aa)
    JAB = np.einsum('kajb,jkab->jab',      OVOV, t2bb)
    jAb =-np.einsum('jbKA,jKbA->jAb', eris_ovOV, t2ab)
    JaB =-np.einsum('kaJB,kJaB->JaB', eris_ovOV, t2ab)
    jaB =-np.einsum('jaKB,jKaB->jaB', eris_ovOV, t2ab)
    eris_ovov = eris_ovOV = eris_OVOV = ovov = OVOV = None
    Hr2aa = lib.direct_sum('ijb+a->ijba', ijb, Fva)
    Hr2bb = lib.direct_sum('ijb+a->ijba', IJB, Fvb)
    Hr2ab = lib.direct_sum('iJb+A->iJbA', iJb, Fvb)
    Hr2ab+= lib.direct_sum('iJB+a->iJaB', iJB, Fva)
    Hr2aa+= lib.direct_sum('-i+jab->ijab', Foa, jab)
    Hr2bb+= lib.direct_sum('-i+jab->ijab', Fob, JAB)
    Hr2ab+= lib.direct_sum('-i+JaB->iJaB', Foa, JaB)
    Hr2ab+= lib.direct_sum('-I+jaB->jIaB', Fob, jaB)
    Hr2aa = Hr2aa + Hr2aa.transpose(0,1,3,2)
    Hr2aa = Hr2aa + Hr2aa.transpose(1,0,2,3)
    Hr2bb = Hr2bb + Hr2bb.transpose(0,1,3,2)
    Hr2bb = Hr2bb + Hr2bb.transpose(1,0,2,3)
    Hr2aa *= .5
    Hr2bb *= .5
    Hr2baaa = lib.direct_sum('Ijb+a->Ijba', Ijb, Fva)
    Hr2aaba = lib.direct_sum('ijb+A->ijAb', ijb, Fvb)
    Hr2aaba+= Fva.reshape(1,1,1,-1)
    Hr2abbb = lib.direct_sum('iJB+A->iJBA', iJB, Fvb)
    Hr2bbab = lib.direct_sum('IJB+a->IJaB', IJB, Fva)
    Hr2bbab+= Fvb.reshape(1,1,1,-1)
    Hr2baaa = Hr2baaa + Hr2baaa.transpose(0,1,3,2)
    Hr2abbb = Hr2abbb + Hr2abbb.transpose(0,1,3,2)
    Hr2baaa+= lib.direct_sum('-I+jab->Ijab', Fob, jab)
    Hr2baaa-= Foa.reshape(1,-1,1,1)
    tmpaaba = lib.direct_sum('-i+jAb->ijAb', Foa, jAb)
    Hr2abbb+= lib.direct_sum('-i+JAB->iJAB', Foa, JAB)
    Hr2abbb-= Fob.reshape(1,-1,1,1)
    tmpbbab = lib.direct_sum('-I+JaB->IJaB', Fob, JaB)
    Hr2aaba+= tmpaaba + tmpaaba.transpose(1,0,2,3)
    Hr2bbab+= tmpbbab + tmpbbab.transpose(1,0,2,3)
    tmpaaba = tmpbbab = None
    Hr2aa += Wovaa.reshape(1,nocca,1,nvira)
    Hr2aa += Wovaa.reshape(nocca,1,1,nvira)
    Hr2aa += Wovaa.reshape(nocca,1,nvira,1)
    Hr2aa += Wovaa.reshape(1,nocca,nvira,1)
    Hr2ab += Wovbb.reshape(1,noccb,1,nvirb)
    Hr2ab += Wovab.reshape(nocca,1,1,nvirb)
    Hr2ab += Wovaa.reshape(nocca,1,nvira,1)
    Hr2ab += Wovba.reshape(1,noccb,nvira,1)
    Hr2bb += Wovbb.reshape(1,noccb,1,nvirb)
    Hr2bb += Wovbb.reshape(noccb,1,1,nvirb)
    Hr2bb += Wovbb.reshape(noccb,1,nvirb,1)
    Hr2bb += Wovbb.reshape(1,noccb,nvirb,1)
    Hr2baaa += Wovaa.reshape(1,nocca,1,nvira)
    Hr2baaa += Wovba.reshape(noccb,1,1,nvira)
    Hr2baaa += Wovba.reshape(noccb,1,nvira,1)
    Hr2baaa += Wovaa.reshape(1,nocca,nvira,1)
    Hr2aaba += Wovaa.reshape(1,nocca,1,nvira)
    Hr2aaba += Wovaa.reshape(nocca,1,1,nvira)
    Hr2aaba += Wovab.reshape(nocca,1,nvirb,1)
    Hr2aaba += Wovab.reshape(1,nocca,nvirb,1)
    Hr2abbb += Wovbb.reshape(1,noccb,1,nvirb)
    Hr2abbb += Wovab.reshape(nocca,1,1,nvirb)
    Hr2abbb += Wovab.reshape(nocca,1,nvirb,1)
    Hr2abbb += Wovbb.reshape(1,noccb,nvirb,1)
    Hr2bbab += Wovbb.reshape(1,noccb,1,nvirb)
    Hr2bbab += Wovbb.reshape(noccb,1,1,nvirb)
    Hr2bbab += Wovba.reshape(noccb,1,nvira,1)
    Hr2bbab += Wovba.reshape(1,noccb,nvira,1)

    Wooaa  = np.einsum('iijj->ij', imds.woooo).copy()
    Wooaa -= np.einsum('ijji->ij', imds.woooo)
    Woobb  = np.einsum('iijj->ij', imds.wOOOO).copy()
    Woobb -= np.einsum('ijji->ij', imds.wOOOO)
    Wooab = np.einsum('iijj->ij', imds.wooOO)
    Wooba = Wooab.T
    Wooaa *= .5
    Woobb *= .5
    Hr2aa += Wooaa.reshape(nocca,nocca,1,1)
    Hr2ab += Wooab.reshape(nocca,noccb,1,1)
    Hr2bb += Woobb.reshape(noccb,noccb,1,1)
    Hr2baaa += Wooba.reshape(noccb,nocca,1,1)
    Hr2aaba += Wooaa.reshape(nocca,nocca,1,1)
    Hr2abbb += Wooab.reshape(nocca,noccb,1,1)
    Hr2bbab += Woobb.reshape(noccb,noccb,1,1)

    #:eris_ovvv = lib.unpack_tril(np.asarray(eris.ovvv).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvira,nvira)
    #:Wvvaa += np.einsum('mb,maab->ab', t1a, eris_ovvv)
    #:Wvvaa -= np.einsum('mb,mbaa->ab', t1a, eris_ovvv)
    mem_now = lib.current_memory()[0]
    max_memory = max(0, eom.max_memory - mem_now)
    blksize = min(nocca, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvira**3*3))))
    for p0,p1 in lib.prange(0, nocca, blksize):
        ovvv = eris.get_ovvv(slice(p0,p1))  # ovvv = eris.ovvv[p0:p1]
        Wvvaa += np.einsum('mb,maab->ab', t1a[p0:p1], ovvv)
        Wvvaa -= np.einsum('mb,mbaa->ab', t1a[p0:p1], ovvv)
        ovvv = None
    #:eris_OVVV = lib.unpack_tril(np.asarray(eris.OVVV).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvirb,nvirb)
    #:Wvvbb += np.einsum('mb,maab->ab', t1b, eris_OVVV)
    #:Wvvbb -= np.einsum('mb,mbaa->ab', t1b, eris_OVVV)
    blksize = min(noccb, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvirb**3*3))))
    for p0, p1 in lib.prange(0, noccb, blksize):
        OVVV = eris.get_OVVV(slice(p0,p1))  # OVVV = eris.OVVV[p0:p1]
        Wvvbb += np.einsum('mb,maab->ab', t1b[p0:p1], OVVV)
        Wvvbb -= np.einsum('mb,mbaa->ab', t1b[p0:p1], OVVV)
        OVVV = None
    #:eris_ovVV = lib.unpack_tril(np.asarray(eris.ovVV).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvirb,nvirb)
    #:Wvvab -= np.einsum('mb,mbaa->ba', t1a, eris_ovVV)
    blksize = min(nocca, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvira*nvirb**2*3))))
    for p0,p1 in lib.prange(0, nocca, blksize):
        ovVV = eris.get_ovVV(slice(p0,p1))  # ovVV = eris.ovVV[p0:p1]
        Wvvab -= np.einsum('mb,mbaa->ba', t1a[p0:p1], ovVV)
        ovVV = None
    blksize = min(noccb, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvirb*nvira**2*3))))
    #:eris_OVvv = lib.unpack_tril(np.asarray(eris.OVvv).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvira,nvira)
    #:Wvvab -= np.einsum('mb,mbaa->ab', t1b, eris_OVvv)
    #idxa = np.arange(nvira)
    #idxa = idxa*(idxa+1)//2+idxa
    #for p0, p1 in lib.prange(0, noccb, blksize):
    #    OVvv = np.asarray(eris.OVvv[p0:p1])
    #    Wvvab -= np.einsum('mb,mba->ab', t1b[p0:p1], OVvv[:,:,idxa])
    #    OVvv = None
    for p0, p1 in lib.prange(0, noccb, blksize):
        OVvv = eris.get_OVvv(slice(p0,p1))  # OVvv = eris.OVvv[p0:p1]
        Wvvab -= np.einsum('mb,mbaa->ab', t1b[p0:p1], OVvv)
        OVvv = None
    Wvvaa = Wvvaa + Wvvaa.T
    Wvvbb = Wvvbb + Wvvbb.T
    #:eris_vvvv = ao2mo.restore(1, np.asarray(eris.vvvv), nvirb)
    #:eris_VVVV = ao2mo.restore(1, np.asarray(eris.VVVV), nvirb)
    #:eris_vvVV = _restore(np.asarray(eris.vvVV), nvira, nvirb)
    #:Wvvaa += np.einsum('aabb->ab', eris_vvvv) - np.einsum('abba->ab', eris_vvvv)
    #:Wvvbb += np.einsum('aabb->ab', eris_VVVV) - np.einsum('abba->ab', eris_VVVV)
    #:Wvvab += np.einsum('aabb->ab', eris_vvVV)
    if eris.vvvv is not None:
        for i in range(nvira):
            i0 = i*(i+1)//2
            vvv = lib.unpack_tril(np.asarray(eris.vvvv[i0:i0+i+1]))
            tmp = np.einsum('bb->b', vvv[i])
            Wvvaa[i] += tmp
            tmp = np.einsum('bb->b', vvv[:,:i+1,i])
            Wvvaa[i,:i+1] -= tmp
            Wvvaa[:i  ,i] -= tmp[:i]
            vvv = lib.unpack_tril(np.asarray(eris.vvVV[i0:i0+i+1]))
            Wvvab[i] += np.einsum('bb->b', vvv[i])
            vvv = None
        for i in range(nvirb):
            i0 = i*(i+1)//2
            vvv = lib.unpack_tril(np.asarray(eris.VVVV[i0:i0+i+1]))
            tmp = np.einsum('bb->b', vvv[i])
            Wvvbb[i] += tmp
            tmp = np.einsum('bb->b', vvv[:,:i+1,i])
            Wvvbb[i,:i+1] -= tmp
            Wvvbb[:i  ,i] -= tmp[:i]
            vvv = None
    Wvvba = Wvvab.T
    Hr2aa += Wvvaa.reshape(1,1,nvira,nvira)
    Hr2ab += Wvvab.reshape(1,1,nvira,nvirb)
    Hr2bb += Wvvbb.reshape(1,1,nvirb,nvirb)
    Hr2baaa += Wvvaa.reshape(1,1,nvira,nvira)
    Hr2aaba += Wvvba.reshape(1,1,nvirb,nvira)
    Hr2abbb += Wvvbb.reshape(1,1,nvirb,nvirb)
    Hr2bbab += Wvvab.reshape(1,1,nvira,nvirb)

    vec_ee = amplitudes_to_vector_ee((Hr1aa,Hr1bb), (Hr2aa,Hr2ab,Hr2bb))
    vec_sf = amplitudes_to_vector_eomsf((Hr1ab,Hr1ba), (Hr2baaa,Hr2aaba,Hr2abbb,Hr2bbab))
    return vec_ee, vec_sf

class EOMEE(eom_rccsd.EOMEE):
    def __init__(self, cc):
        eom_rccsd.EOMEE.__init__(self, cc)
        self.nocc = cc.get_nocc()
        self.nmo = cc.get_nmo()

    kernel = eeccsd
    eeccsd = eeccsd
    get_diag = eeccsd_diag

    def vector_size(self):
        '''size of the vector based on spin-orbital basis'''
        nocc = np.sum(self.nocc)
        nvir = np.sum(self.nmo) - nocc
        return nocc*nvir + nocc*(nocc-1)//2*nvir*(nvir-1)//2

    def make_imds(self, eris=None):
        imds = _IMDS(self._cc, eris=eris)
        imds.make_ee()
        return imds

class EOMEESpinKeep(EOMEE):
    kernel = eomee_ccsd
    eomee_ccsd = eomee_ccsd
    matvec = eomee_ccsd_matvec
    get_diag = eeccsd_diag

    def get_init_guess(self, nroots=1, koopmans=True, diag=None):
        if koopmans:
            nocca, noccb = self.nocc
            nmoa, nmob = self.nmo
            nvira, nvirb = nmoa-nocca, nmob-noccb
# amplitudes are compressed by the function amplitudes_to_vector_ee. sizea is
# the offset in the compressed vector that points to the amplitudes R1_beta
# The addresses of R1_alpha and R1_beta are not contiguous in the compressed
# vector.
            sizea = nocca * nvira + nocca*(nocca-1)//2*nvira*(nvira-1)//2
            diag = np.append(diag[:nocca*nvira], diag[sizea:sizea+noccb*nvirb])
            addr = np.append(np.arange(nocca*nvira),
                             np.arange(sizea,sizea+noccb*nvirb))
            idx = addr[diag.argsort()]
        else:
            idx = diag.argsort()

        size = self.vector_size()
        dtype = getattr(diag, 'dtype', np.double)
        nroots = min(nroots, size)
        guess = []
        for i in idx[:nroots]:
            g = np.zeros(size, dtype)
            g[i] = 1.0
            guess.append(g)
        return guess

    def gen_matvec(self, imds=None, diag=None, **kwargs):
        if imds is None: imds = self.make_imds()
        if diag is None: diag = self.get_diag(imds)[0]
        matvec = lambda xs: [self.matvec(x, imds) for x in xs]
        return matvec, diag

    def vector_to_amplitudes(self, vector, nmo=None, nocc=None):
        if nmo is None: nmo = self.nmo
        if nocc is None: nocc = self.nocc
        return vector_to_amplitudes_ee(vector, nmo, nocc)

    def amplitudes_to_vector(self, r1, r2):
        return amplitudes_to_vector_ee(r1, r2)

    def vector_size(self):
        '''size of the vector based on spin-orbital basis'''
        nocca, noccb = self.nocc
        nmoa, nmob = self.nmo
        nvira, nvirb = nmoa-nocca, nmob-noccb
        sizea = nocca * nvira + nocca*(nocca-1)//2*nvira*(nvira-1)//2
        sizeb = noccb * nvirb + noccb*(noccb-1)//2*nvirb*(nvirb-1)//2
        sizeab = nocca * noccb * nvira * nvirb
        return sizea+sizeb+sizeab

class EOMEESpinFlip(EOMEE):
    kernel = eomsf_ccsd
    eomsf_ccsd = eomsf_ccsd
    matvec = eomsf_ccsd_matvec

    def get_init_guess(self, nroots=1, koopmans=True, diag=None):
        if koopmans:
            nocca, noccb = self.nocc
            nmoa, nmob = self.nmo
            nvira, nvirb = nmoa-nocca, nmob-noccb
            idx = diag[:nocca*nvirb+noccb*nvira].argsort()
        else:
            idx = diag.argsort()

        size = self.vector_size()
        dtype = getattr(diag, 'dtype', np.double)
        nroots = min(nroots, size)
        guess = []
        for i in idx[:nroots]:
            g = np.zeros(size, dtype)
            g[i] = 1.0
            guess.append(g)
        return guess

    def gen_matvec(self, imds=None, diag=None, **kwargs):
        if imds is None: imds = self.make_imds()
        if diag is None: diag = self.get_diag(imds)[1]
        matvec = lambda xs: [self.matvec(x, imds) for x in xs]
        return matvec, diag

    def vector_to_amplitudes(self, vector, nmo=None, nocc=None):
        if nmo is None: nmo = self.nmo
        if nocc is None: nocc = self.nocc
        return vector_to_amplitudes_eomsf(vector, nmo, nocc)

    def amplitudes_to_vector(self, r1, r2):
        return amplitudes_to_vector_eomsf(r1, r2)

    def vector_size(self):
        '''size of the vector based on spin-orbital basis'''
        nocca, noccb = self.nocc
        nmoa, nmob = self.nmo
        nvira, nvirb = nmoa-nocca, nmob-noccb

        nbaaa = noccb*nocca*nvira*(nvira-1)//2
        naaba = nocca*(nocca-1)//2*nvirb*nvira
        nabbb = nocca*noccb*nvirb*(nvirb-1)//2
        nbbab = noccb*(noccb-1)//2*nvira*nvirb
        return nocca*nvirb + noccb*nvira + nbaaa + naaba + nabbb + nbbab

uccsd.UCCSD.EOMIP         = lib.class_as_method(EOMIP)
uccsd.UCCSD.EOMEA         = lib.class_as_method(EOMEA)
uccsd.UCCSD.EOMEE         = lib.class_as_method(EOMEE)
uccsd.UCCSD.EOMEESpinKeep = lib.class_as_method(EOMEESpinKeep)
uccsd.UCCSD.EOMEESpinFlip = lib.class_as_method(EOMEESpinFlip)


class _IMDS:
    # Exactly the same as RCCSD IMDS except
    # -- rintermediates --> uintermediates
    # -- Loo, Lvv, cc_Fov --> Foo, Fvv, Fov
    # -- One less 2-virtual intermediate
    def __init__(self, cc, eris=None):
        self.verbose = cc.verbose
        self.stdout = cc.stdout
        self.t1 = cc.t1
        self.t2 = cc.t2
        if eris is None:
            eris = cc.ao2mo()
        self.eris = eris
        self._made_shared = False
        self.made_ip_imds = False
        self.made_ea_imds = False
        self.made_ee_imds = False

    def _make_shared(self):
        cput0 = (logger.process_clock(), logger.perf_counter())

        t1, t2, eris = self.t1, self.t2, self.eris
        self.Foo, self.FOO = uintermediates.Foo(t1, t2, eris)
        self.Fvv, self.FVV = uintermediates.Fvv(t1, t2, eris)
        self.Fov, self.FOV = uintermediates.Fov(t1, t2, eris)

        # 2 virtuals
        self.Wovvo, self.WovVO, self.WOVvo, self.WOVVO, self.WoVVo, self.WOvvO = \
                uintermediates.Wovvo(t1, t2, eris)
        Wovov = np.asarray(eris.ovov)
        WOVOV = np.asarray(eris.OVOV)
        Wovov = Wovov - Wovov.transpose(0,3,2,1)
        WOVOV = WOVOV - WOVOV.transpose(0,3,2,1)
        self.Wovov = Wovov
        self.WovOV = eris.ovOV
        self.WOVov = None
        self.WOVOV = WOVOV

        self._made_shared = True
        logger.timer_debug1(self, 'EOM-CCSD shared intermediates', *cput0)
        return self

    def make_ip(self):
        if not self._made_shared:
            self._make_shared()

        cput0 = (logger.process_clock(), logger.perf_counter())

        t1, t2, eris = self.t1, self.t2, self.eris

        # 0 or 1 virtuals
        self.Woooo, self.WooOO, _         , self.WOOOO = uintermediates.Woooo(t1, t2, eris)
        self.Wooov, self.WooOV, self.WOOov, self.WOOOV = uintermediates.Wooov(t1, t2, eris)
        self.Woovo, self.WooVO, self.WOOvo, self.WOOVO = uintermediates.Woovo(t1, t2, eris)

        self.made_ip_imds = True
        logger.timer_debug1(self, 'EOM-UCCSD IP intermediates', *cput0)
        return self

    def make_ea(self):
        if not self._made_shared:
            self._make_shared()

        cput0 = (logger.process_clock(), logger.perf_counter())

        t1, t2, eris = self.t1, self.t2, self.eris

        # 3 or 4 virtuals
        self.Wvvov, self.WvvOV, self.WVVov, self.WVVOV = uintermediates.Wvvov(t1, t2, eris)
        self.Wvvvv = None  # too expensive to hold Wvvvv
        self.Wvvvo, self.WvvVO, self.WVVvo, self.WVVVO = uintermediates.Wvvvo(t1, t2, eris)

        # The contribution of Wvvvv
        t1a, t1b = t1
        # The contraction to eris.vvvv is included in eaccsd_matvec
        #:vvvv = eris.vvvv - eris.vvvv.transpose(0,3,2,1)
        #:VVVV = eris.VVVV - eris.VVVV.transpose(0,3,2,1)
        #:self.Wvvvo += lib.einsum('abef,if->abei',      vvvv, t1a)
        #:self.WvvVO += lib.einsum('abef,if->abei', eris_vvVV, t1b)
        #:self.WVVvo += lib.einsum('efab,if->abei', eris_vvVV, t1a)
        #:self.WVVVO += lib.einsum('abef,if->abei',      VVVV, t1b)

        tauaa, tauab, taubb = uccsd.make_tau(t2, t1, t1)
        eris_ovov = np.asarray(eris.ovov)
        eris_OVOV = np.asarray(eris.OVOV)
        eris_ovOV = np.asarray(eris.ovOV)
        ovov = eris_ovov - eris_ovov.transpose(0,3,2,1)
        OVOV = eris_OVOV - eris_OVOV.transpose(0,3,2,1)
        tmp = lib.einsum('menf,if->meni',      ovov, t1a) * .5
        self.Wvvvo += lib.einsum('meni,mnab->aebi', tmp, tauaa)
        tmp = tauaa = None

        tmp = lib.einsum('menf,if->meni',      OVOV, t1b) * .5
        self.WVVVO += lib.einsum('meni,mnab->aebi', tmp, taubb)
        tmp = taubb = None

        tmp = lib.einsum('menf,if->meni', eris_ovOV, t1b)
        self.WvvVO += lib.einsum('meni,mnab->aebi', tmp, tauab)
        tmp = lib.einsum('nfme,if->meni', eris_ovOV, t1a)
        self.WVVvo += lib.einsum('meni,nmba->aebi', tmp, tauab)
        tauab = None
        ovov = OVOV = eris_ovov = eris_OVOV = eris_ovOV = None

        eris_ovvv = eris.get_ovvv(slice(None))
        ovvv = eris_ovvv - eris_ovvv.transpose(0,3,2,1)
        tmp = lib.einsum('mebf,if->mebi', ovvv, t1a)
        tmp = lib.einsum('mebi,ma->aebi', tmp, t1a)
        self.Wvvvo -= tmp - tmp.transpose(2,1,0,3)
        tmp = eris_ovvv = ovvv = None

        eris_OVVV = eris.get_OVVV(slice(None))
        OVVV = eris_OVVV - eris_OVVV.transpose(0,3,2,1)
        tmp = lib.einsum('mebf,if->mebi', OVVV, t1b)
        tmp = lib.einsum('mebi,ma->aebi', tmp, t1b)
        self.WVVVO -= tmp - tmp.transpose(2,1,0,3)
        tmp = eris_OVVV = OVVV = None

        eris_ovVV = eris.get_ovVV(slice(None))
        eris_OVvv = eris.get_OVvv(slice(None))
        tmpaabb = lib.einsum('mebf,if->mebi', eris_ovVV, t1b)
        tmpbaab = lib.einsum('mebf,ie->mfbi', eris_OVvv, t1b)
        tmp  = lib.einsum('mebi,ma->aebi', tmpaabb, t1a)
        tmp += lib.einsum('mfbi,ma->bfai', tmpbaab, t1b)
        self.WvvVO -= tmp
        tmp = tmpaabb = tmpbaab = None

        tmpbbaa = lib.einsum('mebf,if->mebi', eris_OVvv, t1a)
        tmpabba = lib.einsum('mebf,ie->mfbi', eris_ovVV, t1a)
        tmp  = lib.einsum('mebi,ma->aebi', tmpbbaa, t1b)
        tmp += lib.einsum('mfbi,ma->bfai', tmpabba, t1a)
        self.WVVvo -= tmp
        tmp = tmpbbaa = tmpabba = None
        eris_ovVV = eris_OVvv = None
        # The contribution of Wvvvv end

        self.made_ea_imds = True
        logger.timer_debug1(self, 'EOM-UCCSD EA intermediates', *cput0)
        return self

    def make_ee(self):
        cput0 = (logger.process_clock(), logger.perf_counter())
        log = logger.Logger(self.stdout, self.verbose)

        t1, t2, eris = self.t1, self.t2, self.eris
        t1a, t1b = t1
        t2aa, t2ab, t2bb = t2
        nocca, noccb, nvira, nvirb = t2ab.shape
        dtype = np.result_type(t1a, t1b, t2aa, t2ab, t2bb)

        fooa = eris.focka[:nocca,:nocca]
        foob = eris.fockb[:noccb,:noccb]
        fova = eris.focka[:nocca,nocca:]
        fovb = eris.fockb[:noccb,noccb:]
        fvva = eris.focka[nocca:,nocca:]
        fvvb = eris.fockb[noccb:,noccb:]

        self.Fooa = np.zeros((nocca,nocca), dtype=dtype)
        self.Foob = np.zeros((noccb,noccb), dtype=dtype)
        self.Fvva = np.zeros((nvira,nvira), dtype=dtype)
        self.Fvvb = np.zeros((nvirb,nvirb), dtype=dtype)

        wovvo = np.zeros((nocca,nvira,nvira,nocca), dtype=dtype)
        wOVVO = np.zeros((noccb,nvirb,nvirb,noccb), dtype=dtype)
        woVvO = np.zeros((nocca,nvirb,nvira,noccb), dtype=dtype)
        woVVo = np.zeros((nocca,nvirb,nvirb,nocca), dtype=dtype)
        wOvVo = np.zeros((noccb,nvira,nvirb,nocca), dtype=dtype)
        wOvvO = np.zeros((noccb,nvira,nvira,noccb), dtype=dtype)

        wovoo = np.zeros((nocca,nvira,nocca,nocca), dtype=dtype)
        wOVOO = np.zeros((noccb,nvirb,noccb,noccb), dtype=dtype)
        woVoO = np.zeros((nocca,nvirb,nocca,noccb), dtype=dtype)
        wOvOo = np.zeros((noccb,nvira,noccb,nocca), dtype=dtype)

        tauaa, tauab, taubb = uccsd.make_tau(t2, t1, t1)
        #:eris_ovvv = lib.unpack_tril(np.asarray(eris.ovvv).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvira,nvira)
        #:ovvv = eris_ovvv - eris_ovvv.transpose(0,3,2,1)
        #:self.Fvva  = np.einsum('mf,mfae->ae', t1a, ovvv)
        #:self.wovvo = lib.einsum('jf,mebf->mbej', t1a, ovvv)
        #:self.wovoo  = 0.5 * lib.einsum('mebf,ijef->mbij', eris_ovvv, tauaa)
        #:self.wovoo -= 0.5 * lib.einsum('mfbe,ijef->mbij', eris_ovvv, tauaa)
        mem_now = lib.current_memory()[0]
        max_memory = max(0, lib.param.MAX_MEMORY - mem_now)
        blksize = min(nocca, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvira**3*3))))
        for p0,p1 in lib.prange(0, nocca, blksize):
            ovvv = eris.get_ovvv(slice(p0,p1))  # ovvv = eris.ovvv[p0:p1]
            ovvv = ovvv - ovvv.transpose(0,3,2,1)
            self.Fvva += np.einsum('mf,mfae->ae', t1a[p0:p1], ovvv)
            wovvo[p0:p1] = lib.einsum('jf,mebf->mbej', t1a, ovvv)
            wovoo[p0:p1] = 0.5 * lib.einsum('mebf,ijef->mbij', ovvv, tauaa)
            ovvv = None

        #:eris_OVVV = lib.unpack_tril(np.asarray(eris.OVVV).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvirb,nvirb)
        #:OVVV = eris_OVVV - eris_OVVV.transpose(0,3,2,1)
        #:self.Fvvb  = np.einsum('mf,mfae->ae', t1b, OVVV)
        #:self.wOVVO = lib.einsum('jf,mebf->mbej', t1b, OVVV)
        #:self.wOVOO  = 0.5 * lib.einsum('mebf,ijef->mbij', OVVV, taubb)
        blksize = min(noccb, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvirb**3*3))))
        for p0, p1 in lib.prange(0, noccb, blksize):
            OVVV = eris.get_OVVV(slice(p0,p1))  # OVVV = eris.OVVV[p0:p1]
            OVVV = OVVV - OVVV.transpose(0,3,2,1)
            self.Fvvb += np.einsum('mf,mfae->ae', t1b[p0:p1], OVVV)
            wOVVO[p0:p1] = lib.einsum('jf,mebf->mbej', t1b, OVVV)
            wOVOO[p0:p1] = 0.5 * lib.einsum('mebf,ijef->mbij', OVVV, taubb)
            OVVV = None

        #:eris_ovVV = lib.unpack_tril(np.asarray(eris.ovVV).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvirb,nvirb)
        #:self.Fvvb += np.einsum('mf,mfAE->AE', t1a, eris_ovVV)
        #:self.woVvO = lib.einsum('JF,meBF->mBeJ', t1b, eris_ovVV)
        #:self.woVVo = lib.einsum('jf,mfBE->mBEj',-t1a, eris_ovVV)
        #:self.woVoO  = 0.5 * lib.einsum('meBF,iJeF->mBiJ', eris_ovVV, tauab)
        #:self.woVoO += 0.5 * lib.einsum('mfBE,iJfE->mBiJ', eris_ovVV, tauab)
        blksize = min(nocca, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvira*nvirb**2*3))))
        for p0,p1 in lib.prange(0, nocca, blksize):
            ovVV = eris.get_ovVV(slice(p0,p1))  # ovVV = eris.ovVV[p0:p1]
            self.Fvvb += np.einsum('mf,mfAE->AE', t1a[p0:p1], ovVV)
            woVvO[p0:p1] = lib.einsum('JF,meBF->mBeJ', t1b, ovVV)
            woVVo[p0:p1] = lib.einsum('jf,mfBE->mBEj',-t1a, ovVV)
            woVoO[p0:p1] = 0.5 * lib.einsum('meBF,iJeF->mBiJ', ovVV, tauab)
            woVoO[p0:p1]+= 0.5 * lib.einsum('mfBE,iJfE->mBiJ', ovVV, tauab)
            ovVV = None

        #:eris_OVvv = lib.unpack_tril(np.asarray(eris.OVvv).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvira,nvira)
        #:self.Fvva += np.einsum('MF,MFae->ae', t1b, eris_OVvv)
        #:self.wOvVo = lib.einsum('jf,MEbf->MbEj', t1a, eris_OVvv)
        #:self.wOvvO = lib.einsum('JF,MFbe->MbeJ',-t1b, eris_OVvv)
        #:self.wOvOo  = 0.5 * lib.einsum('MEbf,jIfE->MbIj', eris_OVvv, tauab)
        #:self.wOvOo += 0.5 * lib.einsum('MFbe,jIeF->MbIj', eris_OVvv, tauab)
        blksize = min(noccb, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvirb*nvira**2*3))))
        for p0, p1 in lib.prange(0, noccb, blksize):
            OVvv = eris.get_OVvv(slice(p0,p1))  # OVvv = eris.OVvv[p0:p1]
            self.Fvva += np.einsum('MF,MFae->ae', t1b[p0:p1], OVvv)
            wOvVo[p0:p1] = lib.einsum('jf,MEbf->MbEj', t1a, OVvv)
            wOvvO[p0:p1] = lib.einsum('JF,MFbe->MbeJ',-t1b, OVvv)
            wOvOo[p0:p1] = 0.5 * lib.einsum('MEbf,jIfE->MbIj', OVvv, tauab)
            wOvOo[p0:p1]+= 0.5 * lib.einsum('MFbe,jIeF->MbIj', OVvv, tauab)
            OVvv = None

        eris_ovov = np.asarray(eris.ovov)
        eris_OVOV = np.asarray(eris.OVOV)
        eris_ovOV = np.asarray(eris.ovOV)
        ovov = eris_ovov - eris_ovov.transpose(0,3,2,1)
        OVOV = eris_OVOV - eris_OVOV.transpose(0,3,2,1)
        self.Fova = np.einsum('nf,menf->me', t1a,      ovov)
        self.Fova+= np.einsum('NF,meNF->me', t1b, eris_ovOV)
        self.Fova += fova
        self.Fovb = np.einsum('nf,menf->me', t1b,      OVOV)
        self.Fovb+= np.einsum('nf,nfME->ME', t1a, eris_ovOV)
        self.Fovb += fovb
        tilaa, tilab, tilbb = uccsd.make_tau(t2,t1,t1,fac=0.5)
        self.Fooa  = lib.einsum('inef,menf->mi', tilaa, eris_ovov)
        self.Fooa += lib.einsum('iNeF,meNF->mi', tilab, eris_ovOV)
        self.Foob  = lib.einsum('inef,menf->mi', tilbb, eris_OVOV)
        self.Foob += lib.einsum('nIfE,nfME->MI', tilab, eris_ovOV)
        self.Fvva -= lib.einsum('mnaf,menf->ae', tilaa, eris_ovov)
        self.Fvva -= lib.einsum('mNaF,meNF->ae', tilab, eris_ovOV)
        self.Fvvb -= lib.einsum('mnaf,menf->ae', tilbb, eris_OVOV)
        self.Fvvb -= lib.einsum('nMfA,nfME->AE', tilab, eris_ovOV)
        wovvo -= lib.einsum('jnfb,menf->mbej', t2aa,      ovov)
        wovvo += lib.einsum('jNbF,meNF->mbej', t2ab, eris_ovOV)
        wOVVO -= lib.einsum('jnfb,menf->mbej', t2bb,      OVOV)
        wOVVO += lib.einsum('nJfB,nfME->MBEJ', t2ab, eris_ovOV)
        woVvO += lib.einsum('nJfB,menf->mBeJ', t2ab,      ovov)
        woVvO -= lib.einsum('JNFB,meNF->mBeJ', t2bb, eris_ovOV)
        wOvVo -= lib.einsum('jnfb,nfME->MbEj', t2aa, eris_ovOV)
        wOvVo += lib.einsum('jNbF,MENF->MbEj', t2ab,      OVOV)
        woVVo += lib.einsum('jNfB,mfNE->mBEj', t2ab, eris_ovOV)
        wOvvO += lib.einsum('nJbF,neMF->MbeJ', t2ab, eris_ovOV)

        eris_ovoo = np.asarray(eris.ovoo)
        eris_OVOO = np.asarray(eris.OVOO)
        eris_OVoo = np.asarray(eris.OVoo)
        eris_ovOO = np.asarray(eris.ovOO)
        self.Fooa += np.einsum('ne,nemi->mi', t1a, eris_ovoo)
        self.Fooa -= np.einsum('ne,meni->mi', t1a, eris_ovoo)
        self.Fooa += np.einsum('NE,NEmi->mi', t1b, eris_OVoo)
        self.Foob += np.einsum('ne,nemi->mi', t1b, eris_OVOO)
        self.Foob -= np.einsum('ne,meni->mi', t1b, eris_OVOO)
        self.Foob += np.einsum('ne,neMI->MI', t1a, eris_ovOO)
        eris_ovoo = eris_ovoo + np.einsum('nfme,jf->menj', eris_ovov, t1a)
        eris_OVOO = eris_OVOO + np.einsum('nfme,jf->menj', eris_OVOV, t1b)
        eris_OVoo = eris_OVoo + np.einsum('nfme,jf->menj', eris_ovOV, t1a)
        eris_ovOO = eris_ovOO + np.einsum('menf,jf->menj', eris_ovOV, t1b)
        ovoo = eris_ovoo - eris_ovoo.transpose(2,1,0,3)
        OVOO = eris_OVOO - eris_OVOO.transpose(2,1,0,3)
        wovvo += lib.einsum('nb,nemj->mbej', t1a,      ovoo)
        wOVVO += lib.einsum('nb,nemj->mbej', t1b,      OVOO)
        woVvO -= lib.einsum('NB,meNJ->mBeJ', t1b, eris_ovOO)
        wOvVo -= lib.einsum('nb,MEnj->MbEj', t1a, eris_OVoo)
        woVVo += lib.einsum('NB,NEmj->mBEj', t1b, eris_OVoo)
        wOvvO += lib.einsum('nb,neMJ->MbeJ', t1a, eris_ovOO)

        self.Fooa += fooa + 0.5*lib.einsum('me,ie->mi', self.Fova+fova, t1a)
        self.Foob += foob + 0.5*lib.einsum('me,ie->mi', self.Fovb+fovb, t1b)
        self.Fvva += fvva - 0.5*lib.einsum('me,ma->ae', self.Fova+fova, t1a)
        self.Fvvb += fvvb - 0.5*lib.einsum('me,ma->ae', self.Fovb+fovb, t1b)

        # 0 or 1 virtuals
        eris_ovoo = np.asarray(eris.ovoo)
        eris_OVOO = np.asarray(eris.OVOO)
        eris_OVoo = np.asarray(eris.OVoo)
        eris_ovOO = np.asarray(eris.ovOO)
        ovoo = eris_ovoo - eris_ovoo.transpose(2,1,0,3)
        OVOO = eris_OVOO - eris_OVOO.transpose(2,1,0,3)
        woooo = lib.einsum('je,nemi->minj', t1a,      ovoo)
        wOOOO = lib.einsum('je,nemi->minj', t1b,      OVOO)
        wooOO = lib.einsum('JE,NEmi->miNJ', t1b, eris_OVoo)
        woOOo = lib.einsum('je,meNI->mINj',-t1a, eris_ovOO)
        tmpaa = lib.einsum('nemi,jnbe->mbij',      ovoo, t2aa)
        tmpaa+= lib.einsum('NEmi,jNbE->mbij', eris_OVoo, t2ab)
        tmpbb = lib.einsum('nemi,jnbe->mbij',      OVOO, t2bb)
        tmpbb+= lib.einsum('neMI,nJeB->MBIJ', eris_ovOO, t2ab)
        woVoO += lib.einsum('nemi,nJeB->mBiJ',      ovoo, t2ab)
        woVoO += lib.einsum('NEmi,JNBE->mBiJ', eris_OVoo, t2bb)
        woVoO -= lib.einsum('meNI,jNeB->mBjI', eris_ovOO, t2ab)
        wOvOo += lib.einsum('NEMI,jNbE->MbIj',      OVOO, t2ab)
        wOvOo += lib.einsum('neMI,jnbe->MbIj', eris_ovOO, t2aa)
        wOvOo -= lib.einsum('MEni,nJbE->MbJi', eris_OVoo, t2ab)
        wovoo += tmpaa - tmpaa.transpose(0,1,3,2)
        wOVOO += tmpbb - tmpbb.transpose(0,1,3,2)
        self.wooov = np.array(     ovoo.transpose(2,3,0,1), dtype=dtype)
        self.wOOOV = np.array(     OVOO.transpose(2,3,0,1), dtype=dtype)
        self.wooOV = np.array(eris_OVoo.transpose(2,3,0,1), dtype=dtype)
        self.wOOov = np.array(eris_ovOO.transpose(2,3,0,1), dtype=dtype)
#X        self.wOooV =-np.array(eris_OVoo.transpose(0,3,2,1), dtype=dtype)
#X        self.woOOv =-np.array(eris_ovOO.transpose(0,3,2,1), dtype=dtype)
        eris_ovoo = eris_OVOO = eris_ovOO = eris_OVoo = None

        woooo += np.asarray(eris.oooo)
        wOOOO += np.asarray(eris.OOOO)
        wooOO += np.asarray(eris.ooOO)
        self.woooo = woooo - woooo.transpose(0,3,2,1)
        self.wOOOO = wOOOO - wOOOO.transpose(0,3,2,1)
        self.wooOO = wooOO - woOOo.transpose(0,3,2,1)

        eris_ovov = np.asarray(eris.ovov)
        eris_OVOV = np.asarray(eris.OVOV)
        eris_ovOV = np.asarray(eris.ovOV)
        ovov = eris_ovov - eris_ovov.transpose(0,3,2,1)
        OVOV = eris_OVOV - eris_OVOV.transpose(0,3,2,1)
        tauaa, tauab, taubb = uccsd.make_tau(t2,t1,t1)
        self.woooo += 0.5*lib.einsum('ijef,menf->minj', tauaa,      ovov)
        self.wOOOO += 0.5*lib.einsum('ijef,menf->minj', taubb,      OVOV)
        self.wooOO +=     lib.einsum('iJeF,meNF->miNJ', tauab, eris_ovOV)

        self.wooov += lib.einsum('if,mfne->mine', t1a,      ovov)
        self.wOOOV += lib.einsum('if,mfne->mine', t1b,      OVOV)
        self.wooOV += lib.einsum('if,mfNE->miNE', t1a, eris_ovOV)
        self.wOOov += lib.einsum('IF,neMF->MIne', t1b, eris_ovOV)
#X        self.wOooV -= lib.einsum('if,nfME->MinE', t1a, eris_ovOV)
#X        self.woOOv -= lib.einsum('IF,meNF->mINe', t1b, eris_ovOV)

        tmp1aa = lib.einsum('njbf,menf->mbej', t2aa,      ovov)
        tmp1aa-= lib.einsum('jNbF,meNF->mbej', t2ab, eris_ovOV)
        tmp1bb = lib.einsum('njbf,menf->mbej', t2bb,      OVOV)
        tmp1bb-= lib.einsum('nJfB,nfME->MBEJ', t2ab, eris_ovOV)
        tmp1ab = lib.einsum('NJBF,meNF->mBeJ', t2bb, eris_ovOV)
        tmp1ab-= lib.einsum('nJfB,menf->mBeJ', t2ab,      ovov)
        tmp1ba = lib.einsum('njbf,nfME->MbEj', t2aa, eris_ovOV)
        tmp1ba-= lib.einsum('jNbF,MENF->MbEj', t2ab,      OVOV)
        tmp1abba =-lib.einsum('jNfB,mfNE->mBEj', t2ab, eris_ovOV)
        tmp1baab =-lib.einsum('nJbF,neMF->MbeJ', t2ab, eris_ovOV)
        tmpaa = lib.einsum('ie,mbej->mbij', t1a, tmp1aa)
        tmpbb = lib.einsum('ie,mbej->mbij', t1b, tmp1bb)
        tmpab = lib.einsum('ie,mBeJ->mBiJ', t1a, tmp1ab)
        tmpab-= lib.einsum('IE,mBEj->mBjI', t1b, tmp1abba)
        tmpba = lib.einsum('IE,MbEj->MbIj', t1b, tmp1ba)
        tmpba-= lib.einsum('ie,MbeJ->MbJi', t1a, tmp1baab)
        wovoo -= tmpaa - tmpaa.transpose(0,1,3,2)
        wOVOO -= tmpbb - tmpbb.transpose(0,1,3,2)
        woVoO -= tmpab
        wOvOo -= tmpba
        eris_ovov = eris_OVOV = eris_ovOV = None
        eris_ovoo = np.asarray(eris.ovoo)
        eris_OVOO = np.asarray(eris.OVOO)
        eris_ovOO = np.asarray(eris.ovOO)
        eris_OVoo = np.asarray(eris.OVoo)
        wovoo += eris_ovoo.transpose(3,1,2,0) - eris_ovoo.transpose(2,1,0,3)
        wOVOO += eris_OVOO.transpose(3,1,2,0) - eris_OVOO.transpose(2,1,0,3)
        woVoO += eris_OVoo.transpose(3,1,2,0)
        wOvOo += eris_ovOO.transpose(3,1,2,0)
        eris_ovoo = eris_OVOO = eris_ovOO = eris_OVoo = None

        eris_ovvo = np.asarray(eris.ovvo)
        eris_OVVO = np.asarray(eris.OVVO)
        eris_OVvo = np.asarray(eris.OVvo)
        eris_ovVO = np.asarray(eris.ovVO)
        eris_oovv = np.asarray(eris.oovv)
        eris_OOVV = np.asarray(eris.OOVV)
        eris_OOvv = np.asarray(eris.OOvv)
        eris_ooVV = np.asarray(eris.ooVV)
        wovvo += eris_ovvo.transpose(0,2,1,3)
        wOVVO += eris_OVVO.transpose(0,2,1,3)
        woVvO += eris_ovVO.transpose(0,2,1,3)
        wOvVo += eris_OVvo.transpose(0,2,1,3)
        wovvo -= eris_oovv.transpose(0,2,3,1)
        wOVVO -= eris_OOVV.transpose(0,2,3,1)
        woVVo -= eris_ooVV.transpose(0,2,3,1)
        wOvvO -= eris_OOvv.transpose(0,2,3,1)

        tmpaa = lib.einsum('ie,mebj->mbij', t1a, eris_ovvo)
        tmpbb = lib.einsum('ie,mebj->mbij', t1b, eris_OVVO)
        tmpaa-= lib.einsum('ie,mjbe->mbij', t1a, eris_oovv)
        tmpbb-= lib.einsum('ie,mjbe->mbij', t1b, eris_OOVV)
        woVoO += lib.einsum('ie,meBJ->mBiJ', t1a, eris_ovVO)
        woVoO -= lib.einsum('IE,mjBE->mBjI',-t1b, eris_ooVV)
        wOvOo += lib.einsum('IE,MEbj->MbIj', t1b, eris_OVvo)
        wOvOo -= lib.einsum('ie,MJbe->MbJi',-t1a, eris_OOvv)
        wovoo += tmpaa - tmpaa.transpose(0,1,3,2)
        wOVOO += tmpbb - tmpbb.transpose(0,1,3,2)
        wovoo -= lib.einsum('me,ijbe->mbij', self.Fova, t2aa)
        wOVOO -= lib.einsum('me,ijbe->mbij', self.Fovb, t2bb)
        woVoO += lib.einsum('me,iJeB->mBiJ', self.Fova, t2ab)
        wOvOo += lib.einsum('ME,jIbE->MbIj', self.Fovb, t2ab)
        wovoo -= lib.einsum('nb,minj->mbij', t1a, self.woooo)
        wOVOO -= lib.einsum('nb,minj->mbij', t1b, self.wOOOO)
        woVoO -= lib.einsum('NB,miNJ->mBiJ', t1b, self.wooOO)
        wOvOo -= lib.einsum('nb,njMI->MbIj', t1a, self.wooOO)
        eris_ovvo = eris_OVVO = eris_OVvo = eris_ovVO = None
        eris_oovv = eris_OOVV = eris_OOvv = eris_ooVV = None

        self.saved = lib.H5TmpFile()
        self.saved['ovvo'] = wovvo
        self.saved['OVVO'] = wOVVO
        self.saved['oVvO'] = woVvO
        self.saved['OvVo'] = wOvVo
        self.saved['oVVo'] = woVVo
        self.saved['OvvO'] = wOvvO
        self.wovvo = self.saved['ovvo']
        self.wOVVO = self.saved['OVVO']
        self.woVvO = self.saved['oVvO']
        self.wOvVo = self.saved['OvVo']
        self.woVVo = self.saved['oVVo']
        self.wOvvO = self.saved['OvvO']
        self.saved['ovoo'] = wovoo
        self.saved['OVOO'] = wOVOO
        self.saved['oVoO'] = woVoO
        self.saved['OvOo'] = wOvOo
        self.wovoo = self.saved['ovoo']
        self.wOVOO = self.saved['OVOO']
        self.woVoO = self.saved['oVoO']
        self.wOvOo = self.saved['OvOo']

        self.wvovv = self.saved.create_dataset('vovv', (nvira,nocca,nvira,nvira), t1a.dtype.char)
        self.wVOVV = self.saved.create_dataset('VOVV', (nvirb,noccb,nvirb,nvirb), t1a.dtype.char)
        self.wvOvV = self.saved.create_dataset('vOvV', (nvira,noccb,nvira,nvirb), t1a.dtype.char)
        self.wVoVv = self.saved.create_dataset('VoVv', (nvirb,nocca,nvirb,nvira), t1a.dtype.char)

        # 3 or 4 virtuals
        eris_ovoo = np.asarray(eris.ovoo)
        eris_ovov = np.asarray(eris.ovov)
        eris_ovOV = np.asarray(eris.ovOV)
        ovov = eris_ovov - eris_ovov.transpose(0,3,2,1)
        eris_oovv = np.asarray(eris.oovv)
        eris_ovvo = np.asarray(eris.ovvo)
        oovv = eris_oovv - eris_ovvo.transpose(0,3,2,1)
        eris_oovv = eris_ovvo = None
        #:wvovv  = .5 * lib.einsum('meni,mnab->eiab', eris_ovoo, tauaa)
        #:wvovv -= .5 * lib.einsum('me,miab->eiab', self.Fova, t2aa)
        #:tmp1aa = lib.einsum('nibf,menf->mbei', t2aa,      ovov)
        #:tmp1aa-= lib.einsum('iNbF,meNF->mbei', t2ab, eris_ovOV)
        #:wvovv+= lib.einsum('ma,mbei->eiab', t1a, tmp1aa)
        #:wvovv+= lib.einsum('ma,mibe->eiab', t1a,      oovv)
        for p0, p1 in lib.prange(0, nvira, nocca):
            wvovv  = .5*lib.einsum('meni,mnab->eiab', eris_ovoo[:,p0:p1], tauaa)
            wvovv -= .5*lib.einsum('me,miab->eiab', self.Fova[:,p0:p1], t2aa)

            tmp1aa = lib.einsum('nibf,menf->mbei', t2aa, ovov[:,p0:p1])
            tmp1aa-= lib.einsum('iNbF,meNF->mbei', t2ab, eris_ovOV[:,p0:p1])
            wvovv += lib.einsum('ma,mbei->eiab', t1a, tmp1aa)
            wvovv += lib.einsum('ma,mibe->eiab', t1a, oovv[:,:,:,p0:p1])
            self.wvovv[p0:p1] = wvovv
            tmp1aa = None
        eris_ovov = eris_ovoo = eris_ovOV = None

        #:eris_ovvv = lib.unpack_tril(np.asarray(eris.ovvv).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvira,nvira)
        #:ovvv = eris_ovvv - eris_ovvv.transpose(0,3,2,1)
        #:wvovv += lib.einsum('mebf,miaf->eiab',      ovvv, t2aa)
        #:eris_OVvv = lib.unpack_tril(np.asarray(eris.OVvv).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvira,nvira)
        #:wvovv += lib.einsum('MFbe,iMaF->eiab', eris_OVvv, t2ab)
        #:wvovv += eris_ovvv.transpose(2,0,3,1).conj()
        #:self.wvovv -= wvovv - wvovv.transpose(0,1,3,2)
        mem_now = lib.current_memory()[0]
        max_memory = max(0, lib.param.MAX_MEMORY - mem_now)
        blksize = min(nocca, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvira**3*6))))
        for i0,i1 in lib.prange(0, nocca, blksize):
            wvovv = self.wvovv[:,i0:i1]
            for p0,p1 in lib.prange(0, noccb, blksize):
                OVvv = eris.get_OVvv(slice(p0,p1))  # OVvv = eris.OVvv[p0:p1]
                wvovv -= lib.einsum('MFbe,iMaF->eiab', OVvv, t2ab[i0:i1,p0:p1])
                OVvv = None
            for p0,p1 in lib.prange(0, nocca, blksize):
                ovvv = eris.get_ovvv(slice(p0,p1))  # ovvv = eris.ovvv[p0:p1]
                if p0 == i0:
                    wvovv += ovvv.transpose(2,0,3,1).conj()
                ovvv = ovvv - ovvv.transpose(0,3,2,1)
                wvovv -= lib.einsum('mebf,miaf->eiab', ovvv, t2aa[p0:p1,i0:i1])
                ovvv = None
            wvovv = wvovv - wvovv.transpose(0,1,3,2)
            self.wvovv[:,i0:i1] = wvovv

        eris_OVOO = np.asarray(eris.OVOO)
        eris_OVOV = np.asarray(eris.OVOV)
        eris_ovOV = np.asarray(eris.ovOV)
        OVOV = eris_OVOV - eris_OVOV.transpose(0,3,2,1)
        eris_OOVV = np.asarray(eris.OOVV)
        eris_OVVO = np.asarray(eris.OVVO)
        OOVV = eris_OOVV - eris_OVVO.transpose(0,3,2,1)
        eris_OOVV = eris_OVVO = None
        #:wVOVV  = .5*lib.einsum('meni,mnab->eiab', eris_OVOO, taubb)
        #:wVOVV -= .5*lib.einsum('me,miab->eiab', self.Fovb, t2bb)
        #:tmp1bb = lib.einsum('nibf,menf->mbei', t2bb,      OVOV)
        #:tmp1bb-= lib.einsum('nIfB,nfME->MBEI', t2ab, eris_ovOV)
        #:wVOVV += lib.einsum('ma,mbei->eiab', t1b, tmp1bb)
        #:wVOVV += lib.einsum('ma,mibe->eiab', t1b,      OOVV)
        for p0, p1 in lib.prange(0, nvirb, noccb):
            wVOVV  = .5*lib.einsum('meni,mnab->eiab', eris_OVOO[:,p0:p1], taubb)
            wVOVV -= .5*lib.einsum('me,miab->eiab', self.Fovb[:,p0:p1], t2bb)

            tmp1bb = lib.einsum('nibf,menf->mbei', t2bb, OVOV[:,p0:p1])
            tmp1bb-= lib.einsum('nIfB,nfME->MBEI', t2ab, eris_ovOV[:,:,:,p0:p1])
            wVOVV += lib.einsum('ma,mbei->eiab', t1b, tmp1bb)
            wVOVV += lib.einsum('ma,mibe->eiab', t1b, OOVV[:,:,:,p0:p1])
            self.wVOVV[p0:p1] = wVOVV
            tmp1bb = None
        eris_OVOV = eris_OVOO = eris_ovOV = None

        #:eris_OVVV = lib.unpack_tril(np.asarray(eris.OVVV).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvirb,nvirb)
        #:OVVV = eris_OVVV - eris_OVVV.transpose(0,3,2,1)
        #:wVOVV -= lib.einsum('MEBF,MIAF->EIAB',      OVVV, t2bb)
        #:eris_ovVV = lib.unpack_tril(np.asarray(eris.ovVV).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvirb,nvirb)
        #:wVOVV -= lib.einsum('mfBE,mIfA->EIAB', eris_ovVV, t2ab)
        #:wVOVV += eris_OVVV.transpose(2,0,3,1).conj()
        #:self.wVOVV += wVOVV - wVOVV.transpose(0,1,3,2)
        blksize = min(noccb, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvirb**3*6))))
        for i0,i1 in lib.prange(0, noccb, blksize):
            wVOVV = self.wVOVV[:,i0:i1]
            for p0,p1 in lib.prange(0, nocca, blksize):
                ovVV = eris.get_ovVV(slice(p0,p1))  # ovVV = eris.ovVV[p0:p1]
                wVOVV -= lib.einsum('mfBE,mIfA->EIAB', ovVV, t2ab[p0:p1,i0:i1])
                ovVV = None
            for p0,p1 in lib.prange(0, noccb, blksize):
                OVVV = eris.get_OVVV(slice(p0,p1))  # OVVV = eris.OVVV[p0:p1]
                if p0 == i0:
                    wVOVV += OVVV.transpose(2,0,3,1).conj()
                OVVV = OVVV - OVVV.transpose(0,3,2,1)
                wVOVV -= lib.einsum('mebf,miaf->eiab', OVVV, t2bb[p0:p1,i0:i1])
                OVVV = None
            wVOVV = wVOVV - wVOVV.transpose(0,1,3,2)
            self.wVOVV[:,i0:i1] = wVOVV

        eris_ovOV = np.asarray(eris.ovOV)
        eris_ovOO = np.asarray(eris.ovOO)
        eris_OOvv = np.asarray(eris.OOvv)
        eris_ovVO = np.asarray(eris.ovVO)
        #:self.wvOvV = lib.einsum('meNI,mNaB->eIaB', eris_ovOO, tauab)
        #:self.wvOvV -= lib.einsum('me,mIaB->eIaB', self.Fova, t2ab)
        #:tmp1ab = lib.einsum('NIBF,meNF->mBeI', t2bb, eris_ovOV)
        #:tmp1ab-= lib.einsum('nIfB,menf->mBeI', t2ab,      ovov)
        #:tmp1baab = lib.einsum('nIbF,neMF->MbeI', t2ab, eris_ovOV)
        #:tmpab = lib.einsum('ma,mBeI->eIaB', t1a, tmp1ab)
        #:tmpab+= lib.einsum('MA,MbeI->eIbA', t1b, tmp1baab)
        #:tmpab-= lib.einsum('MA,MIbe->eIbA', t1b, eris_OOvv)
        #:tmpab-= lib.einsum('ma,meBI->eIaB', t1a, eris_ovVO)
        #:self.wvOvV += tmpab
        for p0, p1 in lib.prange(0, nvira, nocca):
            wvOvV  = lib.einsum('meNI,mNaB->eIaB', eris_ovOO[:,p0:p1], tauab)
            wvOvV -= lib.einsum('me,mIaB->eIaB', self.Fova[:,p0:p1], t2ab)
            tmp1ab = lib.einsum('NIBF,meNF->mBeI', t2bb, eris_ovOV[:,p0:p1])
            tmp1ab-= lib.einsum('nIfB,menf->mBeI', t2ab, ovov[:,p0:p1])
            wvOvV+= lib.einsum('ma,mBeI->eIaB', t1a, tmp1ab)
            tmp1ab = None
            tmp1baab = lib.einsum('nIbF,neMF->MbeI', t2ab, eris_ovOV[:,p0:p1])
            wvOvV+= lib.einsum('MA,MbeI->eIbA', t1b, tmp1baab)
            tmp1baab = None
            wvOvV-= lib.einsum('MA,MIbe->eIbA', t1b, eris_OOvv[:,:,:,p0:p1])
            wvOvV-= lib.einsum('ma,meBI->eIaB', t1a, eris_ovVO[:,p0:p1])
            self.wvOvV[p0:p1] = wvOvV
        eris_ovOV = eris_ovOO = eris_OOvv = eris_ovVO = None

        #:eris_ovvv = lib.unpack_tril(np.asarray(eris.ovvv).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvira,nvira)
        #:ovvv = eris_ovvv - eris_ovvv.transpose(0,3,2,1)
        #:self.wvOvV -= lib.einsum('mebf,mIfA->eIbA',      ovvv, t2ab)
        #:eris_ovVV = lib.unpack_tril(np.asarray(eris.ovVV).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvirb,nvirb)
        #:self.wvOvV -= lib.einsum('meBF,mIaF->eIaB', eris_ovVV, t2ab)
        #:eris_OVvv = lib.unpack_tril(np.asarray(eris.OVvv).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvira,nvira)
        #:self.wvOvV -= lib.einsum('MFbe,MIAF->eIbA', eris_OVvv, t2bb)
        #:self.wvOvV += eris_OVvv.transpose(2,0,3,1).conj()
        blksize = min(noccb, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvira**3*6))))
        for i0,i1 in lib.prange(0, noccb, blksize):
            wvOvV = self.wvOvV[:,i0:i1]
            for p0,p1 in lib.prange(0, nocca, blksize):
                ovVV = eris.get_ovVV(slice(p0,p1))  # ovVV = eris.ovVV[p0:p1]
                wvOvV -= lib.einsum('meBF,mIaF->eIaB', ovVV, t2ab[p0:p1,i0:i1])
                ovVV = None
            for p0,p1 in lib.prange(0, nocca, blksize):
                ovvv = eris.get_ovvv(slice(p0,p1))  # ovvv = eris.ovvv[p0:p1]
                ovvv = ovvv - ovvv.transpose(0,3,2,1)
                wvOvV -= lib.einsum('mebf,mIfA->eIbA',ovvv, t2ab[p0:p1,i0:i1])
                ovvv = None
            self.wvOvV[:,i0:i1] = wvOvV

        blksize = min(noccb, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvirb*nvira**2*3))))
        for i0,i1 in lib.prange(0, noccb, blksize):
            wvOvV = self.wvOvV[:,i0:i1]
            for p0,p1 in lib.prange(0, noccb, blksize):
                OVvv = eris.get_OVvv(slice(p0,p1))  # OVvv = eris.OVvv[p0:p1]
                if p0 == i0:
                    wvOvV += OVvv.transpose(2,0,3,1).conj()
                wvOvV -= lib.einsum('MFbe,MIAF->eIbA', OVvv, t2bb[p0:p1,i0:i1])
                OVvv = None
            self.wvOvV[:,i0:i1] = wvOvV

        eris_ovOV = np.asarray(eris.ovOV)
        eris_OVoo = np.asarray(eris.OVoo)
        eris_ooVV = np.asarray(eris.ooVV)
        eris_OVvo = np.asarray(eris.OVvo)
        #:self.wVoVv = lib.einsum('MEni,nMbA->EiAb', eris_OVoo, tauab)
        #:self.wVoVv -= lib.einsum('ME,iMbA->EiAb', self.Fovb, t2ab)
        #:tmp1ba = lib.einsum('nibf,nfME->MbEi', t2aa, eris_ovOV)
        #:tmp1ba-= lib.einsum('iNbF,MENF->MbEi', t2ab,      OVOV)
        #:tmp1abba = lib.einsum('iNfB,mfNE->mBEi', t2ab, eris_ovOV)
        #:tmpba = lib.einsum('MA,MbEi->EiAb', t1b, tmp1ba)
        #:tmpba+= lib.einsum('ma,mBEi->EiBa', t1a, tmp1abba)
        #:tmpba-= lib.einsum('ma,miBE->EiBa', t1a, eris_ooVV)
        #:tmpba-= lib.einsum('MA,MEbi->EiAb', t1b, eris_OVvo)
        #:self.wVoVv += tmpba
        for p0, p1 in lib.prange(0, nvirb, noccb):
            wVoVv  = lib.einsum('MEni,nMbA->EiAb', eris_OVoo[:,p0:p1], tauab)
            wVoVv -= lib.einsum('ME,iMbA->EiAb', self.Fovb[:,p0:p1], t2ab)
            tmp1ba = lib.einsum('nibf,nfME->MbEi', t2aa, eris_ovOV[:,:,:,p0:p1])
            tmp1ba-= lib.einsum('iNbF,MENF->MbEi', t2ab, OVOV[:,p0:p1])
            wVoVv += lib.einsum('MA,MbEi->EiAb', t1b, tmp1ba)
            tmp1ba = None
            tmp1abba = lib.einsum('iNfB,mfNE->mBEi', t2ab, eris_ovOV[:,:,:,p0:p1])
            wVoVv += lib.einsum('ma,mBEi->EiBa', t1a, tmp1abba)
            tmp1abba = None
            wVoVv -= lib.einsum('ma,miBE->EiBa', t1a, eris_ooVV[:,:,:,p0:p1])
            wVoVv -= lib.einsum('MA,MEbi->EiAb', t1b, eris_OVvo[:,p0:p1])
            self.wVoVv[p0:p1] = wVoVv
        eris_ovOV = eris_OVoo = eris_ooVV = eris_OVvo = None

        #:eris_OVVV = lib.unpack_tril(np.asarray(eris.OVVV).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvirb,nvirb)
        #:OVVV = eris_OVVV - eris_OVVV.transpose(0,3,2,1)
        #:self.wVoVv -= lib.einsum('MEBF,iMaF->EiBa',      OVVV, t2ab)
        #:eris_OVvv = lib.unpack_tril(np.asarray(eris.OVvv).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvira,nvira)
        #:self.wVoVv -= lib.einsum('MEbf,iMfA->EiAb', eris_OVvv, t2ab)
        #:eris_ovVV = lib.unpack_tril(np.asarray(eris.ovVV).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvirb,nvirb)
        #:self.wVoVv -= lib.einsum('mfBE,miaf->EiBa', eris_ovVV, t2aa)
        #:self.wVoVv += eris_ovVV.transpose(2,0,3,1).conj()
        blksize = min(noccb, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvirb**3*6))))
        for i0,i1 in lib.prange(0, nocca, blksize):
            wVoVv = self.wVoVv[:,i0:i1]
            for p0,p1 in lib.prange(0, noccb, blksize):
                OVvv = eris.get_OVvv(slice(p0,p1))  # OVvv = eris.OVvv[p0:p1]
                wVoVv -= lib.einsum('MEbf,iMfA->EiAb', OVvv, t2ab[i0:i1,p0:p1])
                OVvv = None
            for p0,p1 in lib.prange(0, noccb, blksize):
                OVVV = eris.get_OVVV(slice(p0,p1))  # OVVV = eris.OVVV[p0:p1]
                OVVV = OVVV - OVVV.transpose(0,3,2,1)
                wVoVv -= lib.einsum('MEBF,iMaF->EiBa', OVVV, t2ab[i0:i1,p0:p1])
                OVVV = None
            self.wVoVv[:,i0:i1] = wVoVv

        blksize = min(nocca, max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvira*nvirb**2*3))))
        for i0,i1 in lib.prange(0, nocca, blksize):
            wVoVv = self.wVoVv[:,i0:i1]
            for p0,p1 in lib.prange(0, nocca, blksize):
                ovVV = eris.get_ovVV(slice(p0,p1))  # ovVV = eris.ovVV[p0:p1]
                if p0 == i0:
                    wVoVv += ovVV.transpose(2,0,3,1).conj()
                wVoVv -= lib.einsum('mfBE,miaf->EiBa', ovVV, t2aa[p0:p1,i0:i1])
                ovVV = None
            self.wVoVv[:,i0:i1] = wVoVv

        self.made_ee_imds = True
        log.timer('EOM-UCCSD EE intermediates', *cput0)

def rand_mf(mol, seed=1):
    from pyscf import scf
    from pyscf import gto
    from pyscf import lo
    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = 'sto-3g'
    mol.verbose = 0
    mol.spin = 0
    mol.build()

    np.random.seed(seed)

    mf = scf.UHF(mol).run(conv_tol=1e-14)
    nmo = mol.nao_nr()
    mf.mo_occ = np.zeros((2,nmo))
    mf.mo_occ[0,:4] = 1
    mf.mo_occ[1,:2] = 1
    mf.mo_energy = np.arange(nmo) + np.random.random((2,nmo)) * .3
    mf.mo_energy[mf.mo_occ == 0] += 2

    mo = np.random.random((2,nmo,nmo))
    s = mf.get_ovlp()
    mf.mo_coeff = np.empty_like(mo)
    mf.mo_coeff[0] = lo.orth.vec_lowdin(mo[0], s)
    mf.mo_coeff[1] = lo.orth.vec_lowdin(mo[1], s)
    return mf

def rand_cc_t1_t2(mf, seed=1):
    from pyscf import ao2mo
    from pyscf.cc import uccsd
    mycc = uccsd.UCCSD(mf)

    nocca, noccb = mycc.nocc
    nmoa, nmob = mycc.nmo
    nvira, nvirb = nmoa - nocca, nmob - noccb
    def my_ao2mo(mo):
        eris = ao2mo.kernel(mycc._scf._eri, mo, compact=False)
        eris = ao2mo.restore(1, eris, mf.mol.nao_nr())

        eris = eris + np.cos(eris)*1j
        eris = eris + eris.transpose(1, 0, 3, 2)
        eris = eris + eris.conj().transpose(2, 3, 0, 1)
        return eris
    eris = uccsd._make_eris_incore(mycc)#, ao2mofn=my_ao2mo)

    np.random.seed(seed)
    t1a = (np.random.random((nocca,nvira)) +
           np.random.random((nocca,nvira))*1j - .5 - .5j)
    t1b = (np.random.random((noccb,nvirb)) +
           np.random.random((noccb,nvirb))*1j - .5 - .5j)
    t2aa = (np.random.random((nocca,nocca,nvira,nvira)) +
            np.random.random((nocca,nocca,nvira,nvira))*1j - .5 - .5j)
    t2aa = t2aa - t2aa.transpose(1, 0, 2, 3)
    t2aa = t2aa - t2aa.transpose(0, 1, 3, 2)

    t2ab = (np.random.random((nocca,noccb,nvira,nvirb)) +
            np.random.random((nocca,noccb,nvira,nvirb))*1j - .5 - .5j)

    t2bb = (np.random.random((noccb,noccb,nvirb,nvirb)) +
            np.random.random((noccb,noccb,nvirb,nvirb))*1j - .5 - .5j)
    t2bb = t2bb - t2bb.transpose(1, 0, 2, 3)
    t2bb = t2bb - t2bb.transpose(0, 1, 3, 2)

    t1 = (t1a, t1b)
    t2 = (t2aa, t2ab, t2bb)
    return mycc, eris, t1, t2

def enforce_symm_2p_spin(r1, r2, orbspin, excitation):
    assert(excitation in ['ip', 'ea'])
    if excitation == 'ip':
        nocc, nvir = r2.shape[1:]
    elif excitation == 'ea':
        nocc, nvir = r2.shape[:2]
    else:
        raise NotImplementedError

    idxoa = np.where(orbspin[:nocc] == 0)[0]
    idxob = np.where(orbspin[:nocc] == 1)[0]
    idxva = np.where(orbspin[nocc:] == 0)[0]
    idxvb = np.where(orbspin[nocc:] == 1)[0]

    idxoaa = idxoa[:,None] * nocc + idxoa
    idxobb = idxob[:,None] * nocc + idxob
    idxvaa = idxva[:,None] * nvir + idxva
    idxvbb = idxvb[:,None] * nvir + idxvb

    if excitation == 'ip':
        r2 = r2 - r2.transpose(1, 0, 2)

        r2 = r2.reshape(nocc**2, nvir)
        r2[idxobb.ravel()[:, None], idxva.ravel()] = 0.0
        r2[idxoaa.ravel()[:, None], idxvb.ravel()] = 0.0
        r2 = r2.reshape(nocc, nocc, nvir)

    if excitation == 'ea':
        r2 = r2 - r2.transpose(0, 2, 1)

        r2 = r2.reshape(nocc, nvir**2)
        r2[idxoa.ravel(), idxvbb.ravel()[:, None]] = 0.0
        r2[idxob.ravel(), idxvaa.ravel()[:, None]] = 0.0
        r2 = r2.reshape(nocc, nvir, nvir)

    return r1, r2

def enforce_symm_2p_spin_ip(r1, r2, orbspin):
    return enforce_symm_2p_spin(r1, r2, orbspin, 'ip')

def enforce_symm_2p_spin_ea(r1, r2, orbspin):
    return enforce_symm_2p_spin(r1, r2, orbspin, 'ea')

if __name__ == '__main__':
    from pyscf import gto
    #from pyscf import scf
    #from pyscf.cc import rccsd

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = 'sto-3g'
    mol.verbose = 0
    mol.spin = 0
    mol.build()

    mf = rand_mf(mol)
    mycc, eris, t1, t2 = rand_cc_t1_t2(mf)
    mycc.t1 = t1
    mycc.t2 = t2

    nocca, noccb = mycc.nocc
    nmoa, nmob = mycc.nmo
    nvira, nvirb = nmoa - nocca, nmob - noccb
    nocc = nocca + noccb
    nvir = nvira + nvirb
    nmo = nocc + nvir

    def my_ao2mo(mo):
        nao, nmo = mo.shape
        orbspin = mo.orbspin

#        eris = ao2mo.kernel(mygcc._scf._eri, mo_a + mo_b)
#        sym_forbid = (orbspin[:,None] != orbspin)[np.tril_indices(nmo)]
#        eris[sym_forbid,:] = 0
#        eris[:,sym_forbid] = 0
#        eris = ao2mo.restore(1, eris, nao)
#        return eris
        eris =(np.random.random((nmo,nmo,nmo,nmo)) +
               np.random.random((nmo,nmo,nmo,nmo)) * 1j)

        eris = eris + np.cos(eris)*1j
        eris = eris + eris.transpose(1, 0, 3, 2)
        eris = eris + eris.conj().transpose(2, 3, 0, 1)
        eris[orbspin[:,None] != orbspin] = 0
        eris[:,:,orbspin[:,None] != orbspin] = 0
        return eris

    import pyscf.cc.addons
    from pyscf.cc import gccsd
    mygcc = pyscf.cc.addons.convert_to_gccsd(mycc)
    mygcc._ucc = mycc
    mygcc._ucc_eris = eris
    eris = gccsd._make_eris_incore(mygcc)#, ao2mofn=my_ao2mo)
    orbspin = eris.orbspin

    ## EOM-IP
    myeom = EOMIP(mycc)
    imds = myeom.make_imds()

    np.random.seed(1)
    r1 = np.random.rand(nocc)*1j + np.random.rand(nocc) - 0.5 - 0.5*1j
    r2 = np.random.rand(nocc**2 * nvir)*1j + np.random.rand(nocc**2 * nvir) - 0.5 - 0.5*1j
    r2 = r2.reshape(nocc, nocc, nvir)
    r1, r2 = enforce_symm_2p_spin_ip(r1, r2, orbspin)
    r1, r2 = spin2spatial_ip(r1, r2, orbspin)

    vector = myeom.amplitudes_to_vector(r1, r2)
    r1x, r2x = myeom.vector_to_amplitudes(vector)
    print(abs(r1[0]-r1x[0]).max() < 1e-13 and
          abs(r1[1]-r1x[1]).max() < 1e-13 and
          abs(r2[0]-r2x[0]).max() < 1e-13 and
          abs(r2[1]-r2x[1]).max() < 1e-13 and
          abs(r2[2]-r2x[2]).max() < 1e-13 and
          abs(r2[3]-r2x[3]).max() < 1e-13)
    Hvector = myeom.matvec(vector, imds=imds)
    print('ip', lib.finger(Hvector) - (21.67127462317093-19.068987454261908j))
    print('diag', lib.finger(myeom.get_diag()) - (-9.6676217223549763+9.325219825942975j))

    # EOM-EA
    myeom = EOMEA(mycc)
    imds = myeom.make_imds()

    np.random.seed(1)
    r1 = np.random.rand(nvir)*1j + np.random.rand(nvir) - 0.5 - 0.5*1j
    r2 = np.random.rand(nocc * nvir**2)*1j + np.random.rand(nocc * nvir**2) - 0.5 - 0.5*1j
    r2 = r2.reshape(nocc, nvir, nvir)
    r1, r2 = enforce_symm_2p_spin_ea(r1, r2, orbspin)
    r1, r2 = spin2spatial_ea(r1, r2, orbspin)

    vector = myeom.amplitudes_to_vector(r1, r2)
    r1x, r2x = myeom.vector_to_amplitudes(vector)
    print(abs(r1[0]-r1x[0]).max() < 1e-13 and
          abs(r1[1]-r1x[1]).max() < 1e-13 and
          abs(r2[0]-r2x[0]).max() < 1e-13 and
          abs(r2[1]-r2x[1]).max() < 1e-13 and
          abs(r2[2]-r2x[2]).max() < 1e-13 and
          abs(r2[3]-r2x[3]).max() < 1e-13)
    Hvector = myeom.matvec(vector, imds=imds)
    print('ea', lib.finger(Hvector) - (6.5543877287461187-13.175055314063574j))
    print('diag', lib.finger(myeom.get_diag()) - (-57.353207240857785+1.4052857730841204j))

    mycc = uccsd.UCCSD(mol.UHF().run())
    ecc, t1, t2 = mycc.kernel()
    print(ecc - -0.04946750711013597)
    e,v = mycc.ipccsd(nroots=6)
    print(e[0] - 0.3092874511803249)
    print(e[1] - 0.3092874511803249)
    print(e[2] - 0.4011171373779585)
    print(e[3] - 0.4011171373779585)
    print(e[4] - 0.6107409208314764)
    print(e[5] - 0.6107409208314764)


# pyscf/cc/gccsd.py

#!/usr/bin/env python
# Copyright 2014-2021 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.


import numpy as np
from functools import reduce

from pyscf import lib
from pyscf import ao2mo
from pyscf import scf
from pyscf.lib import logger
from pyscf.cc import ccsd
from pyscf.cc import gintermediates as imd
from pyscf.cc.addons import spatial2spin, spin2spatial
from pyscf import __config__

MEMORYMIN = getattr(__config__, 'cc_ccsd_memorymin', 2000)

#einsum = np.einsum
einsum = lib.einsum

# This is unrestricted (U)CCSD in spin-orbital form.

def update_amps(cc, t1, t2, eris):
    assert(isinstance(eris, _PhysicistsERIs))
    nocc, nvir = t1.shape
    fock = eris.fock

    fov = fock[:nocc,nocc:]
    mo_e_o = eris.mo_energy[:nocc]
    mo_e_v = eris.mo_energy[nocc:] + cc.level_shift

    tau = imd.make_tau(t2, t1, t1)

    Fvv = imd.cc_Fvv(t1, t2, eris)
    Foo = imd.cc_Foo(t1, t2, eris)
    Fov = imd.cc_Fov(t1, t2, eris)
    Woooo = imd.cc_Woooo(t1, t2, eris)
    Wvvvv = imd.cc_Wvvvv(t1, t2, eris)
    Wovvo = imd.cc_Wovvo(t1, t2, eris)

    # Move energy terms to the other side
    Fvv[np.diag_indices(nvir)] -= mo_e_v
    Foo[np.diag_indices(nocc)] -= mo_e_o

    # T1 equation
    t1new  =  einsum('ie,ae->ia', t1, Fvv)
    t1new += -einsum('ma,mi->ia', t1, Foo)
    t1new +=  einsum('imae,me->ia', t2, Fov)
    t1new += -einsum('nf,naif->ia', t1, eris.ovov)
    t1new += -0.5*einsum('imef,maef->ia', t2, eris.ovvv)
    t1new += -0.5*einsum('mnae,mnie->ia', t2, eris.ooov)
    t1new += fov.conj()

    # T2 equation
    Ftmp = Fvv - 0.5*einsum('mb,me->be', t1, Fov)
    tmp = einsum('ijae,be->ijab', t2, Ftmp)
    t2new = tmp - tmp.transpose(0,1,3,2)
    Ftmp = Foo + 0.5*einsum('je,me->mj', t1, Fov)
    tmp = einsum('imab,mj->ijab', t2, Ftmp)
    t2new -= tmp - tmp.transpose(1,0,2,3)
    t2new += np.asarray(eris.oovv).conj()
    t2new += 0.5*einsum('mnab,mnij->ijab', tau, Woooo)
    t2new += 0.5*einsum('ijef,abef->ijab', tau, Wvvvv)
    tmp = einsum('imae,mbej->ijab', t2, Wovvo)
    tmp -= -einsum('ie,ma,mbje->ijab', t1, t1, eris.ovov)
    tmp = tmp - tmp.transpose(1,0,2,3)
    tmp = tmp - tmp.transpose(0,1,3,2)
    t2new += tmp
    tmp = einsum('ie,jeba->ijab', t1, np.array(eris.ovvv).conj())
    t2new += (tmp - tmp.transpose(1,0,2,3))
    tmp = einsum('ma,ijmb->ijab', t1, np.asarray(eris.ooov).conj())
    t2new -= (tmp - tmp.transpose(0,1,3,2))

    eia = mo_e_o[:,None] - mo_e_v
    eijab = lib.direct_sum('ia,jb->ijab', eia, eia)
    t1new /= eia
    t2new /= eijab

    return t1new, t2new


def energy(cc, t1, t2, eris):
    nocc, nvir = t1.shape
    fock = eris.fock
    e = einsum('ia,ia', fock[:nocc,nocc:], t1)
    eris_oovv = np.array(eris.oovv)
    e += 0.25*np.einsum('ijab,ijab', t2, eris_oovv)
    e += 0.5 *np.einsum('ia,jb,ijab', t1, t1, eris_oovv)
    if abs(e.imag) > 1e-4:
        logger.warn(cc, 'Non-zero imaginary part found in GCCSD energy %s', e)
    return e.real

vector_to_amplitudes = ccsd.vector_to_amplitudes_s4
amplitudes_to_vector = ccsd.amplitudes_to_vector_s4

def amplitudes_from_rccsd(t1, t2, orbspin=None):
    '''Convert spatial orbital T1,T2 to spin-orbital T1,T2'''
    return spatial2spin(t1, orbspin), spatial2spin(t2, orbspin)


class GCCSD(ccsd.CCSD):

    conv_tol = getattr(__config__, 'cc_gccsd_GCCSD_conv_tol', 1e-7)
    conv_tol_normt = getattr(__config__, 'cc_gccsd_GCCSD_conv_tol_normt', 1e-6)

    def __init__(self, mf, frozen=None, mo_coeff=None, mo_occ=None):
        assert(isinstance(mf, scf.ghf.GHF))
        ccsd.CCSD.__init__(self, mf, frozen, mo_coeff, mo_occ)

    def init_amps(self, eris=None):
        if eris is None:
            eris = self.ao2mo(self.mo_coeff)
        mo_e = eris.mo_energy
        nocc = self.nocc
        eia = mo_e[:nocc,None] - mo_e[None,nocc:]
        eijab = lib.direct_sum('ia,jb->ijab', eia, eia)
        t1 = eris.fock[:nocc,nocc:] / eia
        eris_oovv = np.array(eris.oovv)
        t2 = eris_oovv / eijab
        self.emp2 = 0.25*einsum('ijab,ijab', t2, eris_oovv.conj()).real
        logger.info(self, 'Init t2, MP2 energy = %.15g', self.emp2)
        return self.emp2, t1, t2

    energy = energy
    update_amps = update_amps

    def kernel(self, t1=None, t2=None, eris=None, mbpt2=False):
        return self.ccsd(t1, t2, eris, mbpt2=mbpt2)
    def ccsd(self, t1=None, t2=None, eris=None, mbpt2=False):
        '''Ground-state unrestricted (U)CCSD.

        Kwargs:
            mbpt2 : bool
                Use one-shot MBPT2 approximation to CCSD.
        '''
        if mbpt2:
            from pyscf.mp import gmp2
            pt = gmp2.GMP2(self._scf, self.frozen, self.mo_coeff, self.mo_occ)
            self.e_corr, self.t2 = pt.kernel(eris=eris)
            nocc, nvir = self.t2.shape[1:3]
            self.t1 = np.zeros((nocc,nvir))
            return self.e_corr, self.t1, self.t2

        # Initialize orbspin so that we can attach it to t1, t2
        if getattr(self.mo_coeff, 'orbspin', None) is None:
            orbspin = scf.ghf.guess_orbspin(self.mo_coeff)
            if not np.any(orbspin == -1):
                self.mo_coeff = lib.tag_array(self.mo_coeff, orbspin=orbspin)

        e_corr, self.t1, self.t2 = ccsd.CCSD.ccsd(self, t1, t2, eris)
        if getattr(eris, 'orbspin', None) is not None:
            self.t1 = lib.tag_array(self.t1, orbspin=eris.orbspin)
            self.t2 = lib.tag_array(self.t2, orbspin=eris.orbspin)
        return e_corr, self.t1, self.t2

    def amplitudes_to_vector(self, t1, t2, out=None):
        return amplitudes_to_vector(t1, t2, out)

    def vector_to_amplitudes(self, vec, nmo=None, nocc=None):
        if nocc is None: nocc = self.nocc
        if nmo is None: nmo = self.nmo
        return vector_to_amplitudes(vec, nmo, nocc)

    def vector_size(self, nmo=None, nocc=None):
        if nocc is None: nocc = self.nocc
        if nmo is None: nmo = self.nmo
        nvir = nmo - nocc
        return nocc * nvir + nocc*(nocc-1)//2*nvir*(nvir-1)//2

    def amplitudes_from_rccsd(self, t1, t2, orbspin=None):
        return amplitudes_from_rccsd(t1, t2, orbspin)


    def solve_lambda(self, t1=None, t2=None, l1=None, l2=None,
                     eris=None):
        from pyscf.cc import gccsd_lambda
        if t1 is None: t1 = self.t1
        if t2 is None: t2 = self.t2
        if eris is None: eris = self.ao2mo(self.mo_coeff)
        self.converged_lambda, self.l1, self.l2 = \
                gccsd_lambda.kernel(self, eris, t1, t2, l1, l2,
                                    max_cycle=self.max_cycle,
                                    tol=self.conv_tol_normt,
                                    verbose=self.verbose)
        return self.l1, self.l2

    def ccsd_t(self, t1=None, t2=None, eris=None):
        from pyscf.cc import gccsd_t
        if t1 is None: t1 = self.t1
        if t2 is None: t2 = self.t2
        if eris is None: eris = self.ao2mo(self.mo_coeff)
        return gccsd_t.kernel(self, eris, t1, t2, self.verbose)

    def ipccsd(self, nroots=1, left=False, koopmans=False, guess=None,
               partition=None, eris=None):
        from pyscf.cc import eom_gccsd
        return eom_gccsd.EOMIP(self).kernel(nroots, left, koopmans, guess,
                                            partition, eris)

    def eaccsd(self, nroots=1, left=False, koopmans=False, guess=None,
               partition=None, eris=None):
        from pyscf.cc import eom_gccsd
        return eom_gccsd.EOMEA(self).kernel(nroots, left, koopmans, guess,
                                            partition, eris)

    def eeccsd(self, nroots=1, koopmans=False, guess=None, eris=None):
        from pyscf.cc import eom_gccsd
        return eom_gccsd.EOMEE(self).kernel(nroots, koopmans, guess, eris)

    def eomip_method(self):
        from pyscf.cc import eom_gccsd
        return eom_gccsd.EOMIP(self)

    def eomea_method(self):
        from pyscf.cc import eom_gccsd
        return eom_gccsd.EOMEA(self)

    def eomee_method(self):
        from pyscf.cc import eom_gccsd
        return eom_gccsd.EOMEE(self)

    def make_rdm1(self, t1=None, t2=None, l1=None, l2=None, ao_repr=False,
                  with_frozen=True, with_mf=True):
        '''Un-relaxed 1-particle density matrix in MO space'''
        from pyscf.cc import gccsd_rdm
        if t1 is None: t1 = self.t1
        if t2 is None: t2 = self.t2
        if l1 is None: l1 = self.l1
        if l2 is None: l2 = self.l2
        if l1 is None: l1, l2 = self.solve_lambda(t1, t2)
        return gccsd_rdm.make_rdm1(self, t1, t2, l1, l2, ao_repr=ao_repr,
                                   with_frozen=with_frozen, with_mf=with_mf)

    def make_rdm2(self, t1=None, t2=None, l1=None, l2=None, ao_repr=False,
                  with_frozen=True, with_dm1=True):
        '''2-particle density matrix in MO space.  The density matrix is
        stored as

        dm2[p,r,q,s] = <p^+ q^+ s r>
        '''
        from pyscf.cc import gccsd_rdm
        if t1 is None: t1 = self.t1
        if t2 is None: t2 = self.t2
        if l1 is None: l1 = self.l1
        if l2 is None: l2 = self.l2
        if l1 is None: l1, l2 = self.solve_lambda(t1, t2)
        return gccsd_rdm.make_rdm2(self, t1, t2, l1, l2, ao_repr=ao_repr,
                                   with_frozen=with_frozen, with_dm1=with_dm1)

    def ao2mo(self, mo_coeff=None):
        nmo = self.nmo
        mem_incore = nmo**4*2 * 8/1e6
        mem_now = lib.current_memory()[0]
        if (self._scf._eri is not None and
            (mem_incore+mem_now < self.max_memory) or self.mol.incore_anyway):
            return _make_eris_incore(self, mo_coeff)
        elif getattr(self._scf, 'with_df', None):
            raise NotImplementedError
        else:
            return _make_eris_outcore(self, mo_coeff)

    def amplitudes_from_ccsd(self, t1, t2):
        '''Convert spatial orbital T1,T2 to spin-orbital T1,T2'''
        return self.spatial2spin(t1), self.spatial2spin(t2)

    def spatial2spin(self, tx, orbspin=None):
        if orbspin is None:
            orbspin = getattr(self.mo_coeff, 'orbspin', None)
            if orbspin is not None:
                orbspin = orbspin[self.get_frozen_mask()]
        return spatial2spin(tx, orbspin)

    def spin2spatial(self, tx, orbspin=None):
        if orbspin is None:
            orbspin = getattr(self.mo_coeff, 'orbspin', None)
            if orbspin is not None:
                orbspin = orbspin[self.get_frozen_mask()]
        return spin2spatial(tx, orbspin)

    def density_fit(self):
        raise NotImplementedError

    def nuc_grad_method(self):
        raise NotImplementedError

    def get_t1_diagnostic(self, t1=None):
        if t1 is None: t1 = self.t1
        raise NotImplementedError
        #return get_t1_diagnostic(t1)

    def get_d1_diagnostic(self, t1=None):
        if t1 is None: t1 = self.t1
        raise NotImplementedError
        #return get_d1_diagnostic(t1)

    def get_d2_diagnostic(self, t2=None):
        if t2 is None: t2 = self.t2
        raise NotImplementedError
        #return get_d2_diagnostic(self.spin2spatial(t2))

CCSD = GCCSD


class _PhysicistsERIs:
    '''<pq||rs> = <pq|rs> - <pq|sr>'''
    def __init__(self, mol=None):
        self.mol = mol
        self.mo_coeff = None
        self.nocc = None
        self.fock = None
        self.e_hf = None
        self.orbspin = None

        self.oooo = None
        self.ooov = None
        self.oovv = None
        self.ovvo = None
        self.ovov = None
        self.ovvv = None
        self.vvvv = None

    def _common_init_(self, mycc, mo_coeff=None):
        if mo_coeff is None:
            mo_coeff = mycc.mo_coeff
        mo_idx = ccsd.get_frozen_mask(mycc)
        if getattr(mo_coeff, 'orbspin', None) is not None:
            self.orbspin = mo_coeff.orbspin[mo_idx]
            mo_coeff = lib.tag_array(mo_coeff[:,mo_idx], orbspin=self.orbspin)
        else:
            orbspin = scf.ghf.guess_orbspin(mo_coeff)
            mo_coeff = mo_coeff[:,mo_idx]
            if not np.any(orbspin == -1):
                self.orbspin = orbspin[mo_idx]
                mo_coeff = lib.tag_array(mo_coeff, orbspin=self.orbspin)
        self.mo_coeff = mo_coeff

        # Note: Recomputed fock matrix since SCF may not be fully converged.
        dm = mycc._scf.make_rdm1(mycc.mo_coeff, mycc.mo_occ)
        vhf = mycc._scf.get_veff(mycc.mol, dm)
        fockao = mycc._scf.get_fock(vhf=vhf, dm=dm)
        self.fock = reduce(np.dot, (mo_coeff.conj().T, fockao, mo_coeff))
        self.e_hf = mycc._scf.energy_tot(dm=dm, vhf=vhf)
        self.nocc = mycc.nocc
        self.mol = mycc.mol

        mo_e = self.mo_energy = self.fock.diagonal().real
        gap = abs(mo_e[:self.nocc,None] - mo_e[None,self.nocc:]).min()
        if gap < 1e-5:
            logger.warn(mycc, 'HOMO-LUMO gap %s too small for GCCSD', gap)
        return self

def _make_eris_incore(mycc, mo_coeff=None, ao2mofn=None):
    eris = _PhysicistsERIs()
    eris._common_init_(mycc, mo_coeff)
    nocc = eris.nocc
    nao, nmo = eris.mo_coeff.shape

    if callable(ao2mofn):
        eri = ao2mofn(eris.mo_coeff).reshape([nmo]*4)
    else:
        assert(eris.mo_coeff.dtype == np.double)
        mo_a = eris.mo_coeff[:nao//2]
        mo_b = eris.mo_coeff[nao//2:]
        orbspin = eris.orbspin
        if orbspin is None:
            eri  = ao2mo.kernel(mycc._scf._eri, mo_a)
            eri += ao2mo.kernel(mycc._scf._eri, mo_b)
            eri1 = ao2mo.kernel(mycc._scf._eri, (mo_a,mo_a,mo_b,mo_b))
            eri += eri1
            eri += eri1.T
        else:
            mo = mo_a + mo_b
            eri = ao2mo.kernel(mycc._scf._eri, mo)
            if eri.size == nmo**4:  # if mycc._scf._eri is a complex array
                sym_forbid = (orbspin[:,None] != orbspin).ravel()
            else:  # 4-fold symmetry
                sym_forbid = (orbspin[:,None] != orbspin)[np.tril_indices(nmo)]
            eri[sym_forbid,:] = 0
            eri[:,sym_forbid] = 0

        if eri.dtype == np.double:
            eri = ao2mo.restore(1, eri, nmo)

    eri = eri.reshape(nmo,nmo,nmo,nmo)
    eri = eri.transpose(0,2,1,3) - eri.transpose(0,2,3,1)

    eris.oooo = eri[:nocc,:nocc,:nocc,:nocc].copy()
    eris.ooov = eri[:nocc,:nocc,:nocc,nocc:].copy()
    eris.oovv = eri[:nocc,:nocc,nocc:,nocc:].copy()
    eris.ovov = eri[:nocc,nocc:,:nocc,nocc:].copy()
    eris.ovvo = eri[:nocc,nocc:,nocc:,:nocc].copy()
    eris.ovvv = eri[:nocc,nocc:,nocc:,nocc:].copy()
    eris.vvvv = eri[nocc:,nocc:,nocc:,nocc:].copy()
    return eris

def _make_eris_outcore(mycc, mo_coeff=None):
    cput0 = (logger.process_clock(), logger.perf_counter())
    log = logger.Logger(mycc.stdout, mycc.verbose)

    eris = _PhysicistsERIs()
    eris._common_init_(mycc, mo_coeff)
    nocc = eris.nocc
    nao, nmo = eris.mo_coeff.shape
    nvir = nmo - nocc
    assert(eris.mo_coeff.dtype == np.double)
    mo_a = eris.mo_coeff[:nao//2]
    mo_b = eris.mo_coeff[nao//2:]
    orbspin = eris.orbspin

    feri = eris.feri = lib.H5TmpFile()
    dtype = np.result_type(eris.mo_coeff).char
    eris.oooo = feri.create_dataset('oooo', (nocc,nocc,nocc,nocc), dtype)
    eris.ooov = feri.create_dataset('ooov', (nocc,nocc,nocc,nvir), dtype)
    eris.oovv = feri.create_dataset('oovv', (nocc,nocc,nvir,nvir), dtype)
    eris.ovov = feri.create_dataset('ovov', (nocc,nvir,nocc,nvir), dtype)
    eris.ovvo = feri.create_dataset('ovvo', (nocc,nvir,nvir,nocc), dtype)
    eris.ovvv = feri.create_dataset('ovvv', (nocc,nvir,nvir,nvir), dtype)

    if orbspin is None:
        orbo_a = mo_a[:,:nocc]
        orbv_a = mo_a[:,nocc:]
        orbo_b = mo_b[:,:nocc]
        orbv_b = mo_b[:,nocc:]

        max_memory = mycc.max_memory-lib.current_memory()[0]
        blksize = min(nocc, max(2, int(max_memory*1e6/8/(nmo**3*2))))
        max_memory = max(MEMORYMIN, max_memory)

        fswap = lib.H5TmpFile()
        ao2mo.kernel(mycc.mol, (orbo_a,mo_a,mo_a,mo_a), fswap, 'aaaa',
                     max_memory=max_memory, verbose=log)
        ao2mo.kernel(mycc.mol, (orbo_a,mo_a,mo_b,mo_b), fswap, 'aabb',
                     max_memory=max_memory, verbose=log)
        ao2mo.kernel(mycc.mol, (orbo_b,mo_b,mo_a,mo_a), fswap, 'bbaa',
                     max_memory=max_memory, verbose=log)
        ao2mo.kernel(mycc.mol, (orbo_b,mo_b,mo_b,mo_b), fswap, 'bbbb',
                     max_memory=max_memory, verbose=log)

        for p0, p1 in lib.prange(0, nocc, blksize):
            tmp  = np.asarray(fswap['aaaa'][p0*nmo:p1*nmo])
            tmp += np.asarray(fswap['aabb'][p0*nmo:p1*nmo])
            tmp += np.asarray(fswap['bbaa'][p0*nmo:p1*nmo])
            tmp += np.asarray(fswap['bbbb'][p0*nmo:p1*nmo])
            tmp = lib.unpack_tril(tmp).reshape(p1-p0,nmo,nmo,nmo)
            eris.oooo[p0:p1] = (tmp[:,:nocc,:nocc,:nocc].transpose(0,2,1,3) -
                                tmp[:,:nocc,:nocc,:nocc].transpose(0,2,3,1))
            eris.ooov[p0:p1] = (tmp[:,:nocc,:nocc,nocc:].transpose(0,2,1,3) -
                                tmp[:,nocc:,:nocc,:nocc].transpose(0,2,3,1))
            eris.ovvv[p0:p1] = (tmp[:,nocc:,nocc:,nocc:].transpose(0,2,1,3) -
                                tmp[:,nocc:,nocc:,nocc:].transpose(0,2,3,1))
            eris.oovv[p0:p1] = (tmp[:,nocc:,:nocc,nocc:].transpose(0,2,1,3) -
                                tmp[:,nocc:,:nocc,nocc:].transpose(0,2,3,1))
            eris.ovov[p0:p1] = (tmp[:,:nocc,nocc:,nocc:].transpose(0,2,1,3) -
                                tmp[:,nocc:,nocc:,:nocc].transpose(0,2,3,1))
            eris.ovvo[p0:p1] = (tmp[:,nocc:,nocc:,:nocc].transpose(0,2,1,3) -
                                tmp[:,:nocc,nocc:,nocc:].transpose(0,2,3,1))
            tmp = None
        cput0 = log.timer_debug1('transforming ovvv', *cput0)

        eris.vvvv = feri.create_dataset('vvvv', (nvir,nvir,nvir,nvir), dtype)
        tril2sq = lib.square_mat_in_trilu_indices(nvir)
        fswap = lib.H5TmpFile()
        ao2mo.kernel(mycc.mol, (orbv_a,orbv_a,orbv_a,orbv_a), fswap, 'aaaa',
                     max_memory=max_memory, verbose=log)
        ao2mo.kernel(mycc.mol, (orbv_a,orbv_a,orbv_b,orbv_b), fswap, 'aabb',
                     max_memory=max_memory, verbose=log)
        ao2mo.kernel(mycc.mol, (orbv_b,orbv_b,orbv_a,orbv_a), fswap, 'bbaa',
                     max_memory=max_memory, verbose=log)
        ao2mo.kernel(mycc.mol, (orbv_b,orbv_b,orbv_b,orbv_b), fswap, 'bbbb',
                     max_memory=max_memory, verbose=log)
        for p0, p1 in lib.prange(0, nvir, blksize):
            off0 = p0*(p0+1)//2
            off1 = p1*(p1+1)//2
            tmp  = np.asarray(fswap['aaaa'][off0:off1])
            tmp += np.asarray(fswap['aabb'][off0:off1])
            tmp += np.asarray(fswap['bbaa'][off0:off1])
            tmp += np.asarray(fswap['bbbb'][off0:off1])

            if p0 > 0:
                c = tmp[ tril2sq[p0:p1,:p0] - off0 ]
                c = lib.unpack_tril(c.reshape((p1-p0)*p0,-1))
                eris.vvvv[p0:p1,:p0] = c.reshape(p1-p0,p0,nvir,nvir)
            c = tmp[ tril2sq[:p1,p0:p1] - off0 ]
            c = lib.unpack_tril(c.reshape((p1-p0)*p1,-1))
            eris.vvvv[:p1,p0:p1] = c.reshape(p1,p1-p0,nvir,nvir)
            tmp = None

        for p0, p1 in lib.prange(0, nvir, blksize):
            tmp = np.asarray(eris.vvvv[p0:p1])
            eris.vvvv[p0:p1] = tmp.transpose(0,2,1,3) - tmp.transpose(0,2,3,1)
        cput0 = log.timer_debug1('transforming vvvv', *cput0)

    else:  # with orbspin
        mo = mo_a + mo_b
        orbo = mo[:,:nocc]
        orbv = mo[:,nocc:]

        max_memory = mycc.max_memory-lib.current_memory()[0]
        blksize = min(nocc, max(2, int(max_memory*1e6/8/(nmo**3*2))))
        max_memory = max(MEMORYMIN, max_memory)

        fswap = lib.H5TmpFile()
        ao2mo.kernel(mycc.mol, (orbo,mo,mo,mo), fswap,
                     max_memory=max_memory, verbose=log)
        sym_forbid = orbspin[:,None] != orbspin

        for p0, p1 in lib.prange(0, nocc, blksize):
            tmp = np.asarray(fswap['eri_mo'][p0*nmo:p1*nmo])
            tmp = lib.unpack_tril(tmp).reshape(p1-p0,nmo,nmo,nmo)
            tmp[sym_forbid[p0:p1]] = 0
            tmp[:,:,sym_forbid] = 0

            eris.oooo[p0:p1] = (tmp[:,:nocc,:nocc,:nocc].transpose(0,2,1,3) -
                                tmp[:,:nocc,:nocc,:nocc].transpose(0,2,3,1))
            eris.ooov[p0:p1] = (tmp[:,:nocc,:nocc,nocc:].transpose(0,2,1,3) -
                                tmp[:,nocc:,:nocc,:nocc].transpose(0,2,3,1))
            eris.ovvv[p0:p1] = (tmp[:,nocc:,nocc:,nocc:].transpose(0,2,1,3) -
                                tmp[:,nocc:,nocc:,nocc:].transpose(0,2,3,1))
            eris.oovv[p0:p1] = (tmp[:,nocc:,:nocc,nocc:].transpose(0,2,1,3) -
                                tmp[:,nocc:,:nocc,nocc:].transpose(0,2,3,1))
            eris.ovov[p0:p1] = (tmp[:,:nocc,nocc:,nocc:].transpose(0,2,1,3) -
                                tmp[:,nocc:,nocc:,:nocc].transpose(0,2,3,1))
            eris.ovvo[p0:p1] = (tmp[:,nocc:,nocc:,:nocc].transpose(0,2,1,3) -
                                tmp[:,:nocc,nocc:,nocc:].transpose(0,2,3,1))
            tmp = None
        cput0 = log.timer_debug1('transforming ovvv', *cput0)

        eris.vvvv = feri.create_dataset('vvvv', (nvir,nvir,nvir,nvir), dtype)
        sym_forbid = (orbspin[nocc:,None]!=orbspin[nocc:])[np.tril_indices(nvir)]
        tril2sq = lib.square_mat_in_trilu_indices(nvir)

        fswap = lib.H5TmpFile()
        ao2mo.kernel(mycc.mol, orbv, fswap, max_memory=max_memory, verbose=log)
        for p0, p1 in lib.prange(0, nvir, blksize):
            off0 = p0*(p0+1)//2
            off1 = p1*(p1+1)//2
            tmp = np.asarray(fswap['eri_mo'][off0:off1])
            tmp[sym_forbid[off0:off1]] = 0
            tmp[:,sym_forbid] = 0

            if p0 > 0:
                c = tmp[ tril2sq[p0:p1,:p0] - off0 ]
                c = lib.unpack_tril(c.reshape((p1-p0)*p0,-1))
                eris.vvvv[p0:p1,:p0] = c.reshape(p1-p0,p0,nvir,nvir)
            c = tmp[ tril2sq[:p1,p0:p1] - off0 ]
            c = lib.unpack_tril(c.reshape((p1-p0)*p1,-1))
            eris.vvvv[:p1,p0:p1] = c.reshape(p1,p1-p0,nvir,nvir)
            tmp = None

        for p0, p1 in lib.prange(0, nvir, blksize):
            tmp = np.asarray(eris.vvvv[p0:p1])
            eris.vvvv[p0:p1] = tmp.transpose(0,2,1,3) - tmp.transpose(0,2,3,1)
        cput0 = log.timer_debug1('transforming vvvv', *cput0)

    return eris


if __name__ == '__main__':
    from pyscf import gto
    mol = gto.Mole()
    mol.atom = [['O', (0.,   0., 0.)],
                ['O', (1.21, 0., 0.)]]
    mol.basis = 'cc-pvdz'
    mol.spin = 2
    mol.build()
    mf = scf.UHF(mol).run()
    mf = scf.addons.convert_to_ghf(mf)

    # Freeze 1s electrons
    frozen = [0,1,2,3]
    gcc = GCCSD(mf, frozen=frozen)
    ecc, t1, t2 = gcc.kernel()
    print(ecc - -0.3486987472235819)

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = 'cc-pvdz'
    mol.spin = 0
    mol.build()
    mf = scf.UHF(mol).run()
    mf = scf.addons.convert_to_ghf(mf)

    mycc = GCCSD(mf)
    ecc, t1, t2 = mycc.kernel()
    print(ecc - -0.2133432712431435)
    e,v = mycc.ipccsd(nroots=8)
    print(e[0] - 0.4335604332073799)
    print(e[2] - 0.5187659896045407)
    print(e[4] - 0.6782876002229172)

    e,v = mycc.eaccsd(nroots=8)
    print(e[0] - 0.16737886338859731)
    print(e[2] - 0.24027613852009164)
    print(e[4] - 0.51006797826488071)

    e,v = mycc.eeccsd(nroots=4)
    print(e[0] - 0.2757159395886167)
    print(e[1] - 0.2757159395886167)
    print(e[2] - 0.2757159395886167)
    print(e[3] - 0.3005716731825082)


# pyscf/cc/gccsd_lambda.py

#!/usr/bin/env python
# Copyright 2014-2020 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: Qiming Sun <osirpt.sun@gmail.com>
#


import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf.cc import ccsd_lambda

einsum = lib.einsum

def kernel(mycc, eris=None, t1=None, t2=None, l1=None, l2=None,
           max_cycle=50, tol=1e-8, verbose=logger.INFO):
    if eris is None: eris = mycc.ao2mo()
    return ccsd_lambda.kernel(mycc, eris, t1, t2, l1, l2, max_cycle, tol,
                              verbose, make_intermediates, update_lambda)

# l2, t2 as ijab
def make_intermediates(mycc, t1, t2, eris):
    nocc, nvir = t1.shape
    foo = eris.fock[:nocc,:nocc]
    fov = eris.fock[:nocc,nocc:]
    fvo = eris.fock[nocc:,:nocc]
    fvv = eris.fock[nocc:,nocc:]

    tau = t2 + einsum('ia,jb->ijab', t1, t1) * 2

    v1 = fvv - einsum('ja,jb->ba', fov, t1)
    v1-= numpy.einsum('jbac,jc->ba', eris.ovvv, t1)
    v1+= einsum('jkca,jkbc->ba', eris.oovv, tau) * .5

    v2 = foo + einsum('ib,jb->ij', fov, t1)
    v2-= numpy.einsum('kijb,kb->ij', eris.ooov, t1)
    v2+= einsum('ikbc,jkbc->ij', eris.oovv, tau) * .5

    v3 = einsum('ijcd,klcd->ijkl', eris.oovv, tau)
    v4 = einsum('ljdb,klcd->jcbk', eris.oovv, t2)
    v4+= numpy.asarray(eris.ovvo)

    v5 = fvo + numpy.einsum('kc,jkbc->bj', fov, t2)
    tmp = fov - numpy.einsum('kldc,ld->kc', eris.oovv, t1)
    v5+= numpy.einsum('kc,kb,jc->bj', tmp, t1, t1)
    v5-= einsum('kljc,klbc->bj', eris.ooov, t2) * .5
    v5+= einsum('kbdc,jkcd->bj', eris.ovvv, t2) * .5

    w3 = v5 + numpy.einsum('jcbk,jb->ck', v4, t1)
    w3 += numpy.einsum('cb,jb->cj', v1, t1)
    w3 -= numpy.einsum('jk,jb->bk', v2, t1)

    woooo = numpy.asarray(eris.oooo) * .5
    woooo+= v3 * .25
    woooo+= einsum('jilc,kc->jilk', eris.ooov, t1)

    wovvo = v4 - numpy.einsum('ljdb,lc,kd->jcbk', eris.oovv, t1, t1)
    wovvo-= einsum('ljkb,lc->jcbk', eris.ooov, t1)
    wovvo+= einsum('jcbd,kd->jcbk', eris.ovvv, t1)

    wovoo = einsum('icdb,jkdb->icjk', eris.ovvv, tau) * .25
    wovoo+= numpy.einsum('jkic->icjk', numpy.asarray(eris.ooov).conj()) * .5
    wovoo+= einsum('icbk,jb->icjk', v4, t1)
    wovoo-= einsum('lijb,klcb->icjk', eris.ooov, t2)

    wvvvo = einsum('jcak,jb->bcak', v4, t1)
    wvvvo+= einsum('jlka,jlbc->bcak', eris.ooov, tau) * .25
    wvvvo-= numpy.einsum('jacb->bcaj', numpy.asarray(eris.ovvv).conj()) * .5
    wvvvo+= einsum('kbad,jkcd->bcaj', eris.ovvv, t2)

    class _IMDS: pass
    imds = _IMDS()
    imds.ftmp = lib.H5TmpFile()
    dtype = numpy.result_type(t2, eris.vvvv).char
    imds.woooo = imds.ftmp.create_dataset('woooo', (nocc,nocc,nocc,nocc), dtype)
    imds.wovvo = imds.ftmp.create_dataset('wovvo', (nocc,nvir,nvir,nocc), dtype)
    imds.wovoo = imds.ftmp.create_dataset('wovoo', (nocc,nvir,nocc,nocc), dtype)
    imds.wvvvo = imds.ftmp.create_dataset('wvvvo', (nvir,nvir,nvir,nocc), dtype)
    imds.woooo[:] = woooo
    imds.wovvo[:] = wovvo
    imds.wovoo[:] = wovoo
    imds.wvvvo[:] = wvvvo
    imds.v1 = v1
    imds.v2 = v2
    imds.w3 = w3
    imds.ftmp.flush()
    return imds


# update L1, L2
def update_lambda(mycc, t1, t2, l1, l2, eris, imds):
    time0 = logger.process_clock(), logger.perf_counter()
    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocc, nvir = t1.shape
    fov = eris.fock[:nocc,nocc:]
    mo_e_o = eris.mo_energy[:nocc]
    mo_e_v = eris.mo_energy[nocc:] + mycc.level_shift
    v1 = imds.v1 - numpy.diag(mo_e_v)
    v2 = imds.v2 - numpy.diag(mo_e_o)

    l1new = numpy.zeros_like(l1)
    l2new = numpy.zeros_like(l2)

    mba = einsum('klca,klcb->ba', l2, t2) * .5
    mij = einsum('kicd,kjcd->ij', l2, t2) * .5
    m3 = einsum('klab,ijkl->ijab', l2, numpy.asarray(imds.woooo))
    tau = t2 + einsum('ia,jb->ijab', t1, t1) * 2
    tmp = einsum('ijcd,klcd->ijkl', l2, tau)
    oovv = numpy.asarray(eris.oovv)
    m3 += einsum('klab,ijkl->ijab', oovv, tmp) * .25
    tmp = einsum('ijcd,kd->ijck', l2, t1)
    m3 -= einsum('kcba,ijck->ijab', eris.ovvv, tmp)
    m3 += einsum('ijcd,cdab->ijab', l2, eris.vvvv) * .5

    l2new += oovv
    l2new += m3
    fov1 = fov + einsum('kjcb,kc->jb', oovv, t1)
    tmp = einsum('ia,jb->ijab', l1, fov1)
    tmp+= einsum('kica,jcbk->ijab', l2, numpy.asarray(imds.wovvo))
    tmp = tmp - tmp.transpose(1,0,2,3)
    l2new += tmp - tmp.transpose(0,1,3,2)
    tmp = einsum('ka,ijkb->ijab', l1, eris.ooov)
    tmp+= einsum('ijca,cb->ijab', l2, v1)
    tmp1vv = mba + einsum('ka,kb->ba', l1, t1)
    tmp+= einsum('ca,ijcb->ijab', tmp1vv, oovv)
    l2new -= tmp - tmp.transpose(0,1,3,2)
    tmp = einsum('ic,jcba->jiba', l1, eris.ovvv)
    tmp+= einsum('kiab,jk->ijab', l2, v2)
    tmp1oo = mij + einsum('ic,kc->ik', l1, t1)
    tmp-= einsum('ik,kjab->ijab', tmp1oo, oovv)
    l2new += tmp - tmp.transpose(1,0,2,3)

    l1new += fov
    l1new += einsum('jb,ibaj->ia', l1, eris.ovvo)
    l1new += einsum('ib,ba->ia', l1, v1)
    l1new -= einsum('ja,ij->ia', l1, v2)
    l1new -= einsum('kjca,icjk->ia', l2, imds.wovoo)
    l1new -= einsum('ikbc,bcak->ia', l2, imds.wvvvo)
    l1new += einsum('ijab,jb->ia', m3, t1)
    l1new += einsum('jiba,bj->ia', l2, imds.w3)
    tmp =(t1 + einsum('kc,kjcb->jb', l1, t2)
          - einsum('bd,jd->jb', tmp1vv, t1)
          - einsum('lj,lb->jb', mij, t1))
    l1new += numpy.einsum('jiba,jb->ia', oovv, tmp)
    l1new += numpy.einsum('icab,bc->ia', eris.ovvv, tmp1vv)
    l1new -= numpy.einsum('jika,kj->ia', eris.ooov, tmp1oo)
    tmp = fov - einsum('kjba,jb->ka', oovv, t1)
    l1new -= numpy.einsum('ik,ka->ia', mij, tmp)
    l1new -= numpy.einsum('ca,ic->ia', mba, tmp)

    eia = lib.direct_sum('i-j->ij', mo_e_o, mo_e_v)
    l1new /= eia
    l2new /= lib.direct_sum('ia+jb->ijab', eia, eia)

    time0 = log.timer_debug1('update l1 l2', *time0)
    return l1new, l2new


if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf
    from pyscf.cc import gccsd

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = '631g'
    mol.spin = 0
    mol.build()
    mf = scf.RHF(mol).run()
    mf0 = mf
    mf = scf.addons.convert_to_ghf(mf)
    mycc = gccsd.GCCSD(mf)
    eris = mycc.ao2mo()
    mycc.kernel(eris=eris)
    l1, l2 = mycc.solve_lambda(mycc.t1, mycc.t2, eris=eris)
    l1 = mycc.spin2spatial(l1, mycc.mo_coeff.orbspin)
    l2 = mycc.spin2spatial(l2, mycc.mo_coeff.orbspin)
    print(lib.finger(l1[0]) --0.0030030170069977758)
    print(lib.finger(l1[1]) --0.0030030170069977758)
    print(lib.finger(l2[0]) --0.041444910588788492 )
    print(lib.finger(l2[1]) - 0.1077575086912813   )
    print(lib.finger(l2[2]) --0.041444910588788492 )
    print(abs(l2[1]-l2[1].transpose(1,0,2,3)-l2[0]).max())
    print(abs(l2[1]-l2[1].transpose(0,1,3,2)-l2[0]).max())

    from pyscf.cc import ccsd
    mycc0 = ccsd.CCSD(mf0)
    eris0 = mycc0.ao2mo()
    mycc0.kernel(eris=eris0)
    t1 = mycc0.t1
    t2 = mycc0.t2
    imds = ccsd_lambda.make_intermediates(mycc0, t1, t2, eris0)
    l1, l2 = ccsd_lambda.update_lambda(mycc0, t1, t2, t1, t2, eris0, imds)
    l1ref, l2ref = ccsd_lambda.update_lambda(mycc0, t1, t2, l1, l2, eris0, imds)
    t1 = mycc.spatial2spin(t1, mycc.mo_coeff.orbspin)
    t2 = mycc.spatial2spin(t2, mycc.mo_coeff.orbspin)
    l1 = mycc.spatial2spin(l1, mycc.mo_coeff.orbspin)
    l2 = mycc.spatial2spin(l2, mycc.mo_coeff.orbspin)
    imds = make_intermediates(mycc, t1, t2, eris)
    l1, l2 = update_lambda(mycc, t1, t2, l1, l2, eris, imds)
    l1 = mycc.spin2spatial(l1, mycc.mo_coeff.orbspin)
    l2 = mycc.spin2spatial(l2, mycc.mo_coeff.orbspin)
    print(abs(l1[0]-l1ref).max())
    print(abs(l2[1]-l2ref).max())


# pyscf/cc/gccsd_rdm.py

#!/usr/bin/env python
# Copyright 2014-2021 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: Qiming Sun <osirpt.sun@gmail.com>
#         Jun Yang
#

import numpy
from pyscf import lib

#einsum = numpy.einsum
einsum = lib.einsum

def _gamma1_intermediates(mycc, t1, t2, l1, l2):
    doo  =-einsum('ie,je->ij', l1, t1)
    doo -= einsum('imef,jmef->ij', l2, t2) * .5

    dvv  = einsum('ma,mb->ab', t1, l1)
    dvv += einsum('mnea,mneb->ab', t2, l2) * .5

    xt1  = einsum('mnef,inef->mi', l2, t2) * .5
    xt2  = einsum('mnfa,mnfe->ae', t2, l2) * .5
    xt2 += einsum('ma,me->ae', t1, l1)
    dvo  = einsum('imae,me->ai', t2, l1)
    dvo -= einsum('mi,ma->ai', xt1, t1)
    dvo -= einsum('ie,ae->ai', t1, xt2)
    dvo += t1.T

    dov = l1

    return doo, dov, dvo, dvv

# gamma2 intermediates in Chemist's notation
# When computing intermediates, the convention
# dm2[q,p,s,r] = <p^\dagger r^\dagger s q> is assumed in this function.
# It changes to dm2[p,q,r,s] = <p^\dagger r^\dagger s q> in _make_rdm2
def _gamma2_intermediates(mycc, t1, t2, l1, l2):
    tau = t2 + einsum('ia,jb->ijab', t1, t1) * 2
    miajb = einsum('ikac,kjcb->iajb', l2, t2)

    goovv = 0.25 * (l2.conj() + tau)
    tmp = einsum('kc,kica->ia', l1, t2)
    goovv += einsum('ia,jb->ijab', tmp, t1)
    tmp = einsum('kc,kb->cb', l1, t1)
    goovv += einsum('cb,ijca->ijab', tmp, t2) * .5
    tmp = einsum('kc,jc->kj', l1, t1)
    goovv += einsum('kiab,kj->ijab', tau, tmp) * .5
    tmp = numpy.einsum('ldjd->lj', miajb)
    goovv -= einsum('lj,liba->ijab', tmp, tau) * .25
    tmp = numpy.einsum('ldlb->db', miajb)
    goovv -= einsum('db,jida->ijab', tmp, tau) * .25
    goovv -= einsum('ldia,ljbd->ijab', miajb, tau) * .5
    tmp = einsum('klcd,ijcd->ijkl', l2, tau) * .25**2
    goovv += einsum('ijkl,klab->ijab', tmp, tau)
    goovv = goovv.conj()

    gvvvv = einsum('ijab,ijcd->abcd', tau, l2) * 0.125
    goooo = einsum('klab,ijab->klij', l2, tau) * 0.125

    gooov  = einsum('jkba,ib->jkia', tau, l1) * -0.25
    gooov += einsum('iljk,la->jkia', goooo, t1)
    tmp = numpy.einsum('icjc->ij', miajb) * .25
    gooov -= einsum('ij,ka->jkia', tmp, t1)
    gooov += einsum('icja,kc->jkia', miajb, t1) * .5
    gooov = gooov.conj()
    gooov += einsum('jkab,ib->jkia', l2, t1) * .25

    govvo  = einsum('ia,jb->ibaj', l1, t1)
    govvo += numpy.einsum('iajb->ibaj', miajb)
    govvo -= einsum('ikac,jc,kb->ibaj', l2, t1, t1)

    govvv  = einsum('ja,ijcb->iacb', l1, tau) * .25
    govvv += einsum('bcad,id->iabc', gvvvv, t1)
    tmp = numpy.einsum('kakb->ab', miajb) * .25
    govvv += einsum('ab,ic->iacb', tmp, t1)
    govvv += einsum('kaib,kc->iabc', miajb, t1) * .5
    govvv = govvv.conj()
    govvv += einsum('ijbc,ja->iabc', l2, t1) * .25

    dovov = goovv.transpose(0,2,1,3) - goovv.transpose(0,3,1,2)
    dvvvv = gvvvv.transpose(0,2,1,3) - gvvvv.transpose(0,3,1,2)
    doooo = goooo.transpose(0,2,1,3) - goooo.transpose(0,3,1,2)
    dovvv = govvv.transpose(0,2,1,3) - govvv.transpose(0,3,1,2)
    dooov = gooov.transpose(0,2,1,3) - gooov.transpose(1,2,0,3)
    dovvo = govvo.transpose(0,2,1,3)
    dovov =(dovov + dovov.transpose(2,3,0,1)) * .5
    dvvvv = dvvvv + dvvvv.transpose(1,0,3,2).conj()
    doooo = doooo + doooo.transpose(1,0,3,2).conj()
    dovvo =(dovvo + dovvo.transpose(3,2,1,0).conj()) * .5
    doovv = None # = -dovvo.transpose(0,3,2,1)
    dvvov = None
    return (dovov, dvvvv, doooo, doovv, dovvo, dvvov, dovvv, dooov)

def make_rdm1(mycc, t1, t2, l1, l2, ao_repr=False, with_frozen=True, with_mf=True):
    r'''
    One-particle density matrix in the molecular spin-orbital representation
    (the occupied-virtual blocks from the orbital response contribution are
    not included).

    dm1[p,q] = <q^\dagger p>  (p,q are spin-orbitals)

    The convention of 1-pdm is based on McWeeney's book, Eq (5.4.20).
    The contraction between 1-particle Hamiltonian and rdm1 is
    E = einsum('pq,qp', h1, rdm1)
    '''
    d1 = _gamma1_intermediates(mycc, t1, t2, l1, l2)
    return _make_rdm1(mycc, d1, with_frozen=with_frozen, ao_repr=ao_repr,
                      with_mf=with_mf)

def make_rdm2(mycc, t1, t2, l1, l2, ao_repr=False, with_frozen=True, with_dm1=True):
    r'''
    Two-particle density matrix in the molecular spin-orbital representation

    dm2[p,q,r,s] = <p^\dagger r^\dagger s q>

    where p,q,r,s are spin-orbitals. p,q correspond to one particle and r,s
    correspond to another particle.  The contraction between ERIs (in
    Chemist's notation) and rdm2 is
    E = einsum('pqrs,pqrs', eri, rdm2)
    '''
    d1 = _gamma1_intermediates(mycc, t1, t2, l1, l2)
    d2 = _gamma2_intermediates(mycc, t1, t2, l1, l2)
    return _make_rdm2(mycc, d1, d2, with_dm1=with_dm1, with_frozen=with_frozen,
                      ao_repr=ao_repr)

def _make_rdm1(mycc, d1, with_frozen=True, ao_repr=False, with_mf=True):
    r'''
    One-particle density matrix in the molecular spin-orbital representation
    (the occupied-virtual blocks from the orbital response contribution are
    not included).

    dm1[p,q] = <q^\dagger p>  (p,q are spin-orbitals)

    The convention of 1-pdm is based on McWeeney's book, Eq (5.4.20).
    The contraction between 1-particle Hamiltonian and rdm1 is
    E = einsum('pq,qp', h1, rdm1)
    '''
    doo, dov, dvo, dvv = d1
    nocc, nvir = dov.shape
    nmo = nocc + nvir

    dm1 = numpy.empty((nmo,nmo), dtype=doo.dtype)
    dm1[:nocc,:nocc] = doo + doo.conj().T
    dm1[:nocc,nocc:] = dov + dvo.conj().T
    dm1[nocc:,:nocc] = dm1[:nocc,nocc:].conj().T
    dm1[nocc:,nocc:] = dvv + dvv.conj().T
    dm1 *= .5
    if with_mf:
        dm1[numpy.diag_indices(nocc)] += 1

    if with_frozen and mycc.frozen is not None:
        nmo = mycc.mo_occ.size
        nocc = numpy.count_nonzero(mycc.mo_occ > 0)
        rdm1 = numpy.zeros((nmo,nmo), dtype=dm1.dtype)
        if with_mf:
            rdm1[numpy.diag_indices(nocc)] = 1
        moidx = numpy.where(mycc.get_frozen_mask())[0]
        rdm1[moidx[:,None],moidx] = dm1
        dm1 = rdm1

    if ao_repr:
        mo = mycc.mo_coeff
        dm1 = lib.einsum('pi,ij,qj->pq', mo, dm1, mo.conj())
    return dm1

def _make_rdm2(mycc, d1, d2, with_dm1=True, with_frozen=True, ao_repr=False):
    r'''
    dm2[p,q,r,s] = <p^\dagger r^\dagger s q>

    Note the contraction between ERIs (in Chemist's notation) and rdm2 is
    E = einsum('pqrs,pqrs', eri, rdm2)
    '''
    dovov, dvvvv, doooo, doovv, dovvo, dvvov, dovvv, dooov = d2
    nocc, nvir = dovov.shape[:2]
    nmo = nocc + nvir

    dm2 = numpy.empty((nmo,nmo,nmo,nmo), dtype=doooo.dtype)

    dovov = numpy.asarray(dovov)
    dm2[:nocc,nocc:,:nocc,nocc:] = dovov
    dm2[nocc:,:nocc,nocc:,:nocc] = dm2[:nocc,nocc:,:nocc,nocc:].transpose(1,0,3,2).conj()
    dovov = None

    dovvo = numpy.asarray(dovvo)
    dm2[:nocc,:nocc,nocc:,nocc:] =-dovvo.transpose(0,3,2,1)
    dm2[nocc:,nocc:,:nocc,:nocc] =-dovvo.transpose(2,1,0,3)
    dm2[:nocc,nocc:,nocc:,:nocc] = dovvo
    dm2[nocc:,:nocc,:nocc,nocc:] = dovvo.transpose(1,0,3,2).conj()
    dovvo = None

    dm2[nocc:,nocc:,nocc:,nocc:] = dvvvv
    dm2[:nocc,:nocc,:nocc,:nocc] = doooo

    dovvv = numpy.asarray(dovvv)
    dm2[:nocc,nocc:,nocc:,nocc:] = dovvv
    dm2[nocc:,nocc:,:nocc,nocc:] = dovvv.transpose(2,3,0,1)
    dm2[nocc:,nocc:,nocc:,:nocc] = dovvv.transpose(3,2,1,0).conj()
    dm2[nocc:,:nocc,nocc:,nocc:] = dovvv.transpose(1,0,3,2).conj()
    dovvv = None

    dooov = numpy.asarray(dooov)
    dm2[:nocc,:nocc,:nocc,nocc:] = dooov
    dm2[:nocc,nocc:,:nocc,:nocc] = dooov.transpose(2,3,0,1)
    dm2[:nocc,:nocc,nocc:,:nocc] = dooov.transpose(1,0,3,2).conj()
    dm2[nocc:,:nocc,:nocc,:nocc] = dooov.transpose(3,2,1,0).conj()

    if with_frozen and mycc.frozen is not None:
        nmo, nmo0 = mycc.mo_occ.size, nmo
        nocc = numpy.count_nonzero(mycc.mo_occ > 0)
        rdm2 = numpy.zeros((nmo,nmo,nmo,nmo), dtype=dm2.dtype)
        moidx = numpy.where(mycc.get_frozen_mask())[0]
        idx = (moidx.reshape(-1,1) * nmo + moidx).ravel()
        lib.takebak_2d(rdm2.reshape(nmo**2,nmo**2),
                       dm2.reshape(nmo0**2,nmo0**2), idx, idx)
        dm2 = rdm2

    if with_dm1:
        dm1 = _make_rdm1(mycc, d1, with_frozen)
        dm1[numpy.diag_indices(nocc)] -= 1

        for i in range(nocc):
            # Be careful with the convention of dm1 and dm2.transpose
            # at the end
            dm2[i,i,:,:] += dm1
            dm2[:,:,i,i] += dm1
            dm2[:,i,i,:] -= dm1
            dm2[i,:,:,i] -= dm1.T

        for i in range(nocc):
            for j in range(nocc):
                dm2[i,i,j,j] += 1
                dm2[i,j,j,i] -= 1

    # dm2 was computed as dm2[p,q,r,s] = < p^\dagger r^\dagger s q > in the
    # above. Transposing it so that it be contracted with ERIs (in Chemist's
    # notation):
    #   E = einsum('pqrs,pqrs', eri, rdm2)
    dm2 = dm2.transpose(1,0,3,2)
    if ao_repr:
        from pyscf.cc import ccsd_rdm
        dm2 = ccsd_rdm._rdm2_mo2ao(dm2, mycc.mo_coeff)
    return dm2


if __name__ == '__main__':
    from functools import reduce
    from pyscf import gto
    from pyscf import scf
    from pyscf import ao2mo
    from pyscf.cc import gccsd

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = '631g'
    mol.spin = 2
    mol.build()
    mf = scf.UHF(mol).run(conv_tol=1.)
    mf = scf.addons.convert_to_ghf(mf)

    mycc = gccsd.GCCSD(mf)
    ecc, t1, t2 = mycc.kernel()
    l1, l2 = mycc.solve_lambda()
    dm1 = make_rdm1(mycc, t1, t2, l1, l2)
    dm2 = make_rdm2(mycc, t1, t2, l1, l2)
    nao = mol.nao_nr()
    mo_a = mf.mo_coeff[:nao]
    mo_b = mf.mo_coeff[nao:]
    nmo = mo_a.shape[1]
    eri = ao2mo.kernel(mf._eri, mo_a+mo_b, compact=False).reshape([nmo]*4)
    orbspin = mf.mo_coeff.orbspin
    sym_forbid = (orbspin[:,None] != orbspin)
    eri[sym_forbid,:,:] = 0
    eri[:,:,sym_forbid] = 0
    hcore = scf.RHF(mol).get_hcore()
    h1 = reduce(numpy.dot, (mo_a.T.conj(), hcore, mo_a))
    h1+= reduce(numpy.dot, (mo_b.T.conj(), hcore, mo_b))
    e1 = numpy.einsum('ij,ji', h1, dm1)
    e1+= numpy.einsum('ijkl,ijkl', eri, dm2) * .5
    e1+= mol.energy_nuc()
    print(e1 - mycc.e_tot)

    #TODO: test 1pdm, 2pdm against FCI


# pyscf/cc/gccsd_t.py

#!/usr/bin/env python
# Copyright 2014-2020 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: Qiming Sun <osirpt.sun@gmail.com>
#

'''
GHF-CCSD(T) with spin-orbital integrals
'''

import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf.cc import gccsd

# spin-orbital formula
# JCP 98, 8718 (1993); DOI:10.1063/1.464480
def kernel(cc, eris, t1=None, t2=None, verbose=logger.INFO):
    assert(isinstance(eris, gccsd._PhysicistsERIs))
    if t1 is None or t2 is None:
        t1, t2 = cc.t1, cc.t2
    nocc, nvir = t1.shape

    bcei = numpy.asarray(eris.ovvv).conj().transpose(3,2,1,0)
    majk = numpy.asarray(eris.ooov).conj().transpose(2,3,0,1)
    bcjk = numpy.asarray(eris.oovv).conj().transpose(2,3,0,1)
    fvo = eris.fock[nocc:,:nocc]
    mo_e = eris.mo_energy
    eijk = lib.direct_sum('i+j+k->ijk', mo_e[:nocc], mo_e[:nocc], mo_e[:nocc])
    eabc = lib.direct_sum('a+b+c->abc', mo_e[nocc:], mo_e[nocc:], mo_e[nocc:])

    t2T = t2.transpose(2,3,0,1)
    t1T = t1.T
    def get_wv(a, b, c):
        w  = numpy.einsum('ejk,ei->ijk', t2T[a,:], bcei[b,c])
        w -= numpy.einsum('im,mjk->ijk', t2T[b,c], majk[:,a])
        v  = numpy.einsum('i,jk->ijk', t1T[a], bcjk[b,c])
        v += numpy.einsum('i,jk->ijk', fvo[a], t2T [b,c])
        v += w
        w = w + w.transpose(2,0,1) + w.transpose(1,2,0)
        return w, v

    et = 0
    for a in range(nvir):
        for b in range(a):
            for c in range(b):
                wabc, vabc = get_wv(a, b, c)
                wcab, vcab = get_wv(c, a, b)
                wbac, vbac = get_wv(b, a, c)

                w = wabc + wcab - wbac
                v = vabc + vcab - vbac
                w /= eijk - eabc[a,b,c]
                et += numpy.einsum('ijk,ijk', w, v.conj())
    et /= 2
    return et


if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf
    from pyscf import cc

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -.957 , .587)],
        [1 , (0.2,  .757 , .487)]]

    mol.basis = '631g'
    mol.build()
    mf = scf.RHF(mol).run(conv_tol=1e-1)
    mycc = cc.CCSD(mf).set(conv_tol=1e-11).run()
    et = mycc.ccsd_t()

    mycc = cc.GCCSD(scf.addons.convert_to_ghf(mf)).set(conv_tol=1e-11).run()
    eris = mycc.ao2mo()
    print(kernel(mycc, eris) - et)

    numpy.random.seed(1)
    mf.mo_coeff = numpy.random.random(mf.mo_coeff.shape) - .9
    mycc = cc.GCCSD(scf.addons.convert_to_ghf(mf))
    eris = mycc.ao2mo()
    nocc = 10
    nvir = mol.nao_nr() * 2 - nocc
    t1 = numpy.random.random((nocc,nvir))*.1 - .2
    t2 = numpy.random.random((nocc,nocc,nvir,nvir))*.1 - .2
    t2 = t2 - t2.transpose(1,0,2,3)
    t2 = t2 - t2.transpose(0,1,3,2)
    print(kernel(mycc, eris, t1, t2) - 263713.3945021223)


# pyscf/cc/gccsd_t_lambda.py

#!/usr/bin/env python
# Copyright 2014-2021 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: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Lambda equation of GHF-CCSD(T) with spin-orbital integrals

Ref:
JCP 98, 8718 (1993); DOI:10.1063/1.464480
JCP 147, 044104 (2017); DOI:10.1063/1.4994918
'''

import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf.cc import ccsd_lambda
from pyscf.cc import gccsd_lambda


def kernel(mycc, eris=None, t1=None, t2=None, l1=None, l2=None,
           max_cycle=50, tol=1e-8, verbose=logger.INFO):
    return ccsd_lambda.kernel(mycc, eris, t1, t2, l1, l2, max_cycle, tol,
                              verbose, make_intermediates, update_lambda)

def make_intermediates(mycc, t1, t2, eris):
    imds = gccsd_lambda.make_intermediates(mycc, t1, t2, eris)

    nocc, nvir = t1.shape
    bcei = numpy.asarray(eris.ovvv).conj().transpose(3,2,1,0)
    majk = numpy.asarray(eris.ooov).conj().transpose(2,3,0,1)
    bcjk = numpy.asarray(eris.oovv).conj().transpose(2,3,0,1)

    mo_e = eris.mo_energy
    eia = mo_e[:nocc,None] - mo_e[nocc:]
    d3 = lib.direct_sum('ia+jb+kc->ijkabc', eia, eia, eia)

    t3c =(numpy.einsum('jkae,bcei->ijkabc', t2, bcei) -
          numpy.einsum('imbc,majk->ijkabc', t2, majk))
    t3c = t3c - t3c.transpose(0,1,2,4,3,5) - t3c.transpose(0,1,2,5,4,3)
    t3c = t3c - t3c.transpose(1,0,2,3,4,5) - t3c.transpose(2,1,0,3,4,5)
    t3c /= d3

    t3d = numpy.einsum('ia,bcjk->ijkabc', t1, bcjk)
    t3d += numpy.einsum('ai,jkbc->ijkabc', eris.fock[nocc:,:nocc], t2)
    t3d = t3d - t3d.transpose(0,1,2,4,3,5) - t3d.transpose(0,1,2,5,4,3)
    t3d = t3d - t3d.transpose(1,0,2,3,4,5) - t3d.transpose(2,1,0,3,4,5)
    t3d /= d3

    l1_t = numpy.einsum('ijkabc,jkbc->ia', t3c.conj(), eris.oovv) / eia
    imds.l1_t = l1_t * .25

    m3 = t3c * 2 + t3d
    tmp = numpy.einsum('ijkaef,kbfe->ijab', m3.conj(), eris.ovvv) * .5
    l2_t = tmp - tmp.transpose(0,1,3,2)
    tmp = numpy.einsum('imnabc,mnjc->ijab', m3.conj(), eris.ooov) * .5
    l2_t -= tmp - tmp.transpose(1,0,2,3)
    l2_t += numpy.einsum('kc,ijkabc->ijab', eris.fock[:nocc,nocc:], t3c.conj())
    imds.l2_t = l2_t / lib.direct_sum('ia+jb->ijab', eia, eia)

    return imds


def update_lambda(mycc, t1, t2, l1, l2, eris=None, imds=None):
    if eris is None: eris = mycc.ao2mo()
    if imds is None: imds = make_intermediates(mycc, t1, t2, eris)
    l1, l2 = gccsd_lambda.update_lambda(mycc, t1, t2, l1, l2, eris, imds)
    l1 += imds.l1_t
    l2 += imds.l2_t
    return l1, l2


if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf
    from pyscf import cc

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = '631g'
    mol.spin = 0
    mol.build()
    mf0 = mf = scf.RHF(mol).run(conv_tol=1.)
    mf = scf.addons.convert_to_ghf(mf)

    from pyscf.cc import ccsd_t_lambda_slow as ccsd_t_lambda
    mycc0 = cc.CCSD(mf0)
    eris0 = mycc0.ao2mo()
    mycc0.kernel(eris=eris0)
    t1 = mycc0.t1
    t2 = mycc0.t2
    imds = ccsd_t_lambda.make_intermediates(mycc0, t1, t2, eris0)
    l1, l2 = ccsd_t_lambda.update_lambda(mycc0, t1, t2, t1, t2, eris0, imds)
    l1ref, l2ref = ccsd_t_lambda.update_lambda(mycc0, t1, t2, l1, l2, eris0, imds)
    mycc = cc.GCCSD(mf)
    eris = mycc.ao2mo()
    t1 = mycc.spatial2spin(t1, mycc.mo_coeff.orbspin)
    t2 = mycc.spatial2spin(t2, mycc.mo_coeff.orbspin)
    l1 = mycc.spatial2spin(l1, mycc.mo_coeff.orbspin)
    l2 = mycc.spatial2spin(l2, mycc.mo_coeff.orbspin)
    imds = make_intermediates(mycc, t1, t2, eris)
    l1, l2 = update_lambda(mycc, t1, t2, l1, l2, eris, imds)
    l1 = mycc.spin2spatial(l1, mycc.mo_coeff.orbspin)
    l2 = mycc.spin2spatial(l2, mycc.mo_coeff.orbspin)
    print(abs(l2[1]-l2[1].transpose(1,0,2,3)-l2[0]).max())
    print(abs(l2[1]-l2[1].transpose(0,1,3,2)-l2[0]).max())
    print(abs(l1[0]-l1ref).max())
    print(abs(l2[1]-l2ref).max())

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = '631g'
    mol.spin = 2
    mol.build()
    mf0 = mf = scf.UHF(mol).run(conv_tol=1)
    mf = scf.addons.convert_to_ghf(mf)

    from pyscf.cc import uccsd_t_lambda
    mycc0 = cc.CCSD(mf0)
    eris0 = mycc0.ao2mo()
    mycc0.kernel(eris=eris0)
    t1 = mycc0.t1
    t2 = mycc0.t2
    imds = uccsd_t_lambda.make_intermediates(mycc0, t1, t2, eris0)
    l1, l2 = uccsd_t_lambda.update_lambda(mycc0, t1, t2, t1, t2, eris0, imds)
    l1ref, l2ref = uccsd_t_lambda.update_lambda(mycc0, t1, t2, l1, l2, eris0, imds)
    mycc = cc.GCCSD(mf)
    eris = mycc.ao2mo()
    t1 = mycc.spatial2spin(t1, mycc.mo_coeff.orbspin)
    t2 = mycc.spatial2spin(t2, mycc.mo_coeff.orbspin)
    l1 = mycc.spatial2spin(l1, mycc.mo_coeff.orbspin)
    l2 = mycc.spatial2spin(l2, mycc.mo_coeff.orbspin)
    imds = make_intermediates(mycc, t1, t2, eris)
    l1, l2 = update_lambda(mycc, t1, t2, l1, l2, eris, imds)
    l1 = mycc.spin2spatial(l1, mycc.mo_coeff.orbspin)
    l2 = mycc.spin2spatial(l2, mycc.mo_coeff.orbspin)
    print(abs(l1[0]-l1ref[0]).max())
    print(abs(l1[1]-l1ref[1]).max())
    print(abs(l2[0]-l2ref[0]).max())
    print(abs(l2[1]-l2ref[1]).max())
    print(abs(l2[2]-l2ref[2]).max())


# pyscf/cc/gccsd_t_rdm.py

#!/usr/bin/env python
# Copyright 2014-2021 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: Qiming Sun <osirpt.sun@gmail.com>
#

import numpy
from pyscf import lib
from pyscf.cc import gccsd_rdm

def _gamma1_intermediates(mycc, t1, t2, l1, l2, eris=None):
    doo, dov, dvo, dvv = gccsd_rdm._gamma1_intermediates(mycc, t1, t2, l1, l2)

    if eris is None: eris = mycc.ao2mo()

    nocc, nvir = t1.shape
    bcei = numpy.asarray(eris.ovvv).conj().transpose(3,2,1,0)
    majk = numpy.asarray(eris.ooov).conj().transpose(2,3,0,1)
    bcjk = numpy.asarray(eris.oovv).conj().transpose(2,3,0,1)

    mo_e = eris.mo_energy
    eia = mo_e[:nocc,None] - mo_e[nocc:]
    d3 = lib.direct_sum('ia+jb+kc->ijkabc', eia, eia, eia)

    t3c =(numpy.einsum('jkae,bcei->ijkabc', t2, bcei) -
          numpy.einsum('imbc,majk->ijkabc', t2, majk))
    t3c = t3c - t3c.transpose(0,1,2,4,3,5) - t3c.transpose(0,1,2,5,4,3)
    t3c = t3c - t3c.transpose(1,0,2,3,4,5) - t3c.transpose(2,1,0,3,4,5)
    t3c /= d3

    t3d = numpy.einsum('ia,bcjk->ijkabc', t1, bcjk)
    t3d += numpy.einsum('ai,jkbc->ijkabc', eris.fock[nocc:,:nocc], t2)
    t3d = t3d - t3d.transpose(0,1,2,4,3,5) - t3d.transpose(0,1,2,5,4,3)
    t3d = t3d - t3d.transpose(1,0,2,3,4,5) - t3d.transpose(2,1,0,3,4,5)
    t3d /= d3

    goo = numpy.einsum('iklabc,jklabc->ij', (t3c+t3d).conj(), t3c) * (1./12)
    gvv = numpy.einsum('ijkacd,ijkbcd->ab', t3c+t3d, t3c.conj()) * (1./12)
    doo[numpy.diag_indices(nocc)] -= goo.diagonal()
    dvv[numpy.diag_indices(nvir)] += gvv.diagonal()
    dvo += numpy.einsum('ijab,ijkabc->ck', t2.conj(), t3c) * (1./4)

    return doo, dov, dvo, dvv

# gamma2 intermediates in Chemist's notation
def _gamma2_intermediates(mycc, t1, t2, l1, l2, eris=None):
    dovov, dvvvv, doooo, doovv, dovvo, dvvov, dovvv, dooov = \
            gccsd_rdm._gamma2_intermediates(mycc, t1, t2, l1, l2)
    if eris is None: eris = mycc.ao2mo()

    nocc, nvir = t1.shape
    bcei = numpy.asarray(eris.ovvv).conj().transpose(3,2,1,0)
    majk = numpy.asarray(eris.ooov).conj().transpose(2,3,0,1)
    bcjk = numpy.asarray(eris.oovv).conj().transpose(2,3,0,1)

    mo_e = eris.mo_energy
    eia = mo_e[:nocc,None] - mo_e[nocc:]
    d3 = lib.direct_sum('ia+jb+kc->ijkabc', eia, eia, eia)

    t3c =(numpy.einsum('jkae,bcei->ijkabc', t2, bcei) -
          numpy.einsum('imbc,majk->ijkabc', t2, majk))
    t3c = t3c - t3c.transpose(0,1,2,4,3,5) - t3c.transpose(0,1,2,5,4,3)
    t3c = t3c - t3c.transpose(1,0,2,3,4,5) - t3c.transpose(2,1,0,3,4,5)
    t3c /= d3

    t3d = numpy.einsum('ia,bcjk->ijkabc', t1, bcjk)
    t3d += numpy.einsum('ai,jkbc->ijkabc', eris.fock[nocc:,:nocc], t2)
    t3d = t3d - t3d.transpose(0,1,2,4,3,5) - t3d.transpose(0,1,2,5,4,3)
    t3d = t3d - t3d.transpose(1,0,2,3,4,5) - t3d.transpose(2,1,0,3,4,5)
    t3d /= d3

    goovv = numpy.einsum('kc,ijkabc->ijab', t1.conj(), t3c).conj() * (1./4)
    dovov += goovv.transpose(0,2,1,3) - goovv.transpose(0,3,1,2)

    m3 = t3c * 2 + t3d
# *(1/8) instead of (1/4) because ooov appears 4 times in the 2pdm tensor due
# to symmetrization, and its contribution is scaled by 1/2 in Tr(H,2pdm)
    gooov  = numpy.einsum('imbc,ijkabc->jkma', t2, m3.conj()) * (1./8)
    dooov -= gooov.transpose(0,2,1,3) - gooov.transpose(1,2,0,3)

    govvv  = numpy.einsum('jkae,ijkabc->iecb', t2, m3.conj()) * (1./8)
    dovvv += govvv.transpose(0,2,1,3) - govvv.transpose(0,3,1,2)
    return dovov, dvvvv, doooo, doovv, dovvo, dvvov, dovvv, dooov

def make_rdm1(mycc, t1, t2, l1, l2, eris=None, ao_repr=False):
    d1 = _gamma1_intermediates(mycc, t1, t2, l1, l2, eris)
    return gccsd_rdm._make_rdm1(mycc, d1, True, ao_repr=ao_repr)

# rdm2 in Chemist's notation
def make_rdm2(mycc, t1, t2, l1, l2, eris=None):
    d1 = _gamma1_intermediates(mycc, t1, t2, l1, l2, eris)
    d2 = _gamma2_intermediates(mycc, t1, t2, l1, l2, eris)
    return gccsd_rdm._make_rdm2(mycc, d1, d2, True, True)


if __name__ == '__main__':
    from functools import reduce
    from pyscf import gto
    from pyscf import scf
    from pyscf import ao2mo
    from pyscf import cc

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -.957 , .587)],
        [1 , (0.2,  .757 , .487)]]
    mol.basis = '631g'
    mol.build()
    mf0 = mf = scf.RHF(mol).run(conv_tol=1.)
    mf = scf.addons.convert_to_ghf(mf)

    from pyscf.cc import ccsd_t_lambda_slow as ccsd_t_lambda
    from pyscf.cc import ccsd_t_rdm_slow as ccsd_t_rdm
    mycc0 = cc.CCSD(mf0)
    eris0 = mycc0.ao2mo()
    mycc0.kernel(eris=eris0)
    t1 = mycc0.t1
    t2 = mycc0.t2
    imds = ccsd_t_lambda.make_intermediates(mycc0, t1, t2, eris0)
    l1, l2 = ccsd_t_lambda.update_lambda(mycc0, t1, t2, t1, t2, eris0, imds)
    dm1ref = ccsd_t_rdm.make_rdm1(mycc0, t1, t2, l1, l2, eris0)
    dm2ref = ccsd_t_rdm.make_rdm2(mycc0, t1, t2, l1, l2, eris0)

    mycc = cc.GCCSD(mf)
    eris = mycc.ao2mo()
    t1 = mycc.spatial2spin(t1, mycc.mo_coeff.orbspin)
    t2 = mycc.spatial2spin(t2, mycc.mo_coeff.orbspin)
    l1 = mycc.spatial2spin(l1, mycc.mo_coeff.orbspin)
    l2 = mycc.spatial2spin(l2, mycc.mo_coeff.orbspin)
    gdm1 = make_rdm1(mycc, t1, t2, l1, l2, eris)
    gdm2 = make_rdm2(mycc, t1, t2, l1, l2, eris)
    idxa = numpy.where(mycc.mo_coeff.orbspin == 0)[0]
    idxb = numpy.where(mycc.mo_coeff.orbspin == 1)[0]

    trdm1 = gdm1[idxa[:,None],idxa]
    trdm1+= gdm1[idxb[:,None],idxb]
    trdm2 = gdm2[idxa[:,None,None,None],idxa[:,None,None],idxa[:,None],idxa]
    trdm2+= gdm2[idxb[:,None,None,None],idxb[:,None,None],idxb[:,None],idxb]
    dm2ab = gdm2[idxa[:,None,None,None],idxa[:,None,None],idxb[:,None],idxb]
    trdm2+= dm2ab
    trdm2+= dm2ab.transpose(2,3,0,1)
    print(abs(trdm1 - dm1ref).max())
    print(abs(trdm2 - dm2ref).max())

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = '631g'
    mol.spin = 2
    mol.charge = 2
    mol.build()
    mf0 = mf = scf.UHF(mol).run(conv_tol=1)
    mf = scf.addons.convert_to_ghf(mf)

    from pyscf.cc import uccsd_t_lambda
    from pyscf.cc import uccsd_t_rdm
    mycc0 = cc.UCCSD(mf0)
    eris0 = mycc0.ao2mo()
    mycc0.kernel(eris=eris0)
    t1 = mycc0.t1
    t2 = mycc0.t2
    imds = uccsd_t_lambda.make_intermediates(mycc0, t1, t2, eris0)
    l1, l2 = uccsd_t_lambda.update_lambda(mycc0, t1, t2, t1, t2, eris0, imds)
    dm1ref = uccsd_t_rdm.make_rdm1(mycc0, t1, t2, l1, l2, eris0)
    dm2ref = uccsd_t_rdm.make_rdm2(mycc0, t1, t2, l1, l2, eris0)

    mycc = cc.GCCSD(mf)
    eris = mycc.ao2mo()
    t1 = mycc.spatial2spin(t1, mycc.mo_coeff.orbspin)
    t2 = mycc.spatial2spin(t2, mycc.mo_coeff.orbspin)
    l1 = mycc.spatial2spin(l1, mycc.mo_coeff.orbspin)
    l2 = mycc.spatial2spin(l2, mycc.mo_coeff.orbspin)
    gdm1 = make_rdm1(mycc, t1, t2, l1, l2, eris)
    gdm2 = make_rdm2(mycc, t1, t2, l1, l2, eris)
    idxa = numpy.where(mycc.mo_coeff.orbspin == 0)[0]
    idxb = numpy.where(mycc.mo_coeff.orbspin == 1)[0]

    print(abs(dm1ref[0] - gdm1[idxa[:,None],idxa]).max())
    print(abs(dm1ref[1] - gdm1[idxb[:,None],idxb]).max())
    print(abs(dm2ref[0] - gdm2[idxa[:,None,None,None],idxa[:,None,None],idxa[:,None],idxa]).max())
    print(abs(dm2ref[1] - gdm2[idxa[:,None,None,None],idxa[:,None,None],idxb[:,None],idxb]).max())
    print(abs(dm2ref[2] - gdm2[idxb[:,None,None,None],idxb[:,None,None],idxb[:,None],idxb]).max())

    ecc, t1, t2 = mycc.kernel(eris=eris)
    e3ref = mycc.e_tot + mycc.ccsd_t()
    l1, l2 = mycc.solve_lambda(eris=eris)
    dm1 = make_rdm1(mycc, t1, t2, l1, l2, eris=eris)
    dm2 = make_rdm2(mycc, t1, t2, l1, l2, eris=eris)
    nao = mol.nao_nr()
    mo_a = mf.mo_coeff[:nao]
    mo_b = mf.mo_coeff[nao:]
    nmo = mo_a.shape[1]
    eri = ao2mo.kernel(mf._eri, mo_a+mo_b, compact=False).reshape([nmo]*4)
    orbspin = mf.mo_coeff.orbspin
    sym_forbid = (orbspin[:,None] != orbspin)
    eri[sym_forbid,:,:] = 0
    eri[:,:,sym_forbid] = 0
    hcore = scf.RHF(mol).get_hcore()
    h1 = reduce(numpy.dot, (mo_a.T.conj(), hcore, mo_a))
    h1+= reduce(numpy.dot, (mo_b.T.conj(), hcore, mo_b))
    e3 = numpy.einsum('ij,ji', h1, dm1)
    e3+= numpy.einsum('ijkl,ijkl', eri, dm2) * .5
    e3+= mol.energy_nuc()
    print(e3 - e3ref)


# pyscf/cc/gccsd_t_slow.py

#!/usr/bin/env python
# Copyright 2014-2021 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: Qiming Sun <osirpt.sun@gmail.com>
#

'''
GHF-CCSD(T) with spin-orbital integrals
'''

import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf.cc import gccsd

# spin-orbital formula
# JCP 98, 8718 (1993); DOI:10.1063/1.464480
def kernel(cc, eris, t1=None, t2=None, max_memory=2000, verbose=logger.INFO):
    assert(isinstance(eris, gccsd._PhysicistsERIs))
    if t1 is None or t2 is None:
        t1, t2 = cc.t1, cc.t2

    nocc, nvir = t1.shape
    bcei = numpy.asarray(eris.ovvv).conj().transpose(3,2,1,0)
    majk = numpy.asarray(eris.ooov).conj().transpose(2,3,0,1)
    bcjk = numpy.asarray(eris.oovv).conj().transpose(2,3,0,1)

    mo_e = eris.mo_energy
    eia = mo_e[:nocc,None] - mo_e[nocc:]
    d3 = lib.direct_sum('ia+jb+kc->ijkabc', eia, eia, eia)

    t3c =(numpy.einsum('jkae,bcei->ijkabc', t2, bcei) -
          numpy.einsum('imbc,majk->ijkabc', t2, majk))
    t3c = t3c - t3c.transpose(0,1,2,4,3,5) - t3c.transpose(0,1,2,5,4,3)
    t3c = t3c - t3c.transpose(1,0,2,3,4,5) - t3c.transpose(2,1,0,3,4,5)
    t3c /= d3
    # e4 = numpy.einsum('ijkabc,ijkabc,ijkabc', t3c.conj(), d3, t3c) / 36
    # sia = numpy.einsum('jkbc,ijkabc->ia', eris.oovv, t3c) * .25
    # e5 = numpy.einsum('ia,ia', sia, t1.conj())
    # et = e4 + e5
    # return et
    t3d = numpy.einsum('ia,bcjk->ijkabc', t1, bcjk)
    t3d += numpy.einsum('ai,jkbc->ijkabc', eris.fock[nocc:,:nocc], t2)
    t3d = t3d - t3d.transpose(0,1,2,4,3,5) - t3d.transpose(0,1,2,5,4,3)
    t3d = t3d - t3d.transpose(1,0,2,3,4,5) - t3d.transpose(2,1,0,3,4,5)
    t3d /= d3
    et = numpy.einsum('ijkabc,ijkabc,ijkabc', (t3c+t3d).conj(), d3, t3c) / 36
    return et


if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf
    from pyscf import cc

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -.957 , .587)],
        [1 , (0.2,  .757 , .487)]]

    mol.basis = '631g'
    mol.build()
    mf = scf.RHF(mol).run(conv_tol=1e-1)
    mycc = cc.CCSD(mf).set(conv_tol=1e-11).run()
    et = mycc.ccsd_t()

    mycc = cc.GCCSD(scf.addons.convert_to_ghf(mf)).set(conv_tol=1e-11).run()
    eris = mycc.ao2mo()
    print(kernel(mycc, eris) - et)


# pyscf/cc/gintermediates.py

#!/usr/bin/env python
# Copyright 2014-2020 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.

import numpy as np
from pyscf import lib
from pyscf.lib import logger

#einsum = np.einsum
einsum = lib.einsum

# Ref: Gauss and Stanton, J. Chem. Phys. 103, 3561 (1995) Table III

# Section (a)

def make_tau(t2, t1a, t1b, fac=1, out=None):
    t1t1 = einsum('ia,jb->ijab', fac*0.5*t1a, t1b)
    t1t1 = t1t1 - t1t1.transpose(1,0,2,3)
    tau1 = t1t1 - t1t1.transpose(0,1,3,2)
    tau1 += t2
    return tau1

def cc_Fvv(t1, t2, eris):
    nocc, nvir = t1.shape
    fov = eris.fock[:nocc,nocc:]
    fvv = eris.fock[nocc:,nocc:]
    eris_vovv = np.asarray(eris.ovvv).transpose(1,0,3,2)
    tau_tilde = make_tau(t2, t1, t1,fac=0.5)
    Fae = fvv - 0.5*einsum('me,ma->ae',fov, t1)
    Fae += einsum('mf,amef->ae', t1, eris_vovv)
    Fae -= 0.5*einsum('mnaf,mnef->ae', tau_tilde, eris.oovv)
    return Fae

def cc_Foo(t1, t2, eris):
    nocc, nvir = t1.shape
    fov = eris.fock[:nocc,nocc:]
    foo = eris.fock[:nocc,:nocc]
    tau_tilde = make_tau(t2, t1, t1,fac=0.5)
    Fmi = ( foo + 0.5*einsum('me,ie->mi',fov, t1)
            + einsum('ne,mnie->mi', t1, eris.ooov)
            + 0.5*einsum('inef,mnef->mi', tau_tilde, eris.oovv) )
    return Fmi

def cc_Fov(t1, t2, eris):
    nocc, nvir = t1.shape
    fov = eris.fock[:nocc,nocc:]
    Fme = fov + einsum('nf,mnef->me', t1, eris.oovv)
    return Fme

def cc_Woooo(t1, t2, eris):
    tau = make_tau(t2, t1, t1)
    tmp = einsum('je,mnie->mnij', t1, eris.ooov)
    Wmnij = eris.oooo + tmp - tmp.transpose(0,1,3,2)
    Wmnij += 0.25*einsum('ijef,mnef->mnij', tau, eris.oovv)
    return Wmnij

def cc_Wvvvv(t1, t2, eris):
    tau = make_tau(t2, t1, t1)
    eris_ovvv = np.asarray(eris.ovvv)
    tmp = einsum('mb,mafe->bafe', t1, eris_ovvv)
    Wabef = np.asarray(eris.vvvv) - tmp + tmp.transpose(1,0,2,3)
    Wabef += einsum('mnab,mnef->abef', tau, 0.25*np.asarray(eris.oovv))
    return Wabef

def cc_Wovvo(t1, t2, eris):
    eris_ovvo = -np.asarray(eris.ovov).transpose(0,1,3,2)
    eris_oovo = -np.asarray(eris.ooov).transpose(0,1,3,2)
    Wmbej  = einsum('jf,mbef->mbej', t1, eris.ovvv)
    Wmbej -= einsum('nb,mnej->mbej', t1, eris_oovo)
    Wmbej -= 0.5*einsum('jnfb,mnef->mbej', t2, eris.oovv)
    Wmbej -= einsum('jf,nb,mnef->mbej', t1, t1, eris.oovv)
    Wmbej += eris_ovvo
    return Wmbej

### Section (b)

def Fvv(t1, t2, eris):
    ccFov = cc_Fov(t1, t2, eris)
    Fae = cc_Fvv(t1, t2, eris) - 0.5*einsum('ma,me->ae', t1,ccFov)
    return Fae

def Foo(t1, t2, eris):
    ccFov = cc_Fov(t1, t2, eris)
    Fmi = cc_Foo(t1, t2, eris) + 0.5*einsum('ie,me->mi', t1,ccFov)
    return Fmi

def Fov(t1, t2, eris):
    Fme = cc_Fov(t1, t2, eris)
    return Fme

def Woooo(t1, t2, eris):
    tau = make_tau(t2, t1, t1)
    Wmnij = 0.25*einsum('ijef,mnef->mnij', tau, eris.oovv)
    Wmnij += cc_Woooo(t1, t2, eris)
    return Wmnij

def Wvvvv(t1, t2, eris):
    tau = make_tau(t2, t1, t1)
    Wabef = cc_Wvvvv(t1, t2, eris)
    Wabef += einsum('mnab,mnef->abef', tau, .25*np.asarray(eris.oovv))
    return Wabef

def Wovvo(t1, t2, eris):
    Wmbej = -0.5*einsum('jnfb,mnef->mbej', t2, eris.oovv)
    Wmbej += cc_Wovvo(t1, t2, eris)
    return Wmbej

def Wooov(t1, t2, eris):
    Wmnie = einsum('if,mnfe->mnie', t1, eris.oovv)
    Wmnie += eris.ooov
    return Wmnie

def Wvovv(t1, t2, eris):
    Wamef = einsum('na,nmef->amef', -t1, eris.oovv)
    Wamef -= np.asarray(eris.ovvv).transpose(1,0,2,3)
    return Wamef

def Wovoo(t1, t2, eris):
    eris_ovvo = -np.asarray(eris.ovov).transpose(0,1,3,2)
    tmp1 = einsum('mnie,jnbe->mbij', eris.ooov, t2)
    tmp2 = einsum('ie,mbej->mbij', t1, eris_ovvo)
    tmp2 -= einsum('ie,njbf,mnef->mbij', t1, t2, eris.oovv)
    FFov = Fov(t1, t2, eris)
    WWoooo = Woooo(t1, t2, eris)
    tau = make_tau(t2, t1, t1)
    Wmbij = einsum('me,ijbe->mbij', -FFov, t2)
    Wmbij -= einsum('nb,mnij->mbij', t1, WWoooo)
    Wmbij += 0.5 * einsum('mbef,ijef->mbij', eris.ovvv, tau)
    Wmbij += tmp1 - tmp1.transpose(0,1,3,2)
    Wmbij += tmp2 - tmp2.transpose(0,1,3,2)
    Wmbij += np.asarray(eris.ooov).conj().transpose(2,3,0,1)
    return Wmbij

def Wvvvo(t1, t2, eris, _Wvvvv=None):
    eris_ovvo = -np.asarray(eris.ovov).transpose(0,1,3,2)
    eris_vvvo = -np.asarray(eris.ovvv).transpose(2,3,1,0).conj()
    eris_oovo = -np.asarray(eris.ooov).transpose(0,1,3,2)
    tmp1 = einsum('mbef,miaf->abei', eris.ovvv, t2)
    tmp2 = einsum('ma,mbei->abei', t1, eris_ovvo)
    tmp2 -= einsum('ma,nibf,mnef->abei', t1, t2, eris.oovv)
    FFov = Fov(t1, t2, eris)
    tau = make_tau(t2, t1, t1)
    Wabei  = 0.5 * einsum('mnei,mnab->abei', eris_oovo, tau)
    Wabei -= einsum('me,miab->abei', FFov, t2)
    Wabei += eris_vvvo
    Wabei -= tmp1 - tmp1.transpose(1,0,2,3)
    Wabei -= tmp2 - tmp2.transpose(1,0,2,3)
    nocc,nvir = t1.shape
    if _Wvvvv is None:
        _Wvvvv = Wvvvv(t1, t2, eris)
    Wabei += einsum('abef,if->abei', _Wvvvv, t1)
    return Wabei

########################################
# T3[2] related contributions to T1/T2
########################################

def _cp(a):
    return np.array(a, copy=False, order='C')

def get_t3p2_imds_slow(cc, t1, t2, eris=None, t3p2_ip_out=None, t3p2_ea_out=None):
    """Calculates T1, T2 amplitudes corrected by second-order T3 contribution
    and intermediates used in IP/EA-CCSD(T)a

    Args:
        cc (:obj:`GCCSD`):
            Object containing coupled-cluster results.
        t1 (:obj:`ndarray`):
            T1 amplitudes.
        t2 (:obj:`ndarray`):
            T2 amplitudes from which the T3[2] amplitudes are formed.
        eris (:obj:`_PhysicistsERIs`):
            Antisymmetrized electron-repulsion integrals in physicist's notation.
        t3p2_ip_out (:obj:`ndarray`):
            Store results of the intermediate used in IP-EOM-CCSD(T)a.
        t3p2_ea_out (:obj:`ndarray`):
            Store results of the intermediate used in EA-EOM-CCSD(T)a.

    Returns:
        delta_ccsd (float):
            Difference of perturbed and unperturbed CCSD ground-state energy,
                energy(T1 + T1[2], T2 + T2[2]) - energy(T1, T2)
        pt1 (:obj:`ndarray`):
            Perturbatively corrected T1 amplitudes.
        pt2 (:obj:`ndarray`):
            Perturbatively corrected T2 amplitudes.

    Reference:
        D. A. Matthews, J. F. Stanton "A new approach to approximate..."
            JCP 145, 124102 (2016); DOI:10.1063/1.4962910, Equation 14
        Shavitt and Bartlett "Many-body Methods in Physics and Chemistry"
            2009, Equation 10.33
    """
    if eris is None:
        eris = cc.ao2mo()
    fock = eris.fock
    nocc, nvir = t1.shape

    fov = fock[:nocc, nocc:]
    #foo = fock[:nocc, :nocc].diagonal()
    #fvv = fock[nocc:, nocc:].diagonal()

    mo_e_o = eris.mo_energy[:nocc]
    mo_e_v = eris.mo_energy[nocc:]

    oovv = _cp(eris.oovv)
    ovvv = _cp(eris.ovvv)
    ooov = _cp(eris.ooov)
    oovv = _cp(eris.oovv)
    vooo = _cp(ooov).conj().transpose(3, 2, 1, 0)
    vvvo = _cp(ovvv).conj().transpose(3, 2, 1, 0)

    ccsd_energy = cc.energy(t1, t2, eris)
    dtype = np.result_type(t1, t2)
    if np.issubdtype(dtype, np.dtype(complex).type):
        logger.error(cc, 't3p2 imds has not been strictly checked for use with complex integrals')

    if t3p2_ip_out is None:
        t3p2_ip_out = np.zeros((nocc,nvir,nocc,nocc), dtype=dtype)
    Wmcik = t3p2_ip_out

    if t3p2_ea_out is None:
        t3p2_ea_out = np.zeros((nvir,nvir,nvir,nocc), dtype=dtype)
    Wacek = t3p2_ea_out

    tmp_t3 = lib.einsum('bcdk,ijad->ijkabc', vvvo, t2)
    tmp_t3 -= lib.einsum('cmkj,imab->ijkabc', vooo, t2)

    # P(ijk)
    tmp_t3 = (tmp_t3 + tmp_t3.transpose(1, 2, 0, 3, 4, 5) +
                       tmp_t3.transpose(2, 0, 1, 3, 4, 5))
    # P(abc)
    tmp_t3 = (tmp_t3 + tmp_t3.transpose(0, 1, 2, 4, 5, 3) +
                       tmp_t3.transpose(0, 1, 2, 5, 3, 4))

    eia = mo_e_o[:, None] - mo_e_v[None, :]
    eijab = eia[:, None, :, None] + eia[None, :, None, :]
    eijkabc = eijab[:, :, None, :, :, None] + eia[None, None, :, None, None, :]
    tmp_t3 /= eijkabc

    pt1 = 0.25 * lib.einsum('mnef,imnaef->ia', oovv, tmp_t3)

    pt2 = lib.einsum('ijmabe,me->ijab', tmp_t3, fov)
    tmp2 = ovvv - 0.0*lib.einsum('nmef,nb->mbfe', oovv, t1)
    tmp = lib.einsum('ijmaef,mbfe->ijab', tmp_t3, tmp2)
    tmp = tmp - tmp.transpose(0, 1, 3, 2)  # P(ab)
    tmp *= 0.5
    pt2 = pt2 + tmp
    tmp2 = ooov + 0.0*lib.einsum('mnef,jf->nmje', oovv, t1)
    tmp = lib.einsum('inmabe,nmje->ijab', tmp_t3, tmp2)
    tmp *= -0.5
    tmp = tmp - tmp.transpose(1, 0, 2, 3)  # P(ij)
    pt2 = pt2 + tmp

    eia = mo_e_o[:, None] - mo_e_v[None, :]
    eijab = eia[:, None, :, None] + eia[None, :, None, :]

    pt1 /= eia
    pt2 /= eijab

    pt1 = pt1 + t1
    pt2 = pt2 + t2

    Wmcik += 0.5*lib.einsum('ijkabc,mjab->mcik', tmp_t3, oovv)
    Wacek += -0.5*lib.einsum('ijkabc,ijeb->acek', tmp_t3, oovv)

    delta_ccsd_energy = cc.energy(pt1, pt2, eris) - ccsd_energy
    logger.info(cc, 'CCSD energy T3[2] correction : %16.12e', delta_ccsd_energy)
    return delta_ccsd_energy, pt1, pt2, Wmcik, Wacek


# pyscf/cc/qcisd.py

#!/usr/bin/env python
# Copyright 2014-2021 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: Qiming Sun <osirpt.sun@gmail.com>
#         Timothy Berkelbach <tim.berkelbach@gmail.com>
#

'''
QCISD for real integrals
8-fold permutation symmetry has been used
(ij|kl) = (ji|kl) = (kl|ij) = ...
'''


import numpy
from pyscf import gto
from pyscf import lib
from pyscf.lib import logger
from pyscf.cc.ccsd import CCSD, _add_vvvv, _flops, _ChemistsERIs
from pyscf import __config__

BLKMIN = getattr(__config__, 'cc_ccsd_blkmin', 4)
MEMORYMIN = getattr(__config__, 'cc_ccsd_memorymin', 2000)


# t1: ia
# t2: ijab
def kernel(mycc, eris=None, t1=None, t2=None, max_cycle=50, tol=1e-8,
           tolnormt=1e-6, verbose=None):
    log = logger.new_logger(mycc, verbose)
    if eris is None:
        eris = mycc.ao2mo(mycc.mo_coeff)
    if t1 is None and t2 is None:
        t1, t2 = mycc.get_init_guess(eris)
    elif t2 is None:
        t2 = mycc.get_init_guess(eris)[1]

    cput1 = cput0 = (logger.process_clock(), logger.perf_counter())
    eold = 0
    eccsd = mycc.energy(t1, t2, eris)
    log.info('Init E_corr(QCISD) = %.15g', eccsd)

    if isinstance(mycc.diis, lib.diis.DIIS):
        adiis = mycc.diis
    elif mycc.diis:
        adiis = lib.diis.DIIS(mycc, mycc.diis_file, incore=mycc.incore_complete)
        adiis.space = mycc.diis_space
    else:
        adiis = None

    conv = False
    for istep in range(max_cycle):
        t1new, t2new = mycc.update_amps(t1, t2, eris)
        tmpvec = mycc.amplitudes_to_vector(t1new, t2new)
        tmpvec -= mycc.amplitudes_to_vector(t1, t2)
        normt = numpy.linalg.norm(tmpvec)
        tmpvec = None
        if mycc.iterative_damping < 1.0:
            alpha = mycc.iterative_damping
            t1new = (1-alpha) * t1 + alpha * t1new
            t2new *= alpha
            t2new += (1-alpha) * t2
        t1, t2 = t1new, t2new
        t1new = t2new = None
        t1, t2 = mycc.run_diis(t1, t2, istep, normt, eccsd-eold, adiis)
        eold, eccsd = eccsd, mycc.energy(t1, t2, eris)
        log.info('cycle = %d  E_corr(QCISD) = %.15g  dE = %.9g  norm(t1,t2) = %.6g',
                 istep+1, eccsd, eccsd - eold, normt)
        cput1 = log.timer('QCISD iter', *cput1)
        if abs(eccsd-eold) < tol and normt < tolnormt:
            conv = True
            break
    log.timer('QCISD', *cput0)
    return conv, eccsd, t1, t2


def update_amps(mycc, t1, t2, eris):
    assert(isinstance(eris, _ChemistsERIs))

    time0 = logger.process_clock(), logger.perf_counter()
    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocc, nvir = t1.shape
    fock = eris.fock
    mo_e_o = eris.mo_energy[:nocc]
    mo_e_v = eris.mo_energy[nocc:] + mycc.level_shift

    t1new = numpy.zeros_like(t1)
    t2new = _add_vvvv(mycc, 0*t1, t2, eris, t2sym='jiba')
    t2new *= .5  # *.5 because t2+t2.transpose(1,0,3,2) in the end
    time1 = log.timer_debug1('vvvv', *time0)

#** make_inter_F
    fov = fock[:nocc,nocc:].copy()
    t1new += fov

    foo = fock[:nocc,:nocc] - numpy.diag(mo_e_o)
    fvv = fock[nocc:,nocc:] - numpy.diag(mo_e_v)

    if mycc.incore_complete:
        fswap = None
    else:
        fswap = lib.H5TmpFile()
    fwVOov, fwVooV = _add_ovvv_(mycc, t1, t2, eris, fvv, t1new, t2new, fswap)
    time1 = log.timer_debug1('ovvv', *time1)

    woooo = numpy.asarray(eris.oooo).transpose(0,2,1,3).copy()

    unit = nocc**2*nvir*7 + nocc**3 + nocc*nvir**2
    mem_now = lib.current_memory()[0]
    max_memory = max(0, mycc.max_memory - mem_now)
    blksize = min(nvir, max(BLKMIN, int((max_memory*.9e6/8-nocc**4)/unit)))
    log.debug1('max_memory %d MB,  nocc,nvir = %d,%d  blksize = %d',
               max_memory, nocc, nvir, blksize)

    for p0, p1 in lib.prange(0, nvir, blksize):
        wVOov = fwVOov[p0:p1]
        wVooV = fwVooV[p0:p1]
        eris_ovoo = eris.ovoo[:,p0:p1]
        eris_oovv = numpy.empty((nocc,nocc,p1-p0,nvir))
        def load_oovv(p0, p1):
            eris_oovv[:] = eris.oovv[:,:,p0:p1]
        with lib.call_in_background(load_oovv, sync=not mycc.async_io) as prefetch_oovv:
            #:eris_oovv = eris.oovv[:,:,p0:p1]
            prefetch_oovv(p0, p1)
            wVOov -= lib.einsum('jbik,ka->bjia', eris_ovoo, t1)
            t2new[:,:,p0:p1] += wVOov.transpose(1,2,0,3)
            eris_ovoo = None
        load_oovv = prefetch_oovv = None

        wVOov *= 0 # QCI

        eris_ovvo = numpy.empty((nocc,p1-p0,nvir,nocc))
        def load_ovvo(p0, p1):
            eris_ovvo[:] = eris.ovvo[:,p0:p1]
        with lib.call_in_background(load_ovvo, sync=not mycc.async_io) as prefetch_ovvo:
            #:eris_ovvo = eris.ovvo[:,p0:p1]
            prefetch_ovvo(p0, p1)
            t1new[:,p0:p1] -= numpy.einsum('jb,jiab->ia', t1, eris_oovv)
            wVooV -= eris_oovv.transpose(2,0,1,3)
            wVOov += wVooV*.5  #: bjia + bija*.5
        load_ovvo = prefetch_ovvo = None

        t2new[:,:,p0:p1] += (eris_ovvo*0.5).transpose(0,3,1,2)
        eris_voov = eris_ovvo.conj().transpose(1,0,3,2)
        t1new[:,p0:p1] += 2*numpy.einsum('jb,aijb->ia', t1, eris_voov)
        eris_ovvo = None
        eris_oovv = tmp = None

        fov[:,p0:p1] += numpy.einsum('kc,aikc->ia', t1, eris_voov) * 2
        fov[:,p0:p1] -= numpy.einsum('kc,akic->ia', t1, eris_voov)

        tau = t2[:,:,p0:p1]
        theta  = tau.transpose(1,0,2,3) * 2
        theta -= tau
        fvv -= lib.einsum('cjia,cjib->ab', theta.transpose(2,1,0,3), eris_voov)
        foo += lib.einsum('aikb,kjab->ij', eris_voov, theta)
        theta = None
        woooo += lib.einsum('ijab,aklb->ijkl', tau, eris_voov)
        tau = None

        def update_wVooV(q0, q1, tau):
            wVooV[:] += lib.einsum('bkic,jkca->bija', eris_voov[:,:,:,q0:q1], tau)
        with lib.call_in_background(update_wVooV, sync=not mycc.async_io) as update_wVooV:
            for q0, q1 in lib.prange(0, nvir, blksize):
                tau  = t2[:,:,q0:q1] * .5
                #:wVooV += lib.einsum('bkic,jkca->bija', eris_voov[:,:,:,q0:q1], tau)
                update_wVooV(q0, q1, tau)
        tau = update_wVooV = None
        def update_t2(q0, q1, tmp):
            t2new[:,:,q0:q1] += tmp.transpose(2,0,1,3)
            tmp *= .5
            t2new[:,:,q0:q1] += tmp.transpose(0,2,1,3)
        with lib.call_in_background(update_t2, sync=not mycc.async_io) as update_t2:
            for q0, q1 in lib.prange(0, nvir, blksize):
                tmp = lib.einsum('jkca,ckib->jaib', t2[:,:,p0:p1,q0:q1], wVooV)
                #:t2new[:,:,q0:q1] += tmp.transpose(2,0,1,3)
                #:tmp *= .5
                #:t2new[:,:,q0:q1] += tmp.transpose(0,2,1,3)
                update_t2(q0, q1, tmp)
                tmp = None

        wVOov += eris_voov
        eris_VOov = -.5 * eris_voov.transpose(0,2,1,3)
        eris_VOov += eris_voov
        eris_voov = None
        def update_wVOov(q0, q1, tau):
            wVOov[:,:,:,q0:q1] += .5 * lib.einsum('aikc,kcjb->aijb', eris_VOov, tau)
        with lib.call_in_background(update_wVOov, sync=not mycc.async_io) as update_wVOov:
            for q0, q1 in lib.prange(0, nvir, blksize):
                tau  = t2[:,:,q0:q1].transpose(1,3,0,2) * 2
                tau -= t2[:,:,q0:q1].transpose(0,3,1,2)
                #:wVOov[:,:,:,q0:q1] += .5 * lib.einsum('aikc,kcjb->aijb', eris_VOov, tau)
                update_wVOov(q0, q1, tau)
                tau = None
        def update_t2(q0, q1, theta):
            t2new[:,:,q0:q1] += lib.einsum('kica,ckjb->ijab', theta, wVOov)
        with lib.call_in_background(update_t2, sync=not mycc.async_io) as update_t2:
            for q0, q1 in lib.prange(0, nvir, blksize):
                theta  = t2[:,:,p0:p1,q0:q1] * 2
                theta -= t2[:,:,p0:p1,q0:q1].transpose(1,0,2,3)
                #:t2new[:,:,q0:q1] += lib.einsum('kica,ckjb->ijab', theta, wVOov)
                update_t2(q0, q1, theta)
                theta = None
        eris_VOov = wVOov = wVooV = update_wVOov = None
        time1 = log.timer_debug1('voov [%d:%d]'%(p0, p1), *time1)
    fwVOov = fwVooV = fswap = None

    for p0, p1 in lib.prange(0, nvir, blksize):
        theta = t2[:,:,p0:p1].transpose(1,0,2,3) * 2 - t2[:,:,p0:p1]
        t1new += numpy.einsum('jb,ijba->ia', fov[:,p0:p1], theta)
        t1new -= lib.einsum('jbki,kjba->ia', eris.ovoo[:,p0:p1], theta)

        tau = t2[:,:,p0:p1]
        t2new[:,:,p0:p1] += .5 * lib.einsum('ijkl,klab->ijab', woooo, tau)
        theta = tau = None

    t2new += lib.einsum('ijac,bc->ijab', t2, fvv)
    t2new -= lib.einsum('ki,kjab->ijab', foo, t2)

    eia = mo_e_o[:,None] - mo_e_v
    t1new += numpy.einsum('ib,ab->ia', t1, fvv)
    t1new -= numpy.einsum('ja,ji->ia', t1, foo)
    t1new /= eia

    #: t2new = t2new + t2new.transpose(1,0,3,2)
    for i in range(nocc):
        if i > 0:
            t2new[i,:i] += t2new[:i,i].transpose(0,2,1)
            t2new[i,:i] /= lib.direct_sum('a,jb->jab', eia[i], eia[:i])
            t2new[:i,i] = t2new[i,:i].transpose(0,2,1)
        t2new[i,i] = t2new[i,i] + t2new[i,i].T
        t2new[i,i] /= lib.direct_sum('a,b->ab', eia[i], eia[i])

    time0 = log.timer_debug1('update t1 t2', *time0)
    return t1new, t2new

def _add_ovvv_(mycc, t1, t2, eris, fvv, t1new, t2new, fswap):
    time1 = logger.process_clock(), logger.perf_counter()
    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocc, nvir = t1.shape
    nvir_pair = nvir * (nvir+1) // 2

    if fswap is None:
        wVOov = numpy.zeros((nvir,nocc,nocc,nvir))
    else:
        wVOov = fswap.create_dataset('wVOov', (nvir,nocc,nocc,nvir), 'f8')
    wooVV = numpy.zeros((nocc,nocc*nvir_pair))

    max_memory = mycc.max_memory - lib.current_memory()[0]
    unit = nocc*nvir**2*3 + nocc**2*nvir + 2
    blksize = min(nvir, max(BLKMIN, int((max_memory*.95e6/8-wooVV.size)/unit)))
    if not mycc.direct:
        unit = nocc*nvir**2*3 + nocc**2*nvir + 2 + nocc*nvir**2 + nocc*nvir
        blksize = min(nvir, max(BLKMIN, int((max_memory*.95e6/8-wooVV.size-nocc**2*nvir)/unit)))
    log.debug1('max_memory %d MB,  nocc,nvir = %d,%d  blksize = %d',
               max_memory, nocc, nvir, blksize)

    def load_ovvv(buf, p0):
        if p0 < nvir:
            p1 = min(nvir, p0+blksize)
            buf[:p1-p0] = eris.ovvv[:,p0:p1].transpose(1,0,2)

    with lib.call_in_background(load_ovvv, sync=not mycc.async_io) as prefetch:
        buf = numpy.empty((blksize,nocc,nvir_pair))
        buf_prefetch = numpy.empty((blksize,nocc,nvir_pair))

        load_ovvv(buf_prefetch, 0)
        for p0, p1 in lib.prange(0, nvir, blksize):
            buf, buf_prefetch = buf_prefetch, buf
            prefetch(buf_prefetch, p1)

            eris_vovv = buf[:p1-p0]
            eris_vovv = lib.unpack_tril(eris_vovv.reshape((p1-p0)*nocc,nvir_pair))
            eris_vovv = eris_vovv.reshape(p1-p0,nocc,nvir,nvir)

            wVOov[p0:p1] = lib.einsum('biac,jc->bija', eris_vovv, t1)

            theta = t2[:,:,p0:p1].transpose(1,2,0,3) * 2
            theta -= t2[:,:,p0:p1].transpose(0,2,1,3)
            t1new += lib.einsum('icjb,cjba->ia', theta, eris_vovv)
            theta = None
            time1 = log.timer_debug1('vovv [%d:%d]'%(p0, p1), *time1)

    if fswap is None:
        wooVV = lib.unpack_tril(wooVV.reshape(nocc**2,nvir_pair))
        return wVOov, wooVV.reshape(nocc,nocc,nvir,nvir).transpose(2,1,0,3)
    else:
        fswap.create_dataset('wVooV', (nvir,nocc,nocc,nvir), 'f8')
        wooVV = wooVV.reshape(nocc,nocc,nvir_pair)
        tril2sq = lib.square_mat_in_trilu_indices(nvir)
        for p0, p1 in lib.prange(0, nvir, blksize):
            fswap['wVooV'][p0:p1] = wooVV[:,:,tril2sq[p0:p1]].transpose(2,1,0,3)
        return fswap['wVOov'], fswap['wVooV']

def energy(mycc, t1=None, t2=None, eris=None):
    '''CCSD correlation energy'''
    if t1 is None: t1 = mycc.t1
    if t2 is None: t2 = mycc.t2
    if eris is None: eris = mycc.ao2mo()

    nocc, nvir = t1.shape
    fock = eris.fock
    e = numpy.einsum('ia,ia', fock[:nocc,nocc:], t1) * 2
    max_memory = mycc.max_memory - lib.current_memory()[0]
    blksize = int(min(nvir, max(BLKMIN, max_memory*.3e6/8/(nocc**2*nvir+1))))
    for p0, p1 in lib.prange(0, nvir, blksize):
        eris_ovvo = eris.ovvo[:,p0:p1]
        tau = t2[:,:,p0:p1]
        e += 2 * numpy.einsum('ijab,iabj', tau, eris_ovvo)
        e -=     numpy.einsum('jiab,iabj', tau, eris_ovvo)
    if abs(e.imag) > 1e-4:
        logger.warn(mycc, 'Non-zero imaginary part found in QCISD energy %s', e)
    return e.real

def as_scanner(cc):
    '''Generating a scanner/solver for QCISD PES.

    The returned solver is a function. This function requires one argument
    "mol" as input and returns total QCISD energy.

    '''
    if isinstance(cc, lib.SinglePointScanner):
        return cc

    logger.info(cc, 'Set %s as a scanner', cc.__class__)

    class QCISD_Scanner(cc.__class__, lib.SinglePointScanner):
        def __init__(self, cc):
            self.__dict__.update(cc.__dict__)
            self._scf = cc._scf.as_scanner()
        def __call__(self, mol_or_geom, **kwargs):
            if isinstance(mol_or_geom, gto.Mole):
                mol = mol_or_geom
            else:
                mol = self.mol.set_geom_(mol_or_geom, inplace=False)

            if self.t2 is not None:
                last_size = self.vector_size()
            else:
                last_size = 0

            self.reset(mol)

            mf_scanner = self._scf
            mf_scanner(mol)
            self.mo_coeff = mf_scanner.mo_coeff
            self.mo_occ = mf_scanner.mo_occ
            if last_size != self.vector_size():
                self.t1 = self.t2 = None
            self.kernel(self.t1, self.t2, **kwargs)
            return self.e_tot
    return QCISD_Scanner(cc)


class QCISD(CCSD):
    '''restricted QCISD
    '''

    def dump_flags(self, verbose=None):
        log = logger.new_logger(self, verbose)
        log.info('')
        log.info('******** %s ********', self.__class__)
        log.info('QCISD nocc = %s, nmo = %s', self.nocc, self.nmo)
        if self.frozen is not None:
            log.info('frozen orbitals %s', self.frozen)
        log.info('max_cycle = %d', self.max_cycle)
        log.info('direct = %d', self.direct)
        log.info('conv_tol = %g', self.conv_tol)
        log.info('conv_tol_normt = %s', self.conv_tol_normt)
        log.info('diis_space = %d', self.diis_space)
        #log.info('diis_file = %s', self.diis_file)
        log.info('diis_start_cycle = %d', self.diis_start_cycle)
        log.info('diis_start_energy_diff = %g', self.diis_start_energy_diff)
        log.info('max_memory %d MB (current use %d MB)',
                 self.max_memory, lib.current_memory()[0])
        if (log.verbose >= logger.DEBUG1 and
            self.__class__ == QCISD):
            nocc = self.nocc
            nvir = self.nmo - self.nocc
            flops = _flops(nocc, nvir)
            log.debug1('total FLOPs %s', flops)
        return self

    energy = energy
    _add_vvvv = _add_vvvv
    update_amps = update_amps

    def kernel(self, t1=None, t2=None, eris=None):
        return self.qcisd(t1, t2, eris)
    def qcisd(self, t1=None, t2=None, eris=None):
        assert(self.mo_coeff is not None)
        assert(self.mo_occ is not None)

        if self.verbose >= logger.WARN:
            self.check_sanity()
        self.dump_flags()

        if eris is None:
            eris = self.ao2mo(self.mo_coeff)

        self.e_hf = getattr(eris, 'e_hf', None)
        if self.e_hf is None:
            self.e_hf = self._scf.e_tot

        self.converged, self.e_corr, self.t1, self.t2 = \
                kernel(self, eris, t1, t2, max_cycle=self.max_cycle,
                       tol=self.conv_tol, tolnormt=self.conv_tol_normt,
                       verbose=self.verbose)
        self._finalize()
        return self.e_corr, self.t1, self.t2

    as_scanner = as_scanner

    def qcisd_t(self, t1=None, t2=None, eris=None):
        from pyscf.cc import qcisd_t
        if t1 is None: t1 = self.t1
        if t2 is None: t2 = self.t2
        if eris is None: eris = self.ao2mo(self.mo_coeff)
        return qcisd_t.kernel(self, eris, t1, t2, self.verbose)

RQCISD = QCISD


if __name__ == '__main__':
    from pyscf import scf

    mol = gto.Mole()
    #mol.atom = [
    #    [8 , (0. , 0.     , 0.)],
    #    [1 , (0. , -0.757 , 0.587)],
    #    [1 , (0. , 0.757  , 0.587)]]
    mol.atom = [['Ne', (0,0,0)]]
    mol.basis = 'cc-pvdz'
    mol.verbose = 7
    mol.spin = 0
    mol.build()
    mf = scf.RHF(mol).run(conv_tol=1e-14)

    mycc = QCISD(mf, frozen=1)
    ecc, t1, t2 = mycc.kernel()
    et = mycc.qcisd_t()
    print("QCISD(T) =", mycc.e_tot+et)

    mol = gto.Mole()
    mol.atom = """C  0.000  0.000  0.000
                  H  0.637  0.637  0.637
                  H -0.637 -0.637  0.637
                  H -0.637  0.637 -0.637
                  H  0.637 -0.637 -0.637"""
    mol.basis = 'cc-pvdz'
    mol.verbose = 7
    mol.spin = 0
    mol.build()
    mf = scf.RHF(mol).run(conv_tol=1e-14)

    mycc = QCISD(mf, frozen=1)
    ecc, t1, t2 = mycc.kernel()
    print(mycc.e_tot - -40.383989)
    et = mycc.qcisd_t()
    print(mycc.e_tot+et - -40.387679)


# pyscf/cc/qcisd_slow.py

#!/usr/bin/env python
# Copyright 2014-2021 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: Timothy Berkelbach <tim.berkelbach@gmail.com>
#

'''
Restricted QCISD implementation
The 4-index integrals are saved on disk entirely (without using any symmetry).

Note MO integrals are treated in chemist's notation

Ref:
'''


import numpy
import numpy as np

from pyscf import lib
from pyscf.lib import logger
from pyscf.cc import rccsd_slow as rccsd
from pyscf.cc import rintermediates as imd
from pyscf import __config__

BLKMIN = getattr(__config__, 'cc_ccsd_blkmin', 4)
MEMORYMIN = getattr(__config__, 'cc_ccsd_memorymin', 2000)


def kernel(mycc, eris=None, t1=None, t2=None, max_cycle=50, tol=1e-8,
           tolnormt=1e-6, verbose=None):
    '''Same as ccsd.kernel with strings modified to correct the method name'''
    log = logger.new_logger(mycc, verbose)
    if eris is None:
        eris = mycc.ao2mo(mycc.mo_coeff)
    if t1 is None and t2 is None:
        t1, t2 = mycc.get_init_guess(eris)
    elif t2 is None:
        t2 = mycc.get_init_guess(eris)[1]

    cput1 = cput0 = (logger.process_clock(), logger.perf_counter())
    eold = 0
    eccsd = mycc.energy(t1, t2, eris)
    log.info('Init E_corr(QCISD) = %.15g', eccsd)

    if isinstance(mycc.diis, lib.diis.DIIS):
        adiis = mycc.diis
    elif mycc.diis:
        adiis = lib.diis.DIIS(mycc, mycc.diis_file, incore=mycc.incore_complete)
        adiis.space = mycc.diis_space
    else:
        adiis = None

    conv = False
    for istep in range(max_cycle):
        t1new, t2new = mycc.update_amps(t1, t2, eris)
        tmpvec = mycc.amplitudes_to_vector(t1new, t2new)
        tmpvec -= mycc.amplitudes_to_vector(t1, t2)
        normt = numpy.linalg.norm(tmpvec)
        tmpvec = None
        if mycc.iterative_damping < 1.0:
            alpha = mycc.iterative_damping
            t1new = (1-alpha) * t1 + alpha * t1new
            t2new *= alpha
            t2new += (1-alpha) * t2
        t1, t2 = t1new, t2new
        t1new = t2new = None
        t1, t2 = mycc.run_diis(t1, t2, istep, normt, eccsd-eold, adiis)
        eold, eccsd = eccsd, mycc.energy(t1, t2, eris)
        log.info('cycle = %d  E_corr(QCISD) = %.15g  dE = %.9g  norm(t1,t2) = %.6g',
                 istep+1, eccsd, eccsd - eold, normt)
        cput1 = log.timer('QCISD iter', *cput1)
        if abs(eccsd-eold) < tol and normt < tolnormt:
            conv = True
            break
    log.timer('QCISD', *cput0)
    return conv, eccsd, t1, t2


def update_amps(cc, t1, t2, eris):
    # Ref: Hirata et al., J. Chem. Phys. 120, 2581 (2004) Eqs.(35)-(36)
    nocc, nvir = t1.shape
    fock = eris.fock

    fov = fock[:nocc,nocc:].copy()
    foo = fock[:nocc,:nocc].copy()
    fvv = fock[nocc:,nocc:].copy()

    Foo = imd.cc_Foo(0*t1,t2,eris)
    Fvv = imd.cc_Fvv(0*t1,t2,eris)
    Fov = imd.cc_Fov(t1,t2,eris)

    Foo -= np.diag(np.diag(foo))
    Fvv -= np.diag(np.diag(fvv))

    # T1 equation
    t1new = np.asarray(fov).conj().copy()
    t1new +=   lib.einsum('ac,ic->ia', Fvv, t1)
    t1new +=  -lib.einsum('ki,ka->ia', Foo, t1)
    t1new += 2*lib.einsum('kc,kica->ia', Fov, t2)
    t1new +=  -lib.einsum('kc,ikca->ia', Fov, t2)
    t1new += 2*lib.einsum('kcai,kc->ia', eris.ovvo, t1)
    t1new +=  -lib.einsum('kiac,kc->ia', eris.oovv, t1)
    eris_ovvv = np.asarray(eris.ovvv)
    t1new += 2*lib.einsum('kdac,ikcd->ia', eris_ovvv, t2)
    t1new +=  -lib.einsum('kcad,ikcd->ia', eris_ovvv, t2)
    t1new +=-2*lib.einsum('kilc,klac->ia', eris.ooov, t2)
    t1new +=   lib.einsum('likc,klac->ia', eris.ooov, t2)

    # T2 equation
    t2new = np.asarray(eris.ovov).conj().transpose(0,2,1,3).copy()
    Loo = imd.Loo(0*t1, t2, eris)
    Lvv = imd.Lvv(0*t1, t2, eris)
    Loo -= np.diag(np.diag(foo))
    Lvv -= np.diag(np.diag(fvv))
    Woooo = imd.cc_Woooo(0*t1, t2, eris)
    Wvoov = imd.cc_Wvoov(0*t1, t2, eris)
    Wvovo = imd.cc_Wvovo(0*t1, t2, eris)
    Wvvvv = imd.cc_Wvvvv(0*t1, t2, eris)
    t2new += lib.einsum('klij,klab->ijab', Woooo, t2)
    t2new += lib.einsum('abcd,ijcd->ijab', Wvvvv, t2)
    tmp = lib.einsum('ac,ijcb->ijab', Lvv, t2)
    t2new += (tmp + tmp.transpose(1,0,3,2))
    tmp = lib.einsum('ki,kjab->ijab', Loo, t2)
    t2new -= (tmp + tmp.transpose(1,0,3,2))
    tmp  = 2*lib.einsum('akic,kjcb->ijab', Wvoov, t2)
    tmp -=   lib.einsum('akci,kjcb->ijab', Wvovo, t2)
    t2new += (tmp + tmp.transpose(1,0,3,2))
    tmp = lib.einsum('akic,kjbc->ijab', Wvoov, t2)
    t2new -= (tmp + tmp.transpose(1,0,3,2))
    tmp = lib.einsum('bkci,kjac->ijab', Wvovo, t2)
    t2new -= (tmp + tmp.transpose(1,0,3,2))

    tmp2 = np.asarray(eris.ovvv).conj().transpose(1,3,0,2)
    tmp = lib.einsum('abic,jc->ijab', tmp2, t1)
    t2new += (tmp + tmp.transpose(1,0,3,2))
    tmp2 = np.asarray(eris.ooov).transpose(3,1,2,0).conj()
    tmp = lib.einsum('akij,kb->ijab', tmp2, t1)
    t2new -= (tmp + tmp.transpose(1,0,3,2))

    mo_e = eris.fock.diagonal().real
    eia = mo_e[:nocc,None] - mo_e[None,nocc:]
    eijab = lib.direct_sum('ia,jb->ijab',eia,eia)
    t1new /= eia
    t2new /= eijab

    return t1new, t2new


class QCISD(rccsd.RCCSD):
    '''restricted QCISD
    '''

    def kernel(self, t1=None, t2=None, eris=None):
        return self.qcisd(t1, t2, eris)
    def qcisd(self, t1=None, t2=None, eris=None):
        assert(self.mo_coeff is not None)
        assert(self.mo_occ is not None)

        if self.verbose >= logger.WARN:
            self.check_sanity()
        self.dump_flags()

        if eris is None:
            eris = self.ao2mo(self.mo_coeff)

        self.e_hf = getattr(eris, 'e_hf', None)
        if self.e_hf is None:
            self.e_hf = self._scf.e_tot

        self.converged, self.e_corr, self.t1, self.t2 = \
                kernel(self, eris, t1, t2, max_cycle=self.max_cycle,
                       tol=self.conv_tol, tolnormt=self.conv_tol_normt,
                       verbose=self.verbose)
        self._finalize()
        return self.e_corr, self.t1, self.t2

    def energy(self, t1=None, t2=None, eris=None):
        return rccsd.energy(self, t1*0, t2, eris)

    update_amps = update_amps

    def qcisd_t(self, t1=None, t2=None, eris=None):
        from pyscf.cc import qcisd_t_slow as qcisd_t
        if t1 is None: t1 = self.t1
        if t2 is None: t2 = self.t2
        if eris is None: eris = self.ao2mo(self.mo_coeff)
        return qcisd_t.kernel(self, eris, t1, t2, self.verbose)

    def density_fit(self, auxbasis=None, with_df=None):
        raise NotImplementedError


if __name__ == '__main__':
    from pyscf import gto, scf

    mol = gto.Mole()
    mol.atom = """C  0.000  0.000  0.000
                  H  0.637  0.637  0.637
                  H -0.637 -0.637  0.637
                  H -0.637  0.637 -0.637
                  H  0.637 -0.637 -0.637"""
    mol.basis = 'cc-pvdz'
    mol.verbose = 7
    mol.spin = 0
    mol.build()
    mf = scf.RHF(mol).run(conv_tol=1e-14)

    mycc = QCISD(mf, frozen=1)
    ecc, t1, t2 = mycc.kernel()
    print(mycc.e_tot - -40.383989)
    et = mycc.qcisd_t()
    print(mycc.e_tot+et - -40.387679)


# pyscf/cc/qcisd_t.py

#!/usr/bin/env python
# Copyright 2014-2020 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: Qiming Sun <osirpt.sun@gmail.com>
#         Timothy Berkelbach <tim.berkelbach@gmail.com>
#

'''
RHF-QCISD(T) for real integrals
'''


import ctypes
import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf.cc import _ccsd
from pyscf.cc.ccsd_t import _sort_eri, _sort_t2_vooo_

# t3 as ijkabc

# JCP 94, 442 (1991); DOI:10.1063/1.460359.  Error in Eq (1), should be [ia] >= [jb] >= [kc]
def kernel(mycc, eris, t1=None, t2=None, verbose=logger.NOTE):
    cpu1 = cpu0 = (logger.process_clock(), logger.perf_counter())
    log = logger.new_logger(mycc, verbose)
    if t1 is None: t1 = mycc.t1
    if t2 is None: t2 = mycc.t2

    nocc, nvir = t1.shape
    nmo = nocc + nvir

    dtype = numpy.result_type(t1, t2, eris.ovoo.dtype)
    if mycc.incore_complete:
        ftmp = None
        eris_vvop = numpy.zeros((nvir,nvir,nocc,nmo), dtype)
    else:
        ftmp = lib.H5TmpFile()
        eris_vvop = ftmp.create_dataset('vvop', (nvir,nvir,nocc,nmo), dtype)

    orbsym = _sort_eri(mycc, eris, nocc, nvir, eris_vvop, log)

    mo_energy, t1T, t2T, vooo, fvo, restore_t2_inplace = \
            _sort_t2_vooo_(mycc, orbsym, t1, t2, eris)
    cpu1 = log.timer_debug1('QCISD(T) sort_eri', *cpu1)

    cpu2 = list(cpu1)
    orbsym = numpy.hstack((numpy.sort(orbsym[:nocc]),numpy.sort(orbsym[nocc:])))
    o_ir_loc = numpy.append(0, numpy.cumsum(numpy.bincount(orbsym[:nocc], minlength=8)))
    v_ir_loc = numpy.append(0, numpy.cumsum(numpy.bincount(orbsym[nocc:], minlength=8)))
    o_sym = orbsym[:nocc]
    oo_sym = (o_sym[:,None] ^ o_sym).ravel()
    oo_ir_loc = numpy.append(0, numpy.cumsum(numpy.bincount(oo_sym, minlength=8)))
    nirrep = max(oo_sym) + 1

    orbsym   = orbsym.astype(numpy.int32)
    o_ir_loc = o_ir_loc.astype(numpy.int32)
    v_ir_loc = v_ir_loc.astype(numpy.int32)
    oo_ir_loc = oo_ir_loc.astype(numpy.int32)
    if dtype == numpy.complex128:
        drv = _ccsd.libcc.QCIsd_t_zcontract
    else:
        drv = _ccsd.libcc.QCIsd_t_contract
    et_sum = numpy.zeros(1, dtype=dtype)
    def contract(a0, a1, b0, b1, cache):
        cache_row_a, cache_col_a, cache_row_b, cache_col_b = cache
        drv(et_sum.ctypes.data_as(ctypes.c_void_p),
            mo_energy.ctypes.data_as(ctypes.c_void_p),
            t1T.ctypes.data_as(ctypes.c_void_p),
            t2T.ctypes.data_as(ctypes.c_void_p),
            vooo.ctypes.data_as(ctypes.c_void_p),
            fvo.ctypes.data_as(ctypes.c_void_p),
            ctypes.c_int(nocc), ctypes.c_int(nvir),
            ctypes.c_int(a0), ctypes.c_int(a1),
            ctypes.c_int(b0), ctypes.c_int(b1),
            ctypes.c_int(nirrep),
            o_ir_loc.ctypes.data_as(ctypes.c_void_p),
            v_ir_loc.ctypes.data_as(ctypes.c_void_p),
            oo_ir_loc.ctypes.data_as(ctypes.c_void_p),
            orbsym.ctypes.data_as(ctypes.c_void_p),
            cache_row_a.ctypes.data_as(ctypes.c_void_p),
            cache_col_a.ctypes.data_as(ctypes.c_void_p),
            cache_row_b.ctypes.data_as(ctypes.c_void_p),
            cache_col_b.ctypes.data_as(ctypes.c_void_p))
        cpu2[:] = log.timer_debug1('contract %d:%d,%d:%d'%(a0,a1,b0,b1), *cpu2)

    # The rest 20% memory for cache b
    mem_now = lib.current_memory()[0]
    max_memory = max(0, mycc.max_memory - mem_now)
    bufsize = (max_memory*.5e6/8-nocc**3*3*lib.num_threads())/(nocc*nmo)  #*.5 for async_io
    bufsize *= .5  #*.5 upper triangular part is loaded
    bufsize *= .8  #*.8 for [a0:a1]/[b0:b1] partition
    bufsize = max(8, bufsize)
    log.debug('max_memory %d MB (%d MB in use)', max_memory, mem_now)
    with lib.call_in_background(contract, sync=not mycc.async_io) as async_contract:
        for a0, a1 in reversed(list(lib.prange_tril(0, nvir, bufsize))):
            cache_row_a = numpy.asarray(eris_vvop[a0:a1,:a1], order='C')
            if a0 == 0:
                cache_col_a = cache_row_a
            else:
                cache_col_a = numpy.asarray(eris_vvop[:a0,a0:a1], order='C')
            async_contract(a0, a1, a0, a1, (cache_row_a,cache_col_a,
                                            cache_row_a,cache_col_a))

            for b0, b1 in lib.prange_tril(0, a0, bufsize/8):
                cache_row_b = numpy.asarray(eris_vvop[b0:b1,:b1], order='C')
                if b0 == 0:
                    cache_col_b = cache_row_b
                else:
                    cache_col_b = numpy.asarray(eris_vvop[:b0,b0:b1], order='C')
                async_contract(a0, a1, b0, b1, (cache_row_a,cache_col_a,
                                                cache_row_b,cache_col_b))

    t2 = restore_t2_inplace(t2T)
    et_sum *= 2
    if abs(et_sum[0].imag) > 1e-4:
        logger.warn(mycc, 'Non-zero imaginary part of QCISD(T) energy was found %s',
                    et_sum[0])
    et = et_sum[0].real
    log.timer('QCISD(T)', *cpu0)
    log.note('QCISD(T) correction = %.15g', et)
    return et


# pyscf/cc/qcisd_t_slow.py

#!/usr/bin/env python
# Copyright 2014-2020 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: Qiming Sun <osirpt.sun@gmail.com>
#         Timothy Berkelbach <tim.berkelbach@gmail.com>
#

import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf.cc.ccsd_t_slow import r3

'''
QCISD(T)
'''

# t3 as ijkabc

# JCP 94, 442 (1991); DOI:10.1063/1.460359.  Error in Eq (1), should be [ia] >= [jb] >= [kc]
def kernel(mycc, eris, t1=None, t2=None, verbose=logger.NOTE):
    if isinstance(verbose, logger.Logger):
        log = verbose
    else:
        log = logger.Logger(mycc.stdout, verbose)

    if t1 is None: t1 = mycc.t1
    if t2 is None: t2 = mycc.t2

    t1T = t1.T
    t2T = t2.transpose(2,3,0,1)

    nocc, nvir = t1.shape
    mo_e = eris.fock.diagonal().real
    e_occ, e_vir = mo_e[:nocc], mo_e[nocc:]
    eijk = lib.direct_sum('i,j,k->ijk', e_occ, e_occ, e_occ)

    eris_vvov = eris.get_ovvv().conj().transpose(1,3,0,2)
    eris_vooo = numpy.asarray(eris.ovoo).conj().transpose(1,0,3,2)
    eris_vvoo = numpy.asarray(eris.ovov).conj().transpose(1,3,0,2)
    fvo = eris.fock[nocc:,:nocc]
    def get_w(a, b, c):
        w = numpy.einsum('if,fkj->ijk', eris_vvov[a,b], t2T[c,:])
        w-= numpy.einsum('ijm,mk->ijk', eris_vooo[a,:], t2T[b,c])
        return w
    def get_v(a, b, c):
        v = numpy.einsum('ij,k->ijk', eris_vvoo[a,b], t1T[c])
        v+= numpy.einsum('ij,k->ijk', t2T[a,b], fvo[c])
        return v

    et = 0
    for a in range(nvir):
        for b in range(a+1):
            for c in range(b+1):
                d3 = eijk - e_vir[a] - e_vir[b] - e_vir[c]
                if a == c:  # a == b == c
                    d3 *= 6
                elif a == b or b == c:
                    d3 *= 2

                wabc = get_w(a, b, c)
                wacb = get_w(a, c, b)
                wbac = get_w(b, a, c)
                wbca = get_w(b, c, a)
                wcab = get_w(c, a, b)
                wcba = get_w(c, b, a)
                vabc = get_v(a, b, c)
                vacb = get_v(a, c, b)
                vbac = get_v(b, a, c)
                vbca = get_v(b, c, a)
                vcab = get_v(c, a, b)
                vcba = get_v(c, b, a)
                zabc = r3(wabc + vabc) / d3
                zacb = r3(wacb + vacb) / d3
                zbac = r3(wbac + vbac) / d3
                zbca = r3(wbca + vbca) / d3
                zcab = r3(wcab + vcab) / d3
                zcba = r3(wcba + vcba) / d3

                et+= numpy.einsum('ijk,ijk', wabc, zabc.conj())
                et+= numpy.einsum('ikj,ijk', wacb, zabc.conj())
                et+= numpy.einsum('jik,ijk', wbac, zabc.conj())
                et+= numpy.einsum('jki,ijk', wbca, zabc.conj())
                et+= numpy.einsum('kij,ijk', wcab, zabc.conj())
                et+= numpy.einsum('kji,ijk', wcba, zabc.conj())

                et+= numpy.einsum('ijk,ijk', wacb, zacb.conj())
                et+= numpy.einsum('ikj,ijk', wabc, zacb.conj())
                et+= numpy.einsum('jik,ijk', wcab, zacb.conj())
                et+= numpy.einsum('jki,ijk', wcba, zacb.conj())
                et+= numpy.einsum('kij,ijk', wbac, zacb.conj())
                et+= numpy.einsum('kji,ijk', wbca, zacb.conj())

                et+= numpy.einsum('ijk,ijk', wbac, zbac.conj())
                et+= numpy.einsum('ikj,ijk', wbca, zbac.conj())
                et+= numpy.einsum('jik,ijk', wabc, zbac.conj())
                et+= numpy.einsum('jki,ijk', wacb, zbac.conj())
                et+= numpy.einsum('kij,ijk', wcba, zbac.conj())
                et+= numpy.einsum('kji,ijk', wcab, zbac.conj())

                et+= numpy.einsum('ijk,ijk', wbca, zbca.conj())
                et+= numpy.einsum('ikj,ijk', wbac, zbca.conj())
                et+= numpy.einsum('jik,ijk', wcba, zbca.conj())
                et+= numpy.einsum('jki,ijk', wcab, zbca.conj())
                et+= numpy.einsum('kij,ijk', wabc, zbca.conj())
                et+= numpy.einsum('kji,ijk', wacb, zbca.conj())

                et+= numpy.einsum('ijk,ijk', wcab, zcab.conj())
                et+= numpy.einsum('ikj,ijk', wcba, zcab.conj())
                et+= numpy.einsum('jik,ijk', wacb, zcab.conj())
                et+= numpy.einsum('jki,ijk', wabc, zcab.conj())
                et+= numpy.einsum('kij,ijk', wbca, zcab.conj())
                et+= numpy.einsum('kji,ijk', wbac, zcab.conj())

                et+= numpy.einsum('ijk,ijk', wcba, zcba.conj())
                et+= numpy.einsum('ikj,ijk', wcab, zcba.conj())
                et+= numpy.einsum('jik,ijk', wbca, zcba.conj())
                et+= numpy.einsum('jki,ijk', wbac, zcba.conj())
                et+= numpy.einsum('kij,ijk', wacb, zcba.conj())
                et+= numpy.einsum('kji,ijk', wabc, zcba.conj())
    et *= 2
    log.info('QCISD(T) correction = %.15g', et)
    return et


# pyscf/cc/rccsd.py

#!/usr/bin/env python
# Copyright 2014-2021 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: Timothy Berkelbach <tim.berkelbach@gmail.com>
#

'''
Restricted CCSD implementation which supports both real and complex integrals.
The 4-index integrals are saved on disk entirely (without using any symmetry).
This code is slower than the pyscf.cc.ccsd implementation.

Note MO integrals are treated in chemist's notation

Ref: Hirata et al., J. Chem. Phys. 120, 2581 (2004)
'''


import numpy as np

from pyscf import lib
from pyscf import ao2mo
from pyscf.lib import logger
from pyscf.cc import ccsd
from pyscf.cc import rintermediates as imd
from pyscf.mp import mp2
from pyscf import __config__

BLKMIN = getattr(__config__, 'cc_ccsd_blkmin', 4)
MEMORYMIN = getattr(__config__, 'cc_ccsd_memorymin', 2000)

def update_amps(cc, t1, t2, eris):
    # Ref: Hirata et al., J. Chem. Phys. 120, 2581 (2004) Eqs.(35)-(36)
    assert(isinstance(eris, ccsd._ChemistsERIs))
    nocc, nvir = t1.shape
    fock = eris.fock
    mo_e_o = eris.mo_energy[:nocc]
    mo_e_v = eris.mo_energy[nocc:] + cc.level_shift

    fov = fock[:nocc,nocc:].copy()
    foo = fock[:nocc,:nocc].copy()
    fvv = fock[nocc:,nocc:].copy()

    Foo = imd.cc_Foo(t1,t2,eris)
    Fvv = imd.cc_Fvv(t1,t2,eris)
    Fov = imd.cc_Fov(t1,t2,eris)

    # Move energy terms to the other side
    Foo[np.diag_indices(nocc)] -= mo_e_o
    Fvv[np.diag_indices(nvir)] -= mo_e_v

    # T1 equation
    t1new  =-2*np.einsum('kc,ka,ic->ia', fov, t1, t1)
    t1new +=   np.einsum('ac,ic->ia', Fvv, t1)
    t1new +=  -np.einsum('ki,ka->ia', Foo, t1)
    t1new += 2*np.einsum('kc,kica->ia', Fov, t2)
    t1new +=  -np.einsum('kc,ikca->ia', Fov, t2)
    t1new +=   np.einsum('kc,ic,ka->ia', Fov, t1, t1)
    t1new += fov.conj()
    t1new += 2*np.einsum('kcai,kc->ia', eris.ovvo, t1)
    t1new +=  -np.einsum('kiac,kc->ia', eris.oovv, t1)
    eris_ovvv = np.asarray(eris.get_ovvv())
    t1new += 2*lib.einsum('kdac,ikcd->ia', eris_ovvv, t2)
    t1new +=  -lib.einsum('kcad,ikcd->ia', eris_ovvv, t2)
    t1new += 2*lib.einsum('kdac,kd,ic->ia', eris_ovvv, t1, t1)
    t1new +=  -lib.einsum('kcad,kd,ic->ia', eris_ovvv, t1, t1)
    eris_ovoo = np.asarray(eris.ovoo, order='C')
    t1new +=-2*lib.einsum('lcki,klac->ia', eris_ovoo, t2)
    t1new +=   lib.einsum('kcli,klac->ia', eris_ovoo, t2)
    t1new +=-2*lib.einsum('lcki,lc,ka->ia', eris_ovoo, t1, t1)
    t1new +=   lib.einsum('kcli,lc,ka->ia', eris_ovoo, t1, t1)

    # T2 equation
    tmp2  = lib.einsum('kibc,ka->abic', eris.oovv, -t1)
    tmp2 += np.asarray(eris_ovvv).conj().transpose(1,3,0,2)
    tmp = lib.einsum('abic,jc->ijab', tmp2, t1)
    t2new = tmp + tmp.transpose(1,0,3,2)
    tmp2  = lib.einsum('kcai,jc->akij', eris.ovvo, t1)
    tmp2 += eris_ovoo.transpose(1,3,0,2).conj()
    tmp = lib.einsum('akij,kb->ijab', tmp2, t1)
    t2new -= tmp + tmp.transpose(1,0,3,2)
    t2new += np.asarray(eris.ovov).conj().transpose(0,2,1,3)
    if cc.cc2:
        Woooo2 = np.asarray(eris.oooo).transpose(0,2,1,3).copy()
        Woooo2 += lib.einsum('lcki,jc->klij', eris_ovoo, t1)
        Woooo2 += lib.einsum('kclj,ic->klij', eris_ovoo, t1)
        Woooo2 += lib.einsum('kcld,ic,jd->klij', eris.ovov, t1, t1)
        t2new += lib.einsum('klij,ka,lb->ijab', Woooo2, t1, t1)
        Wvvvv = lib.einsum('kcbd,ka->abcd', eris_ovvv, -t1)
        Wvvvv = Wvvvv + Wvvvv.transpose(1,0,3,2)
        Wvvvv += np.asarray(eris.vvvv).transpose(0,2,1,3)
        t2new += lib.einsum('abcd,ic,jd->ijab', Wvvvv, t1, t1)
        Lvv2 = fvv - np.einsum('kc,ka->ac', fov, t1)
        Lvv2 -= np.diag(np.diag(fvv))
        tmp = lib.einsum('ac,ijcb->ijab', Lvv2, t2)
        t2new += (tmp + tmp.transpose(1,0,3,2))
        Loo2 = foo + np.einsum('kc,ic->ki', fov, t1)
        Loo2 -= np.diag(np.diag(foo))
        tmp = lib.einsum('ki,kjab->ijab', Loo2, t2)
        t2new -= (tmp + tmp.transpose(1,0,3,2))
    else:
        Loo = imd.Loo(t1, t2, eris)
        Lvv = imd.Lvv(t1, t2, eris)
        Loo[np.diag_indices(nocc)] -= mo_e_o
        Lvv[np.diag_indices(nvir)] -= mo_e_v

        Woooo = imd.cc_Woooo(t1, t2, eris)
        Wvoov = imd.cc_Wvoov(t1, t2, eris)
        Wvovo = imd.cc_Wvovo(t1, t2, eris)
        Wvvvv = imd.cc_Wvvvv(t1, t2, eris)

        tau = t2 + np.einsum('ia,jb->ijab', t1, t1)
        t2new += lib.einsum('klij,klab->ijab', Woooo, tau)
        t2new += lib.einsum('abcd,ijcd->ijab', Wvvvv, tau)
        tmp = lib.einsum('ac,ijcb->ijab', Lvv, t2)
        t2new += (tmp + tmp.transpose(1,0,3,2))
        tmp = lib.einsum('ki,kjab->ijab', Loo, t2)
        t2new -= (tmp + tmp.transpose(1,0,3,2))
        tmp  = 2*lib.einsum('akic,kjcb->ijab', Wvoov, t2)
        tmp -=   lib.einsum('akci,kjcb->ijab', Wvovo, t2)
        t2new += (tmp + tmp.transpose(1,0,3,2))
        tmp = lib.einsum('akic,kjbc->ijab', Wvoov, t2)
        t2new -= (tmp + tmp.transpose(1,0,3,2))
        tmp = lib.einsum('bkci,kjac->ijab', Wvovo, t2)
        t2new -= (tmp + tmp.transpose(1,0,3,2))

    eia = mo_e_o[:,None] - mo_e_v
    eijab = lib.direct_sum('ia,jb->ijab',eia,eia)
    t1new /= eia
    t2new /= eijab

    return t1new, t2new


def energy(cc, t1=None, t2=None, eris=None):
    '''RCCSD correlation energy'''
    if t1 is None: t1 = cc.t1
    if t2 is None: t2 = cc.t2
    if eris is None: eris = cc.ao2mo()

    nocc, nvir = t1.shape
    fock = eris.fock
    e = 2*np.einsum('ia,ia', fock[:nocc,nocc:], t1)
    tau = np.einsum('ia,jb->ijab',t1,t1)
    tau += t2
    eris_ovov = np.asarray(eris.ovov)
    e += 2*np.einsum('ijab,iajb', tau, eris_ovov)
    e +=  -np.einsum('ijab,ibja', tau, eris_ovov)
    if abs(e.imag) > 1e-4:
        logger.warn(cc, 'Non-zero imaginary part found in RCCSD energy %s', e)
    return e.real


class RCCSD(ccsd.CCSD):
    '''restricted CCSD with IP-EOM, EA-EOM, EE-EOM, and SF-EOM capabilities

    Ground-state CCSD is performed in optimized ccsd.CCSD and EOM is performed here.
    '''

    def kernel(self, t1=None, t2=None, eris=None, mbpt2=False):
        return self.ccsd(t1, t2, eris, mbpt2)
    def ccsd(self, t1=None, t2=None, eris=None, mbpt2=False):
        '''Ground-state CCSD.

        Kwargs:
            mbpt2 : bool
                Use one-shot MBPT2 approximation to CCSD.
        '''
        if mbpt2:
            pt = mp2.MP2(self._scf, self.frozen, self.mo_coeff, self.mo_occ)
            self.e_corr, self.t2 = pt.kernel(eris=eris)
            nocc, nvir = self.t2.shape[1:3]
            self.t1 = np.zeros((nocc,nvir))
            return self.e_corr, self.t1, self.t2

        if eris is None:
            eris = self.ao2mo(self.mo_coeff)
        return ccsd.CCSD.ccsd(self, t1, t2, eris)

    def ao2mo(self, mo_coeff=None):
        nmo = self.nmo
        nao = self.mo_coeff.shape[0]
        nmo_pair = nmo * (nmo+1) // 2
        nao_pair = nao * (nao+1) // 2
        mem_incore = (max(nao_pair**2, nmo**4) + nmo_pair**2) * 8/1e6
        mem_now = lib.current_memory()[0]
        if (self._scf._eri is not None and
            (mem_incore+mem_now < self.max_memory) or self.mol.incore_anyway):
            return _make_eris_incore(self, mo_coeff)

        elif getattr(self._scf, 'with_df', None):
            logger.warn(self, 'CCSD detected DF being used in the HF object. '
                        'MO integrals are computed based on the DF 3-index tensors.\n'
                        'It\'s recommended to use dfccsd.CCSD for the '
                        'DF-CCSD calculations')
            raise NotImplementedError
            #return _make_df_eris_outcore(self, mo_coeff)

        else:
            return _make_eris_outcore(self, mo_coeff)

    energy = energy
    update_amps = update_amps

    def solve_lambda(self, t1=None, t2=None, l1=None, l2=None,
                     eris=None):
        from pyscf.cc import rccsd_lambda
        if t1 is None: t1 = self.t1
        if t2 is None: t2 = self.t2
        if eris is None: eris = self.ao2mo(self.mo_coeff)
        self.converged_lambda, self.l1, self.l2 = \
                rccsd_lambda.kernel(self, eris, t1, t2, l1, l2,
                                    max_cycle=self.max_cycle,
                                    tol=self.conv_tol_normt,
                                    verbose=self.verbose)
        return self.l1, self.l2

    def ccsd_t(self, t1=None, t2=None, eris=None):
        return ccsd.CCSD.ccsd_t(self, t1, t2, eris)

    def density_fit(self, auxbasis=None, with_df=None):
        raise NotImplementedError


class _ChemistsERIs(ccsd._ChemistsERIs):

    def get_ovvv(self, *slices):
        '''To access a subblock of ovvv tensor'''
        if slices:
            return self.ovvv[slices]
        else:
            return self.ovvv

def _make_eris_incore(mycc, mo_coeff=None, ao2mofn=None):
    cput0 = (logger.process_clock(), logger.perf_counter())
    eris = _ChemistsERIs()
    eris._common_init_(mycc, mo_coeff)
    nocc = eris.nocc
    nmo = eris.fock.shape[0]

    if callable(ao2mofn):
        eri1 = ao2mofn(eris.mo_coeff).reshape([nmo]*4)
    else:
        eri1 = ao2mo.incore.full(mycc._scf._eri, eris.mo_coeff)
        eri1 = ao2mo.restore(1, eri1, nmo)
    eris.oooo = eri1[:nocc,:nocc,:nocc,:nocc].copy()
    eris.ovoo = eri1[:nocc,nocc:,:nocc,:nocc].copy()
    eris.ovov = eri1[:nocc,nocc:,:nocc,nocc:].copy()
    eris.oovv = eri1[:nocc,:nocc,nocc:,nocc:].copy()
    eris.ovvo = eri1[:nocc,nocc:,nocc:,:nocc].copy()
    eris.ovvv = eri1[:nocc,nocc:,nocc:,nocc:].copy()
    eris.vvvv = eri1[nocc:,nocc:,nocc:,nocc:].copy()
    logger.timer(mycc, 'CCSD integral transformation', *cput0)
    return eris

def _make_eris_outcore(mycc, mo_coeff=None):
    cput0 = (logger.process_clock(), logger.perf_counter())
    log = logger.Logger(mycc.stdout, mycc.verbose)
    eris = _ChemistsERIs()
    eris._common_init_(mycc, mo_coeff)

    mol = mycc.mol
    mo_coeff = eris.mo_coeff
    nocc = eris.nocc
    nao, nmo = mo_coeff.shape
    nvir = nmo - nocc
    eris.feri1 = lib.H5TmpFile()
    eris.oooo = eris.feri1.create_dataset('oooo', (nocc,nocc,nocc,nocc), 'f8')
    eris.ovoo = eris.feri1.create_dataset('ovoo', (nocc,nvir,nocc,nocc), 'f8', chunks=(nocc,1,nocc,nocc))
    eris.ovov = eris.feri1.create_dataset('ovov', (nocc,nvir,nocc,nvir), 'f8', chunks=(nocc,1,nocc,nvir))
    eris.ovvo = eris.feri1.create_dataset('ovvo', (nocc,nvir,nvir,nocc), 'f8', chunks=(nocc,1,nvir,nocc))
    eris.ovvv = eris.feri1.create_dataset('ovvv', (nocc,nvir,nvir,nvir), 'f8')
    eris.oovv = eris.feri1.create_dataset('oovv', (nocc,nocc,nvir,nvir), 'f8', chunks=(nocc,nocc,1,nvir))
    eris.vvvv = eris.feri1.create_dataset('vvvv', (nvir,nvir,nvir,nvir), 'f8')
    max_memory = max(MEMORYMIN, mycc.max_memory-lib.current_memory()[0])

    ftmp = lib.H5TmpFile()
    ao2mo.full(mol, mo_coeff, ftmp, max_memory=max_memory, verbose=log)
    eri = ftmp['eri_mo']

    nocc_pair = nocc*(nocc+1)//2
    tril2sq = lib.square_mat_in_trilu_indices(nmo)
    oo = eri[:nocc_pair]
    eris.oooo[:] = ao2mo.restore(1, oo[:,:nocc_pair], nocc)
    oovv = lib.take_2d(oo, tril2sq[:nocc,:nocc].ravel(), tril2sq[nocc:,nocc:].ravel())
    eris.oovv[:] = oovv.reshape(nocc,nocc,nvir,nvir)
    oo = oovv = None

    tril2sq = lib.square_mat_in_trilu_indices(nmo)
    blksize = min(nvir, max(BLKMIN, int(max_memory*1e6/8/nmo**3/2)))
    for p0, p1 in lib.prange(0, nvir, blksize):
        q0, q1 = p0+nocc, p1+nocc
        off0 = q0*(q0+1)//2
        off1 = q1*(q1+1)//2
        buf = lib.unpack_tril(eri[off0:off1])

        tmp = buf[ tril2sq[q0:q1,:nocc] - off0 ]
        eris.ovoo[:,p0:p1] = tmp[:,:,:nocc,:nocc].transpose(1,0,2,3)
        eris.ovvo[:,p0:p1] = tmp[:,:,nocc:,:nocc].transpose(1,0,2,3)
        eris.ovov[:,p0:p1] = tmp[:,:,:nocc,nocc:].transpose(1,0,2,3)
        eris.ovvv[:,p0:p1] = tmp[:,:,nocc:,nocc:].transpose(1,0,2,3)

        tmp = buf[ tril2sq[q0:q1,nocc:q1] - off0 ]
        eris.vvvv[p0:p1,:p1] = tmp[:,:,nocc:,nocc:]
        if p0 > 0:
            eris.vvvv[:p0,p0:p1] = tmp[:,:p0,nocc:,nocc:].transpose(1,0,2,3)
        buf = tmp = None
    log.timer('CCSD integral transformation', *cput0)
    return eris


if __name__ == '__main__':
    from pyscf import scf
    from pyscf import gto
    from pyscf.cc import gccsd

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = 'cc-pvdz'
    #mol.basis = '3-21G'
    mol.verbose = 0
    mol.spin = 0
    mol.build()
    mf = scf.RHF(mol).run(conv_tol=1e-14)

    mycc = RCCSD(mf)
    mycc.max_memory = 0
    eris = mycc.ao2mo()
    emp2, t1, t2 = mycc.init_amps(eris)
    print(lib.finger(t2) - 0.044540097905897198)
    np.random.seed(1)
    t1 = np.random.random(t1.shape)*.1
    t2 = np.random.random(t2.shape)*.1
    t2 = t2 + t2.transpose(1,0,3,2)
    t1, t2 = update_amps(mycc, t1, t2, eris)
    print(lib.finger(t1) - 0.25118555558133576)
    print(lib.finger(t2) - 0.02352137419932243)

    ecc, t1, t2 = mycc.kernel()
    print(ecc - -0.21334326214236796)

    e, v = mycc.ipccsd(nroots=3)
    print(e[0] - 0.43356041409195489)
    print(e[1] - 0.51876598058509493)
    print(e[2] - 0.6782879569941862 )

    e, v = mycc.eeccsd(nroots=4)
    print(e[0] - 0.2757159395886167)
    print(e[1] - 0.2757159395886167)
    print(e[2] - 0.2757159395886167)
    print(e[3] - 0.3005716731825082)


    mol = gto.Mole()
    mol.verbose = 0
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = '631g'
    mol.build()
    mf = scf.RHF(mol)
    mf.conv_tol = 1e-16
    mf.scf()
    mo_coeff = mf.mo_coeff + np.sin(mf.mo_coeff) * .01j
    nao = mo_coeff.shape[0]
    eri = ao2mo.restore(1, mf._eri, nao)
    eri0 = lib.einsum('pqrs,pi,qj,rk,sl->ijkl', eri, mo_coeff.conj(), mo_coeff,
                      mo_coeff.conj(), mo_coeff)

    nocc, nvir = 5, nao-5
    eris = _ChemistsERIs(mol)
    eris.oooo = eri0[:nocc,:nocc,:nocc,:nocc].copy()
    eris.ovoo = eri0[:nocc,nocc:,:nocc,:nocc].copy()
    eris.oovv = eri0[:nocc,:nocc,nocc:,nocc:].copy()
    eris.ovvo = eri0[:nocc,nocc:,nocc:,:nocc].copy()
    eris.ovov = eri0[:nocc,nocc:,:nocc,nocc:].copy()
    eris.ovvv = eri0[:nocc,nocc:,nocc:,nocc:].copy()
    eris.vvvv = eri0[nocc:,nocc:,nocc:,nocc:].copy()
    eris.fock = np.diag(mf.mo_energy)
    eris.mo_energy = mf.mo_energy

    np.random.seed(1)
    t1 = np.random.random((nocc,nvir)) + np.random.random((nocc,nvir))*.1j - .5
    t2 = np.random.random((nocc,nocc,nvir,nvir)) - .5
    t2 = t2 + np.sin(t2) * .1j
    t2 = t2 + t2.transpose(1,0,3,2)

    mycc = RCCSD(mf)
    t1new_ref, t2new_ref = update_amps(mycc, t1, t2, eris)

    orbspin = np.zeros(nao*2, dtype=int)
    orbspin[1::2] = 1
    eri1 = np.zeros([nao*2]*4, dtype=np.complex128)
    eri1[0::2,0::2,0::2,0::2] = eri1[0::2,0::2,1::2,1::2] = eri0
    eri1[1::2,1::2,0::2,0::2] = eri1[1::2,1::2,1::2,1::2] = eri0
    eri1 = eri1.transpose(0,2,1,3) - eri1.transpose(0,2,3,1)
    erig = gccsd._PhysicistsERIs(mol)
    nocc *= 2
    nvir *= 2
    erig.oooo = eri1[:nocc,:nocc,:nocc,:nocc].copy()
    erig.ooov = eri1[:nocc,:nocc,:nocc,nocc:].copy()
    erig.ovov = eri1[:nocc,nocc:,:nocc,nocc:].copy()
    erig.ovvo = eri1[:nocc,nocc:,nocc:,:nocc].copy()
    erig.oovv = eri1[:nocc,:nocc,nocc:,nocc:].copy()
    erig.ovvv = eri1[:nocc,nocc:,nocc:,nocc:].copy()
    erig.vvvv = eri1[nocc:,nocc:,nocc:,nocc:].copy()
    mo_e = np.array([mf.mo_energy]*2)
    erig.fock = np.diag(mo_e.T.ravel())

    myccg = gccsd.GCCSD(scf.addons.convert_to_ghf(mf))
    t1, t2 = myccg.amplitudes_from_ccsd(t1, t2)
    t1new, t2new = gccsd.update_amps(myccg, t1, t2, erig)
    print(abs(t1new[0::2,0::2]-t1new_ref).max())
    t2aa = t2new[0::2,0::2,0::2,0::2]
    t2ab = t2new[0::2,1::2,0::2,1::2]
    print(abs(t2ab-t2new_ref).max())
    print(abs(t2ab-t2ab.transpose(1,0,2,3) - t2aa).max())



# pyscf/cc/rccsd_lambda.py

#!/usr/bin/env python
# Copyright 2014-2020 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: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Restricted CCSD lambda equation solver which supports both real and complex
integrals.  This code is slower than the pyscf.cc.ccsd_lambda implementation.

Note MO integrals are treated in chemist's notation
'''



import numpy as np
from pyscf import lib
from pyscf.lib import logger
from pyscf.cc import _ccsd
from pyscf.cc import ccsd_lambda

einsum = lib.einsum

def kernel(mycc, eris=None, t1=None, t2=None, l1=None, l2=None,
           max_cycle=50, tol=1e-8, verbose=logger.INFO):
    if eris is None: eris = mycc.ao2mo()
    return ccsd_lambda.kernel(mycc, eris, t1, t2, l1, l2, max_cycle, tol,
                              verbose, make_intermediates, update_lambda)


def make_intermediates(mycc, t1, t2, eris):
    nocc, nvir = t1.shape
    foo = eris.fock[:nocc,:nocc]
    fov = eris.fock[:nocc,nocc:]
    fvo = eris.fock[nocc:,:nocc]
    fvv = eris.fock[nocc:,nocc:]

    tau = _ccsd.make_tau(t2, t1, t1)
    ovov = np.asarray(eris.ovov)
    ovoo = np.asarray(eris.ovoo)
    ovov1 = ovov * 2 - ovov.transpose(0,3,2,1)
    ovoo1 = ovoo * 2 - ovoo.transpose(2,1,0,3)

    v1  = fvv - lib.einsum('ja,jb->ba', fov, t1)
    v1 -= lib.einsum('jakc,jkbc->ba', ovov1, tau)
    v2  = foo + lib.einsum('ib,jb->ij', fov, t1)
    v2 += lib.einsum('ibkc,jkbc->ij', ovov1, tau)
    v2 += np.einsum('kbij,kb->ij', ovoo1, t1)
    v4 = fov + np.einsum('jbkc,kc->jb', ovov1, t1)

    v5  = np.einsum('kc,jkbc->bj', fov, t2) * 2
    v5 -= np.einsum('kc,jkcb->bj', fov, t2)
    v5 += fvo
    v5 += lib.einsum('kc,kb,jc->bj', v4, t1, t1)
    v5 -= lib.einsum('lckj,klbc->bj', ovoo1, t2)

    oooo = np.asarray(eris.oooo)
    woooo  = lib.einsum('icjl,kc->ikjl', ovoo, t1)
    woooo += lib.einsum('jcil,kc->iljk', ovoo, t1)
    woooo += oooo.copy()
    woooo += lib.einsum('icjd,klcd->ikjl', ovov, tau)

    theta = t2*2 - t2.transpose(0,1,3,2)
    v4OVvo  = lib.einsum('ldjb,klcd->jbck', ovov1, t2)
    v4OVvo -= lib.einsum('ldjb,kldc->jbck', ovov, t2)
    v4OVvo += np.asarray(eris.ovvo)

    v4oVVo  = lib.einsum('jdlb,kldc->jbck', ovov, t2)
    v4oVVo -= np.asarray(eris.oovv).transpose(0,3,2,1)

    v4ovvo = v4OVvo*2 + v4oVVo
    w3 = np.einsum('jbck,jb->ck', v4ovvo, t1)

    woovo  = lib.einsum('ibck,jb->ijck', v4ovvo, t1)
    woovo = woovo - woovo.transpose(0,3,2,1)
    woovo += lib.einsum('ibck,jb->ikcj', v4OVvo-v4oVVo, t1)
    woovo += ovoo1.conj().transpose(3,2,1,0)

    woovo += lib.einsum('lcik,jlbc->ikbj', ovoo1, theta)
    woovo -= lib.einsum('lcik,jlbc->ijbk', ovoo1, t2)
    woovo -= lib.einsum('iclk,ljbc->ijbk', ovoo1, t2)

    wvvvo  = lib.einsum('jack,jb->back', v4ovvo, t1)
    wvvvo = wvvvo - wvvvo.transpose(2,1,0,3)
    wvvvo += lib.einsum('jack,jb->cabk', v4OVvo-v4oVVo, t1)
    wvvvo -= lib.einsum('lajk,jlbc->cabk', ovoo1, tau)

    wOVvo  = v4OVvo
    woVVo  = v4oVVo
    wOVvo -= np.einsum('jbld,kd,lc->jbck', ovov, t1, t1)
    woVVo += np.einsum('jdlb,kd,lc->jbck', ovov, t1, t1)
    wOVvo -= lib.einsum('jblk,lc->jbck', ovoo, t1)
    woVVo += lib.einsum('lbjk,lc->jbck', ovoo, t1)
    v4ovvo = v4OVvo = v4oVVo = None

    ovvv = np.asarray(eris.get_ovvv())
    wvvvo += lib.einsum('kacd,kjbd->bacj', ovvv, t2) * 1.5

    wOVvo += lib.einsum('jbcd,kd->jbck', ovvv, t1)
    woVVo -= lib.einsum('jdcb,kd->jbck', ovvv, t1)

    ovvv = ovvv*2 - ovvv.transpose(0,3,2,1)
    v1 += np.einsum('jcba,jc->ba', ovvv, t1)
    v5 += lib.einsum('kdbc,jkcd->bj', ovvv, t2)
    woovo += lib.einsum('idcb,jkdb->ijck', ovvv, tau)

    tmp = lib.einsum('kdca,jkbd->cabj', ovvv, theta)
    wvvvo -= tmp
    wvvvo += tmp.transpose(2,1,0,3) * .5
    wvvvo -= ovvv.conj().transpose(3,2,1,0)
    ovvv = tmp = None

    w3 += v5
    w3 += np.einsum('cb,jb->cj', v1, t1)
    w3 -= np.einsum('jk,jb->bk', v2, t1)

    class _IMDS: pass
    imds = _IMDS()
    imds.ftmp = lib.H5TmpFile()
    dtype = np.result_type(t2, eris.vvvv).char
    imds.woooo = imds.ftmp.create_dataset('woooo', (nocc,nocc,nocc,nocc), dtype)
    imds.wovvo = imds.ftmp.create_dataset('wovvo', (nocc,nvir,nvir,nocc), dtype)
    imds.woVVo = imds.ftmp.create_dataset('woVVo', (nocc,nvir,nvir,nocc), dtype)
    imds.woovo = imds.ftmp.create_dataset('woovo', (nocc,nocc,nvir,nocc), dtype)
    imds.wvvvo = imds.ftmp.create_dataset('wvvvo', (nvir,nvir,nvir,nocc), dtype)

    imds.woooo[:] = woooo
    imds.wovvo[:] = wOVvo*2 + woVVo
    imds.woVVo[:] = woVVo
    imds.woovo[:] = woovo
    imds.wvvvo[:] = wvvvo
    imds.v1 = v1
    imds.v2 = v2
    imds.w3 = w3
    imds.v4 = v4
    imds.ftmp.flush()
    return imds


# update L1, L2
def update_lambda(mycc, t1, t2, l1, l2, eris, imds):
    time0 = logger.process_clock(), logger.perf_counter()
    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocc, nvir = t1.shape
    l1new = np.zeros_like(l1)

    foo = eris.fock[:nocc,:nocc]
    fov = eris.fock[:nocc,nocc:]
    fvv = eris.fock[nocc:,nocc:]

    tau = _ccsd.make_tau(t2, t1, t1)

    theta = t2*2 - t2.transpose(0,1,3,2)
    mvv = lib.einsum('klca,klcb->ba', l2, theta)
    moo = lib.einsum('kicd,kjcd->ij', l2, theta)
    mvv1 = lib.einsum('jc,jb->bc', l1, t1) + mvv
    moo1 = lib.einsum('ic,kc->ik', l1, t1) + moo

    # m3 = einsum('ijab,acbd->ijcd', l2, vvvv)
    #    = einsum('ijab,cadb->ijcd', l2.conj(), vvvv).conj()
    m3 = mycc._add_vvvv(None, l2.conj(), eris, with_ovvv=False).conj()
    m3 += lib.einsum('klab,ikjl->ijab', l2, imds.woooo)
    m3 *= .5

    ovov = np.asarray(eris.ovov)
    l2tau = np.einsum('ijcd,klcd->ijkl', l2, tau)
    m3 += np.einsum('kalb,ijkl->ijab', ovov, l2tau) * .5
    l2tau = None

    l2new = ovov.transpose(0,2,1,3) * .5
    l2new += lib.einsum('ijac,cb->ijab', l2, imds.v1)
    l2new -= lib.einsum('ikab,jk->ijab', l2, imds.v2)
    l2new -= lib.einsum('ca,icjb->ijab', mvv1, ovov)
    l2new -= lib.einsum('ik,kajb->ijab', moo1, ovov)

    ovov = ovov * 2 - ovov.transpose(0,3,2,1)
    l1new -= np.einsum('ik,ka->ia', moo, imds.v4)
    l1new -= np.einsum('ca,ic->ia', mvv, imds.v4)
    l2new += np.einsum('ia,jb->ijab', l1, imds.v4)

    tmp  = t1 + np.einsum('kc,kjcb->jb', l1, theta)
    tmp -= lib.einsum('bd,jd->jb', mvv1, t1)
    tmp -= lib.einsum('lj,lb->jb', moo, t1)
    l1new += np.einsum('jbia,jb->ia', ovov, tmp)
    ovov = tmp = None

    ovvv = np.asarray(eris.get_ovvv())
    l1new += np.einsum('iacb,bc->ia', ovvv, mvv1) * 2
    l1new -= np.einsum('ibca,bc->ia', ovvv, mvv1)
    l2new += lib.einsum('ic,jbca->jiba', l1, ovvv)
    l2t1 = np.einsum('ijcd,kd->ijck', l2, t1)
    m3 -= np.einsum('kbca,ijck->ijab', ovvv, l2t1)
    l2t1 = ovvv = None

    l2new += m3
    l1new += np.einsum('ijab,jb->ia', m3, t1) * 2
    l1new += np.einsum('jiba,jb->ia', m3, t1) * 2
    l1new -= np.einsum('ijba,jb->ia', m3, t1)
    l1new -= np.einsum('jiab,jb->ia', m3, t1)

    ovoo = np.asarray(eris.ovoo)
    l1new -= np.einsum('iajk,kj->ia', ovoo, moo1) * 2
    l1new += np.einsum('jaik,kj->ia', ovoo, moo1)
    l2new -= lib.einsum('ka,jbik->ijab', l1, ovoo)
    ovoo = None

    l2theta = l2*2 - l2.transpose(0,1,3,2)
    l2new += lib.einsum('ikac,jbck->ijab', l2theta, imds.wovvo) * .5
    tmp = lib.einsum('ikca,jbck->ijab', l2, imds.woVVo)
    l2new += tmp * .5
    l2new += tmp.transpose(1,0,2,3)
    l2theta = None

    l1new += fov
    l1new += lib.einsum('ib,ba->ia', l1, imds.v1)
    l1new -= lib.einsum('ja,ij->ia', l1, imds.v2)

    l1new += np.einsum('jb,iabj->ia', l1, eris.ovvo) * 2
    l1new -= np.einsum('jb,ijba->ia', l1, eris.oovv)

    l1new -= lib.einsum('ijbc,bacj->ia', l2, imds.wvvvo)
    l1new -= lib.einsum('kjca,ijck->ia', l2, imds.woovo)

    l1new += np.einsum('ijab,bj->ia', l2, imds.w3) * 2
    l1new -= np.einsum('ijba,bj->ia', l2, imds.w3)

    eia = lib.direct_sum('i-j->ij', foo.diagonal(), fvv.diagonal() + mycc.level_shift)
    l1new /= eia
    l1new += l1

    l2new = l2new + l2new.transpose(1,0,3,2)
    l2new /= lib.direct_sum('ia+jb->ijab', eia, eia)
    l2new += l2

    time0 = log.timer_debug1('update l1 l2', *time0)
    return l1new, l2new


if __name__ == '__main__':
    from functools import reduce
    from pyscf import gto
    from pyscf import scf
    from pyscf import ao2mo
    from pyscf.cc import rccsd

    mol = gto.Mole()
    mol.verbose = 0
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]

    mol.basis = 'cc-pvdz'
    mol.build()
    mf = scf.RHF(mol)
    mf.conv_tol = 1e-16
    mf.scf()

    mcc = rccsd.RCCSD(mf)
    mcc.conv_tol = 1e-12
    ecc, t1, t2 = mcc.kernel()

    eris = mcc.ao2mo()
    l1, l2 = mcc.solve_lambda(t1, t2, eris=eris)
    print(np.linalg.norm(l1)-0.0132626841292)
    print(np.linalg.norm(l2)-0.212575609057)

    from pyscf.cc import ccsd_rdm
    dm1 = ccsd_rdm.make_rdm1(mcc, t1, t2, l1, l2)
    dm2 = ccsd_rdm.make_rdm2(mcc, t1, t2, l1, l2)
    h1 = reduce(np.dot, (mf.mo_coeff.T, mf.get_hcore(), mf.mo_coeff))
    nmo = h1.shape[0]
    eri = ao2mo.full(mf._eri, mf.mo_coeff)
    eri = ao2mo.restore(1, eri, nmo).reshape((nmo,)*4)
    e1 = np.einsum('pq,pq', h1, dm1)
    e2 = np.einsum('pqrs,pqrs', eri, dm2) * .5
    print(e1+e2+mol.energy_nuc() - mf.e_tot - ecc)



# pyscf/cc/rccsd_slow.py

#!/usr/bin/env python
# Copyright 2014-2021 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.

'''
Restricted CCSD

Ref: Stanton et al., J. Chem. Phys. 94, 4334 (1990)
Ref: Hirata et al., J. Chem. Phys. 120, 2581 (2004)
'''

from functools import reduce

import numpy as np

from pyscf import lib
from pyscf import ao2mo
from pyscf.lib import logger
from pyscf.cc import ccsd
from pyscf.cc import rintermediates as imd
from pyscf.lib import linalg_helper

#einsum = np.einsum
einsum = lib.einsum

# note MO integrals are treated in chemist's notation

def update_amps(cc, t1, t2, eris):
    # Ref: Hirata et al., J. Chem. Phys. 120, 2581 (2004) Eqs.(35)-(36)
    nocc, nvir = t1.shape
    fock = eris.fock

    fov = fock[:nocc,nocc:].copy()
    foo = fock[:nocc,:nocc].copy()
    fvv = fock[nocc:,nocc:].copy()

    Foo = imd.cc_Foo(t1,t2,eris)
    Fvv = imd.cc_Fvv(t1,t2,eris)
    Fov = imd.cc_Fov(t1,t2,eris)

    # Move energy terms to the other side
    Foo -= np.diag(np.diag(foo))
    Fvv -= np.diag(np.diag(fvv))

    # T1 equation
    t1new = np.asarray(fov).conj().copy()
    t1new +=-2*einsum('kc,ka,ic->ia', fov, t1, t1)
    t1new +=   einsum('ac,ic->ia', Fvv, t1)
    t1new +=  -einsum('ki,ka->ia', Foo, t1)
    t1new += 2*einsum('kc,kica->ia', Fov, t2)
    t1new +=  -einsum('kc,ikca->ia', Fov, t2)
    t1new +=   einsum('kc,ic,ka->ia', Fov, t1, t1)
    t1new += 2*einsum('kcai,kc->ia', eris.ovvo, t1)
    t1new +=  -einsum('kiac,kc->ia', eris.oovv, t1)
    eris_ovvv = np.asarray(eris.ovvv)
    t1new += 2*einsum('kdac,ikcd->ia', eris_ovvv, t2)
    t1new +=  -einsum('kcad,ikcd->ia', eris_ovvv, t2)
    t1new += 2*einsum('kdac,kd,ic->ia', eris_ovvv, t1, t1)
    t1new +=  -einsum('kcad,kd,ic->ia', eris_ovvv, t1, t1)
    t1new +=-2*einsum('kilc,klac->ia', eris.ooov, t2)
    t1new +=   einsum('likc,klac->ia', eris.ooov, t2)
    t1new +=-2*einsum('kilc,lc,ka->ia', eris.ooov, t1, t1)
    t1new +=   einsum('likc,lc,ka->ia', eris.ooov, t1, t1)

    # T2 equation
    t2new = np.asarray(eris.ovov).conj().transpose(0,2,1,3).copy()
    if cc.cc2:
        Woooo2 = np.asarray(eris.oooo).transpose(0,2,1,3).copy()
        Woooo2 += einsum('kilc,jc->klij', eris.ooov, t1)
        Woooo2 += einsum('ljkc,ic->klij', eris.ooov, t1)
        Woooo2 += einsum('kcld,ic,jd->klij', eris.ovov, t1, t1)
        t2new += einsum('klij,ka,lb->ijab', Woooo2, t1, t1)
        Wvvvv = einsum('kcbd,ka->abcd', eris_ovvv, -t1)
        Wvvvv = Wvvvv + Wvvvv.transpose(1,0,3,2)
        Wvvvv += np.asarray(eris.vvvv).transpose(0,2,1,3)
        t2new += einsum('abcd,ic,jd->ijab', Wvvvv, t1, t1)
        Lvv2 = fvv - einsum('kc,ka->ac', fov, t1)
        Lvv2 -= np.diag(np.diag(fvv))
        tmp = einsum('ac,ijcb->ijab', Lvv2, t2)
        t2new += (tmp + tmp.transpose(1,0,3,2))
        Loo2 = foo + einsum('kc,ic->ki', fov, t1)
        Loo2 -= np.diag(np.diag(foo))
        tmp = einsum('ki,kjab->ijab', Loo2, t2)
        t2new -= (tmp + tmp.transpose(1,0,3,2))
    else:
        Loo = imd.Loo(t1, t2, eris)
        Lvv = imd.Lvv(t1, t2, eris)
        Loo -= np.diag(np.diag(foo))
        Lvv -= np.diag(np.diag(fvv))
        Woooo = imd.cc_Woooo(t1, t2, eris)
        Wvoov = imd.cc_Wvoov(t1, t2, eris)
        Wvovo = imd.cc_Wvovo(t1, t2, eris)
        Wvvvv = imd.cc_Wvvvv(t1, t2, eris)
        tau = t2 + einsum('ia,jb->ijab', t1, t1)
        t2new += einsum('klij,klab->ijab', Woooo, tau)
        t2new += einsum('abcd,ijcd->ijab', Wvvvv, tau)
        tmp = einsum('ac,ijcb->ijab', Lvv, t2)
        t2new += (tmp + tmp.transpose(1,0,3,2))
        tmp = einsum('ki,kjab->ijab', Loo, t2)
        t2new -= (tmp + tmp.transpose(1,0,3,2))
        tmp  = 2*einsum('akic,kjcb->ijab', Wvoov, t2)
        tmp -=   einsum('akci,kjcb->ijab', Wvovo, t2)
        t2new += (tmp + tmp.transpose(1,0,3,2))
        tmp = einsum('akic,kjbc->ijab', Wvoov, t2)
        t2new -= (tmp + tmp.transpose(1,0,3,2))
        tmp = einsum('bkci,kjac->ijab', Wvovo, t2)
        t2new -= (tmp + tmp.transpose(1,0,3,2))

    tmp2  = einsum('kibc,ka->abic', eris.oovv, -t1)
    tmp2 += np.asarray(eris.ovvv).conj().transpose(1,3,0,2)
    tmp = einsum('abic,jc->ijab', tmp2, t1)
    t2new += (tmp + tmp.transpose(1,0,3,2))
    tmp2  = einsum('kcai,jc->akij', eris.ovvo, t1)
    tmp2 += np.asarray(eris.ooov).transpose(3,1,2,0).conj()
    tmp = einsum('akij,kb->ijab', tmp2, t1)
    t2new -= (tmp + tmp.transpose(1,0,3,2))

    mo_e = eris.fock.diagonal().real
    eia = mo_e[:nocc,None] - mo_e[None,nocc:]
    eijab = lib.direct_sum('ia,jb->ijab',eia,eia)
    t1new /= eia
    t2new /= eijab

    return t1new, t2new


def energy(cc, t1, t2, eris):
    nocc, nvir = t1.shape
    fock = eris.fock
    e = 2*einsum('ia,ia', fock[:nocc,nocc:], t1)
    tau = einsum('ia,jb->ijab',t1,t1)
    tau += t2
    eris_ovov = np.asarray(eris.ovov)
    e += 2*einsum('ijab,iajb', tau, eris_ovov)
    e +=  -einsum('ijab,ibja', tau, eris_ovov)
    return e.real


class RCCSD(ccsd.CCSD):

    def __init__(self, mf, frozen=None, mo_coeff=None, mo_occ=None):
        ccsd.CCSD.__init__(self, mf, frozen, mo_coeff, mo_occ)
        self.max_space = 20
        self._keys = self._keys.union(['max_space'])

    def dump_flags(self, verbose=None):
        ccsd.CCSD.dump_flags(self, verbose)

    def init_amps(self, eris):
        nocc = self.nocc
        mo_e = eris.fock.diagonal().real
        eia = mo_e[:nocc,None] - mo_e[None,nocc:]
        eijab = lib.direct_sum('ia,jb->ijab', eia, eia)
        t1 = eris.fock[:nocc,nocc:].conj() / eia
        eris_ovov = np.asarray(eris.ovov)
        t2 = eris_ovov.transpose(0,2,1,3).conj() / eijab
        self.emp2  = 2*einsum('ijab,iajb', t2, eris_ovov)
        self.emp2 -=   einsum('ijab,ibja', t2, eris_ovov)
        lib.logger.info(self, 'Init t2, MP2 energy = %.15g', self.emp2)
        return self.emp2, t1, t2


    def kernel(self, t1=None, t2=None, eris=None, mbpt2=False, cc2=False):
        return self.ccsd(t1, t2, eris, mbpt2, cc2)
    def ccsd(self, t1=None, t2=None, eris=None, mbpt2=False, cc2=False):
        '''Ground-state CCSD.

        Kwargs:
            mbpt2 : bool
                Use one-shot MBPT2 approximation to CCSD.
            cc2 : bool
                Use CC2 approximation to CCSD.
        '''
        if mbpt2 and cc2:
            raise RuntimeError('MBPT2 and CC2 are mutually exclusive approximations to the CCSD ground state.')
        if eris is None: eris = self.ao2mo(self.mo_coeff)
        self.eris = eris
        self.dump_flags()
        if mbpt2:
            cctyp = 'MBPT2'
            self.e_corr, self.t1, self.t2 = self.init_amps(eris)
        else:
            if cc2:
                cctyp = 'CC2'
                self.cc2 = True
            else:
                cctyp = 'CCSD'
                self.cc2 = False
            self.converged, self.e_corr, self.t1, self.t2 = \
                    ccsd.kernel(self, eris, t1, t2, max_cycle=self.max_cycle,
                                tol=self.conv_tol, tolnormt=self.conv_tol_normt,
                                verbose=self.verbose)
            if self.converged:
                logger.info(self, '%s converged', cctyp)
            else:
                logger.info(self, '%s not converged', cctyp)
        if self._scf.e_tot == 0:
            logger.note(self, 'E_corr = %.16g', self.e_corr)
        else:
            logger.note(self, 'E(%s) = %.16g  E_corr = %.16g',
                        cctyp, self.e_tot, self.e_corr)
        return self.e_corr, self.t1, self.t2

    def ao2mo(self, mo_coeff=None):
        return _ChemistsERIs(self, mo_coeff)

    energy = energy
    update_amps = update_amps

    def nip(self):
        nocc = self.nocc
        nvir = self.nmo - nocc
        self._nip = nocc + nocc*nocc*nvir
        return self._nip

    def nea(self):
        nocc = self.nocc
        nvir = self.nmo - nocc
        self._nea = nvir + nocc*nvir*nvir
        return self._nea

    def nee(self):
        nocc = self.nocc
        nvir = self.nmo - nocc
        self._nee = nocc*nvir + nocc*nocc*nvir*nvir
        return self._nee

    def ipccsd(self, nroots=1, left=False, koopmans=False, guess=None, partition=None):
        '''Calculate (N-1)-electron charged excitations via IP-EOM-CCSD.

        Kwargs:
            nroots : int
                Number of roots (eigenvalues) requested
            partition : bool or str
                Use a matrix-partitioning for the doubles-doubles block.
                Can be None, 'mp' (Moller-Plesset, i.e. orbital energies on the diagonal),
                or 'full' (full diagonal elements).
            koopmans : bool
                Calculate Koopmans'-like (quasiparticle) excitations only, targeting via
                overlap.
            guess : list of ndarray
                List of guess vectors to use for targeting via overlap.
        '''
        cput0 = (logger.process_clock(), logger.perf_counter())
        log = logger.Logger(self.stdout, self.verbose)
        size = self.nip()
        nroots = min(nroots,size)
        if partition:
            partition = partition.lower()
            assert partition in ['mp','full']
        self.ip_partition = partition
        if partition == 'full':
            self._ipccsd_diag_matrix2 = self.vector_to_amplitudes_ip(self.ipccsd_diag())[1]

        adiag = self.ipccsd_diag()
        user_guess = False
        if guess:
            user_guess = True
            assert len(guess) == nroots
            for g in guess:
                assert g.size == size
        else:
            guess = []
            if koopmans:
                for n in range(nroots):
                    g = np.zeros(size)
                    g[self.nocc-n-1] = 1.0
                    guess.append(g)
            else:
                idx = adiag.argsort()[:nroots]
                for i in idx:
                    g = np.zeros(size)
                    g[i] = 1.0
                    guess.append(g)

        def precond(r, e0, x0):
            return r/(e0-adiag+1e-12)

        if left:
            matvec = self.lipccsd_matvec
        else:
            matvec = self.ipccsd_matvec
        eig = linalg_helper.eig
        if user_guess or koopmans:
            def pickeig(w, v, nr, envs):
                x0 = linalg_helper._gen_x0(envs['v'], envs['xs'])
                idx = np.argmax( np.abs(np.dot(np.array(guess).conj(),np.array(x0).T)), axis=1 )
                return lib.linalg_helper._eigs_cmplx2real(w, v, idx)
            eip, evecs = eig(matvec, guess, precond, pick=pickeig,
                             tol=self.conv_tol, max_cycle=self.max_cycle,
                             max_space=self.max_space, nroots=nroots, verbose=self.verbose)
        else:
            eip, evecs = eig(matvec, guess, precond,
                             tol=self.conv_tol, max_cycle=self.max_cycle,
                             max_space=self.max_space, nroots=nroots, verbose=self.verbose)

        self.eip = eip.real

        if nroots == 1:
            eip, evecs = [self.eip], [evecs]
        for n, en, vn in zip(range(nroots), eip, evecs):
            logger.info(self, 'IP root %d E = %.16g  qpwt = %0.6g',
                        n, en, np.linalg.norm(vn[:self.nocc])**2)
        log.timer('IP-CCSD', *cput0)
        if nroots == 1:
            return eip[0], evecs[0]
        else:
            return eip, evecs

    def ipccsd_matvec(self, vector):
        # Ref: Nooijen and Snijders, J. Chem. Phys. 102, 1681 (1995) Eqs.(8)-(9)
        if not getattr(self, 'imds', None):
            self.imds = _IMDS(self)
        if not self.imds.made_ip_imds:
            self.imds.make_ip(self.ip_partition)
        imds = self.imds

        r1,r2 = self.vector_to_amplitudes_ip(vector)

        # 1h-1h block
        Hr1 = -einsum('ki,k->i',imds.Loo,r1)
        #1h-2h1p block
        Hr1 += 2*einsum('ld,ild->i',imds.Fov,r2)
        Hr1 +=  -einsum('kd,kid->i',imds.Fov,r2)
        Hr1 += -2*einsum('klid,kld->i',imds.Wooov,r2)
        Hr1 +=    einsum('lkid,kld->i',imds.Wooov,r2)

        # 2h1p-1h block
        Hr2 = -einsum('kbij,k->ijb',imds.Wovoo,r1)
        # 2h1p-2h1p block
        if self.ip_partition == 'mp':
            nocc, nvir = self.t1.shape
            fock = self.eris.fock
            foo = fock[:nocc,:nocc]
            fvv = fock[nocc:,nocc:]
            Hr2 += einsum('bd,ijd->ijb',fvv,r2)
            Hr2 += -einsum('ki,kjb->ijb',foo,r2)
            Hr2 += -einsum('lj,ilb->ijb',foo,r2)
        elif self.ip_partition == 'full':
            Hr2 += self._ipccsd_diag_matrix2*r2
        else:
            Hr2 += einsum('bd,ijd->ijb',imds.Lvv,r2)
            Hr2 += -einsum('ki,kjb->ijb',imds.Loo,r2)
            Hr2 += -einsum('lj,ilb->ijb',imds.Loo,r2)
            Hr2 +=  einsum('klij,klb->ijb',imds.Woooo,r2)
            Hr2 += 2*einsum('lbdj,ild->ijb',imds.Wovvo,r2)
            Hr2 +=  -einsum('kbdj,kid->ijb',imds.Wovvo,r2)
            Hr2 +=  -einsum('lbjd,ild->ijb',imds.Wovov,r2) #typo in Ref
            Hr2 +=  -einsum('kbid,kjd->ijb',imds.Wovov,r2)
            tmp = 2*einsum('lkdc,kld->c',imds.Woovv,r2)
            tmp += -einsum('kldc,kld->c',imds.Woovv,r2)
            Hr2 += -einsum('c,ijcb->ijb',tmp,self.t2)

        vector = self.amplitudes_to_vector_ip(Hr1,Hr2)
        return vector

    def lipccsd_matvec(self, vector):
        if not getattr(self, 'imds', None):
            self.imds = _IMDS(self)
        if not self.imds.made_ip_imds:
            self.imds.make_ip(self.ip_partition)
        imds = self.imds

        r1,r2 = self.vector_to_amplitudes_ip(vector)

        # 1h-1h block
        Hr1 = -einsum('ki,i->k',imds.Loo,r1)
        #1h-2h1p block
        Hr1 += -einsum('kbij,ijb->k',imds.Wovoo,r2)

        # 2h1p-1h block
        Hr2 = -einsum('kd,l->kld',imds.Fov,r1)
        Hr2 += 2.*einsum('ld,k->kld',imds.Fov,r1)
        Hr2 += -2.*einsum('klid,i->kld',imds.Wooov,r1)
        Hr2 += einsum('lkid,i->kld',imds.Wooov,r1)
        # 2h1p-2h1p block
        if self.ip_partition == 'mp':
            nocc, nvir = self.t1.shape
            fock = self.eris.fock
            foo = fock[:nocc,:nocc]
            fvv = fock[nocc:,nocc:]
            Hr2 += einsum('bd,klb->kld',fvv,r2)
            Hr2 += -einsum('ki,ild->kld',foo,r2)
            Hr2 += -einsum('lj,kjd->kld',foo,r2)
        elif self.ip_partition == 'full':
            Hr2 += self._ipccsd_diag_matrix2*r2
        else:
            Hr2 += einsum('bd,klb->kld',imds.Lvv,r2)
            Hr2 += -einsum('ki,ild->kld',imds.Loo,r2)
            Hr2 += -einsum('lj,kjd->kld',imds.Loo,r2)
            Hr2 += 2.*einsum('lbdj,kjb->kld',imds.Wovvo,r2)
            Hr2 += -einsum('kbdj,ljb->kld',imds.Wovvo,r2)
            Hr2 += -einsum('lbjd,kjb->kld',imds.Wovov,r2)
            Hr2 += einsum('klij,ijd->kld',imds.Woooo,r2)
            Hr2 += -einsum('kbid,ilb->kld',imds.Wovov,r2)
            tmp = einsum('ijcb,ijb->c',t2,r2)
            Hr2 += einsum('kldc,c->kld',imds.Woovv,tmp)
            Hr2 += -2.*einsum('lkdc,c->kld',imds.Woovv,tmp)

        vector = self.amplitudes_to_vector_ip(Hr1,Hr2)
        return vector

    def ipccsd_diag(self):
        if not getattr(self, 'imds', None):
            self.imds = _IMDS(self)
        if not self.imds.made_ip_imds:
            self.imds.make_ip(self.ip_partition)
        imds = self.imds

        t1, t2 = self.t1, self.t2
        nocc, nvir = t1.shape
        fock = self.eris.fock
        foo = fock[:nocc,:nocc]
        fvv = fock[nocc:,nocc:]

        Hr1 = -np.diag(imds.Loo)
        Hr2 = np.zeros((nocc,nocc,nvir), t1.dtype)
        for i in range(nocc):
            for j in range(nocc):
                for b in range(nvir):
                    if self.ip_partition == 'mp':
                        Hr2[i,j,b] += fvv[b,b]
                        Hr2[i,j,b] += -foo[i,i]
                        Hr2[i,j,b] += -foo[j,j]
                    else:
                        Hr2[i,j,b] += imds.Lvv[b,b]
                        Hr2[i,j,b] += -imds.Loo[i,i]
                        Hr2[i,j,b] += -imds.Loo[j,j]
                        Hr2[i,j,b] += imds.Woooo[i,j,i,j]
                        Hr2[i,j,b] += 2*imds.Wovvo[j,b,b,j]
                        Hr2[i,j,b] += -imds.Wovvo[i,b,b,i]*(i==j)
                        Hr2[i,j,b] += -imds.Wovov[j,b,j,b]
                        Hr2[i,j,b] += -imds.Wovov[i,b,i,b]
                        Hr2[i,j,b] += -2*np.dot(imds.Woovv[j,i,b,:],t2[i,j,:,b])
                        Hr2[i,j,b] += np.dot(imds.Woovv[i,j,b,:],t2[i,j,:,b])

        vector = self.amplitudes_to_vector_ip(Hr1,Hr2)
        return vector

    def vector_to_amplitudes_ip(self,vector):
        nocc = self.nocc
        nvir = self.nmo - nocc
        r1 = vector[:nocc].copy()
        r2 = vector[nocc:].copy().reshape(nocc,nocc,nvir)
        return [r1,r2]

    def amplitudes_to_vector_ip(self,r1,r2):
        nocc = self.nocc
        nvir = self.nmo - nocc
        size = self.nip()
        vector = np.zeros(size, r1.dtype)
        vector[:nocc] = r1.copy()
        vector[nocc:] = r2.copy().reshape(nocc*nocc*nvir)
        return vector

    def ipccsd_star_contract(self, ipccsd_evals, ipccsd_evecs, lipccsd_evecs):
        assert self.ip_partition is None
        t2, eris = self.t2, self.eris
        fock = eris.fock
        nocc = self.nocc
        nvir = self.nmo - nocc

        #fov = fock[:nocc,nocc:]
        foo = fock[:nocc,:nocc]
        fvv = fock[nocc:,nocc:]

        ovov = _cp(eris.ovov)
        ovvv = _cp(eris.ovvv)
        vovv = ovvv.conj().transpose(1,0,3,2)
        oovv = _cp(eris.oovv)
        ovvo = _cp(eris.ovvo)
        ooov = _cp(eris.ooov)
        vooo = ooov.conj().transpose(3,2,1,0)
        oooo = _cp(eris.oooo)

        eijkab = np.zeros((nocc,nocc,nocc,nvir,nvir))
        for i,j,k in lib.cartesian_prod([range(nocc),range(nocc),range(nocc)]):
            for a,b in lib.cartesian_prod([range(nvir),range(nvir)]):
                eijkab[i,j,k,a,b] = foo[i,i] + foo[j,j] + foo[k,k] - fvv[a,a] - fvv[b,b]

        ipccsd_evecs  = np.array(ipccsd_evecs)
        lipccsd_evecs = np.array(lipccsd_evecs)
        e = []
        for _eval, _evec, _levec in zip(ipccsd_evals, ipccsd_evecs, lipccsd_evecs):
            l1,l2 = self.vector_to_amplitudes_ip(_levec)
            r1,r2 = self.vector_to_amplitudes_ip(_evec)
            ldotr = np.dot(l1.conj(),r1) + np.dot(l2.ravel(),r2.ravel())
            l1 /= ldotr
            l2 /= ldotr
            l2 = 1./3*(l2 + 2.*l2.transpose(1,0,2))

            _eijkab = eijkab + _eval
            _eijkab = 1./_eijkab

            lijkab = 0.5*einsum('iajb,k->ijkab', ovov, l1)
            lijkab += einsum('iaeb,jke->ijkab', ovvv, l2)
            lijkab += -einsum('kmjb,ima->ijkab', ooov, l2)
            lijkab += -einsum('imjb,mka->ijkab', ooov, l2)
            lijkab = lijkab + lijkab.transpose(1,0,2,4,3)

            rijkab = -einsum('mkbe,m,ijae->ijkab', oovv, r1, t2)
            rijkab -= einsum('mebj,m,ikae->ijkab', ovvo, r1, t2)
            rijkab += einsum('mjnk,n,imab->ijkab', oooo, r1, t2)
            rijkab +=  einsum('aibe,kje->ijkab', vovv, r2)
            rijkab += -einsum('bjmk,mia->ijkab', vooo, r2)
            rijkab += -einsum('bjmi,kma->ijkab', vooo, r2)
            rijkab = rijkab + rijkab.transpose(1,0,2,4,3)

            lijkab = 4.*lijkab \
                   - 2.*lijkab.transpose(1,0,2,3,4) \
                   - 2.*lijkab.transpose(2,1,0,3,4) \
                   - 2.*lijkab.transpose(0,2,1,3,4) \
                   + 1.*lijkab.transpose(1,2,0,3,4) \
                   + 1.*lijkab.transpose(2,0,1,3,4)

            deltaE = 0.5*einsum('ijkab,ijkab,ijkab',lijkab,rijkab,_eijkab)
            deltaE = deltaE.real
            logger.note(self, "Exc. energy, delta energy = %16.12f, %16.12f",
                        _eval+deltaE, deltaE)
            e.append(_eval+deltaE)
        return e

    def eaccsd(self, nroots=1, left=False, koopmans=False, guess=None, partition=None):
        '''Calculate (N+1)-electron charged excitations via EA-EOM-CCSD.

        Kwargs:
            See ipccd()
        '''
        cput0 = (logger.process_clock(), logger.perf_counter())
        log = logger.Logger(self.stdout, self.verbose)
        size = self.nea()
        nroots = min(nroots,size)
        if partition:
            partition = partition.lower()
            assert partition in ['mp','full']
        self.ea_partition = partition
        if partition == 'full':
            self._eaccsd_diag_matrix2 = self.vector_to_amplitudes_ea(self.eaccsd_diag())[1]

        adiag = self.eaccsd_diag()
        user_guess = False
        if guess:
            user_guess = True
            assert len(guess) == nroots
            for g in guess:
                assert g.size == size
        else:
            guess = []
            if koopmans:
                for n in range(nroots):
                    g = np.zeros(size)
                    g[n] = 1.0
                    guess.append(g)
            else:
                idx = adiag.argsort()[:nroots]
                for i in idx:
                    g = np.zeros(size)
                    g[i] = 1.0
                    guess.append(g)

        def precond(r, e0, x0):
            return r/(e0-adiag+1e-12)

        if left:
            matvec = self.leaccsd_matvec
        else:
            matvec = self.eaccsd_matvec
        eig = linalg_helper.eig
        if user_guess or koopmans:
            def pickeig(w, v, nr, envs):
                x0 = linalg_helper._gen_x0(envs['v'], envs['xs'])
                idx = np.argmax( np.abs(np.dot(np.array(guess).conj(),np.array(x0).T)), axis=1 )
                return lib.linalg_helper._eigs_cmplx2real(w, v, idx)
            eea, evecs = eig(matvec, guess, precond, pick=pickeig,
                             tol=self.conv_tol, max_cycle=self.max_cycle,
                             max_space=self.max_space, nroots=nroots, verbose=self.verbose)
        else:
            eea, evecs = eig(matvec, guess, precond,
                             tol=self.conv_tol, max_cycle=self.max_cycle,
                             max_space=self.max_space, nroots=nroots, verbose=self.verbose)

        self.eea = eea.real

        if nroots == 1:
            eea, evecs = [self.eea], [evecs]
        nvir = self.nmo - self.nocc
        for n, en, vn in zip(range(nroots), eea, evecs):
            logger.info(self, 'EA root %d E = %.16g  qpwt = %0.6g',
                        n, en, np.linalg.norm(vn[:nvir])**2)
        log.timer('EA-CCSD', *cput0)
        if nroots == 1:
            return eea[0], evecs[0]
        else:
            return eea, evecs

    def eaccsd_matvec(self,vector):
        # Ref: Nooijen and Bartlett, J. Chem. Phys. 102, 3629 (1994) Eqs.(30)-(31)
        if not getattr(self, 'imds', None):
            self.imds = _IMDS(self)
        if not self.imds.made_ea_imds:
            self.imds.make_ea(self.ea_partition)
        imds = self.imds

        r1,r2 = self.vector_to_amplitudes_ea(vector)

        # Eq. (30)
        # 1p-1p block
        Hr1 =  einsum('ac,c->a',imds.Lvv,r1)
        # 1p-2p1h block
        Hr1 += einsum('ld,lad->a',2.*imds.Fov,r2)
        Hr1 += einsum('ld,lda->a',  -imds.Fov,r2)
        Hr1 += 2*einsum('alcd,lcd->a',imds.Wvovv,r2)
        Hr1 +=  -einsum('aldc,lcd->a',imds.Wvovv,r2)
        # Eq. (31)
        # 2p1h-1p block
        Hr2 = einsum('abcj,c->jab',imds.Wvvvo,r1)
        # 2p1h-2p1h block
        if self.ea_partition == 'mp':
            nocc, nvir = self.t1.shape
            fock = self.eris.fock
            foo = fock[:nocc,:nocc]
            fvv = fock[nocc:,nocc:]
            Hr2 +=  einsum('ac,jcb->jab',fvv,r2)
            Hr2 +=  einsum('bd,jad->jab',fvv,r2)
            Hr2 += -einsum('lj,lab->jab',foo,r2)
        elif self.ea_partition == 'full':
            Hr2 += self._eaccsd_diag_matrix2*r2
        else:
            Hr2 +=  einsum('ac,jcb->jab',imds.Lvv,r2)
            Hr2 +=  einsum('bd,jad->jab',imds.Lvv,r2)
            Hr2 += -einsum('lj,lab->jab',imds.Loo,r2)
            Hr2 += 2*einsum('lbdj,lad->jab',imds.Wovvo,r2)
            Hr2 +=  -einsum('lbjd,lad->jab',imds.Wovov,r2)
            Hr2 +=  -einsum('lajc,lcb->jab',imds.Wovov,r2)
            Hr2 +=  -einsum('lbcj,lca->jab',imds.Wovvo,r2)
            nvir = self.nmo-self.nocc
            for a in range(nvir):
                Hr2[:,a,:] += einsum('bcd,jcd->jb',imds.Wvvvv[a],r2)
            tmp = (2*einsum('klcd,lcd->k',imds.Woovv,r2)
                    -einsum('kldc,lcd->k',imds.Woovv,r2))
            Hr2 += -einsum('k,kjab->jab',tmp,self.t2)

        vector = self.amplitudes_to_vector_ea(Hr1,Hr2)
        return vector

    def leaccsd_matvec(self,vector):
        # Note this is not the same left EA equations used by Nooijen and Bartlett.
        # Small changes were made so that the same type L2 basis was used for both the
        # left EA and left IP equations.  You will note more similarity for these
        # equations to the left IP equations than for the left EA equations by Nooijen.
        if not getattr(self, 'imds', None):
            self.imds = _IMDS(self)
        if not self.imds.made_ea_imds:
            self.imds.make_ea(self.ea_partition)
        imds = self.imds

        r1,r2 = self.vector_to_amplitudes_ea(vector)

        # Eq. (30)
        # 1p-1p block
        Hr1 = einsum('ac,a->c',imds.Lvv,r1)
        # 1p-2p1h block
        Hr1 += einsum('abcj,jab->c',imds.Wvvvo,r2)
        # Eq. (31)
        # 2p1h-1p block
        Hr2 = 2.*einsum('c,ld->lcd',r1,imds.Fov)
        Hr2 +=   -einsum('d,lc->lcd',r1,imds.Fov)
        Hr2 += 2.*einsum('a,alcd->lcd',r1,imds.Wvovv)
        Hr2 +=   -einsum('a,aldc->lcd',r1,imds.Wvovv)
        # 2p1h-2p1h block
        if self.ea_partition == 'mp':
            nocc, nvir = self.t1.shape
            fock = self.eris.fock
            foo = fock[:nocc,:nocc]
            fvv = fock[nocc:,nocc:]
            Hr2 += einsum('lad,ac->lcd',r2,fvv)
            Hr2 += einsum('lcb,bd->lcd',r2,fvv)
            Hr2 += -einsum('jcd,lj->lcd',r2,foo)
        elif self.ea_partition == 'full':
            Hr2 += self._eaccsd_diag_matrix2*r2
        else:
            Hr2 += einsum('lad,ac->lcd',r2,imds.Lvv)
            Hr2 += einsum('lcb,bd->lcd',r2,imds.Lvv)
            Hr2 += -einsum('jcd,lj->lcd',r2,imds.Loo)
            Hr2 += 2.*einsum('jcb,lbdj->lcd',r2,imds.Wovvo)
            Hr2 +=   -einsum('jcb,lbjd->lcd',r2,imds.Wovov)
            Hr2 +=   -einsum('lajc,jad->lcd',imds.Wovov,r2)
            Hr2 +=   -einsum('lbcj,jdb->lcd',imds.Wovvo,r2)
            nvir = self.nmo-self.nocc
            for a in range(nvir):
                Hr2 += einsum('lb,bcd->lcd',r2[:,a,:],imds.Wvvvv[a])
            tmp = einsum('ijcb,ibc->j',t2,r2)
            Hr2 +=     einsum('kjef,j->kef',imds.Woovv,tmp)
            Hr2 += -2.*einsum('kjfe,j->kef',imds.Woovv,tmp)

        vector = self.amplitudes_to_vector_ea(Hr1,Hr2)
        return vector

    def eaccsd_diag(self):
        if not getattr(self, 'imds', None):
            self.imds = _IMDS(self)
        if not self.imds.made_ea_imds:
            self.imds.make_ea(self.ea_partition)
        imds = self.imds

        t1, t2 = self.t1, self.t2
        nocc, nvir = t1.shape

        fock = self.eris.fock
        foo = fock[:nocc,:nocc]
        fvv = fock[nocc:,nocc:]

        Hr1 = np.diag(imds.Lvv)
        Hr2 = np.zeros((nocc,nvir,nvir), t1.dtype)
        for a in range(nvir):
            if self.ea_partition != 'mp':
                _Wvvvva = np.array(imds.Wvvvv[a])
            for b in range(nvir):
                for j in range(nocc):
                    if self.ea_partition == 'mp':
                        Hr2[j,a,b] += fvv[a,a]
                        Hr2[j,a,b] += fvv[b,b]
                        Hr2[j,a,b] += -foo[j,j]
                    else:
                        Hr2[j,a,b] += imds.Lvv[a,a]
                        Hr2[j,a,b] += imds.Lvv[b,b]
                        Hr2[j,a,b] += -imds.Loo[j,j]
                        Hr2[j,a,b] += 2*imds.Wovvo[j,b,b,j]
                        Hr2[j,a,b] += -imds.Wovov[j,b,j,b]
                        Hr2[j,a,b] += -imds.Wovov[j,a,j,a]
                        Hr2[j,a,b] += -imds.Wovvo[j,b,b,j]*(a==b)
                        Hr2[j,a,b] += _Wvvvva[b,a,b]
                        Hr2[j,a,b] += -2*np.dot(imds.Woovv[:,j,a,b],t2[:,j,a,b])
                        Hr2[j,a,b] += np.dot(imds.Woovv[:,j,b,a],t2[:,j,a,b])

        vector = self.amplitudes_to_vector_ea(Hr1,Hr2)
        return vector

    def vector_to_amplitudes_ea(self,vector):
        nocc = self.nocc
        nvir = self.nmo - nocc
        r1 = vector[:nvir].copy()
        r2 = vector[nvir:].copy().reshape(nocc,nvir,nvir)
        return [r1,r2]

    def amplitudes_to_vector_ea(self,r1,r2):
        nocc = self.nocc
        nvir = self.nmo - nocc
        size = self.nea()
        vector = np.zeros(size, r1.dtype)
        vector[:nvir] = r1.copy()
        vector[nvir:] = r2.copy().reshape(nocc*nvir*nvir)
        return vector

    def eaccsd_star_contract(self, eaccsd_evals, eaccsd_evecs, leaccsd_evecs):
        assert self.ea_partition is None
        t2, eris = self.t2, self.eris
        fock = eris.fock
        nocc = self.nocc
        nvir = self.nmo - nocc

        #fov = fock[:nocc,nocc:]
        foo = fock[:nocc,:nocc]
        fvv = fock[nocc:,nocc:]

        ovov = _cp(eris.ovov)
        ovvv = _cp(eris.ovvv)
        vovv = ovvv.conj().transpose(1,0,3,2)
        ooov = _cp(eris.ooov)
        vooo = ooov.conj().transpose(3,2,1,0)
        oovv = _cp(eris.oovv)
        #oooo = _cp(eris.oooo)
        vvvv = _cp(eris.vvvv)
        ovvo = _cp(eris.ovvo)

        eijabc = np.zeros((nocc,nocc,nvir,nvir,nvir))
        for i,j in lib.cartesian_prod([range(nocc),range(nocc)]):
            for a,b,c in lib.cartesian_prod([range(nvir),range(nvir),range(nvir)]):
                eijabc[i,j,a,b,c] = foo[i,i] + foo[j,j] - fvv[a,a] - fvv[b,b] - fvv[c,c]

        eaccsd_evecs  = np.array(eaccsd_evecs)
        leaccsd_evecs = np.array(leaccsd_evecs)
        e = []
        for _eval, _evec, _levec in zip(eaccsd_evals, eaccsd_evecs, leaccsd_evecs):
            l1,l2 = self.vector_to_amplitudes_ea(_levec)
            r1,r2 = self.vector_to_amplitudes_ea(_evec)
            ldotr = np.dot(l1.conj(),r1) + np.dot(l2.ravel(),r2.ravel())
            l1 /= ldotr
            l2 /= ldotr
            l2 = 1./3*(1.*l2 + 2.*l2.transpose(0,2,1))
            r2 = r2.transpose(0,2,1)

            _eijabc = eijabc + _eval
            _eijabc = 1./_eijabc

            lijabc = -0.5*einsum('c,iajb->ijabc', l1, ovov)
            lijabc += einsum('jmia,mbc->ijabc', ooov, l2)
            lijabc -= einsum('iaeb,jec->ijabc', ovvv, l2)
            lijabc -= einsum('jbec,iae->ijabc', ovvv, l2)
            lijabc = lijabc + lijabc.transpose(1,0,3,2,4)

            rijabc = -einsum('becf,f,ijae->ijabc', vvvv, r1, t2)
            rijabc += einsum('mjce,e,imab->ijabc', oovv, r1, t2)
            rijabc += einsum('mebj,e,imac->ijabc', ovvo, r1, t2)
            rijabc += einsum('aimj,mbc->ijabc', vooo, r2)
            rijabc += -einsum('bjce,iae->ijabc', vovv, r2)
            rijabc += -einsum('aibe,jec->ijabc', vovv, r2)
            rijabc = rijabc + rijabc.transpose(1,0,3,2,4)

            lijabc =  4.*lijabc \
                    - 2.*lijabc.transpose(0,1,3,2,4) \
                    - 2.*lijabc.transpose(0,1,4,3,2) \
                    - 2.*lijabc.transpose(0,1,2,4,3) \
                    + 1.*lijabc.transpose(0,1,3,4,2) \
                    + 1.*lijabc.transpose(0,1,4,2,3)
            deltaE = 0.5*einsum('ijabc,ijabc,ijabc',lijabc,rijabc,_eijabc)
            deltaE = deltaE.real
            logger.note(self, "Exc. energy, delta energy = %16.12f, %16.12f",
                        _eval+deltaE, deltaE)
            e.append(_eval+deltaE)
        return e


    def eeccsd(self, nroots=1, koopmans=False, guess=None, partition=None):
        '''Calculate N-electron neutral excitations via EE-EOM-CCSD.

        Kwargs:
            nroots : int
                Number of roots (eigenvalues) requested
            partition : bool or str
                Use a matrix-partitioning for the doubles-doubles block.
                Can be None, 'mp' (Moller-Plesset, i.e. orbital energies on the diagonal),
                or 'full' (full diagonal elements).
            koopmans : bool
                Calculate Koopmans'-like (1p1h) excitations only, targeting via
                overlap.
            guess : list of ndarray
                List of guess vectors to use for targeting via overlap.
        '''
        cput0 = (logger.process_clock(), logger.perf_counter())
        log = logger.Logger(self.stdout, self.verbose)
        size = self.nee()
        nroots = min(nroots,size)
        if partition:
            partition = partition.lower()
            assert partition in ['mp','full']
        self.ee_partition = partition
        if partition == 'full':
            self._eeccsd_diag_matrix2 = self.vector_to_amplitudes_ee(self.eeccsd_diag())[1]

        nvir = self.nmo - self.nocc
        adiag = self.eeccsd_diag()
        user_guess = False
        if guess:
            user_guess = True
            assert len(guess) == nroots
            for g in guess:
                assert g.size == size
        else:
            guess = []
            idx = adiag.argsort()
            n = 0
            for i in idx:
                g = np.zeros(size)
                g[i] = 1.0
                if koopmans:
                    if np.linalg.norm(g[:self.nocc*nvir])**2 > 0.8:
                        guess.append(g)
                        n += 1
                else:
                    guess.append(g)
                    n += 1
                if n == nroots:
                    break

        def precond(r, e0, x0):
            return r/(e0-adiag+1e-12)

        eig = linalg_helper.eig
        if user_guess or koopmans:
            def pickeig(w, v, nr, envs):
                x0 = linalg_helper._gen_x0(envs['v'], envs['xs'])
                idx = np.argmax( np.abs(np.dot(np.array(guess).conj(),np.array(x0).T)), axis=1 )
                return lib.linalg_helper._eigs_cmplx2real(w, v, idx)
            eee, evecs = eig(self.eeccsd_matvec, guess, precond, pick=pickeig,
                             tol=self.conv_tol, max_cycle=self.max_cycle,
                             max_space=self.max_space, nroots=nroots, verbose=self.verbose)
        else:
            eee, evecs = eig(self.eeccsd_matvec, guess, precond,
                             tol=self.conv_tol, max_cycle=self.max_cycle,
                             max_space=self.max_space, nroots=nroots, verbose=self.verbose)

        self.eee = eee.real

        if nroots == 1:
            eee, evecs = [self.eee], [evecs]
        for n, en, vn in zip(range(nroots), eee, evecs):
            logger.info(self, 'EE root %d E = %.16g  qpwt = %0.6g',
                        n, en, np.linalg.norm(vn[:self.nocc*nvir])**2)
        log.timer('EE-CCSD', *cput0)
        if nroots == 1:
            return eee[0], evecs[0]
        else:
            return eee, evecs

    def eeccsd_matvec(self,vector):
        raise NotImplementedError

    def eeccsd_diag(self):
        raise NotImplementedError

    def vector_to_amplitudes_ee(self,vector):
        nocc = self.nocc
        nvir = self.nmo - nocc
        r1 = vector[:nocc*nvir].copy().reshape((nocc,nvir))
        r2 = vector[nocc*nvir:].copy().reshape((nocc,nocc,nvir,nvir))
        return [r1,r2]

    def amplitudes_to_vector_ee(self,r1,r2):
        nocc = self.nocc
        nvir = self.nmo - nocc
        size = self.nee()
        vector = np.zeros(size, r1.dtype)
        vector[:nocc*nvir] = r1.copy().reshape(nocc*nvir)
        vector[nocc*nvir:] = r2.copy().reshape(nocc*nocc*nvir*nvir)
        return vector


class _ChemistsERIs:
    def __init__(self, cc, mo_coeff=None, method='incore',
                 ao2mofn=ao2mo.outcore.general_iofree):
        if mo_coeff is None:
            self.mo_coeff = mo_coeff = ccsd._mo_without_core(cc, cc.mo_coeff)
        else:
            self.mo_coeff = mo_coeff = ccsd._mo_without_core(cc, mo_coeff)
        dm = cc._scf.make_rdm1(cc.mo_coeff, cc.mo_occ)
        fockao = cc._scf.get_hcore() + cc._scf.get_veff(cc.mol, dm)
        self.fock = reduce(np.dot, (mo_coeff.T, fockao, mo_coeff))

        nocc = cc.nocc
        nmo = cc.nmo
        eri1 = ao2mo.incore.full(cc._scf._eri, mo_coeff)
        eri1 = ao2mo.restore(1, eri1, nmo)
        self.oooo = eri1[:nocc,:nocc,:nocc,:nocc].copy()
        self.ooov = eri1[:nocc,:nocc,:nocc,nocc:].copy()
        self.ovoo = eri1[:nocc,nocc:,:nocc,:nocc].copy()
        self.ovov = eri1[:nocc,nocc:,:nocc,nocc:].copy()
        self.oovv = eri1[:nocc,:nocc,nocc:,nocc:].copy()
        self.ovvo = eri1[:nocc,nocc:,nocc:,:nocc].copy()
        self.ovvv = eri1[:nocc,nocc:,nocc:,nocc:].copy()
        self.vvvv = eri1[nocc:,nocc:,nocc:,nocc:].copy()

    def get_ovvv(self, *slices):
        '''To access a subblock of ovvv tensor'''
        if slices:
            return self.ovvv[slices]
        else:
            return self.ovvv

class _IMDS:
    def __init__(self, cc):
        self.verbose = cc.verbose
        self.stdout = cc.stdout
        self.t1 = cc.t1
        self.t2 = cc.t2
        self.eris = cc.eris
        self.made_ip_imds = False
        self.made_ea_imds = False
        self._made_shared_2e = False

    def _make_shared_1e(self):
        cput0 = (logger.process_clock(), logger.perf_counter())
        log = logger.Logger(self.stdout, self.verbose)

        t1,t2,eris = self.t1, self.t2, self.eris
        self.Loo = imd.Loo(t1,t2,eris)
        self.Lvv = imd.Lvv(t1,t2,eris)
        self.Fov = imd.cc_Fov(t1,t2,eris)

        log.timer('EOM-CCSD shared one-electron intermediates', *cput0)

    def _make_shared_2e(self):
        cput0 = (logger.process_clock(), logger.perf_counter())
        log = logger.Logger(self.stdout, self.verbose)

        t1,t2,eris = self.t1, self.t2, self.eris
        # 2 virtuals
        self.Wovov = imd.Wovov(t1,t2,eris)
        self.Wovvo = imd.Wovvo(t1,t2,eris)
        self.Woovv = np.asarray(eris.ovov).transpose(0,2,1,3)

        log.timer('EOM-CCSD shared two-electron intermediates', *cput0)

    def make_ip(self, ip_partition=None):
        self._make_shared_1e()
        if self._made_shared_2e is False and ip_partition != 'mp':
            self._make_shared_2e()
            self._made_shared_2e = True

        cput0 = (logger.process_clock(), logger.perf_counter())
        log = logger.Logger(self.stdout, self.verbose)

        t1,t2,eris = self.t1, self.t2, self.eris

        # 0 or 1 virtuals
        if ip_partition != 'mp':
            self.Woooo = imd.Woooo(t1,t2,eris)
        self.Wooov = imd.Wooov(t1,t2,eris)
        self.Wovoo = imd.Wovoo(t1,t2,eris)
        self.made_ip_imds = True
        log.timer('EOM-CCSD IP intermediates', *cput0)

    def make_ea(self, ea_partition=None):
        self._make_shared_1e()
        if self._made_shared_2e is False and ea_partition != 'mp':
            self._make_shared_2e()
            self._made_shared_2e = True

        cput0 = (logger.process_clock(), logger.perf_counter())
        log = logger.Logger(self.stdout, self.verbose)

        t1,t2,eris = self.t1, self.t2, self.eris

        # 3 or 4 virtuals
        self.Wvovv = imd.Wvovv(t1,t2,eris)
        if ea_partition == 'mp':
            self.Wvvvo = imd.Wvvvo(t1,t2,eris)
        else:
            self.Wvvvv = imd.Wvvvv(t1,t2,eris)
            self.Wvvvo = imd.Wvvvo(t1,t2,eris,self.Wvvvv)
        self.made_ea_imds = True
        log.timer('EOM-CCSD EA intermediates', *cput0)

    def make_ee(self):
        raise NotImplementedError

def _cp(a):
    return np.array(a, copy=False, order='C')


if __name__ == '__main__':
    from pyscf import scf
    from pyscf import gto

    mol = gto.M()
    nocc, nvir = 5, 12
    nmo = nocc + nvir
    nmo_pair = nmo*(nmo+1)//2
    mf = scf.RHF(mol)
    np.random.seed(12)
    mf._eri = np.random.random(nmo_pair*(nmo_pair+1)//2)
    mf.mo_coeff = np.random.random((nmo,nmo))
    mf.mo_energy = np.arange(0., nmo)
    mf.mo_occ = np.zeros(nmo)
    mf.mo_occ[:nocc] = 2
    vhf = mf.get_veff(mol, mf.make_rdm1())
    cinv = np.linalg.inv(mf.mo_coeff)
    mf.get_hcore = lambda *args: (reduce(np.dot, (cinv.T*mf.mo_energy, cinv)) - vhf)
    mycc = RCCSD(mf)
    eris = mycc.ao2mo()
    a = np.random.random((nmo,nmo)) * .1
    eris.fock += a + a.T.conj()
    t1 = np.random.random((nocc,nvir)) * .1
    t2 = np.random.random((nocc,nocc,nvir,nvir)) * .1
    t2 = t2 + t2.transpose(1,0,3,2)

    mycc.cc2 = False
    t1a, t2a = update_amps(mycc, t1, t2, eris)
    print(lib.finger(t1a) - -106360.5276951083)
    print(lib.finger(t2a) - 66540.100267798145)
    mycc.cc2 = True
    t1a, t2a = update_amps(mycc, t1, t2, eris)
    print(lib.finger(t1a) - -106360.5276951083)
    print(lib.finger(t2a) - -1517.9391800662809)

    eri1 = np.random.random((nmo,nmo,nmo,nmo)) + np.random.random((nmo,nmo,nmo,nmo))*1j
    eri1 = eri1.transpose(0,2,1,3)
    eri1 = eri1 + eri1.transpose(1,0,3,2).conj()
    eri1 = eri1 + eri1.transpose(2,3,0,1)
    eri1 *= .1
    eris.oooo = eri1[:nocc,:nocc,:nocc,:nocc].copy()
    eris.ooov = eri1[:nocc,:nocc,:nocc,nocc:].copy()
    eris.ovoo = eri1[:nocc,nocc:,:nocc,:nocc].copy()
    eris.ovov = eri1[:nocc,nocc:,:nocc,nocc:].copy()
    eris.oovv = eri1[:nocc,:nocc,nocc:,nocc:].copy()
    eris.ovvo = eri1[:nocc,nocc:,nocc:,:nocc].copy()
    eris.ovvv = eri1[:nocc,nocc:,nocc:,nocc:].copy()
    eris.vvvv = eri1[nocc:,nocc:,nocc:,nocc:].copy()
    a = np.random.random((nmo,nmo)) * .1j
    eris.fock = eris.fock + a + a.T.conj()

    t1 = t1 + np.random.random((nocc,nvir)) * .1j
    t2 = t2 + np.random.random((nocc,nocc,nvir,nvir)) * .1j
    t2 = t2 + t2.transpose(1,0,3,2)
    mycc.cc2 = False
    t1a, t2a = update_amps(mycc, t1, t2, eris)
    print(lib.finger(t1a) - (-13.32050019680894-1.8825765910430254j))
    print(lib.finger(t2a) - (9.2521062044785189+29.999480274811873j))
    mycc.cc2 = True
    t1a, t2a = update_amps(mycc, t1, t2, eris)
    print(lib.finger(t1a) - (-13.32050019680894-1.8825765910430254j))
    print(lib.finger(t2a) - (-0.056223856104895858+0.025472249329733986j))

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = 'cc-pvdz'
    #mol.basis = '3-21G'
    mol.verbose = 0
    mol.spin = 0
    mol.build()
    mf = scf.RHF(mol).run(conv_tol=1e-14)

    mycc = RCCSD(mf)
    eris = mycc.ao2mo()
    emp2, t1, t2 = mycc.init_amps(eris)
    print(lib.finger(t2) - 0.044540097905897198)
    np.random.seed(1)
    t1 = np.random.random(t1.shape)*.1
    t2 = np.random.random(t2.shape)*.1
    t2 = t2 + t2.transpose(1,0,3,2)
    t1, t2 = update_amps(mycc, t1, t2, eris)
    print(lib.finger(t1) - 0.25118555558133576)
    print(lib.finger(t2) - 0.02352137419932243)

    ecc, t1, t2 = mycc.kernel()
    print(ecc - -0.21334326214236796)

    part = None
    print("IP energies... (right eigenvector)")
    e,v = mycc.ipccsd(nroots=3)
    print(e[0] - 0.43356041409195489)
    print(e[1] - 0.51876598058509493)
    print(e[2] - 0.6782879569941862 )

    print("IP energies... (left eigenvector)")
    le,lv = mycc.ipccsd(nroots=3,left=True)
    print(le[0] - 0.43356040428879794)
    print(le[1] - 0.51876597800180335)
    print(le[2] - 0.67828755013874864)

    e = mycc.ipccsd_star_contract(e,v,lv)
    print(e[0] - 0.43793202073189047)
    print(e[1] - 0.52287073446559729)
    print(e[2] - 0.67994597948852287)

    print("EA energies... (right eigenvector)")
    e,v = mycc.eaccsd(nroots=3)
    print(e[0] - 0.16737886282063008)
    print(e[1] - 0.24027622989542635)
    print(e[2] - 0.51006796667905585)

    print("EA energies... (left eigenvector)")
    le,lv = mycc.eaccsd(nroots=3,left=True)
    print(le[0] - 0.16737896537079733)
    print(le[1] - 0.24027634198123343)
    print(le[2] - 0.51006809015066612)

    e = mycc.eaccsd_star_contract(e,v,lv)
    print(e[0] - 0.16656250953550664)
    print(e[1] - 0.23944144521387614)
    print(e[2] - 0.41399436888830721)

    # Note: Not implemented
    #e,v = mycc.eeccsd(nroots=4)


# pyscf/cc/rintermediates.py

#!/usr/bin/env python
# Copyright 2014-2021 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.

'''
Intermediates for restricted CCSD.  Complex integrals are supported.
'''

import numpy as np
from pyscf import lib
from pyscf.lib import logger
from pyscf import ao2mo

# This is restricted (R)CCSD
# Ref: Hirata et al., J. Chem. Phys. 120, 2581 (2004)

### Eqs. (37)-(39) "kappa"

def cc_Foo(t1, t2, eris):
    nocc, nvir = t1.shape
    foo = eris.fock[:nocc,:nocc]
    eris_ovov = np.asarray(eris.ovov)
    Fki  = 2*lib.einsum('kcld,ilcd->ki', eris_ovov, t2)
    Fki -=   lib.einsum('kdlc,ilcd->ki', eris_ovov, t2)
    Fki += 2*lib.einsum('kcld,ic,ld->ki', eris_ovov, t1, t1)
    Fki -=   lib.einsum('kdlc,ic,ld->ki', eris_ovov, t1, t1)
    Fki += foo
    return Fki

def cc_Fvv(t1, t2, eris):
    nocc, nvir = t1.shape
    fvv = eris.fock[nocc:,nocc:]
    eris_ovov = np.asarray(eris.ovov)
    Fac  =-2*lib.einsum('kcld,klad->ac', eris_ovov, t2)
    Fac +=   lib.einsum('kdlc,klad->ac', eris_ovov, t2)
    Fac -= 2*lib.einsum('kcld,ka,ld->ac', eris_ovov, t1, t1)
    Fac +=   lib.einsum('kdlc,ka,ld->ac', eris_ovov, t1, t1)
    Fac += fvv
    return Fac

def cc_Fov(t1, t2, eris):
    nocc, nvir = t1.shape
    fov = eris.fock[:nocc,nocc:]
    eris_ovov = np.asarray(eris.ovov)
    Fkc  = 2*np.einsum('kcld,ld->kc', eris_ovov, t1)
    Fkc -=   np.einsum('kdlc,ld->kc', eris_ovov, t1)
    Fkc += fov
    return Fkc

### Eqs. (40)-(41) "lambda"

def Loo(t1, t2, eris):
    nocc, nvir = t1.shape
    fov = eris.fock[:nocc,nocc:]
    Lki = cc_Foo(t1, t2, eris) + np.einsum('kc,ic->ki',fov, t1)
    eris_ovoo = np.asarray(eris.ovoo)
    Lki += 2*np.einsum('lcki,lc->ki', eris_ovoo, t1)
    Lki -=   np.einsum('kcli,lc->ki', eris_ovoo, t1)
    return Lki

def Lvv(t1, t2, eris):
    nocc, nvir = t1.shape
    fov = eris.fock[:nocc,nocc:]
    Lac = cc_Fvv(t1, t2, eris) - np.einsum('kc,ka->ac',fov, t1)
    eris_ovvv = np.asarray(eris.get_ovvv())
    Lac += 2*np.einsum('kdac,kd->ac', eris_ovvv, t1)
    Lac -=   np.einsum('kcad,kd->ac', eris_ovvv, t1)
    return Lac

### Eqs. (42)-(45) "chi"

def cc_Woooo(t1, t2, eris):
    eris_ovoo = np.asarray(eris.ovoo)
    Wklij  = lib.einsum('lcki,jc->klij', eris_ovoo, t1)
    Wklij += lib.einsum('kclj,ic->klij', eris_ovoo, t1)
    eris_ovov = np.asarray(eris.ovov)
    Wklij += lib.einsum('kcld,ijcd->klij', eris_ovov, t2)
    Wklij += lib.einsum('kcld,ic,jd->klij', eris_ovov, t1, t1)
    Wklij += np.asarray(eris.oooo).transpose(0,2,1,3)
    return Wklij

def cc_Wvvvv(t1, t2, eris):
    # Incore
    eris_ovvv = np.asarray(eris.get_ovvv())
    Wabcd  = lib.einsum('kdac,kb->abcd', eris_ovvv,-t1)
    Wabcd -= lib.einsum('kcbd,ka->abcd', eris_ovvv, t1)
    Wabcd += np.asarray(_get_vvvv(eris)).transpose(0,2,1,3)
    return Wabcd

def cc_Wvoov(t1, t2, eris):
    eris_ovvv = np.asarray(eris.get_ovvv())
    eris_ovoo = np.asarray(eris.ovoo)
    Wakic  = lib.einsum('kcad,id->akic', eris_ovvv, t1)
    Wakic -= lib.einsum('kcli,la->akic', eris_ovoo, t1)
    Wakic += np.asarray(eris.ovvo).transpose(2,0,3,1)
    eris_ovov = np.asarray(eris.ovov)
    Wakic -= 0.5*lib.einsum('ldkc,ilda->akic', eris_ovov, t2)
    Wakic -= 0.5*lib.einsum('lckd,ilad->akic', eris_ovov, t2)
    Wakic -= lib.einsum('ldkc,id,la->akic', eris_ovov, t1, t1)
    Wakic += lib.einsum('ldkc,ilad->akic', eris_ovov, t2)
    return Wakic

def cc_Wvovo(t1, t2, eris):
    eris_ovvv = np.asarray(eris.get_ovvv())
    eris_ovoo = np.asarray(eris.ovoo)
    Wakci  = lib.einsum('kdac,id->akci', eris_ovvv, t1)
    Wakci -= lib.einsum('lcki,la->akci', eris_ovoo, t1)
    Wakci += np.asarray(eris.oovv).transpose(2,0,3,1)
    eris_ovov = np.asarray(eris.ovov)
    Wakci -= 0.5*lib.einsum('lckd,ilda->akci', eris_ovov, t2)
    Wakci -= lib.einsum('lckd,id,la->akci', eris_ovov, t1, t1)
    return Wakci

def Wooov(t1, t2, eris):
    eris_ovov = np.asarray(eris.ovov)
    Wklid  = lib.einsum('ic,kcld->klid', t1, eris_ovov)
    Wklid += np.asarray(eris.ovoo).transpose(2,0,3,1)
    return Wklid

def Wvovv(t1, t2, eris):
    eris_ovov = np.asarray(eris.ovov)
    Walcd  = lib.einsum('ka,kcld->alcd',-t1, eris_ovov)
    Walcd += np.asarray(eris.get_ovvv()).transpose(2,0,3,1)
    return Walcd

def W1ovvo(t1, t2, eris):
    eris_ovov = np.asarray(eris.ovov)
    Wkaci  = 2*lib.einsum('kcld,ilad->kaci', eris_ovov, t2)
    Wkaci +=  -lib.einsum('kcld,liad->kaci', eris_ovov, t2)
    Wkaci +=  -lib.einsum('kdlc,ilad->kaci', eris_ovov, t2)
    Wkaci += np.asarray(eris.ovvo).transpose(0,2,1,3)
    return Wkaci

def W2ovvo(t1, t2, eris):
    Wkaci = lib.einsum('la,lkic->kaci',-t1, Wooov(t1, t2, eris))
    eris_ovvv = np.asarray(eris.get_ovvv())
    Wkaci += lib.einsum('kcad,id->kaci', eris_ovvv, t1)
    return Wkaci

def Wovvo(t1, t2, eris):
    Wkaci = W1ovvo(t1, t2, eris) + W2ovvo(t1, t2, eris)
    return Wkaci

def W1ovov(t1, t2, eris):
    eris_ovov = np.asarray(eris.ovov)
    Wkbid = -lib.einsum('kcld,ilcb->kbid', eris_ovov, t2)
    Wkbid += np.asarray(eris.oovv).transpose(0,2,1,3)
    return Wkbid

def W2ovov(t1, t2, eris):
    Wkbid = lib.einsum('klid,lb->kbid', Wooov(t1, t2, eris),-t1)
    eris_ovvv = np.asarray(eris.get_ovvv())
    Wkbid += lib.einsum('kcbd,ic->kbid', eris_ovvv, t1)
    return Wkbid

def Wovov(t1, t2, eris):
    return W1ovov(t1, t2, eris) + W2ovov(t1, t2, eris)

def Woooo(t1, t2, eris):
    eris_ovov = np.asarray(eris.ovov)
    Wklij  = lib.einsum('kcld,ijcd->klij', eris_ovov, t2)
    Wklij += lib.einsum('kcld,ic,jd->klij', eris_ovov, t1, t1)
    eris_ovoo = np.asarray(eris.ovoo)
    Wklij += lib.einsum('ldki,jd->klij', eris_ovoo, t1)
    Wklij += lib.einsum('kclj,ic->klij', eris_ovoo, t1)
    Wklij += np.asarray(eris.oooo).transpose(0,2,1,3)
    return Wklij

def Wvvvv(t1, t2, eris):
    eris_ovov = np.asarray(eris.ovov)
    Wabcd  = lib.einsum('kcld,klab->abcd', eris_ovov, t2)
    Wabcd += lib.einsum('kcld,ka,lb->abcd', eris_ovov, t1, t1)
    Wabcd += np.asarray(_get_vvvv(eris)).transpose(0,2,1,3)
    eris_ovvv = np.asarray(eris.get_ovvv())
    Wabcd -= lib.einsum('ldac,lb->abcd', eris_ovvv, t1)
    Wabcd -= lib.einsum('kcbd,ka->abcd', eris_ovvv, t1)
    return Wabcd

def Wvvvo(t1, t2, eris, _Wvvvv=None):
    nocc,nvir = t1.shape
    eris_ovvv = np.asarray(eris.get_ovvv())
    # Check if t1=0 (HF+MBPT(2))
    # don't make vvvv if you can avoid it!
    Wabcj  =  -lib.einsum('alcj,lb->abcj', W1ovov(t1, t2, eris).transpose(1,0,3,2), t1)
    Wabcj +=  -lib.einsum('kbcj,ka->abcj', W1ovvo(t1, t2, eris), t1)
    Wabcj += 2*lib.einsum('ldac,ljdb->abcj', eris_ovvv, t2)
    Wabcj +=  -lib.einsum('ldac,ljbd->abcj', eris_ovvv, t2)
    Wabcj +=  -lib.einsum('lcad,ljdb->abcj', eris_ovvv, t2)
    Wabcj +=  -lib.einsum('kcbd,jkda->abcj', eris_ovvv, t2)
    eris_ovoo = np.asarray(eris.ovoo)
    Wabcj +=   lib.einsum('kclj,lkba->abcj', eris_ovoo, t2)
    Wabcj +=   lib.einsum('kclj,lb,ka->abcj', eris_ovoo, t1, t1)
    Wabcj +=  -lib.einsum('kc,kjab->abcj', cc_Fov(t1, t2, eris), t2)
    Wabcj += np.asarray(eris_ovvv).transpose(3,1,2,0).conj()
    if np.any(t1):
        if _Wvvvv is None:
            _Wvvvv = Wvvvv(t1, t2, eris)
        Wabcj += lib.einsum('abcd,jd->abcj', _Wvvvv, t1)
    return Wabcj

def Wovoo(t1, t2, eris):
    eris_ovoo = np.asarray(eris.ovoo)
    eris_ovvv = np.asarray(eris.get_ovvv())
    Wkbij  =   lib.einsum('kbid,jd->kbij', W1ovov(t1, t2, eris), t1)
    Wkbij +=  -lib.einsum('klij,lb->kbij', Woooo(t1, t2, eris), t1)
    Wkbij +=   lib.einsum('kbcj,ic->kbij', W1ovvo(t1, t2, eris), t1)
    Wkbij += 2*lib.einsum('ldki,ljdb->kbij', eris_ovoo, t2)
    Wkbij +=  -lib.einsum('ldki,jldb->kbij', eris_ovoo, t2)
    Wkbij +=  -lib.einsum('kdli,ljdb->kbij', eris_ovoo, t2)
    Wkbij +=   lib.einsum('kcbd,jidc->kbij', eris_ovvv, t2)
    Wkbij +=   lib.einsum('kcbd,jd,ic->kbij', eris_ovvv, t1, t1)
    Wkbij +=  -lib.einsum('kclj,libc->kbij', eris_ovoo, t2)
    Wkbij +=   lib.einsum('kc,ijcb->kbij', cc_Fov(t1, t2, eris), t2)
    Wkbij += np.asarray(eris_ovoo).transpose(3,1,2,0).conj()
    return Wkbij

def _get_vvvv(eris):
    if eris.vvvv is None and getattr(eris, 'vvL', None) is not None:  # DF eris
        vvL = np.asarray(eris.vvL)
        nvir = int(np.sqrt(eris.vvL.shape[0]*2))
        return ao2mo.restore(1, lib.dot(vvL, vvL.T), nvir)
    elif eris.vvvv.ndim == 2:
        nvir = int(np.sqrt(eris.vvvv.shape[0]*2))
        return ao2mo.restore(1, np.asarray(eris.vvvv), nvir)
    else:
        return eris.vvvv

def _cp(a):
    return np.array(a, copy=False, order='C')

def get_t3p2_imds_slow(cc, t1, t2, eris=None, t3p2_ip_out=None, t3p2_ea_out=None):
    """Calculates T1, T2 amplitudes corrected by second-order T3 contribution
    and intermediates used in IP/EA-CCSD(T)a

    For description of arguments, see `get_t3p2_imds_slow` in `gintermediates.py`.
    """
    if eris is None:
        eris = cc.ao2mo()
    fock = eris.fock
    nocc, nvir = t1.shape

    fov = fock[:nocc, nocc:]
    #foo = fock[:nocc, :nocc].diagonal()
    #fvv = fock[nocc:, nocc:].diagonal()

    dtype = np.result_type(t1, t2)
    if np.issubdtype(dtype, np.dtype(complex).type):
        logger.error(cc, 't3p2 imds has not been strictly checked for use with complex integrals')

    mo_e_o = eris.mo_energy[:nocc]
    mo_e_v = eris.mo_energy[nocc:]

    ovov = _cp(eris.ovov)
    ovvv = _cp(eris.get_ovvv())
    eris_vvov = eris.get_ovvv().conj().transpose(1,3,0,2)  # Physicist notation
    eris_vooo = eris.ovoo[:].conj().transpose(1,0,3,2)  # Chemist notation

    ccsd_energy = cc.energy(t1, t2, eris)
    dtype = np.result_type(t1, t2)

    if t3p2_ip_out is None:
        t3p2_ip_out = np.zeros((nocc,nvir,nocc,nocc), dtype=dtype)
    Wmbkj = t3p2_ip_out

    if t3p2_ea_out is None:
        t3p2_ea_out = np.zeros((nvir,nvir,nvir,nocc), dtype=dtype)
    Wcbej = t3p2_ea_out

    tmp_t3  = np.einsum('abif,kjcf->ijkabc', eris_vvov, t2)
    tmp_t3 -= np.einsum('aimj,mkbc->ijkabc', eris_vooo, t2)

    tmp_t3 = (tmp_t3 + tmp_t3.transpose(0,2,1,3,5,4)
                     + tmp_t3.transpose(1,0,2,4,3,5)
                     + tmp_t3.transpose(1,2,0,4,5,3)
                     + tmp_t3.transpose(2,0,1,5,3,4)
                     + tmp_t3.transpose(2,1,0,5,4,3))

    eia = mo_e_o[:, None] - mo_e_v[None, :]
    eijab = eia[:, None, :, None] + eia[None, :, None, :]
    eijkabc = eijab[:, :, None, :, :, None] + eia[None, None, :, None, None, :]
    tmp_t3 /= eijkabc

    Ptmp_t3 = 2.*tmp_t3 - tmp_t3.transpose(1,0,2,3,4,5) - tmp_t3.transpose(2,1,0,3,4,5)
    pt1 =  0.5 * lib.einsum('jbkc,ijkabc->ia', 2.*ovov - ovov.transpose(0,3,2,1), Ptmp_t3)

    tmp =  0.5 * lib.einsum('ijkabc,ia->jkbc', Ptmp_t3, fov)
    pt2 = tmp + tmp.transpose(1,0,3,2)

    #     b\    / \    /
    #  /\---\  /   \  /
    # i\/a  d\/j   k\/c
    # ------------------
    tmp =  lib.einsum('ijkabc,iadb->jkdc', Ptmp_t3, ovvv)
    pt2 += (tmp + tmp.transpose(1,0,3,2))

    #     m\    / \    /
    #  /\---\  /   \  /
    # i\/a  j\/b   k\/c
    # ------------------
    tmp =  lib.einsum('ijkabc,iajm->mkbc', Ptmp_t3, eris.ovoo)
    pt2 -= (tmp + tmp.transpose(1,0,3,2))

    eia = mo_e_o[:, None] - mo_e_v[None, :]
    eijab = eia[:, None, :, None] + eia[None, :, None, :]

    pt1 /= eia
    pt2 /= eijab

    pt1 += t1
    pt2 += t2

    Wmbkj +=  lib.einsum('ijkabc,mcia->mbkj', Ptmp_t3, ovov)

    Wcbej += -1.0*lib.einsum('ijkabc,iake->cbej', Ptmp_t3, ovov)

    delta_ccsd_energy = cc.energy(pt1, pt2, eris) - ccsd_energy
    logger.info(cc, 'CCSD energy T3[2] correction : %14.8e', delta_ccsd_energy)
    return delta_ccsd_energy, pt1, pt2, Wmbkj, Wcbej


# pyscf/cc/uccsd.py

#!/usr/bin/env python
# Copyright 2014-2021 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: Timothy Berkelbach <tim.berkelbach@gmail.com>
#         Qiming Sun <osirpt.sun@gmail.com>
#

'''
UCCSD with spatial integrals
'''


from functools import reduce
import numpy as np

from pyscf import lib
from pyscf import ao2mo
from pyscf.lib import logger
from pyscf.cc import ccsd
from pyscf.ao2mo import _ao2mo
from pyscf.mp import ump2
from pyscf import scf
from pyscf import __config__

MEMORYMIN = getattr(__config__, 'cc_ccsd_memorymin', 2000)

# This is unrestricted (U)CCSD, in spatial-orbital form.

def update_amps(cc, t1, t2, eris):
    time0 = logger.process_clock(), logger.perf_counter()
    log = logger.Logger(cc.stdout, cc.verbose)
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nocca, noccb, nvira, nvirb = t2ab.shape
    mo_ea_o = eris.mo_energy[0][:nocca]
    mo_ea_v = eris.mo_energy[0][nocca:] + cc.level_shift
    mo_eb_o = eris.mo_energy[1][:noccb]
    mo_eb_v = eris.mo_energy[1][noccb:] + cc.level_shift
    fova = eris.focka[:nocca,nocca:]
    fovb = eris.fockb[:noccb,noccb:]

    u1a = np.zeros_like(t1a)
    u1b = np.zeros_like(t1b)
    #:eris_vvvv = ao2mo.restore(1, np.asarray(eris.vvvv), nvirb)
    #:eris_VVVV = ao2mo.restore(1, np.asarray(eris.VVVV), nvirb)
    #:eris_vvVV = _restore(np.asarray(eris.vvVV), nvira, nvirb)
    #:u2aa += lib.einsum('ijef,aebf->ijab', tauaa, eris_vvvv) * .5
    #:u2bb += lib.einsum('ijef,aebf->ijab', taubb, eris_VVVV) * .5
    #:u2ab += lib.einsum('iJeF,aeBF->iJaB', tauab, eris_vvVV)
    tauaa, tauab, taubb = make_tau(t2, t1, t1)
    u2aa, u2ab, u2bb = cc._add_vvvv(None, (tauaa,tauab,taubb), eris)
    u2aa *= .5
    u2bb *= .5

    Fooa =  .5 * lib.einsum('me,ie->mi', fova, t1a)
    Foob =  .5 * lib.einsum('me,ie->mi', fovb, t1b)
    Fvva = -.5 * lib.einsum('me,ma->ae', fova, t1a)
    Fvvb = -.5 * lib.einsum('me,ma->ae', fovb, t1b)
    Fooa += eris.focka[:nocca,:nocca] - np.diag(mo_ea_o)
    Foob += eris.fockb[:noccb,:noccb] - np.diag(mo_eb_o)
    Fvva += eris.focka[nocca:,nocca:] - np.diag(mo_ea_v)
    Fvvb += eris.fockb[noccb:,noccb:] - np.diag(mo_eb_v)
    dtype = u2aa.dtype
    wovvo = np.zeros((nocca,nvira,nvira,nocca), dtype=dtype)
    wOVVO = np.zeros((noccb,nvirb,nvirb,noccb), dtype=dtype)
    woVvO = np.zeros((nocca,nvirb,nvira,noccb), dtype=dtype)
    woVVo = np.zeros((nocca,nvirb,nvirb,nocca), dtype=dtype)
    wOvVo = np.zeros((noccb,nvira,nvirb,nocca), dtype=dtype)
    wOvvO = np.zeros((noccb,nvira,nvira,noccb), dtype=dtype)

    mem_now = lib.current_memory()[0]
    max_memory = max(0, cc.max_memory - mem_now - u2aa.size*8e-6)
    if nvira > 0 and nocca > 0:
        blksize = max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvira**3*3+1)))
        for p0,p1 in lib.prange(0, nocca, blksize):
            ovvv = eris.get_ovvv(slice(p0,p1))  # ovvv = eris.ovvv[p0:p1]
            ovvv = ovvv - ovvv.transpose(0,3,2,1)
            Fvva += np.einsum('mf,mfae->ae', t1a[p0:p1], ovvv)
            wovvo[p0:p1] += lib.einsum('jf,mebf->mbej', t1a, ovvv)
            u1a += 0.5*lib.einsum('mief,meaf->ia', t2aa[p0:p1], ovvv)
            u2aa[:,p0:p1] += lib.einsum('ie,mbea->imab', t1a, ovvv.conj())
            tmp1aa = lib.einsum('ijef,mebf->ijmb', tauaa, ovvv)
            u2aa -= lib.einsum('ijmb,ma->ijab', tmp1aa, t1a[p0:p1]*.5)
            ovvv = tmp1aa = None

    if nvirb > 0 and noccb > 0:
        blksize = max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvirb**3*3+1)))
        for p0,p1 in lib.prange(0, noccb, blksize):
            OVVV = eris.get_OVVV(slice(p0,p1))  # OVVV = eris.OVVV[p0:p1]
            OVVV = OVVV - OVVV.transpose(0,3,2,1)
            Fvvb += np.einsum('mf,mfae->ae', t1b[p0:p1], OVVV)
            wOVVO[p0:p1] = lib.einsum('jf,mebf->mbej', t1b, OVVV)
            u1b += 0.5*lib.einsum('MIEF,MEAF->IA', t2bb[p0:p1], OVVV)
            u2bb[:,p0:p1] += lib.einsum('ie,mbea->imab', t1b, OVVV.conj())
            tmp1bb = lib.einsum('ijef,mebf->ijmb', taubb, OVVV)
            u2bb -= lib.einsum('ijmb,ma->ijab', tmp1bb, t1b[p0:p1]*.5)
            OVVV = tmp1bb = None

    if nvirb > 0 and nocca > 0:
        blksize = max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvira*nvirb**2*3+1)))
        for p0,p1 in lib.prange(0, nocca, blksize):
            ovVV = eris.get_ovVV(slice(p0,p1))  # ovVV = eris.ovVV[p0:p1]
            Fvvb += np.einsum('mf,mfAE->AE', t1a[p0:p1], ovVV)
            woVvO[p0:p1] = lib.einsum('JF,meBF->mBeJ', t1b, ovVV)
            woVVo[p0:p1] = lib.einsum('jf,mfBE->mBEj',-t1a, ovVV)
            u1b += lib.einsum('mIeF,meAF->IA', t2ab[p0:p1], ovVV)
            u2ab[p0:p1] += lib.einsum('IE,maEB->mIaB', t1b, ovVV.conj())
            tmp1ab = lib.einsum('iJeF,meBF->iJmB', tauab, ovVV)
            u2ab -= lib.einsum('iJmB,ma->iJaB', tmp1ab, t1a[p0:p1])
            ovVV = tmp1ab = None

    if nvira > 0 and noccb > 0:
        blksize = max(ccsd.BLKMIN, int(max_memory*1e6/8/(nvirb*nvira**2*3+1)))
        for p0,p1 in lib.prange(0, noccb, blksize):
            OVvv = eris.get_OVvv(slice(p0,p1))  # OVvv = eris.OVvv[p0:p1]
            Fvva += np.einsum('MF,MFae->ae', t1b[p0:p1], OVvv)
            wOvVo[p0:p1] = lib.einsum('jf,MEbf->MbEj', t1a, OVvv)
            wOvvO[p0:p1] = lib.einsum('JF,MFbe->MbeJ',-t1b, OVvv)
            u1a += lib.einsum('iMfE,MEaf->ia', t2ab[:,p0:p1], OVvv)
            u2ab[:,p0:p1] += lib.einsum('ie,MBea->iMaB', t1a, OVvv.conj())
            tmp1abba = lib.einsum('iJeF,MFbe->iJbM', tauab, OVvv)
            u2ab -= lib.einsum('iJbM,MA->iJbA', tmp1abba, t1b[p0:p1])
            OVvv = tmp1abba = None

    eris_ovov = np.asarray(eris.ovov)
    eris_ovoo = np.asarray(eris.ovoo)
    Woooo = lib.einsum('je,nemi->mnij', t1a, eris_ovoo)
    Woooo = Woooo - Woooo.transpose(0,1,3,2)
    Woooo += np.asarray(eris.oooo).transpose(0,2,1,3)
    Woooo += lib.einsum('ijef,menf->mnij', tauaa, eris_ovov) * .5
    u2aa += lib.einsum('mnab,mnij->ijab', tauaa, Woooo*.5)
    Woooo = tauaa = None
    ovoo = eris_ovoo - eris_ovoo.transpose(2,1,0,3)
    Fooa += np.einsum('ne,nemi->mi', t1a, ovoo)
    u1a += 0.5*lib.einsum('mnae,meni->ia', t2aa, ovoo)
    wovvo += lib.einsum('nb,nemj->mbej', t1a, ovoo)
    ovoo = eris_ovoo = None

    tilaa = make_tau_aa(t2[0], t1a, t1a, fac=0.5)
    ovov = eris_ovov - eris_ovov.transpose(0,3,2,1)
    Fvva -= .5 * lib.einsum('mnaf,menf->ae', tilaa, ovov)
    Fooa += .5 * lib.einsum('inef,menf->mi', tilaa, ovov)
    Fova = np.einsum('nf,menf->me',t1a, ovov)
    u2aa += ovov.conj().transpose(0,2,1,3) * .5
    wovvo -= 0.5*lib.einsum('jnfb,menf->mbej', t2aa, ovov)
    woVvO += 0.5*lib.einsum('nJfB,menf->mBeJ', t2ab, ovov)
    tmpaa = lib.einsum('jf,menf->mnej', t1a, ovov)
    wovvo -= lib.einsum('nb,mnej->mbej', t1a, tmpaa)
    eris_ovov = ovov = tmpaa = tilaa = None

    eris_OVOV = np.asarray(eris.OVOV)
    eris_OVOO = np.asarray(eris.OVOO)
    WOOOO = lib.einsum('je,nemi->mnij', t1b, eris_OVOO)
    WOOOO = WOOOO - WOOOO.transpose(0,1,3,2)
    WOOOO += np.asarray(eris.OOOO).transpose(0,2,1,3)
    WOOOO += lib.einsum('ijef,menf->mnij', taubb, eris_OVOV) * .5
    u2bb += lib.einsum('mnab,mnij->ijab', taubb, WOOOO*.5)
    WOOOO = taubb = None
    OVOO = eris_OVOO - eris_OVOO.transpose(2,1,0,3)
    Foob += np.einsum('ne,nemi->mi', t1b, OVOO)
    u1b += 0.5*lib.einsum('mnae,meni->ia', t2bb, OVOO)
    wOVVO += lib.einsum('nb,nemj->mbej', t1b, OVOO)
    OVOO = eris_OVOO = None

    tilbb = make_tau_aa(t2[2], t1b, t1b, fac=0.5)
    OVOV = eris_OVOV - eris_OVOV.transpose(0,3,2,1)
    Fvvb -= .5 * lib.einsum('MNAF,MENF->AE', tilbb, OVOV)
    Foob += .5 * lib.einsum('inef,menf->mi', tilbb, OVOV)
    Fovb = np.einsum('nf,menf->me',t1b, OVOV)
    u2bb += OVOV.conj().transpose(0,2,1,3) * .5
    wOVVO -= 0.5*lib.einsum('jnfb,menf->mbej', t2bb, OVOV)
    wOvVo += 0.5*lib.einsum('jNbF,MENF->MbEj', t2ab, OVOV)
    tmpbb = lib.einsum('jf,menf->mnej', t1b, OVOV)
    wOVVO -= lib.einsum('nb,mnej->mbej', t1b, tmpbb)
    eris_OVOV = OVOV = tmpbb = tilbb = None

    eris_OVoo = np.asarray(eris.OVoo)
    eris_ovOO = np.asarray(eris.ovOO)
    Fooa += np.einsum('NE,NEmi->mi', t1b, eris_OVoo)
    u1a -= lib.einsum('nMaE,MEni->ia', t2ab, eris_OVoo)
    wOvVo -= lib.einsum('nb,MEnj->MbEj', t1a, eris_OVoo)
    woVVo += lib.einsum('NB,NEmj->mBEj', t1b, eris_OVoo)
    Foob += np.einsum('ne,neMI->MI', t1a, eris_ovOO)
    u1b -= lib.einsum('mNeA,meNI->IA', t2ab, eris_ovOO)
    woVvO -= lib.einsum('NB,meNJ->mBeJ', t1b, eris_ovOO)
    wOvvO += lib.einsum('nb,neMJ->MbeJ', t1a, eris_ovOO)
    WoOoO = lib.einsum('JE,NEmi->mNiJ', t1b, eris_OVoo)
    WoOoO+= lib.einsum('je,neMI->nMjI', t1a, eris_ovOO)
    WoOoO += np.asarray(eris.ooOO).transpose(0,2,1,3)
    eris_OVoo = eris_ovOO = None

    eris_ovOV = np.asarray(eris.ovOV)
    WoOoO += lib.einsum('iJeF,meNF->mNiJ', tauab, eris_ovOV)
    u2ab += lib.einsum('mNaB,mNiJ->iJaB', tauab, WoOoO)
    WoOoO = None

    tilab = make_tau_ab(t2[1], t1 , t1 , fac=0.5)
    Fvva -= lib.einsum('mNaF,meNF->ae', tilab, eris_ovOV)
    Fvvb -= lib.einsum('nMfA,nfME->AE', tilab, eris_ovOV)
    Fooa += lib.einsum('iNeF,meNF->mi', tilab, eris_ovOV)
    Foob += lib.einsum('nIfE,nfME->MI', tilab, eris_ovOV)
    Fova += np.einsum('NF,meNF->me',t1b, eris_ovOV)
    Fovb += np.einsum('nf,nfME->ME',t1a, eris_ovOV)
    u2ab += eris_ovOV.conj().transpose(0,2,1,3)
    wovvo += 0.5*lib.einsum('jNbF,meNF->mbej', t2ab, eris_ovOV)
    wOVVO += 0.5*lib.einsum('nJfB,nfME->MBEJ', t2ab, eris_ovOV)
    wOvVo -= 0.5*lib.einsum('jnfb,nfME->MbEj', t2aa, eris_ovOV)
    woVvO -= 0.5*lib.einsum('JNFB,meNF->mBeJ', t2bb, eris_ovOV)
    woVVo += 0.5*lib.einsum('jNfB,mfNE->mBEj', t2ab, eris_ovOV)
    wOvvO += 0.5*lib.einsum('nJbF,neMF->MbeJ', t2ab, eris_ovOV)
    tmpabab = lib.einsum('JF,meNF->mNeJ', t1b, eris_ovOV)
    tmpbaba = lib.einsum('jf,nfME->MnEj', t1a, eris_ovOV)
    woVvO -= lib.einsum('NB,mNeJ->mBeJ', t1b, tmpabab)
    wOvVo -= lib.einsum('nb,MnEj->MbEj', t1a, tmpbaba)
    woVVo += lib.einsum('NB,NmEj->mBEj', t1b, tmpbaba)
    wOvvO += lib.einsum('nb,nMeJ->MbeJ', t1a, tmpabab)
    tmpabab = tmpbaba = tilab = None

    Fova += fova
    Fovb += fovb
    u1a += fova.conj()
    u1a += np.einsum('ie,ae->ia', t1a, Fvva)
    u1a -= np.einsum('ma,mi->ia', t1a, Fooa)
    u1a -= np.einsum('imea,me->ia', t2aa, Fova)
    u1a += np.einsum('iMaE,ME->ia', t2ab, Fovb)
    u1b += fovb.conj()
    u1b += np.einsum('ie,ae->ia',t1b,Fvvb)
    u1b -= np.einsum('ma,mi->ia',t1b,Foob)
    u1b -= np.einsum('imea,me->ia', t2bb, Fovb)
    u1b += np.einsum('mIeA,me->IA', t2ab, Fova)

    eris_oovv = np.asarray(eris.oovv)
    eris_ovvo = np.asarray(eris.ovvo)
    wovvo -= eris_oovv.transpose(0,2,3,1)
    wovvo += eris_ovvo.transpose(0,2,1,3)
    oovv = eris_oovv - eris_ovvo.transpose(0,3,2,1)
    u1a-= np.einsum('nf,niaf->ia', t1a,      oovv)
    tmp1aa = lib.einsum('ie,mjbe->mbij', t1a,      oovv)
    u2aa += 2*lib.einsum('ma,mbij->ijab', t1a, tmp1aa)
    eris_ovvo = eris_oovv = oovv = tmp1aa = None

    eris_OOVV = np.asarray(eris.OOVV)
    eris_OVVO = np.asarray(eris.OVVO)
    wOVVO -= eris_OOVV.transpose(0,2,3,1)
    wOVVO += eris_OVVO.transpose(0,2,1,3)
    OOVV = eris_OOVV - eris_OVVO.transpose(0,3,2,1)
    u1b-= np.einsum('nf,niaf->ia', t1b,      OOVV)
    tmp1bb = lib.einsum('ie,mjbe->mbij', t1b,      OOVV)
    u2bb += 2*lib.einsum('ma,mbij->ijab', t1b, tmp1bb)
    eris_OVVO = eris_OOVV = OOVV = None

    eris_ooVV = np.asarray(eris.ooVV)
    eris_ovVO = np.asarray(eris.ovVO)
    woVVo -= eris_ooVV.transpose(0,2,3,1)
    woVvO += eris_ovVO.transpose(0,2,1,3)
    u1b+= np.einsum('nf,nfAI->IA', t1a, eris_ovVO)
    tmp1ab = lib.einsum('ie,meBJ->mBiJ', t1a, eris_ovVO)
    tmp1ab+= lib.einsum('IE,mjBE->mBjI', t1b, eris_ooVV)
    u2ab -= lib.einsum('ma,mBiJ->iJaB', t1a, tmp1ab)
    eris_ooVV = eris_ovVO = tmp1ab = None

    eris_OOvv = np.asarray(eris.OOvv)
    eris_OVvo = np.asarray(eris.OVvo)
    wOvvO -= eris_OOvv.transpose(0,2,3,1)
    wOvVo += eris_OVvo.transpose(0,2,1,3)
    u1a+= np.einsum('NF,NFai->ia', t1b, eris_OVvo)
    tmp1ba = lib.einsum('IE,MEbj->MbIj', t1b, eris_OVvo)
    tmp1ba+= lib.einsum('ie,MJbe->MbJi', t1a, eris_OOvv)
    u2ab -= lib.einsum('MA,MbIj->jIbA', t1b, tmp1ba)
    eris_OOvv = eris_OVvo = tmp1ba = None

    u2aa += 2*lib.einsum('imae,mbej->ijab', t2aa, wovvo)
    u2aa += 2*lib.einsum('iMaE,MbEj->ijab', t2ab, wOvVo)
    u2bb += 2*lib.einsum('imae,mbej->ijab', t2bb, wOVVO)
    u2bb += 2*lib.einsum('mIeA,mBeJ->IJAB', t2ab, woVvO)
    u2ab += lib.einsum('imae,mBeJ->iJaB', t2aa, woVvO)
    u2ab += lib.einsum('iMaE,MBEJ->iJaB', t2ab, wOVVO)
    u2ab += lib.einsum('iMeA,MbeJ->iJbA', t2ab, wOvvO)
    u2ab += lib.einsum('IMAE,MbEj->jIbA', t2bb, wOvVo)
    u2ab += lib.einsum('mIeA,mbej->jIbA', t2ab, wovvo)
    u2ab += lib.einsum('mIaE,mBEj->jIaB', t2ab, woVVo)
    wovvo = wOVVO = woVvO = wOvVo = woVVo = wOvvO = None

    Ftmpa = Fvva - .5*lib.einsum('mb,me->be', t1a, Fova)
    Ftmpb = Fvvb - .5*lib.einsum('mb,me->be', t1b, Fovb)
    u2aa += lib.einsum('ijae,be->ijab', t2aa, Ftmpa)
    u2bb += lib.einsum('ijae,be->ijab', t2bb, Ftmpb)
    u2ab += lib.einsum('iJaE,BE->iJaB', t2ab, Ftmpb)
    u2ab += lib.einsum('iJeA,be->iJbA', t2ab, Ftmpa)
    Ftmpa = Fooa + 0.5*lib.einsum('je,me->mj', t1a, Fova)
    Ftmpb = Foob + 0.5*lib.einsum('je,me->mj', t1b, Fovb)
    u2aa -= lib.einsum('imab,mj->ijab', t2aa, Ftmpa)
    u2bb -= lib.einsum('imab,mj->ijab', t2bb, Ftmpb)
    u2ab -= lib.einsum('iMaB,MJ->iJaB', t2ab, Ftmpb)
    u2ab -= lib.einsum('mIaB,mj->jIaB', t2ab, Ftmpa)

    eris_ovoo = np.asarray(eris.ovoo).conj()
    eris_OVOO = np.asarray(eris.OVOO).conj()
    eris_OVoo = np.asarray(eris.OVoo).conj()
    eris_ovOO = np.asarray(eris.ovOO).conj()
    ovoo = eris_ovoo - eris_ovoo.transpose(2,1,0,3)
    OVOO = eris_OVOO - eris_OVOO.transpose(2,1,0,3)
    u2aa -= lib.einsum('ma,jbim->ijab', t1a, ovoo)
    u2bb -= lib.einsum('ma,jbim->ijab', t1b, OVOO)
    u2ab -= lib.einsum('ma,JBim->iJaB', t1a, eris_OVoo)
    u2ab -= lib.einsum('MA,ibJM->iJbA', t1b, eris_ovOO)
    eris_ovoo = eris_OVoo = eris_OVOO = eris_ovOO = None

    u2aa *= .5
    u2bb *= .5
    u2aa = u2aa - u2aa.transpose(0,1,3,2)
    u2aa = u2aa - u2aa.transpose(1,0,2,3)
    u2bb = u2bb - u2bb.transpose(0,1,3,2)
    u2bb = u2bb - u2bb.transpose(1,0,2,3)

    eia_a = lib.direct_sum('i-a->ia', mo_ea_o, mo_ea_v)
    eia_b = lib.direct_sum('i-a->ia', mo_eb_o, mo_eb_v)
    u1a /= eia_a
    u1b /= eia_b

    u2aa /= lib.direct_sum('ia+jb->ijab', eia_a, eia_a)
    u2ab /= lib.direct_sum('ia+jb->ijab', eia_a, eia_b)
    u2bb /= lib.direct_sum('ia+jb->ijab', eia_b, eia_b)

    time0 = log.timer_debug1('update t1 t2', *time0)
    t1new = u1a, u1b
    t2new = u2aa, u2ab, u2bb
    return t1new, t2new


def energy(cc, t1=None, t2=None, eris=None):
    '''UCCSD correlation energy'''
    if t1 is None: t1 = cc.t1
    if t2 is None: t2 = cc.t2
    if eris is None: eris = cc.ao2mo()

    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nocca, noccb, nvira, nvirb = t2ab.shape
    eris_ovov = np.asarray(eris.ovov)
    eris_OVOV = np.asarray(eris.OVOV)
    eris_ovOV = np.asarray(eris.ovOV)
    fova = eris.focka[:nocca,nocca:]
    fovb = eris.fockb[:noccb,noccb:]
    e  = np.einsum('ia,ia', fova, t1a)
    e += np.einsum('ia,ia', fovb, t1b)
    e += 0.25*np.einsum('ijab,iajb',t2aa,eris_ovov)
    e -= 0.25*np.einsum('ijab,ibja',t2aa,eris_ovov)
    e += 0.25*np.einsum('ijab,iajb',t2bb,eris_OVOV)
    e -= 0.25*np.einsum('ijab,ibja',t2bb,eris_OVOV)
    e +=      np.einsum('iJaB,iaJB',t2ab,eris_ovOV)
    e += 0.5*np.einsum('ia,jb,iajb',t1a,t1a,eris_ovov)
    e -= 0.5*np.einsum('ia,jb,ibja',t1a,t1a,eris_ovov)
    e += 0.5*np.einsum('ia,jb,iajb',t1b,t1b,eris_OVOV)
    e -= 0.5*np.einsum('ia,jb,ibja',t1b,t1b,eris_OVOV)
    e +=     np.einsum('ia,jb,iajb',t1a,t1b,eris_ovOV)
    if abs(e.imag) > 1e-4:
        logger.warn(cc, 'Non-zero imaginary part found in UCCSD energy %s', e)
    return e.real


get_nocc = ump2.get_nocc
get_nmo = ump2.get_nmo
get_frozen_mask = ump2.get_frozen_mask


def amplitudes_to_vector(t1, t2, out=None):
    nocca, nvira = t1[0].shape
    noccb, nvirb = t1[1].shape
    sizea = nocca * nvira + nocca*(nocca-1)//2*nvira*(nvira-1)//2
    sizeb = noccb * nvirb + noccb*(noccb-1)//2*nvirb*(nvirb-1)//2
    sizeab = nocca * noccb * nvira * nvirb
    vector = np.ndarray(sizea+sizeb+sizeab, t2[0].dtype, buffer=out)
    ccsd.amplitudes_to_vector_s4(t1[0], t2[0], out=vector[:sizea])
    ccsd.amplitudes_to_vector_s4(t1[1], t2[2], out=vector[sizea:])
    vector[sizea+sizeb:] = t2[1].ravel()
    return vector

def vector_to_amplitudes(vector, nmo, nocc):
    nocca, noccb = nocc
    nmoa, nmob = nmo
    nvira, nvirb = nmoa-nocca, nmob-noccb
    nocc = nocca + noccb
    nvir = nvira + nvirb
    nov = nocc * nvir
    size = nov + nocc*(nocc-1)//2*nvir*(nvir-1)//2
    if vector.size == size:
        #return ccsd.vector_to_amplitudes_s4(vector, nmo, nocc)
        raise RuntimeError('Input vector is GCCSD vecotr')
    else:
        sizea = nocca * nvira + nocca*(nocca-1)//2*nvira*(nvira-1)//2
        sizeb = noccb * nvirb + noccb*(noccb-1)//2*nvirb*(nvirb-1)//2
        sections = np.cumsum([sizea, sizeb])
        veca, vecb, t2ab = np.split(vector, sections)
        t1a, t2aa = ccsd.vector_to_amplitudes_s4(veca, nmoa, nocca)
        t1b, t2bb = ccsd.vector_to_amplitudes_s4(vecb, nmob, noccb)
        t2ab = t2ab.copy().reshape(nocca,noccb,nvira,nvirb)
        return (t1a,t1b), (t2aa,t2ab,t2bb)

def amplitudes_from_rccsd(t1, t2):
    t2aa = t2 - t2.transpose(0,1,3,2)
    return (t1,t1), (t2aa,t2,t2aa)


def _add_vvVV(mycc, t1, t2ab, eris, out=None):
    '''Ht2 = np.einsum('iJcD,acBD->iJaB', t2ab, vvVV)
    without using symmetry in t2ab or Ht2
    '''
    time0 = logger.process_clock(), logger.perf_counter()
    if t2ab.size == 0:
        return np.zeros_like(t2ab)
    if t1 is not None:
        t2ab = make_tau_ab(t2ab, t1, t1)

    log = logger.Logger(mycc.stdout, mycc.verbose)
    nocca, noccb, nvira, nvirb = t2ab.shape

    if mycc.direct:  # AO direct CCSD
        if getattr(eris, 'mo_coeff', None) is not None:
            mo_a, mo_b = eris.mo_coeff
        else:
            moidxa, moidxb = mycc.get_frozen_mask()
            mo_a = mycc.mo_coeff[0][:,moidxa]
            mo_b = mycc.mo_coeff[1][:,moidxb]
        # Note tensor t2ab may be t2bbab from eom_uccsd code.  In that
        # particular case, nocca, noccb do not equal to the actual number of
        # alpha/beta occupied orbitals. orbva and orbvb cannot be indexed as
        # mo_a[:,nocca:] and mo_b[:,noccb:]
        orbva = mo_a[:,-nvira:]
        orbvb = mo_b[:,-nvirb:]
        tau = lib.einsum('ijab,pa->ijpb', t2ab, orbva)
        tau = lib.einsum('ijab,pb->ijap', tau, orbvb)
        time0 = logger.timer_debug1(mycc, 'vvvv-tau mo2ao', *time0)
        buf = eris._contract_vvVV_t2(mycc, tau, mycc.direct, out, log)
        mo = np.asarray(np.hstack((orbva, orbvb)), order='F')
        Ht2 = _ao2mo.nr_e2(buf.reshape(nocca*noccb,-1), mo.conj(),
                           (0,nvira,nvira,nvira+nvirb), 's1', 's1')
        return Ht2.reshape(t2ab.shape)
    else:
        return eris._contract_vvVV_t2(mycc, t2ab, mycc.direct, out, log)

def _add_vvvv(mycc, t1, t2, eris, out=None, with_ovvv=False, t2sym=None):
    time0 = logger.process_clock(), logger.perf_counter()
    log = logger.Logger(mycc.stdout, mycc.verbose)
    if t1 is None:
        t2aa, t2ab, t2bb = t2
    else:
        t2aa, t2ab, t2bb = make_tau(t2, t1, t1)
    nocca, nvira = t2aa.shape[1:3]
    noccb, nvirb = t2bb.shape[1:3]

    if mycc.direct:
        assert(t2sym is None)
        if with_ovvv:
            raise NotImplementedError
        if getattr(eris, 'mo_coeff', None) is not None:
            mo_a, mo_b = eris.mo_coeff
        else:
            moidxa, moidxb = mycc.get_frozen_mask()
            mo_a = mycc.mo_coeff[0][:,moidxa]
            mo_b = mycc.mo_coeff[1][:,moidxb]
        nao = mo_a.shape[0]
        otrila = np.tril_indices(nocca,-1)
        otrilb = np.tril_indices(noccb,-1)
        if nocca > 1:
            tauaa = lib.einsum('xab,pa->xpb', t2aa[otrila], mo_a[:,nocca:])
            tauaa = lib.einsum('xab,pb->xap', tauaa, mo_a[:,nocca:])
        else:
            tauaa = np.zeros((0,nao,nao))
        if noccb > 1:
            taubb = lib.einsum('xab,pa->xpb', t2bb[otrilb], mo_b[:,noccb:])
            taubb = lib.einsum('xab,pb->xap', taubb, mo_b[:,noccb:])
        else:
            taubb = np.zeros((0,nao,nao))
        tauab = lib.einsum('ijab,pa->ijpb', t2ab, mo_a[:,nocca:])
        tauab = lib.einsum('ijab,pb->ijap', tauab, mo_b[:,noccb:])
        tau = np.vstack((tauaa, taubb, tauab.reshape(nocca*noccb,nao,nao)))
        tauaa = taubb = tauab = None
        time0 = log.timer_debug1('vvvv-tau', *time0)

        buf = ccsd._contract_vvvv_t2(mycc, mycc.mol, None, tau, out, log)

        mo = np.asarray(np.hstack((mo_a[:,nocca:], mo_b[:,noccb:])), order='F')
        u2aa = np.zeros_like(t2aa)
        if nocca > 1:
            u2tril = buf[:otrila[0].size]
            u2tril = _ao2mo.nr_e2(u2tril.reshape(-1,nao**2), mo.conj(),
                                  (0,nvira,0,nvira), 's1', 's1')
            u2tril = u2tril.reshape(otrila[0].size,nvira,nvira)
            u2aa[otrila[1],otrila[0]] = u2tril.transpose(0,2,1)
            u2aa[otrila] = u2tril

        u2bb = np.zeros_like(t2bb)
        if noccb > 1:
            u2tril = buf[otrila[0].size:otrila[0].size+otrilb[0].size]
            u2tril = _ao2mo.nr_e2(u2tril.reshape(-1,nao**2), mo.conj(),
                                  (nvira,nvira+nvirb,nvira,nvira+nvirb), 's1', 's1')
            u2tril = u2tril.reshape(otrilb[0].size,nvirb,nvirb)
            u2bb[otrilb[1],otrilb[0]] = u2tril.transpose(0,2,1)
            u2bb[otrilb] = u2tril

        u2ab = _ao2mo.nr_e2(buf[-nocca*noccb:].reshape(nocca*noccb,nao**2), mo,
                            (0,nvira,nvira,nvira+nvirb), 's1', 's1')
        u2ab = u2ab.reshape(t2ab.shape)

    else:
        assert(not with_ovvv)
        if t2sym is None:
            tmp = eris._contract_vvvv_t2(mycc, t2aa[np.tril_indices(nocca)],
                                         mycc.direct, None)
            u2aa = ccsd._unpack_t2_tril(tmp, nocca, nvira, None, 'jiba')
            tmp = eris._contract_VVVV_t2(mycc, t2bb[np.tril_indices(noccb)],
                                         mycc.direct, None)
            u2bb = ccsd._unpack_t2_tril(tmp, noccb, nvirb, None, 'jiba')
            u2ab = eris._contract_vvVV_t2(mycc, t2ab, mycc.direct, None)
        else:
            u2aa = eris._contract_vvvv_t2(mycc, t2aa, mycc.direct, None)
            u2bb = eris._contract_VVVV_t2(mycc, t2bb, mycc.direct, None)
            u2ab = eris._contract_vvVV_t2(mycc, t2ab, mycc.direct, None)

    return u2aa,u2ab,u2bb


class UCCSD(ccsd.CCSD):

    conv_tol = getattr(__config__, 'cc_uccsd_UCCSD_conv_tol', 1e-7)
    conv_tol_normt = getattr(__config__, 'cc_uccsd_UCCSD_conv_tol_normt', 1e-6)

# Attribute frozen can be
# * An integer : The same number of inner-most alpha and beta orbitals are frozen
# * One list : Same alpha and beta orbital indices to be frozen
# * A pair of list : First list is the orbital indices to be frozen for alpha
#       orbitals, second list is for beta orbitals
    def __init__(self, mf, frozen=None, mo_coeff=None, mo_occ=None):
        assert isinstance(mf, scf.uhf.UHF)
        ccsd.CCSD.__init__(self, mf, frozen, mo_coeff, mo_occ)

    get_nocc = get_nocc
    get_nmo = get_nmo
    get_frozen_mask = get_frozen_mask

    def init_amps(self, eris=None):
        time0 = logger.process_clock(), logger.perf_counter()
        if eris is None:
            eris = self.ao2mo(self.mo_coeff)
        nocca, noccb = self.nocc

        fova = eris.focka[:nocca,nocca:]
        fovb = eris.fockb[:noccb,noccb:]
        mo_ea_o = eris.mo_energy[0][:nocca]
        mo_ea_v = eris.mo_energy[0][nocca:]
        mo_eb_o = eris.mo_energy[1][:noccb]
        mo_eb_v = eris.mo_energy[1][noccb:]
        eia_a = lib.direct_sum('i-a->ia', mo_ea_o, mo_ea_v)
        eia_b = lib.direct_sum('i-a->ia', mo_eb_o, mo_eb_v)

        t1a = fova.conj() / eia_a
        t1b = fovb.conj() / eia_b

        eris_ovov = np.asarray(eris.ovov)
        eris_OVOV = np.asarray(eris.OVOV)
        eris_ovOV = np.asarray(eris.ovOV)
        t2aa = eris_ovov.transpose(0,2,1,3) / lib.direct_sum('ia+jb->ijab', eia_a, eia_a)
        t2ab = eris_ovOV.transpose(0,2,1,3) / lib.direct_sum('ia+jb->ijab', eia_a, eia_b)
        t2bb = eris_OVOV.transpose(0,2,1,3) / lib.direct_sum('ia+jb->ijab', eia_b, eia_b)
        t2aa = t2aa - t2aa.transpose(0,1,3,2)
        t2bb = t2bb - t2bb.transpose(0,1,3,2)
        e  =      np.einsum('iJaB,iaJB', t2ab, eris_ovOV)
        e += 0.25*np.einsum('ijab,iajb', t2aa, eris_ovov)
        e -= 0.25*np.einsum('ijab,ibja', t2aa, eris_ovov)
        e += 0.25*np.einsum('ijab,iajb', t2bb, eris_OVOV)
        e -= 0.25*np.einsum('ijab,ibja', t2bb, eris_OVOV)
        self.emp2 = e.real
        logger.info(self, 'Init t2, MP2 energy = %.15g', self.emp2)
        logger.timer(self, 'init mp2', *time0)
        return self.emp2, (t1a,t1b), (t2aa,t2ab,t2bb)

    energy = energy
    update_amps = update_amps
    _add_vvvv = _add_vvvv
    _add_vvVV = _add_vvVV

    def kernel(self, t1=None, t2=None, eris=None, mbpt2=False):
        return self.ccsd(t1, t2, eris, mbpt2)
    def ccsd(self, t1=None, t2=None, eris=None, mbpt2=False):
        '''Ground-state unrestricted (U)CCSD.

        Kwargs:
            mbpt2 : bool
                Use one-shot MBPT2 approximation to CCSD.
        '''
        if mbpt2:
            pt = ump2.UMP2(self._scf, self.frozen, self.mo_coeff, self.mo_occ)
            self.e_corr, self.t2 = pt.kernel(eris=eris)
            t2ab = self.t2[1]
            nocca, noccb, nvira, nvirb = t2ab.shape
            self.t1 = (np.zeros((nocca,nvira)), np.zeros((noccb,nvirb)))
            return self.e_corr, self.t1, self.t2

        return ccsd.CCSD.ccsd(self, t1, t2, eris)

    def solve_lambda(self, t1=None, t2=None, l1=None, l2=None,
                     eris=None):
        from pyscf.cc import uccsd_lambda
        if t1 is None: t1 = self.t1
        if t2 is None: t2 = self.t2
        if eris is None: eris = self.ao2mo(self.mo_coeff)
        self.converged_lambda, self.l1, self.l2 = \
                uccsd_lambda.kernel(self, eris, t1, t2, l1, l2,
                                    max_cycle=self.max_cycle,
                                    tol=self.conv_tol_normt,
                                    verbose=self.verbose)
        return self.l1, self.l2

    def ccsd_t(self, t1=None, t2=None, eris=None):
        from pyscf.cc import uccsd_t
        if t1 is None: t1 = self.t1
        if t2 is None: t2 = self.t2
        if eris is None: eris = self.ao2mo(self.mo_coeff)
        return uccsd_t.kernel(self, eris, t1, t2, self.verbose)
    uccsd_t = ccsd_t

    def make_rdm1(self, t1=None, t2=None, l1=None, l2=None, ao_repr=False,
                  with_frozen=True, with_mf=True):
        '''Un-relaxed 1-particle density matrix in MO space

        Returns:
            dm1a, dm1b
        '''
        from pyscf.cc import uccsd_rdm
        if t1 is None: t1 = self.t1
        if t2 is None: t2 = self.t2
        if l1 is None: l1 = self.l1
        if l2 is None: l2 = self.l2
        if l1 is None: l1, l2 = self.solve_lambda(t1, t2)
        return uccsd_rdm.make_rdm1(self, t1, t2, l1, l2, ao_repr=ao_repr,
                                   with_frozen=with_frozen, with_mf=with_mf)

    def make_rdm2(self, t1=None, t2=None, l1=None, l2=None, ao_repr=False,
                  with_frozen=True, with_dm1=True):
        '''2-particle density matrix in spin-orbital basis.
        '''
        from pyscf.cc import uccsd_rdm
        if t1 is None: t1 = self.t1
        if t2 is None: t2 = self.t2
        if l1 is None: l1 = self.l1
        if l2 is None: l2 = self.l2
        if l1 is None: l1, l2 = self.solve_lambda(t1, t2)
        return uccsd_rdm.make_rdm2(self, t1, t2, l1, l2, ao_repr=ao_repr,
                                   with_frozen=with_frozen, with_dm1=with_dm1)

    def spin_square(self, mo_coeff=None, s=None):
        from pyscf.fci.spin_op import spin_square_general
        if mo_coeff is None:
            mo_coeff = self.mo_coeff
        if s is None:
            s = self._scf.get_ovlp()

        dma,dmb        = self.make_rdm1()
        dmaa,dmab,dmbb = self.make_rdm2()

        return spin_square_general(dma,dmb,dmaa,dmab,dmbb,mo_coeff,s)

    def ao2mo(self, mo_coeff=None):
        nmoa, nmob = self.get_nmo()
        nao = self.mo_coeff[0].shape[0]
        nmo_pair = nmoa * (nmoa+1) // 2
        nao_pair = nao * (nao+1) // 2
        mem_incore = (max(nao_pair**2, nmoa**4) + nmo_pair**2) * 8/1e6
        mem_now = lib.current_memory()[0]
        if (self._scf._eri is not None and
            (mem_incore+mem_now < self.max_memory or self.incore_complete)):
            return _make_eris_incore(self, mo_coeff)

        elif getattr(self._scf, 'with_df', None):
            # TODO: Uncomment once there is an unrestricted DF-CCSD implementation
            #logger.warn(self, 'UCCSD detected DF being used in the HF object. '
            #            'MO integrals are computed based on the DF 3-index tensors.\n'
            #            'It\'s recommended to use dfccsd.CCSD for the '
            #            'DF-CCSD calculations')
            return _make_df_eris_outcore(self, mo_coeff)
        else:
            return _make_eris_outcore(self, mo_coeff)

    def ipccsd(self, nroots=1, left=False, koopmans=False, guess=None,
               partition=None, eris=None):
        from pyscf.cc import eom_uccsd
        return eom_uccsd.EOMIP(self).kernel(nroots, left, koopmans, guess,
                                            partition, eris)

    def eaccsd(self, nroots=1, left=False, koopmans=False, guess=None,
               partition=None, eris=None):
        from pyscf.cc import eom_uccsd
        return eom_uccsd.EOMEA(self).kernel(nroots, left, koopmans, guess,
                                            partition, eris)

    def eeccsd(self, nroots=1, koopmans=False, guess=None, eris=None):
        from pyscf.cc import eom_uccsd
        return eom_uccsd.EOMEE(self).kernel(nroots, koopmans, guess, eris)

    def eomee_ccsd(self, nroots=1, koopmans=False, guess=None, eris=None):
        from pyscf.cc import eom_uccsd
        return eom_uccsd.EOMEESpinKeep(self).kernel(nroots, koopmans, guess, eris)

    def eomsf_ccsd(self, nroots=1, koopmans=False, guess=None, eris=None):
        from pyscf.cc import eom_uccsd
        return eom_uccsd.EOMEESpinFlip(self).kernel(nroots, koopmans, guess, eris)

    def eomip_method(self):
        from pyscf.cc import eom_uccsd
        return eom_uccsd.EOMIP(self)

    def eomea_method(self):
        from pyscf.cc import eom_uccsd
        return eom_uccsd.EOMEA(self)

    def eomee_method(self):
        from pyscf.cc import eom_uccsd
        return eom_uccsd.EOMEE(self)

    def density_fit(self):
        raise NotImplementedError

    def nuc_grad_method(self):
        from pyscf.grad import uccsd
        return uccsd.Gradients(self)

    def amplitudes_to_vector(self, t1, t2, out=None):
        return amplitudes_to_vector(t1, t2, out)

    def vector_to_amplitudes(self, vector, nmo=None, nocc=None):
        if nocc is None: nocc = self.nocc
        if nmo is None: nmo = self.nmo
        return vector_to_amplitudes(vector, nmo, nocc)

    def vector_size(self, nmo=None, nocc=None):
        if nocc is None: nocc = self.nocc
        if nmo is None: nmo = self.nmo
        nocca, noccb = nocc
        nmoa, nmob = nmo
        nvira, nvirb = nmoa-nocca, nmob-noccb
        sizea = nocca * nvira + nocca*(nocca-1)//2*nvira*(nvira-1)//2
        sizeb = noccb * nvirb + noccb*(noccb-1)//2*nvirb*(nvirb-1)//2
        sizeab = nocca * noccb * nvira * nvirb
        return sizea + sizeb + sizeab

    def amplitudes_from_rccsd(self, t1, t2):
        return amplitudes_from_rccsd(t1, t2)

    def get_t1_diagnostic(self, t1=None):
        if t1 is None: t1 = self.t1
        raise NotImplementedError
        #return get_t1_diagnostic(t1)

    def get_d1_diagnostic(self, t1=None):
        if t1 is None: t1 = self.t1
        raise NotImplementedError
        #return get_d1_diagnostic(t1)

    def get_d2_diagnostic(self, t2=None):
        if t2 is None: t2 = self.t2
        raise NotImplementedError
        #return get_d2_diagnostic(t2)

CCSD = UCCSD


class _ChemistsERIs(ccsd._ChemistsERIs):
    def __init__(self, mol=None):
        ccsd._ChemistsERIs.__init__(self, mol)
        self.OOOO = None
        self.OVOO = None
        self.OVOV = None
        self.OOVV = None
        self.OVVO = None
        self.OVVV = None
        self.VVVV = None

        self.ooOO = None
        self.ovOO = None
        self.ovOV = None
        self.ooVV = None
        self.ovVO = None
        self.ovVV = None
        self.vvVV = None

        self.OVoo = None
        self.OOvv = None
        self.OVvo = None

    def _common_init_(self, mycc, mo_coeff=None):
        if mo_coeff is None:
            mo_coeff = mycc.mo_coeff
        mo_idx = mycc.get_frozen_mask()
        self.mo_coeff = mo_coeff = \
                (mo_coeff[0][:,mo_idx[0]], mo_coeff[1][:,mo_idx[1]])
        # Note: Recomputed fock matrix since SCF may not be fully converged.
        dm = mycc._scf.make_rdm1(mycc.mo_coeff, mycc.mo_occ)
        vhf = mycc._scf.get_veff(mycc.mol, dm)
        fockao = mycc._scf.get_fock(vhf=vhf, dm=dm)
        self.focka = reduce(np.dot, (mo_coeff[0].conj().T, fockao[0], mo_coeff[0]))
        self.fockb = reduce(np.dot, (mo_coeff[1].conj().T, fockao[1], mo_coeff[1]))
        self.fock = (self.focka, self.fockb)
        self.e_hf = mycc._scf.energy_tot(dm=dm, vhf=vhf)
        nocca, noccb = self.nocc = mycc.nocc
        self.mol = mycc.mol

        mo_ea = self.focka.diagonal().real
        mo_eb = self.fockb.diagonal().real
        self.mo_energy = (mo_ea, mo_eb)
        gap_a = abs(mo_ea[:nocca,None] - mo_ea[None,nocca:])
        gap_b = abs(mo_eb[:noccb,None] - mo_eb[None,noccb:])
        if gap_a.size > 0:
            gap_a = gap_a.min()
        else:
            gap_a = 1e9
        if gap_b.size > 0:
            gap_b = gap_b.min()
        else:
            gap_b = 1e9
        if gap_a < 1e-5 or gap_b < 1e-5:
            logger.warn(mycc, 'HOMO-LUMO gap (%s,%s) too small for UCCSD',
                        gap_a, gap_b)
        return self

    def get_ovvv(self, *slices):
        return _get_ovvv_base(self.ovvv, *slices)

    def get_ovVV(self, *slices):
        return _get_ovvv_base(self.ovVV, *slices)

    def get_OVvv(self, *slices):
        return _get_ovvv_base(self.OVvv, *slices)

    def get_OVVV(self, *slices):
        return _get_ovvv_base(self.OVVV, *slices)

    def _contract_VVVV_t2(self, mycc, t2, vvvv_or_direct=False, out=None, verbose=None):
        if isinstance(vvvv_or_direct, np.ndarray):
            vvvv = vvvv_or_direct
        elif vvvv_or_direct:
            vvvv = None
        else:
            vvvv = self.VVVV
        return ccsd._contract_vvvv_t2(mycc, self.mol, vvvv, t2, out, verbose)

    def _contract_vvVV_t2(self, mycc, t2, vvvv_or_direct=False, out=None, verbose=None):
        if isinstance(vvvv_or_direct, np.ndarray):
            vvvv = vvvv_or_direct
        elif vvvv_or_direct:
            vvvv = None
        else:
            vvvv = self.vvVV
        return ccsd._contract_vvvv_t2(mycc, self.mol, vvvv, t2, out, verbose)

def _get_ovvv_base(ovvv, *slices):
    if ovvv.ndim == 3:
        ovw = np.asarray(ovvv[slices])
        nocc, nvir, nvir_pair = ovw.shape
        ovvv = lib.unpack_tril(ovw.reshape(nocc*nvir,nvir_pair))
        nvir1 = ovvv.shape[2]
        return ovvv.reshape(nocc,nvir,nvir1,nvir1)
    elif slices:
        return ovvv[slices]
    else:
        return ovvv

def _make_eris_incore(mycc, mo_coeff=None, ao2mofn=None):
    eris = _ChemistsERIs()
    eris._common_init_(mycc, mo_coeff)

    nocca, noccb = mycc.nocc
    nmoa, nmob = mycc.nmo
    nvira, nvirb = nmoa-nocca, nmob-noccb

    moa = eris.mo_coeff[0]
    mob = eris.mo_coeff[1]
    nmoa = moa.shape[1]
    nmob = mob.shape[1]

    if callable(ao2mofn):
        eri_aa = ao2mofn(moa).reshape([nmoa]*4)
        eri_bb = ao2mofn(mob).reshape([nmob]*4)
        eri_ab = ao2mofn((moa,moa,mob,mob))
    else:
        eri_aa = ao2mo.restore(1, ao2mo.full(mycc._scf._eri, moa), nmoa)
        eri_bb = ao2mo.restore(1, ao2mo.full(mycc._scf._eri, mob), nmob)
        eri_ab = ao2mo.general(mycc._scf._eri, (moa,moa,mob,mob), compact=False)
    eri_ba = eri_ab.reshape(nmoa,nmoa,nmob,nmob).transpose(2,3,0,1)

    eri_aa = eri_aa.reshape(nmoa,nmoa,nmoa,nmoa)
    eri_ab = eri_ab.reshape(nmoa,nmoa,nmob,nmob)
    eri_ba = eri_ba.reshape(nmob,nmob,nmoa,nmoa)
    eri_bb = eri_bb.reshape(nmob,nmob,nmob,nmob)
    eris.oooo = eri_aa[:nocca,:nocca,:nocca,:nocca].copy()
    eris.ovoo = eri_aa[:nocca,nocca:,:nocca,:nocca].copy()
    eris.ovov = eri_aa[:nocca,nocca:,:nocca,nocca:].copy()
    eris.oovv = eri_aa[:nocca,:nocca,nocca:,nocca:].copy()
    eris.ovvo = eri_aa[:nocca,nocca:,nocca:,:nocca].copy()
    eris.ovvv = eri_aa[:nocca,nocca:,nocca:,nocca:].copy()
    eris.vvvv = eri_aa[nocca:,nocca:,nocca:,nocca:].copy()

    eris.OOOO = eri_bb[:noccb,:noccb,:noccb,:noccb].copy()
    eris.OVOO = eri_bb[:noccb,noccb:,:noccb,:noccb].copy()
    eris.OVOV = eri_bb[:noccb,noccb:,:noccb,noccb:].copy()
    eris.OOVV = eri_bb[:noccb,:noccb,noccb:,noccb:].copy()
    eris.OVVO = eri_bb[:noccb,noccb:,noccb:,:noccb].copy()
    eris.OVVV = eri_bb[:noccb,noccb:,noccb:,noccb:].copy()
    eris.VVVV = eri_bb[noccb:,noccb:,noccb:,noccb:].copy()

    eris.ooOO = eri_ab[:nocca,:nocca,:noccb,:noccb].copy()
    eris.ovOO = eri_ab[:nocca,nocca:,:noccb,:noccb].copy()
    eris.ovOV = eri_ab[:nocca,nocca:,:noccb,noccb:].copy()
    eris.ooVV = eri_ab[:nocca,:nocca,noccb:,noccb:].copy()
    eris.ovVO = eri_ab[:nocca,nocca:,noccb:,:noccb].copy()
    eris.ovVV = eri_ab[:nocca,nocca:,noccb:,noccb:].copy()
    eris.vvVV = eri_ab[nocca:,nocca:,noccb:,noccb:].copy()

    #eris.OOoo = eri_ba[:noccb,:noccb,:nocca,:nocca].copy()
    eris.OVoo = eri_ba[:noccb,noccb:,:nocca,:nocca].copy()
    #eris.OVov = eri_ba[:noccb,noccb:,:nocca,nocca:].copy()
    eris.OOvv = eri_ba[:noccb,:noccb,nocca:,nocca:].copy()
    eris.OVvo = eri_ba[:noccb,noccb:,nocca:,:nocca].copy()
    eris.OVvv = eri_ba[:noccb,noccb:,nocca:,nocca:].copy()
    #eris.VVvv = eri_ba[noccb:,noccb:,nocca:,nocca:].copy()

    if not callable(ao2mofn):
        ovvv = eris.ovvv.reshape(nocca*nvira,nvira,nvira)
        eris.ovvv = lib.pack_tril(ovvv).reshape(nocca,nvira,nvira*(nvira+1)//2)
        eris.vvvv = ao2mo.restore(4, eris.vvvv, nvira)

        OVVV = eris.OVVV.reshape(noccb*nvirb,nvirb,nvirb)
        eris.OVVV = lib.pack_tril(OVVV).reshape(noccb,nvirb,nvirb*(nvirb+1)//2)
        eris.VVVV = ao2mo.restore(4, eris.VVVV, nvirb)

        ovVV = eris.ovVV.reshape(nocca*nvira,nvirb,nvirb)
        eris.ovVV = lib.pack_tril(ovVV).reshape(nocca,nvira,nvirb*(nvirb+1)//2)
        vvVV = eris.vvVV.reshape(nvira**2,nvirb**2)
        idxa = np.tril_indices(nvira)
        idxb = np.tril_indices(nvirb)
        eris.vvVV = lib.take_2d(vvVV, idxa[0]*nvira+idxa[1], idxb[0]*nvirb+idxb[1])

        OVvv = eris.OVvv.reshape(noccb*nvirb,nvira,nvira)
        eris.OVvv = lib.pack_tril(OVvv).reshape(noccb,nvirb,nvira*(nvira+1)//2)
    return eris

def _make_df_eris_outcore(mycc, mo_coeff=None):
    cput0 = (logger.process_clock(), logger.perf_counter())
    log = logger.Logger(mycc.stdout, mycc.verbose)
    eris = _ChemistsERIs()
    eris._common_init_(mycc, mo_coeff)

    moa, mob = eris.mo_coeff
    nocca, noccb = eris.nocc
    nao = moa.shape[0]
    nmoa = moa.shape[1]
    nmob = mob.shape[1]
    nvira = nmoa - nocca
    nvirb = nmob - noccb
    nvira_pair = nvira*(nvira+1)//2
    nvirb_pair = nvirb*(nvirb+1)//2
    naux = mycc._scf.with_df.get_naoaux()

    # --- Three-center integrals
    # (L|aa)
    Loo = np.empty((naux,nocca,nocca))
    Lov = np.empty((naux,nocca,nvira))
    Lvo = np.empty((naux,nvira,nocca))
    Lvv = np.empty((naux,nvira_pair))
    # (L|bb)
    LOO = np.empty((naux,noccb,noccb))
    LOV = np.empty((naux,noccb,nvirb))
    LVO = np.empty((naux,nvirb,noccb))
    LVV = np.empty((naux,nvirb_pair))
    p1 = 0
    oa, va = np.s_[:nocca], np.s_[nocca:]
    ob, vb = np.s_[:noccb], np.s_[noccb:]
    # Transform three-center integrals to MO basis
    einsum = lib.einsum
    for eri1 in mycc._scf.with_df.loop():
        eri1 = lib.unpack_tril(eri1).reshape(-1,nao,nao)
        # (L|aa)
        Lpq = einsum('Lab,ap,bq->Lpq', eri1, moa, moa)
        p0, p1 = p1, p1 + Lpq.shape[0]
        blk = np.s_[p0:p1]
        Loo[blk] = Lpq[:,oa,oa]
        Lov[blk] = Lpq[:,oa,va]
        Lvo[blk] = Lpq[:,va,oa]
        Lvv[blk] = lib.pack_tril(Lpq[:,va,va].reshape(-1,nvira,nvira))
        # (L|bb)
        Lpq = einsum('Lab,ap,bq->Lpq', eri1, mob, mob)
        LOO[blk] = Lpq[:,ob,ob]
        LOV[blk] = Lpq[:,ob,vb]
        LVO[blk] = Lpq[:,vb,ob]
        LVV[blk] = lib.pack_tril(Lpq[:,vb,vb].reshape(-1,nvirb,nvirb))
    Loo = Loo.reshape(naux,nocca*nocca)
    Lov = Lov.reshape(naux,nocca*nvira)
    Lvo = Lvo.reshape(naux,nocca*nvira)
    LOO = LOO.reshape(naux,noccb*noccb)
    LOV = LOV.reshape(naux,noccb*nvirb)
    LVO = LVO.reshape(naux,noccb*nvirb)

    # --- Four-center integrals
    dot = lib.ddot
    eris.feri1 = lib.H5TmpFile()
    # (aa|aa)
    eris.oooo = eris.feri1.create_dataset('oooo', (nocca,nocca,nocca,nocca), 'f8')
    eris.oovv = eris.feri1.create_dataset('oovv', (nocca,nocca,nvira,nvira), 'f8', chunks=(nocca,nocca,1,nvira))
    eris.ovoo = eris.feri1.create_dataset('ovoo', (nocca,nvira,nocca,nocca), 'f8', chunks=(nocca,1,nocca,nocca))
    eris.ovvo = eris.feri1.create_dataset('ovvo', (nocca,nvira,nvira,nocca), 'f8', chunks=(nocca,1,nvira,nocca))
    eris.ovov = eris.feri1.create_dataset('ovov', (nocca,nvira,nocca,nvira), 'f8', chunks=(nocca,1,nocca,nvira))
    eris.ovvv = eris.feri1.create_dataset('ovvv', (nocca,nvira,nvira_pair), 'f8')
    eris.vvvv = eris.feri1.create_dataset('vvvv', (nvira_pair,nvira_pair), 'f8')
    eris.oooo[:] = dot(Loo.T, Loo).reshape(nocca,nocca,nocca,nocca)
    eris.ovoo[:] = dot(Lov.T, Loo).reshape(nocca,nvira,nocca,nocca)
    eris.oovv[:] = lib.unpack_tril(dot(Loo.T, Lvv)).reshape(nocca,nocca,nvira,nvira)
    eris.ovvo[:] = dot(Lov.T, Lvo).reshape(nocca,nvira,nvira,nocca)
    eris.ovov[:] = dot(Lov.T, Lov).reshape(nocca,nvira,nocca,nvira)
    eris.ovvv[:] = dot(Lov.T, Lvv).reshape(nocca,nvira,nvira_pair)
    eris.vvvv[:] = dot(Lvv.T, Lvv)
    # (bb|bb)
    eris.OOOO = eris.feri1.create_dataset('OOOO', (noccb,noccb,noccb,noccb), 'f8')
    eris.OOVV = eris.feri1.create_dataset('OOVV', (noccb,noccb,nvirb,nvirb), 'f8', chunks=(noccb,noccb,1,nvirb))
    eris.OVOO = eris.feri1.create_dataset('OVOO', (noccb,nvirb,noccb,noccb), 'f8', chunks=(noccb,1,noccb,noccb))
    eris.OVVO = eris.feri1.create_dataset('OVVO', (noccb,nvirb,nvirb,noccb), 'f8', chunks=(noccb,1,nvirb,noccb))
    eris.OVOV = eris.feri1.create_dataset('OVOV', (noccb,nvirb,noccb,nvirb), 'f8', chunks=(noccb,1,noccb,nvirb))
    eris.OVVV = eris.feri1.create_dataset('OVVV', (noccb,nvirb,nvirb_pair), 'f8')
    eris.VVVV = eris.feri1.create_dataset('VVVV', (nvirb_pair,nvirb_pair), 'f8')
    eris.OOOO[:] = dot(LOO.T, LOO).reshape(noccb,noccb,noccb,noccb)
    eris.OVOO[:] = dot(LOV.T, LOO).reshape(noccb,nvirb,noccb,noccb)
    eris.OOVV[:] = lib.unpack_tril(dot(LOO.T, LVV)).reshape(noccb,noccb,nvirb,nvirb)
    eris.OVVO[:] = dot(LOV.T, LVO).reshape(noccb,nvirb,nvirb,noccb)
    eris.OVOV[:] = dot(LOV.T, LOV).reshape(noccb,nvirb,noccb,nvirb)
    eris.OVVV[:] = dot(LOV.T, LVV).reshape(noccb,nvirb,nvirb_pair)
    eris.VVVV[:] = dot(LVV.T, LVV)
    # (aa|bb)
    eris.ooOO = eris.feri1.create_dataset('ooOO', (nocca,nocca,noccb,noccb), 'f8')
    eris.ooVV = eris.feri1.create_dataset('ooVV', (nocca,nocca,nvirb,nvirb), 'f8', chunks=(nocca,nocca,1,nvirb))
    eris.ovOO = eris.feri1.create_dataset('ovOO', (nocca,nvira,noccb,noccb), 'f8', chunks=(nocca,1,noccb,noccb))
    eris.ovVO = eris.feri1.create_dataset('ovVO', (nocca,nvira,nvirb,noccb), 'f8', chunks=(nocca,1,nvirb,noccb))
    eris.ovOV = eris.feri1.create_dataset('ovOV', (nocca,nvira,noccb,nvirb), 'f8', chunks=(nocca,1,noccb,nvirb))
    eris.ovVV = eris.feri1.create_dataset('ovVV', (nocca,nvira,nvirb_pair), 'f8')
    eris.vvVV = eris.feri1.create_dataset('vvVV', (nvira_pair,nvirb_pair), 'f8')
    eris.ooOO[:] = dot(Loo.T, LOO).reshape(nocca,nocca,noccb,noccb)
    eris.ovOO[:] = dot(Lov.T, LOO).reshape(nocca,nvira,noccb,noccb)
    eris.ooVV[:] = lib.unpack_tril(dot(Loo.T, LVV)).reshape(nocca,nocca,nvirb,nvirb)
    eris.ovVO[:] = dot(Lov.T, LVO).reshape(nocca,nvira,nvirb,noccb)
    eris.ovOV[:] = dot(Lov.T, LOV).reshape(nocca,nvira,noccb,nvirb)
    eris.ovVV[:] = dot(Lov.T, LVV).reshape(nocca,nvira,nvirb_pair)
    eris.vvVV[:] = dot(Lvv.T, LVV)
    # (bb|aa)
    eris.OOvv = eris.feri1.create_dataset('OOvv', (noccb,noccb,nvira,nvira), 'f8', chunks=(noccb,noccb,1,nvira))
    eris.OVoo = eris.feri1.create_dataset('OVoo', (noccb,nvirb,nocca,nocca), 'f8', chunks=(noccb,1,nocca,nocca))
    eris.OVvo = eris.feri1.create_dataset('OVvo', (noccb,nvirb,nvira,nocca), 'f8', chunks=(noccb,1,nvira,nocca))
    eris.OVvv = eris.feri1.create_dataset('OVvv', (noccb,nvirb,nvira_pair), 'f8')
    eris.OVoo[:] = dot(LOV.T, Loo).reshape(noccb,nvirb,nocca,nocca)
    eris.OOvv[:] = lib.unpack_tril(dot(LOO.T, Lvv)).reshape(noccb,noccb,nvira,nvira)
    eris.OVvo[:] = dot(LOV.T, Lvo).reshape(noccb,nvirb,nvira,nocca)
    eris.OVvv[:] = dot(LOV.T, Lvv).reshape(noccb,nvirb,nvira_pair)

    log.timer('CCSD integral transformation', *cput0)
    return eris

def _make_eris_outcore(mycc, mo_coeff=None):
    eris = _ChemistsERIs()
    eris._common_init_(mycc, mo_coeff)

    nocca, noccb = mycc.nocc
    nmoa, nmob = mycc.nmo
    nvira, nvirb = nmoa-nocca, nmob-noccb

    moa = eris.mo_coeff[0]
    mob = eris.mo_coeff[1]
    nmoa = moa.shape[1]
    nmob = mob.shape[1]

    orboa = moa[:,:nocca]
    orbob = mob[:,:noccb]
    orbva = moa[:,nocca:]
    orbvb = mob[:,noccb:]
    eris.feri = lib.H5TmpFile()
    eris.oooo = eris.feri.create_dataset('oooo', (nocca,nocca,nocca,nocca), 'f8')
    eris.ovoo = eris.feri.create_dataset('ovoo', (nocca,nvira,nocca,nocca), 'f8')
    eris.ovov = eris.feri.create_dataset('ovov', (nocca,nvira,nocca,nvira), 'f8')
    eris.oovv = eris.feri.create_dataset('oovv', (nocca,nocca,nvira,nvira), 'f8')
    eris.ovvo = eris.feri.create_dataset('ovvo', (nocca,nvira,nvira,nocca), 'f8')
    eris.ovvv = eris.feri.create_dataset('ovvv', (nocca,nvira,nvira*(nvira+1)//2), 'f8')
    #eris.vvvv = eris.feri.create_dataset('vvvv', (nvira,nvira,nvira,nvira), 'f8')
    eris.OOOO = eris.feri.create_dataset('OOOO', (noccb,noccb,noccb,noccb), 'f8')
    eris.OVOO = eris.feri.create_dataset('OVOO', (noccb,nvirb,noccb,noccb), 'f8')
    eris.OVOV = eris.feri.create_dataset('OVOV', (noccb,nvirb,noccb,nvirb), 'f8')
    eris.OOVV = eris.feri.create_dataset('OOVV', (noccb,noccb,nvirb,nvirb), 'f8')
    eris.OVVO = eris.feri.create_dataset('OVVO', (noccb,nvirb,nvirb,noccb), 'f8')
    eris.OVVV = eris.feri.create_dataset('OVVV', (noccb,nvirb,nvirb*(nvirb+1)//2), 'f8')
    #eris.VVVV = eris.feri.create_dataset('VVVV', (nvirb,nvirb,nvirb,nvirb), 'f8')
    eris.ooOO = eris.feri.create_dataset('ooOO', (nocca,nocca,noccb,noccb), 'f8')
    eris.ovOO = eris.feri.create_dataset('ovOO', (nocca,nvira,noccb,noccb), 'f8')
    eris.ovOV = eris.feri.create_dataset('ovOV', (nocca,nvira,noccb,nvirb), 'f8')
    eris.ooVV = eris.feri.create_dataset('ooVV', (nocca,nocca,nvirb,nvirb), 'f8')
    eris.ovVO = eris.feri.create_dataset('ovVO', (nocca,nvira,nvirb,noccb), 'f8')
    eris.ovVV = eris.feri.create_dataset('ovVV', (nocca,nvira,nvirb*(nvirb+1)//2), 'f8')
    #eris.vvVV = eris.feri.create_dataset('vvVV', (nvira,nvira,nvirb,nvirb), 'f8')
    eris.OVoo = eris.feri.create_dataset('OVoo', (noccb,nvirb,nocca,nocca), 'f8')
    eris.OOvv = eris.feri.create_dataset('OOvv', (noccb,noccb,nvira,nvira), 'f8')
    eris.OVvo = eris.feri.create_dataset('OVvo', (noccb,nvirb,nvira,nocca), 'f8')
    eris.OVvv = eris.feri.create_dataset('OVvv', (noccb,nvirb,nvira*(nvira+1)//2), 'f8')

    cput1 = logger.process_clock(), logger.perf_counter()
    mol = mycc.mol
    # <ij||pq> = <ij|pq> - <ij|qp> = (ip|jq) - (iq|jp)
    tmpf = lib.H5TmpFile()
    if nocca > 0:
        ao2mo.general(mol, (orboa,moa,moa,moa), tmpf, 'aa')
        buf = np.empty((nmoa,nmoa,nmoa))
        for i in range(nocca):
            lib.unpack_tril(tmpf['aa'][i*nmoa:(i+1)*nmoa], out=buf)
            eris.oooo[i] = buf[:nocca,:nocca,:nocca]
            eris.ovoo[i] = buf[nocca:,:nocca,:nocca]
            eris.ovov[i] = buf[nocca:,:nocca,nocca:]
            eris.oovv[i] = buf[:nocca,nocca:,nocca:]
            eris.ovvo[i] = buf[nocca:,nocca:,:nocca]
            eris.ovvv[i] = lib.pack_tril(buf[nocca:,nocca:,nocca:])
        del(tmpf['aa'])

    if noccb > 0:
        buf = np.empty((nmob,nmob,nmob))
        ao2mo.general(mol, (orbob,mob,mob,mob), tmpf, 'bb')
        for i in range(noccb):
            lib.unpack_tril(tmpf['bb'][i*nmob:(i+1)*nmob], out=buf)
            eris.OOOO[i] = buf[:noccb,:noccb,:noccb]
            eris.OVOO[i] = buf[noccb:,:noccb,:noccb]
            eris.OVOV[i] = buf[noccb:,:noccb,noccb:]
            eris.OOVV[i] = buf[:noccb,noccb:,noccb:]
            eris.OVVO[i] = buf[noccb:,noccb:,:noccb]
            eris.OVVV[i] = lib.pack_tril(buf[noccb:,noccb:,noccb:])
        del(tmpf['bb'])

    if nocca > 0:
        buf = np.empty((nmoa,nmob,nmob))
        ao2mo.general(mol, (orboa,moa,mob,mob), tmpf, 'ab')
        for i in range(nocca):
            lib.unpack_tril(tmpf['ab'][i*nmoa:(i+1)*nmoa], out=buf)
            eris.ooOO[i] = buf[:nocca,:noccb,:noccb]
            eris.ovOO[i] = buf[nocca:,:noccb,:noccb]
            eris.ovOV[i] = buf[nocca:,:noccb,noccb:]
            eris.ooVV[i] = buf[:nocca,noccb:,noccb:]
            eris.ovVO[i] = buf[nocca:,noccb:,:noccb]
            eris.ovVV[i] = lib.pack_tril(buf[nocca:,noccb:,noccb:])
        del(tmpf['ab'])

    if noccb > 0:
        buf = np.empty((nmob,nmoa,nmoa))
        ao2mo.general(mol, (orbob,mob,moa,moa), tmpf, 'ba')
        for i in range(noccb):
            lib.unpack_tril(tmpf['ba'][i*nmob:(i+1)*nmob], out=buf)
            eris.OVoo[i] = buf[noccb:,:nocca,:nocca]
            eris.OOvv[i] = buf[:noccb,nocca:,nocca:]
            eris.OVvo[i] = buf[noccb:,nocca:,:nocca]
            eris.OVvv[i] = lib.pack_tril(buf[noccb:,nocca:,nocca:])
        del(tmpf['ba'])
    buf = None
    cput1 = logger.timer_debug1(mycc, 'transforming oopq, ovpq', *cput1)

    if not mycc.direct:
        ao2mo.full(mol, orbva, eris.feri, dataname='vvvv')
        ao2mo.full(mol, orbvb, eris.feri, dataname='VVVV')
        ao2mo.general(mol, (orbva,orbva,orbvb,orbvb), eris.feri, dataname='vvVV')
        eris.vvvv = eris.feri['vvvv']
        eris.VVVV = eris.feri['VVVV']
        eris.vvVV = eris.feri['vvVV']
        cput1 = logger.timer_debug1(mycc, 'transforming vvvv', *cput1)

    return eris


def make_tau(t2, t1, r1, fac=1, out=None):
    t1a, t1b = t1
    r1a, r1b = r1
    tau1aa = make_tau_aa(t2[0], t1a, r1a, fac, out)
    tau1bb = make_tau_aa(t2[2], t1b, r1b, fac, out)
    tau1ab = make_tau_ab(t2[1], t1, r1, fac, out)
    return tau1aa, tau1ab, tau1bb

def make_tau_aa(t2aa, t1a, r1a, fac=1, out=None):
    tau1aa = np.einsum('ia,jb->ijab', t1a, r1a)
    tau1aa-= np.einsum('ia,jb->jiab', t1a, r1a)
    tau1aa = tau1aa - tau1aa.transpose(0,1,3,2)
    tau1aa *= fac * .5
    tau1aa += t2aa
    return tau1aa

def make_tau_ab(t2ab, t1, r1, fac=1, out=None):
    t1a, t1b = t1
    r1a, r1b = r1
    tau1ab = np.einsum('ia,jb->ijab', t1a, r1b)
    tau1ab+= np.einsum('ia,jb->ijab', r1a, t1b)
    tau1ab *= fac * .5
    tau1ab += t2ab
    return tau1ab

def _flops(nocc, nmo):
    nocca, noccb = nocc
    nmoa, nmob = nmo
    nvira, nvirb = nmoa - nocca, nmob - noccb

    # vvvv
    fp  = nocca**2 * nvira**4
    fp += noccb**2 * nvirb**4
    fp += nocca    * nvira**2 * noccb    * nvirb**2 * 2

    # ovvv
    fp += nocca**2 * nvira**3 * 2
    fp += nocca**2 * nvira**3 * 2
    fp += nocca**2 * nvira**3 * 2
    fp += nocca**3 * nvira**3 * 2
    fp += nocca**3 * nvira**2 * 2

    # OVVV
    fp += noccb**2 * nvirb**3 * 2
    fp += noccb**2 * nvirb**3 * 2
    fp += noccb**2 * nvirb**3 * 2
    fp += noccb**3 * nvirb**3 * 2
    fp += noccb**3 * nvirb**2 * 2

    # ovVV
    fp += nocca    * nvira * noccb    * nvirb**2 * 2
    fp += nocca**2 * nvira            * nvirb**2 * 2
    fp += nocca    * nvira * noccb    * nvirb**2 * 2
    fp += nocca    * nvira * noccb    * nvirb**2 * 2
    fp += nocca**2 * nvira * noccb    * nvirb**2 * 2
    fp += nocca**2 * nvira * noccb    * nvirb    * 2

    # OVvv
    fp += nocca    * nvira**2 * noccb    * nvirb * 2
    fp +=            nvira**2 * noccb**2 * nvirb * 2
    fp += nocca    * nvira**2 * noccb    * nvirb * 2
    fp += nocca    * nvira**2 * noccb    * nvirb * 2
    fp += nocca    * nvira**2 * noccb**2 * nvirb * 2
    fp += nocca    * nvira    * noccb**2 * nvirb * 2

    fp += nocca**4 * nvira    * 2
    fp += nocca**4 * nvira**2 * 2
    fp += nocca**4 * nvira**2 * 2
    fp += nocca**3 * nvira**2 * 2
    fp += nocca**3 * nvira**2 * 2
    fp += nocca**2 * nvira**3 * 2
    fp += nocca**3 * nvira**2 * 2
    fp += nocca**3 * nvira**3 * 2
    fp += nocca**2 * nvira**2 * noccb    * nvirb    * 2
    fp += nocca**3 * nvira**2 * 2
    fp += nocca**3 * nvira**2 * 2

    fp += noccb**4 * nvirb    * 2
    fp += noccb**4 * nvirb**2 * 2
    fp += noccb**4 * nvirb**2 * 2
    fp += noccb**3 * nvirb**2 * 2
    fp += noccb**3 * nvirb**2 * 2
    fp += noccb**2 * nvirb**3 * 2
    fp += noccb**3 * nvirb**2 * 2
    fp += noccb**3 * nvirb**3 * 2
    fp += noccb**2 * nvirb**2 * nocca    * nvira    * 2
    fp += noccb**3 * nvirb**2 * 2
    fp += noccb**3 * nvirb**2 * 2

    fp += nocca**2 * nvira    * noccb    * nvirb    * 2
    fp += nocca**2 * nvira    * noccb    * nvirb    * 2
    fp += nocca**2            * noccb    * nvirb**2 * 2
    fp += noccb**2 * nvirb    * nocca    * nvira    * 2
    fp += noccb**2 * nvirb    * nocca    * nvira    * 2
    fp += noccb**2            * nocca    * nvira**2 * 2

    fp += nocca**2            * noccb**2 * nvirb    * 2
    fp += nocca**2 * nvira    * noccb**2            * 2
    fp += nocca**2 * nvira    * noccb**2 * nvirb    * 2
    fp += nocca**2 * nvira    * noccb**2 * nvirb    * 2

    fp += nocca    * nvira**2 * noccb    * nvirb    * 2
    fp += nocca    * nvira    * noccb    * nvirb**2 * 2
    fp += nocca    * nvira**2 * noccb    * nvirb    * 2
    fp += nocca    * nvira    * noccb    * nvirb**2 * 2

    fp += nocca**2 * nvira**2 * noccb    * nvirb    * 2
    fp += nocca    * nvira    * noccb**2 * nvirb**2 * 2
    fp += nocca**2 * nvira**2 * noccb    * nvirb    * 2
    fp += nocca    * nvira    * noccb**2 * nvirb**2 * 2
    fp += nocca**2 * nvira    * noccb    * nvirb**2 * 2
    fp += nocca    * nvira**2 * noccb**2 * nvirb    * 2

    fp += nocca    * nvira    * noccb**2 * nvirb    * 2
    fp += nocca**2 * nvira    * noccb    * nvirb    * 2
    fp += nocca    * nvira    * noccb**2 * nvirb    * 2
    fp += nocca**2 * nvira    * noccb    * nvirb    * 2
    fp += nocca**2            * noccb    * nvirb**2 * 2
    fp += nocca    * nvira**2 * noccb**2            * 2

    fp += nocca**3 * nvira**2 * 2
    fp += nocca**3 * nvira**2 * 2
    fp += noccb**3 * nvirb**2 * 2
    fp += noccb**3 * nvirb**2 * 2

    fp += nocca**2 * nvira    * noccb    * nvirb    * 2
    fp += nocca**2            * noccb    * nvirb**2 * 2
    fp += nocca**2 * nvira    * noccb    * nvirb    * 2
    fp += noccb**2 * nvirb    * nocca    * nvira    * 2
    fp += noccb**2            * nocca    * nvira**2 * 2
    fp += noccb**2 * nvirb    * nocca    * nvira    * 2

    fp += nocca**3 * nvira**3 * 2
    fp += nocca**2 * nvira**2 * noccb    * nvirb    * 2
    fp += noccb**3 * nvirb**3 * 2
    fp += noccb**2 * nvirb**2 * nocca    * nvira    * 2

    fp += nocca**2 * nvira**2 * noccb    * nvirb    * 2
    fp += nocca    * nvira    * noccb**2 * nvirb**2 * 2
    fp += nocca    * nvira**2 * noccb**2 * nvirb    * 2
    fp += nocca**2 * nvira**2 * noccb    * nvirb    * 2
    fp += nocca**2 * nvira**2 * noccb    * nvirb    * 2
    fp += nocca**2 * nvira    * noccb    * nvirb**2 * 2

    fp += nocca**2 * nvira**3 * 2
    fp += noccb**2 * nvirb**3 * 2
    fp += nocca    * nvira    * noccb    * nvirb**2 * 2
    fp += nocca    * nvira**2 * noccb    * nvirb    * 2
    fp += nocca**3 * nvira**2 * 2
    fp += noccb**3 * nvirb**2 * 2
    fp += nocca    * nvira    * noccb**2 * nvirb    * 2
    fp += nocca**2 * nvira    * noccb    * nvirb    * 2

    fp += nocca**2 * nvira**3 * 2
    fp += noccb**2 * nvirb**3 * 2
    fp += nocca**2 * nvira    * noccb    * nvirb    * 2
    fp += nocca    * nvira    * noccb**2 * nvirb    * 2
    return fp


if __name__ == '__main__':
    from pyscf import gto

    mol = gto.Mole()
    mol.atom = [['O', (0.,   0., 0.)],
                ['O', (1.21, 0., 0.)]]
    mol.basis = 'cc-pvdz'
    mol.spin = 2
    mol.build()
    mf = scf.UHF(mol).run()
    # Freeze 1s electrons
    # also acceptable
    #frozen = 4 or [2,2]
    frozen = [[0,1], [0,1]]
    ucc = UCCSD(mf, frozen=frozen)
    eris = ucc.ao2mo()
    ecc, t1, t2 = ucc.kernel(eris=eris)
    print(ecc - -0.3486987472235819)

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = 'cc-pvdz'
    mol.spin = 0
    mol.build()
    mf = scf.UHF(mol).run()

    mycc = UCCSD(mf)
    mycc.direct = True
    ecc, t1, t2 = mycc.kernel()
    print(ecc - -0.2133432712431435)
    print(mycc.ccsd_t() - -0.003060021865720902)

    e,v = mycc.ipccsd(nroots=8)
    print(e[0] - 0.4335604332073799)
    print(e[2] - 0.5187659896045407)
    print(e[4] - 0.6782876002229172)

    e,v = mycc.eaccsd(nroots=8)
    print(e[0] - 0.16737886338859731)
    print(e[2] - 0.24027613852009164)
    print(e[4] - 0.51006797826488071)

    e,v = mycc.eeccsd(nroots=4)
    print(e[0] - 0.2757159395886167)
    print(e[1] - 0.2757159395886167)
    print(e[2] - 0.2757159395886167)
    print(e[3] - 0.3005716731825082)


# pyscf/cc/uccsd_lambda.py

#!/usr/bin/env python
# Copyright 2014-2020 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: Qiming Sun <osirpt.sun@gmail.com>
#


import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf.cc import uccsd
from pyscf.cc import ccsd_lambda

einsum = lib.einsum

def kernel(mycc, eris=None, t1=None, t2=None, l1=None, l2=None,
           max_cycle=50, tol=1e-8, verbose=logger.INFO):
    if eris is None: eris = mycc.ao2mo()
    return ccsd_lambda.kernel(mycc, eris, t1, t2, l1, l2, max_cycle, tol,
                              verbose, make_intermediates, update_lambda)

# l2, t2 as ijab
def make_intermediates(mycc, t1, t2, eris):
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nocca, nvira = t1a.shape
    noccb, nvirb = t1b.shape

    fooa = eris.focka[:nocca,:nocca]
    fova = eris.focka[:nocca,nocca:]
    fvoa = eris.focka[nocca:,:nocca]
    fvva = eris.focka[nocca:,nocca:]
    foob = eris.fockb[:noccb,:noccb]
    fovb = eris.fockb[:noccb,noccb:]
    fvob = eris.fockb[noccb:,:noccb]
    fvvb = eris.fockb[noccb:,noccb:]

    tauaa, tauab, taubb = uccsd.make_tau(t2, t1, t1)

    ovov = numpy.asarray(eris.ovov)
    ovov = ovov - ovov.transpose(0,3,2,1)
    OVOV = numpy.asarray(eris.OVOV)
    OVOV = OVOV - OVOV.transpose(0,3,2,1)
    ovOV = numpy.asarray(eris.ovOV)

    v1a  = fvva - einsum('ja,jb->ba', fova, t1a)
    v1b  = fvvb - einsum('ja,jb->ba', fovb, t1b)
    v1a += einsum('jcka,jkbc->ba', ovov, tauaa) * .5
    v1a -= einsum('jaKC,jKbC->ba', ovOV, tauab) * .5
    v1a -= einsum('kaJC,kJbC->ba', ovOV, tauab) * .5
    v1b += einsum('jcka,jkbc->ba', OVOV, taubb) * .5
    v1b -= einsum('kcJA,kJcB->BA', ovOV, tauab) * .5
    v1b -= einsum('jcKA,jKcB->BA', ovOV, tauab) * .5

    v2a  = fooa + einsum('ib,jb->ij', fova, t1a)
    v2b  = foob + einsum('ib,jb->ij', fovb, t1b)
    v2a += einsum('ibkc,jkbc->ij', ovov, tauaa) * .5
    v2a += einsum('ibKC,jKbC->ij', ovOV, tauab)
    v2b += einsum('ibkc,jkbc->ij', OVOV, taubb) * .5
    v2b += einsum('kcIB,kJcB->IJ', ovOV, tauab)

    ovoo = numpy.asarray(eris.ovoo)
    ovoo = ovoo - ovoo.transpose(2,1,0,3)
    OVOO = numpy.asarray(eris.OVOO)
    OVOO = OVOO - OVOO.transpose(2,1,0,3)
    OVoo = numpy.asarray(eris.OVoo)
    ovOO = numpy.asarray(eris.ovOO)
    v2a -= numpy.einsum('ibkj,kb->ij', ovoo, t1a)
    v2a += numpy.einsum('KBij,KB->ij', OVoo, t1b)
    v2b -= numpy.einsum('ibkj,kb->ij', OVOO, t1b)
    v2b += numpy.einsum('kbIJ,kb->IJ', ovOO, t1a)

    v5a  = fvoa + numpy.einsum('kc,jkbc->bj', fova, t2aa)
    v5a += numpy.einsum('KC,jKbC->bj', fovb, t2ab)
    v5b  = fvob + numpy.einsum('kc,jkbc->bj', fovb, t2bb)
    v5b += numpy.einsum('kc,kJcB->BJ', fova, t2ab)
    tmp  = fova - numpy.einsum('kdlc,ld->kc', ovov, t1a)
    tmp += numpy.einsum('kcLD,LD->kc', ovOV, t1b)
    v5a += einsum('kc,kb,jc->bj', tmp, t1a, t1a)
    tmp  = fovb - numpy.einsum('kdlc,ld->kc', OVOV, t1b)
    tmp += numpy.einsum('ldKC,ld->KC', ovOV, t1a)
    v5b += einsum('kc,kb,jc->bj', tmp, t1b, t1b)
    v5a -= einsum('lckj,klbc->bj', ovoo, t2aa) * .5
    v5a -= einsum('LCkj,kLbC->bj', OVoo, t2ab)
    v5b -= einsum('LCKJ,KLBC->BJ', OVOO, t2bb) * .5
    v5b -= einsum('lcKJ,lKcB->BJ', ovOO, t2ab)

    oooo = numpy.asarray(eris.oooo)
    OOOO = numpy.asarray(eris.OOOO)
    ooOO = numpy.asarray(eris.ooOO)
    woooo  = einsum('icjl,kc->ikjl', ovoo, t1a)
    wOOOO  = einsum('icjl,kc->ikjl', OVOO, t1b)
    wooOO  = einsum('icJL,kc->ikJL', ovOO, t1a)
    wooOO += einsum('JCil,KC->ilJK', OVoo, t1b)
    woooo += (oooo - oooo.transpose(0,3,2,1)) * .5
    wOOOO += (OOOO - OOOO.transpose(0,3,2,1)) * .5
    wooOO += ooOO.copy()
    woooo += einsum('icjd,klcd->ikjl', ovov, tauaa) * .25
    wOOOO += einsum('icjd,klcd->ikjl', OVOV, taubb) * .25
    wooOO += einsum('icJD,kLcD->ikJL', ovOV, tauab)

    v4ovvo  = einsum('jbld,klcd->jbck', ovov, t2aa)
    v4ovvo += einsum('jbLD,kLcD->jbck', ovOV, t2ab)
    v4ovvo += numpy.asarray(eris.ovvo)
    v4ovvo -= numpy.asarray(eris.oovv).transpose(0,3,2,1)
    v4OVVO  = einsum('jbld,klcd->jbck', OVOV, t2bb)
    v4OVVO += einsum('ldJB,lKdC->JBCK', ovOV, t2ab)
    v4OVVO += numpy.asarray(eris.OVVO)
    v4OVVO -= numpy.asarray(eris.OOVV).transpose(0,3,2,1)
    v4OVvo  = einsum('ldJB,klcd->JBck', ovOV, t2aa)
    v4OVvo += einsum('JBLD,kLcD->JBck', OVOV, t2ab)
    v4OVvo += numpy.asarray(eris.OVvo)
    v4ovVO  = einsum('jbLD,KLCD->jbCK', ovOV, t2bb)
    v4ovVO += einsum('jbld,lKdC->jbCK', ovov, t2ab)
    v4ovVO += numpy.asarray(eris.ovVO)
    v4oVVo  = einsum('jdLB,kLdC->jBCk', ovOV, t2ab)
    v4oVVo -= numpy.asarray(eris.ooVV).transpose(0,3,2,1)
    v4OvvO  = einsum('lbJD,lKcD->JbcK', ovOV, t2ab)
    v4OvvO -= numpy.asarray(eris.OOvv).transpose(0,3,2,1)

    woovo  = einsum('ibck,jb->ijck', v4ovvo, t1a)
    wOOVO  = einsum('ibck,jb->ijck', v4OVVO, t1b)
    wOOvo  = einsum('IBck,JB->IJck', v4OVvo, t1b)
    wOOvo -= einsum('IbcK,jb->IKcj', v4OvvO, t1a)
    wooVO  = einsum('ibCK,jb->ijCK', v4ovVO, t1a)
    wooVO -= einsum('iBCk,JB->ikCJ', v4oVVo, t1b)
    woovo += ovoo.conj().transpose(3,2,1,0) * .5
    wOOVO += OVOO.conj().transpose(3,2,1,0) * .5
    wooVO += OVoo.conj().transpose(3,2,1,0)
    wOOvo += ovOO.conj().transpose(3,2,1,0)
    woovo -= einsum('iclk,jlbc->ikbj', ovoo, t2aa)
    woovo += einsum('LCik,jLbC->ikbj', OVoo, t2ab)
    wOOVO -= einsum('iclk,jlbc->ikbj', OVOO, t2bb)
    wOOVO += einsum('lcIK,lJcB->IKBJ', ovOO, t2ab)
    wooVO -= einsum('iclk,lJcB->ikBJ', ovoo, t2ab)
    wooVO += einsum('LCik,JLBC->ikBJ', OVoo, t2bb)
    wooVO -= einsum('icLK,jLcB->ijBK', ovOO, t2ab)
    wOOvo -= einsum('ICLK,jLbC->IKbj', OVOO, t2ab)
    wOOvo += einsum('lcIK,jlbc->IKbj', ovOO, t2aa)
    wOOvo -= einsum('IClk,lJbC->IJbk', OVoo, t2ab)

    wvvvo  = einsum('jack,jb->back', v4ovvo, t1a)
    wVVVO  = einsum('jack,jb->back', v4OVVO, t1b)
    wVVvo  = einsum('JAck,JB->BAck', v4OVvo, t1b)
    wVVvo -= einsum('jACk,jb->CAbk', v4oVVo, t1a)
    wvvVO  = einsum('jaCK,jb->baCK', v4ovVO, t1a)
    wvvVO -= einsum('JacK,JB->caBK', v4OvvO, t1b)
    wvvvo += einsum('lajk,jlbc->back', .25*ovoo, tauaa)
    wVVVO += einsum('lajk,jlbc->back', .25*OVOO, taubb)
    wVVvo -= einsum('LAjk,jLcB->BAck', OVoo, tauab)
    wvvVO -= einsum('laJK,lJbC->baCK', ovOO, tauab)

    w3a  = numpy.einsum('jbck,jb->ck', v4ovvo, t1a)
    w3a += numpy.einsum('JBck,JB->ck', v4OVvo, t1b)
    w3b  = numpy.einsum('jbck,jb->ck', v4OVVO, t1b)
    w3b += numpy.einsum('jbCK,jb->CK', v4ovVO, t1a)

    wovvo  = v4ovvo
    wOVVO  = v4OVVO
    wovVO  = v4ovVO
    wOVvo  = v4OVvo
    woVVo  = v4oVVo
    wOvvO  = v4OvvO
    wovvo += lib.einsum('jbld,kd,lc->jbck', ovov, t1a, -t1a)
    wOVVO += lib.einsum('jbld,kd,lc->jbck', OVOV, t1b, -t1b)
    wovVO += lib.einsum('jbLD,KD,LC->jbCK', ovOV, t1b, -t1b)
    wOVvo += lib.einsum('ldJB,kd,lc->JBck', ovOV, t1a, -t1a)
    woVVo += lib.einsum('jdLB,kd,LC->jBCk', ovOV, t1a,  t1b)
    wOvvO += lib.einsum('lbJD,KD,lc->JbcK', ovOV, t1b,  t1a)
    wovvo -= einsum('jblk,lc->jbck', ovoo, t1a)
    wOVVO -= einsum('jblk,lc->jbck', OVOO, t1b)
    wovVO -= einsum('jbLK,LC->jbCK', ovOO, t1b)
    wOVvo -= einsum('JBlk,lc->JBck', OVoo, t1a)
    woVVo += einsum('LBjk,LC->jBCk', OVoo, t1b)
    wOvvO += einsum('lbJK,lc->JbcK', ovOO, t1a)

    if nvira > 0 and nocca > 0:
        ovvv = numpy.asarray(eris.get_ovvv())
        ovvv = ovvv - ovvv.transpose(0,3,2,1)
        v1a -= numpy.einsum('jabc,jc->ba', ovvv, t1a)
        v5a += einsum('kdbc,jkcd->bj', ovvv, t2aa) * .5
        woovo += einsum('idcb,kjbd->ijck', ovvv, tauaa) * .25
        wovvo += einsum('jbcd,kd->jbck', ovvv, t1a)
        wvvvo -= ovvv.conj().transpose(3,2,1,0) * .5
        wvvvo += einsum('jacd,kjbd->cabk', ovvv, t2aa)
        wvvVO += einsum('jacd,jKdB->caBK', ovvv, t2ab)
        ovvv = tmp = None

    if nvirb > 0 and noccb > 0:
        OVVV = numpy.asarray(eris.get_OVVV())
        OVVV = OVVV - OVVV.transpose(0,3,2,1)
        v1b -= numpy.einsum('jabc,jc->ba', OVVV, t1b)
        v5b += einsum('KDBC,JKCD->BJ', OVVV, t2bb) * .5
        wOOVO += einsum('idcb,kjbd->ijck', OVVV, taubb) * .25
        wOVVO += einsum('jbcd,kd->jbck', OVVV, t1b)
        wVVVO -= OVVV.conj().transpose(3,2,1,0) * .5
        wVVVO += einsum('jacd,kjbd->cabk', OVVV, t2bb)
        wVVvo += einsum('JACD,kJbD->CAbk', OVVV, t2ab)
        OVVV = tmp = None

    if nvirb > 0 and nocca > 0:
        OVvv = numpy.asarray(eris.get_OVvv())
        v1a += numpy.einsum('JCba,JC->ba', OVvv, t1b)
        v5a += einsum('KDbc,jKcD->bj', OVvv, t2ab)
        wOOvo += einsum('IDcb,kJbD->IJck', OVvv, tauab)
        wOVvo += einsum('JBcd,kd->JBck', OVvv, t1a)
        wOvvO -= einsum('JDcb,KD->JbcK', OVvv, t1b)
        wvvVO -= OVvv.conj().transpose(3,2,1,0)
        wvvvo -= einsum('KDca,jKbD->cabj', OVvv, t2ab)
        wvvVO -= einsum('KDca,JKBD->caBJ', OVvv, t2bb)
        wVVvo += einsum('KAcd,jKdB->BAcj', OVvv, t2ab)
        OVvv = tmp = None

    if nvira > 0 and noccb > 0:
        ovVV = numpy.asarray(eris.get_ovVV())
        v1b += numpy.einsum('jcBA,jc->BA', ovVV, t1a)
        v5b += einsum('kdBC,kJdC->BJ', ovVV, t2ab)
        wooVO += einsum('idCB,jKdB->ijCK', ovVV, tauab)
        wovVO += einsum('jbCD,KD->jbCK', ovVV, t1b)
        woVVo -= einsum('jdCB,kd->jBCk', ovVV, t1a)
        wVVvo -= ovVV.conj().transpose(3,2,1,0)
        wVVVO -= einsum('kdCA,kJdB->CABJ', ovVV, t2ab)
        wVVvo -= einsum('kdCA,jkbd->CAbj', ovVV, t2aa)
        wvvVO += einsum('kaCD,kJbD->baCJ', ovVV, t2ab)
        ovVV = tmp = None

    w3a += v5a
    w3b += v5b
    w3a += lib.einsum('cb,jb->cj', v1a, t1a)
    w3b += lib.einsum('cb,jb->cj', v1b, t1b)
    w3a -= lib.einsum('jk,jb->bk', v2a, t1a)
    w3b -= lib.einsum('jk,jb->bk', v2b, t1b)

    class _IMDS: pass
    imds = _IMDS()
    imds.ftmp = lib.H5TmpFile()
    dtype = numpy.result_type(t2ab, eris.vvvv).char
    imds.woooo = imds.ftmp.create_dataset('woooo', (nocca,nocca,nocca,nocca), dtype)
    imds.wooOO = imds.ftmp.create_dataset('wooOO', (nocca,nocca,noccb,noccb), dtype)
    imds.wOOOO = imds.ftmp.create_dataset('wOOOO', (noccb,noccb,noccb,noccb), dtype)
    imds.wovvo = imds.ftmp.create_dataset('wovvo', (nocca,nvira,nvira,nocca), dtype)
    imds.wOVVO = imds.ftmp.create_dataset('wOVVO', (noccb,nvirb,nvirb,noccb), dtype)
    imds.wovVO = imds.ftmp.create_dataset('wovVO', (nocca,nvira,nvirb,noccb), dtype)
    imds.wOVvo = imds.ftmp.create_dataset('wOVvo', (noccb,nvirb,nvira,nocca), dtype)
    imds.woVVo = imds.ftmp.create_dataset('woVVo', (nocca,nvirb,nvirb,nocca), dtype)
    imds.wOvvO = imds.ftmp.create_dataset('wOvvO', (noccb,nvira,nvira,noccb), dtype)
    imds.woovo = imds.ftmp.create_dataset('woovo', (nocca,nocca,nvira,nocca), dtype)
    imds.wOOVO = imds.ftmp.create_dataset('wOOVO', (noccb,noccb,nvirb,noccb), dtype)
    imds.wOOvo = imds.ftmp.create_dataset('wOOvo', (noccb,noccb,nvira,nocca), dtype)
    imds.wooVO = imds.ftmp.create_dataset('wooVO', (nocca,nocca,nvirb,noccb), dtype)
    imds.wvvvo = imds.ftmp.create_dataset('wvvvo', (nvira,nvira,nvira,nocca), dtype)
    imds.wVVVO = imds.ftmp.create_dataset('wVVVO', (nvirb,nvirb,nvirb,noccb), dtype)
    imds.wVVvo = imds.ftmp.create_dataset('wVVvo', (nvirb,nvirb,nvira,nocca), dtype)
    imds.wvvVO = imds.ftmp.create_dataset('wvvVO', (nvira,nvira,nvirb,noccb), dtype)

    imds.woooo[:] = woooo
    imds.wOOOO[:] = wOOOO
    imds.wooOO[:] = wooOO
    imds.wovvo[:] = wovvo
    imds.wOVVO[:] = wOVVO
    imds.wovVO[:] = wovVO
    imds.wOVvo[:] = wOVvo
    imds.woVVo[:] = woVVo
    imds.wOvvO[:] = wOvvO
    imds.woovo[:] = woovo
    imds.wOOVO[:] = wOOVO
    imds.wOOvo[:] = wOOvo
    imds.wooVO[:] = wooVO
    imds.wvvvo[:] = wvvvo
    imds.wVVVO[:] = wVVVO
    imds.wVVvo[:] = wVVvo
    imds.wvvVO[:] = wvvVO
    imds.v1a = v1a
    imds.v1b = v1b
    imds.v2a = v2a
    imds.v2b = v2b
    imds.w3a = w3a
    imds.w3b = w3b
    imds.ftmp.flush()
    return imds


# update L1, L2
def update_lambda(mycc, t1, t2, l1, l2, eris, imds):
    time0 = logger.process_clock(), logger.perf_counter()
    log = logger.Logger(mycc.stdout, mycc.verbose)

    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    l1a, l1b = l1
    l2aa, l2ab, l2bb = l2
    nocca, nvira = t1a.shape
    noccb, nvirb = t1b.shape
    u1a = numpy.zeros_like(l1a)
    u1b = numpy.zeros_like(l1b)
    u2aa = numpy.zeros_like(l2aa)
    u2ab = numpy.zeros_like(l2ab)
    u2bb = numpy.zeros_like(l2bb)
    mo_ea_o = eris.mo_energy[0][:nocca]
    mo_ea_v = eris.mo_energy[0][nocca:] + mycc.level_shift
    mo_eb_o = eris.mo_energy[1][:noccb]
    mo_eb_v = eris.mo_energy[1][noccb:] + mycc.level_shift

    fova = eris.focka[:nocca,nocca:]
    fovb = eris.fockb[:noccb,noccb:]
    v1a = imds.v1a - numpy.diag(mo_ea_v)
    v1b = imds.v1b - numpy.diag(mo_eb_v)
    v2a = imds.v2a - numpy.diag(mo_ea_o)
    v2b = imds.v2b - numpy.diag(mo_eb_o)

    mvv = einsum('klca,klcb->ba', l2aa, t2aa) * .5
    mvv+= einsum('lKaC,lKbC->ba', l2ab, t2ab)
    mVV = einsum('klca,klcb->ba', l2bb, t2bb) * .5
    mVV+= einsum('kLcA,kLcB->BA', l2ab, t2ab)
    moo = einsum('kicd,kjcd->ij', l2aa, t2aa) * .5
    moo+= einsum('iKdC,jKdC->ij', l2ab, t2ab)
    mOO = einsum('kicd,kjcd->ij', l2bb, t2bb) * .5
    mOO+= einsum('kIcD,kJcD->IJ', l2ab, t2ab)

    #m3 = lib.einsum('ijcd,cdab->ijab', l2, eris.vvvv) * .5
    m3aa, m3ab, m3bb = mycc._add_vvvv(None, (l2aa.conj(),l2ab.conj(),l2bb.conj()), eris)
    m3aa = m3aa.conj()
    m3ab = m3ab.conj()
    m3bb = m3bb.conj()
    m3aa += lib.einsum('klab,ikjl->ijab', l2aa, numpy.asarray(imds.woooo))
    m3bb += lib.einsum('klab,ikjl->ijab', l2bb, numpy.asarray(imds.wOOOO))
    m3ab += lib.einsum('kLaB,ikJL->iJaB', l2ab, numpy.asarray(imds.wooOO))

    ovov = numpy.asarray(eris.ovov)
    ovov = ovov - ovov.transpose(0,3,2,1)
    OVOV = numpy.asarray(eris.OVOV)
    OVOV = OVOV - OVOV.transpose(0,3,2,1)
    ovOV = numpy.asarray(eris.ovOV)
    mvv1 = einsum('jc,jb->bc', l1a, t1a) + mvv
    mVV1 = einsum('jc,jb->bc', l1b, t1b) + mVV
    moo1 = einsum('ic,kc->ik', l1a, t1a) + moo
    mOO1 = einsum('ic,kc->ik', l1b, t1b) + mOO
    if nvira > 0 and nocca > 0:
        ovvv = numpy.asarray(eris.get_ovvv())
        ovvv = ovvv - ovvv.transpose(0,3,2,1)
        tmp = lib.einsum('ijcd,kd->ijck', l2aa, t1a)
        m3aa -= lib.einsum('kbca,ijck->ijab', ovvv, tmp)

        tmp = einsum('ic,jbca->jiba', l1a, ovvv)
        tmp+= einsum('kiab,jk->ijab', l2aa, v2a)
        tmp-= einsum('ik,kajb->ijab', moo1, ovov)
        u2aa += tmp - tmp.transpose(1,0,2,3)
        u1a += numpy.einsum('iacb,bc->ia', ovvv, mvv1)
        ovvv = tmp = None

    if nvirb > 0 and noccb > 0:
        OVVV = numpy.asarray(eris.get_OVVV())
        OVVV = OVVV - OVVV.transpose(0,3,2,1)
        tmp = lib.einsum('ijcd,kd->ijck', l2bb, t1b)
        m3bb -= lib.einsum('kbca,ijck->ijab', OVVV, tmp)

        tmp = einsum('ic,jbca->jiba', l1b, OVVV)
        tmp+= einsum('kiab,jk->ijab', l2bb, v2b)
        tmp-= einsum('ik,kajb->ijab', mOO1, OVOV)
        u2bb += tmp - tmp.transpose(1,0,2,3)
        u1b += numpy.einsum('iaCB,BC->ia', OVVV, mVV1)
        OVVV = tmp = None

    if nvirb > 0 and nocca > 0:
        OVvv = numpy.asarray(eris.get_OVvv())
        tmp = lib.einsum('iJcD,KD->iJcK', l2ab, t1b)
        m3ab -= lib.einsum('KBca,iJcK->iJaB', OVvv, tmp)

        tmp = einsum('ic,JAcb->JibA', l1a, OVvv)
        tmp-= einsum('kIaB,jk->IjaB', l2ab, v2a)
        tmp-= einsum('IK,jaKB->IjaB', mOO1, ovOV)
        u2ab += tmp.transpose(1,0,2,3)
        u1b += numpy.einsum('iacb,bc->ia', OVvv, mvv1)
        OVvv = tmp = None

    if nvira > 0 and noccb > 0:
        ovVV = numpy.asarray(eris.get_ovVV())
        tmp = lib.einsum('iJdC,kd->iJCk', l2ab, t1a)
        m3ab -= lib.einsum('kaCB,iJCk->iJaB', ovVV, tmp)

        tmp = einsum('IC,jbCA->jIbA', l1b, ovVV)
        tmp-= einsum('iKaB,JK->iJaB', l2ab, v2b)
        tmp-= einsum('ik,kaJB->iJaB', moo1, ovOV)
        u2ab += tmp
        u1a += numpy.einsum('iaCB,BC->ia', ovVV, mVV1)
        ovVV = tmp = None

    tauaa, tauab, taubb = uccsd.make_tau(t2, t1, t1)
    tmp = lib.einsum('ijcd,klcd->ijkl', l2aa, tauaa)
    ovov = numpy.asarray(eris.ovov)
    ovov = ovov - ovov.transpose(0,3,2,1)
    m3aa += lib.einsum('kalb,ijkl->ijab', ovov, tmp) * .25

    tmp = lib.einsum('ijcd,klcd->ijkl', l2bb, taubb)
    OVOV = numpy.asarray(eris.OVOV)
    OVOV = OVOV - OVOV.transpose(0,3,2,1)
    m3bb += lib.einsum('kalb,ijkl->ijab', OVOV, tmp) * .25

    tmp = lib.einsum('iJcD,kLcD->iJkL', l2ab, tauab)
    ovOV = numpy.asarray(eris.ovOV)
    m3ab += lib.einsum('kaLB,iJkL->iJaB', ovOV, tmp) * .5
    tmp = lib.einsum('iJdC,lKdC->iJKl', l2ab, tauab)
    m3ab += lib.einsum('laKB,iJKl->iJaB', ovOV, tmp) * .5

    u1a += numpy.einsum('ijab,jb->ia', m3aa, t1a)
    u1a += numpy.einsum('iJaB,JB->ia', m3ab, t1b)
    u1b += numpy.einsum('IJAB,JB->IA', m3bb, t1b)
    u1b += numpy.einsum('jIbA,jb->IA', m3ab, t1a)

    u2aa += m3aa
    u2bb += m3bb
    u2ab += m3ab
    u2aa += ovov.transpose(0,2,1,3)
    u2bb += OVOV.transpose(0,2,1,3)
    u2ab += ovOV.transpose(0,2,1,3)

    fov1 = fova + numpy.einsum('kcjb,kc->jb', ovov, t1a)
    fov1+= numpy.einsum('jbKC,KC->jb', ovOV, t1b)
    tmp = numpy.einsum('ia,jb->ijab', l1a, fov1)
    tmp+= einsum('kica,jbck->ijab', l2aa, imds.wovvo)
    tmp+= einsum('iKaC,jbCK->ijab', l2ab, imds.wovVO)
    tmp = tmp - tmp.transpose(1,0,2,3)
    u2aa += tmp - tmp.transpose(0,1,3,2)

    fov1 = fovb + numpy.einsum('kcjb,kc->jb', OVOV, t1b)
    fov1+= numpy.einsum('kcJB,kc->JB', ovOV, t1a)
    tmp = numpy.einsum('ia,jb->ijab', l1b, fov1)
    tmp+= einsum('kica,jbck->ijab', l2bb, imds.wOVVO)
    tmp+= einsum('kIcA,JBck->IJAB', l2ab, imds.wOVvo)
    tmp = tmp - tmp.transpose(1,0,2,3)
    u2bb += tmp - tmp.transpose(0,1,3,2)

    fov1 = fovb + numpy.einsum('kcjb,kc->jb', OVOV, t1b)
    fov1+= numpy.einsum('kcJB,kc->JB', ovOV, t1a)
    u2ab += numpy.einsum('ia,JB->iJaB', l1a, fov1)
    u2ab += einsum('iKaC,JBCK->iJaB', l2ab, imds.wOVVO)
    u2ab += einsum('kica,JBck->iJaB', l2aa, imds.wOVvo)
    u2ab += einsum('kIaC,jBCk->jIaB', l2ab, imds.woVVo)
    u2ab += einsum('iKcA,JbcK->iJbA', l2ab, imds.wOvvO)
    fov1 = fova + numpy.einsum('kcjb,kc->jb', ovov, t1a)
    fov1+= numpy.einsum('jbKC,KC->jb', ovOV, t1b)
    u2ab += numpy.einsum('ia,jb->jiba', l1b, fov1)
    u2ab += einsum('kIcA,jbck->jIbA', l2ab, imds.wovvo)
    u2ab += einsum('KICA,jbCK->jIbA', l2bb, imds.wovVO)

    ovoo = numpy.asarray(eris.ovoo)
    ovoo = ovoo - ovoo.transpose(2,1,0,3)
    OVOO = numpy.asarray(eris.OVOO)
    OVOO = OVOO - OVOO.transpose(2,1,0,3)
    OVoo = numpy.asarray(eris.OVoo)
    ovOO = numpy.asarray(eris.ovOO)
    tmp = einsum('ka,jbik->ijab', l1a, ovoo)
    tmp+= einsum('ijca,cb->ijab', l2aa, v1a)
    tmp+= einsum('ca,icjb->ijab', mvv1, ovov)
    u2aa -= tmp - tmp.transpose(0,1,3,2)
    tmp = einsum('ka,jbik->ijab', l1b, OVOO)
    tmp+= einsum('ijca,cb->ijab', l2bb, v1b)
    tmp+= einsum('ca,icjb->ijab', mVV1, OVOV)
    u2bb -= tmp - tmp.transpose(0,1,3,2)
    u2ab -= einsum('ka,JBik->iJaB', l1a, OVoo)
    u2ab += einsum('iJaC,CB->iJaB', l2ab, v1b)
    u2ab -= einsum('ca,icJB->iJaB', mvv1, ovOV)
    u2ab -= einsum('KA,ibJK->iJbA', l1b, ovOO)
    u2ab += einsum('iJcA,cb->iJbA', l2ab, v1a)
    u2ab -= einsum('CA,ibJC->iJbA', mVV1, ovOV)

    u1a += fova
    u1b += fovb
    u1a += einsum('ib,ba->ia', l1a, v1a)
    u1a -= einsum('ja,ij->ia', l1a, v2a)
    u1b += einsum('ib,ba->ia', l1b, v1b)
    u1b -= einsum('ja,ij->ia', l1b, v2b)

    u1a += numpy.einsum('jb,iabj->ia', l1a, eris.ovvo)
    u1a -= numpy.einsum('jb,ijba->ia', l1a, eris.oovv)
    u1a += numpy.einsum('JB,iaBJ->ia', l1b, eris.ovVO)
    u1b += numpy.einsum('jb,iabj->ia', l1b, eris.OVVO)
    u1b -= numpy.einsum('jb,ijba->ia', l1b, eris.OOVV)
    u1b += numpy.einsum('jb,iabj->ia', l1a, eris.OVvo)

    u1a -= einsum('kjca,ijck->ia', l2aa, imds.woovo)
    u1a -= einsum('jKaC,ijCK->ia', l2ab, imds.wooVO)
    u1b -= einsum('kjca,ijck->ia', l2bb, imds.wOOVO)
    u1b -= einsum('kJcA,IJck->IA', l2ab, imds.wOOvo)

    u1a -= einsum('ikbc,back->ia', l2aa, imds.wvvvo)
    u1a -= einsum('iKbC,baCK->ia', l2ab, imds.wvvVO)
    u1b -= einsum('IKBC,BACK->IA', l2bb, imds.wVVVO)
    u1b -= einsum('kIcB,BAck->IA', l2ab, imds.wVVvo)

    u1a += numpy.einsum('jiba,bj->ia', l2aa, imds.w3a)
    u1a += numpy.einsum('iJaB,BJ->ia', l2ab, imds.w3b)
    u1b += numpy.einsum('JIBA,BJ->IA', l2bb, imds.w3b)
    u1b += numpy.einsum('jIbA,bj->IA', l2ab, imds.w3a)

    tmpa  = t1a + numpy.einsum('kc,kjcb->jb', l1a, t2aa)
    tmpa += numpy.einsum('KC,jKbC->jb', l1b, t2ab)
    tmpa -= einsum('bd,jd->jb', mvv1, t1a)
    tmpa -= einsum('lj,lb->jb', moo, t1a)
    tmpb  = t1b + numpy.einsum('kc,kjcb->jb', l1b, t2bb)
    tmpb += numpy.einsum('kc,kJcB->JB', l1a, t2ab)
    tmpb -= einsum('bd,jd->jb', mVV1, t1b)
    tmpb -= einsum('lj,lb->jb', mOO, t1b)
    u1a += numpy.einsum('jbia,jb->ia', ovov, tmpa)
    u1a += numpy.einsum('iaJB,JB->ia', ovOV, tmpb)
    u1b += numpy.einsum('jbia,jb->ia', OVOV, tmpb)
    u1b += numpy.einsum('jbIA,jb->IA', ovOV, tmpa)

    u1a -= numpy.einsum('iajk,kj->ia', ovoo, moo1)
    u1a -= numpy.einsum('iaJK,KJ->ia', ovOO, mOO1)
    u1b -= numpy.einsum('iajk,kj->ia', OVOO, mOO1)
    u1b -= numpy.einsum('IAjk,kj->IA', OVoo, moo1)

    tmp  = fova - numpy.einsum('kbja,jb->ka', ovov, t1a)
    tmp += numpy.einsum('kaJB,JB->ka', ovOV, t1b)
    u1a -= lib.einsum('ik,ka->ia', moo, tmp)
    u1a -= lib.einsum('ca,ic->ia', mvv, tmp)
    tmp  = fovb - numpy.einsum('kbja,jb->ka', OVOV, t1b)
    tmp += numpy.einsum('jbKA,jb->KA', ovOV, t1a)
    u1b -= lib.einsum('ik,ka->ia', mOO, tmp)
    u1b -= lib.einsum('ca,ic->ia', mVV, tmp)

    eia = lib.direct_sum('i-j->ij', mo_ea_o, mo_ea_v)
    eIA = lib.direct_sum('i-j->ij', mo_eb_o, mo_eb_v)
    u1a /= eia
    u1b /= eIA

    u2aa /= lib.direct_sum('ia+jb->ijab', eia, eia)
    u2ab /= lib.direct_sum('ia+jb->ijab', eia, eIA)
    u2bb /= lib.direct_sum('ia+jb->ijab', eIA, eIA)

    time0 = log.timer_debug1('update l1 l2', *time0)
    return (u1a,u1b), (u2aa,u2ab,u2bb)


if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf
    from pyscf.cc import gccsd

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = '631g'
    mol.spin = 2
    mol.build()
    mf = scf.UHF(mol).run()
    mycc = gccsd.GCCSD(scf.addons.convert_to_ghf(mf))
    eris = mycc.ao2mo()
    mycc.kernel()
    l1, l2 = mycc.solve_lambda(mycc.t1, mycc.t2, eris=eris)
    l1ref = mycc.spin2spatial(l1, mycc.mo_coeff.orbspin)
    l2ref = mycc.spin2spatial(l2, mycc.mo_coeff.orbspin)

    mycc = uccsd.UCCSD(mf)
    eris = mycc.ao2mo()
    mycc.kernel()
    conv, l1, l2 = kernel(mycc, eris, mycc.t1, mycc.t2, tol=1e-8)
    print(abs(l1[0]-l1ref[0]).max())
    print(abs(l1[1]-l1ref[1]).max())
    print(abs(l2[0]-l2ref[0]).max())
    print(abs(l2[1]-l2ref[1]).max())
    print(abs(l2[2]-l2ref[2]).max())


# pyscf/cc/uccsd_rdm.py

#!/usr/bin/env python
# Copyright 2014-2020 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: Qiming Sun <osirpt.sun@gmail.com>
#         Jun Yang
#

import numpy
from pyscf import lib
from pyscf import ao2mo

#einsum = numpy.einsum
einsum = lib.einsum

#TODO: optimize memory use

def _gamma1_intermediates(cc, t1, t2, l1, l2):
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    l1a, l1b = l1
    l2aa, l2ab, l2bb = l2
    nocca, nvira = t1a.shape
    noccb, nvirb = t1b.shape

    tmpA  = einsum('imef,jmef->ij', l2ab, t2ab)
    tmpA += einsum('imef,jmef->ij', l2aa, t2aa) * .5

    tmpB  = einsum('mief,mjef->ij', l2ab, t2ab)
    tmpB += einsum('imef,jmef->ij', l2bb, t2bb) * .5

    tmpC  = einsum('mnae,mnbe->ab', t2ab, l2ab)
    tmpC += einsum('mnae,mnbe->ab', t2aa, l2aa) * .5

    tmpD  = einsum('mnea,mneb->ab', t2ab, l2ab)
    tmpD += einsum('mnae,mnbe->ab', t2bb, l2bb) * .5

    dooa  = -einsum('ie,je->ij', l1a, t1a)
    dooa -= tmpA
    doob  = -einsum('ie,je->ij', l1b, t1b)
    doob -= tmpB

    dvva  = einsum('ma,mb->ab', t1a, l1a)
    dvva += tmpC
    dvvb  = einsum('ma,mb->ab', t1b, l1b)
    dvvb += tmpD

    xt1a = tmpA
    xt2a = tmpC
    xt2a += einsum('ma,me->ae', t1a, l1a)

    xt1b = tmpB
    xt2b = tmpD
    xt2b += einsum('ma,me->ae', t1b, l1b)

    dvoa  = numpy.einsum('imae,me->ai', t2aa, l1a)
    dvoa += numpy.einsum('imae,me->ai', t2ab, l1b)
    dvoa -= einsum('mi,ma->ai', xt1a, t1a)
    dvoa -= einsum('ie,ae->ai', t1a, xt2a)
    dvoa += t1a.T

    dvob  = numpy.einsum('imae,me->ai', t2bb, l1b)
    dvob += numpy.einsum('miea,me->ai', t2ab, l1a)
    dvob -= einsum('mi,ma->ai', xt1b, t1b)
    dvob -= einsum('ie,ae->ai', t1b, xt2b)
    dvob += t1b.T

    dova = l1a
    dovb = l1b

    return ((dooa, doob), (dova, dovb), (dvoa, dvob), (dvva, dvvb))

# gamma2 intermediates in Chemist's notation
#TODO: hold d2 intermediates in h5fobj
def _gamma2_outcore(cc, t1, t2, l1, l2, h5fobj, compress_vvvv=False):
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    l1a, l1b = l1
    l2aa, l2ab, l2bb = l2

    tauaa = t2aa + numpy.einsum('ia,jb->ijab', 2*t1a, t1a)
    tauab = t2ab + numpy.einsum('ia,jb->ijab',   t1a, t1b)
    taubb = t2bb + numpy.einsum('ia,jb->ijab', 2*t1b, t1b)
    miajb = einsum('ikac,kjcb->iajb', l2aa, t2aa)
    miajb+= einsum('ikac,jkbc->iajb', l2ab, t2ab)
    miaJB = einsum('ikac,kjcb->iajb', l2aa, t2ab)
    miaJB+= einsum('ikac,kjcb->iajb', l2ab, t2bb)
    mIAjb = einsum('kica,jkbc->iajb', l2bb, t2ab)
    mIAjb+= einsum('kica,kjcb->iajb', l2ab, t2aa)
    mIAJB = einsum('ikac,kjcb->iajb', l2bb, t2bb)
    mIAJB+= einsum('kica,kjcb->iajb', l2ab, t2ab)
    miAjB = einsum('ikca,jkcb->iajb', l2ab, t2ab)
    mIaJb = einsum('kiac,kjbc->iajb', l2ab, t2ab)

    goovv = (l2aa.conj() + tauaa) * .25
    goOvV = (l2ab.conj() + tauab) * .5
    gOOVV = (l2bb.conj() + taubb) * .25

    tmpa  = einsum('kc,kica->ia', l1a, t2aa)
    tmpa += einsum('kc,ikac->ia', l1b, t2ab)
    tmpb  = einsum('kc,kica->ia', l1b, t2bb)
    tmpb += einsum('kc,kica->ia', l1a, t2ab)
    goovv += einsum('ia,jb->ijab', tmpa, t1a)
    goOvV += einsum('ia,jb->ijab', tmpa, t1b) * .5
    goOvV += einsum('ia,jb->jiba', tmpb, t1a) * .5
    gOOVV += einsum('ia,jb->ijab', tmpb, t1b)

    tmpa = einsum('kc,kb->cb', l1a, t1a)
    tmpb = einsum('kc,kb->cb', l1b, t1b)
    goovv += einsum('cb,ijca->ijab', tmpa, t2aa) * .5
    goOvV -= einsum('cb,ijac->ijab', tmpb, t2ab) * .5
    goOvV -= einsum('cb,jica->jiba', tmpa, t2ab) * .5
    gOOVV += einsum('cb,ijca->ijab', tmpb, t2bb) * .5
    tmpa = einsum('kc,jc->kj', l1a, t1a)
    tmpb = einsum('kc,jc->kj', l1b, t1b)
    goovv += einsum('kiab,kj->ijab', tauaa, tmpa) * .5
    goOvV -= einsum('ikab,kj->ijab', tauab , tmpb) * .5
    goOvV -= einsum('kiba,kj->jiba', tauab , tmpa) * .5
    gOOVV += einsum('kiab,kj->ijab', taubb, tmpb) * .5

    tmpa  = numpy.einsum('ldjd->lj', miajb)
    tmpa += numpy.einsum('ldjd->lj', miAjB)
    tmpb  = numpy.einsum('ldjd->lj', mIAJB)
    tmpb += numpy.einsum('ldjd->lj', mIaJb)
    goovv -= einsum('lj,liba->ijab', tmpa, tauaa) * .25
    goOvV -= einsum('lj,ilab->ijab', tmpb, tauab) * .25
    goOvV -= einsum('lj,liba->jiba', tmpa, tauab) * .25
    gOOVV -= einsum('lj,liba->ijab', tmpb, taubb) * .25
    tmpa  = numpy.einsum('ldlb->db', miajb)
    tmpa += numpy.einsum('ldlb->db', mIaJb)
    tmpb  = numpy.einsum('ldlb->db', mIAJB)
    tmpb += numpy.einsum('ldlb->db', miAjB)
    goovv -= einsum('db,jida->ijab', tmpa, tauaa) * .25
    goOvV -= einsum('db,ijad->ijab', tmpb, tauab) * .25
    goOvV -= einsum('db,jida->jiba', tmpa, tauab) * .25
    gOOVV -= einsum('db,jida->ijab', tmpb, taubb) * .25

    goovv -= einsum('ldia,ljbd->ijab', miajb, tauaa) * .5
    goovv += einsum('LDia,jLbD->ijab', mIAjb, t2ab ) * .5
    gOOVV -= einsum('ldia,ljbd->ijab', mIAJB, taubb) * .5
    gOOVV += einsum('ldia,ljdb->ijab', miaJB, t2ab ) * .5
    goOvV -= einsum('LDia,LJBD->iJaB', mIAjb, taubb) * .25
    goOvV += einsum('ldia,lJdB->iJaB', miajb, t2ab ) * .25
    goOvV -= einsum('ldIA,ljbd->jIbA', miaJB, tauaa) * .25
    goOvV += einsum('LDIA,jLbD->jIbA', mIAJB, t2ab ) * .25
    goOvV += einsum('lDiA,lJbD->iJbA', miAjB, tauab) * .5
    goOvV += einsum('LdIa,jd,LB->jIaB', mIaJb, t1a, t1b) * .5

    tmpaa = einsum('klcd,ijcd->ijkl', l2aa, tauaa) * .25**2
    tmpbb = einsum('klcd,ijcd->ijkl', l2bb, taubb) * .25**2
    tmpabab = einsum('kLcD,iJcD->iJkL', l2ab, tauab) * .5
    goovv += einsum('ijkl,klab->ijab', tmpaa, tauaa)
    goOvV += einsum('ijkl,klab->ijab', tmpabab, tauab)
    gOOVV += einsum('ijkl,klab->ijab', tmpbb, taubb)
    goovv = goovv.conj()
    goOvV = goOvV.conj()
    gOOVV = gOOVV.conj()

    gvvvv = einsum('ijab,ijcd->abcd', tauaa, l2aa) * .125
    gvVvV = einsum('ijab,ijcd->abcd', tauab, l2ab) * .25
    gVVVV = einsum('ijab,ijcd->abcd', taubb, l2bb) * .125

    goooo = einsum('ijab,klab->ijkl', l2aa, tauaa) * .125
    goOoO = einsum('ijab,klab->ijkl', l2ab, tauab) * .25
    gOOOO = einsum('ijab,klab->ijkl', l2bb, taubb) * .125

    gooov = einsum('jkba,ib->jkia', tauaa, -0.25 * l1a)
    goOoV = einsum('jkba,ib->jkia', tauab, -0.5  * l1a)
    gOoOv = einsum('kjab,ib->jkia', tauab, -0.5  * l1b)
    gOOOV = einsum('jkba,ib->jkia', taubb, -0.25 * l1b)

    gooov += einsum('iljk,la->jkia', goooo, t1a)
    goOoV += einsum('iljk,la->jkia', goOoO, t1b) * 2
    gOoOv += einsum('likj,la->jkia', goOoO, t1a) * 2
    gOOOV += einsum('iljk,la->jkia', gOOOO, t1b)

    tmpa  = numpy.einsum('icjc->ij', miajb) * .25
    tmpa += numpy.einsum('icjc->ij', miAjB) * .25
    tmpb  = numpy.einsum('icjc->ij', mIAJB) * .25
    tmpb += numpy.einsum('icjc->ij', mIaJb) * .25
    gooov -= einsum('ij,ka->jkia', tmpa, t1a)
    goOoV -= einsum('ij,ka->jkia', tmpa, t1b)
    gOoOv -= einsum('ij,ka->jkia', tmpb, t1a)
    gOOOV -= einsum('ij,ka->jkia', tmpb, t1b)

    gooov += einsum('icja,kc->jkia', miajb, .5 * t1a)
    goOoV += einsum('icja,kc->jkia', miAjB, .5 * t1b)
    goOoV -= einsum('icJA,kc->kJiA', miaJB, .5 * t1a)
    gOoOv += einsum('icja,kc->jkia', mIaJb, .5 * t1a)
    gOoOv -= einsum('ICja,KC->KjIa', mIAjb, .5 * t1b)
    gOOOV += einsum('icja,kc->jkia', mIAJB, .5 * t1b)

    gooov = gooov.conj()
    goOoV = goOoV.conj()
    gOoOv = gOoOv.conj()
    gOOOV = gOOOV.conj()
    gooov += einsum('jkab,ib->jkia', l2aa, .25*t1a)
    goOoV -= einsum('jkba,ib->jkia', l2ab, .5 *t1a)
    gOoOv -= einsum('kjab,ib->jkia', l2ab, .5 *t1b)
    gOOOV += einsum('jkab,ib->jkia', l2bb, .25*t1b)

    govvv = einsum('ja,ijcb->iacb', .25 * l1a, tauaa)
    goVvV = einsum('ja,ijcb->iacb', .5  * l1b, tauab)
    gOvVv = einsum('ja,jibc->iacb', .5  * l1a, tauab)
    gOVVV = einsum('ja,ijcb->iacb', .25 * l1b, taubb)

    govvv += einsum('bcad,id->iabc', gvvvv, t1a)
    goVvV -= einsum('bcda,id->iabc', gvVvV, t1a) * 2
    gOvVv -= einsum('cbad,id->iabc', gvVvV, t1b) * 2
    gOVVV += einsum('bcad,id->iabc', gVVVV, t1b)

    tmpa  = numpy.einsum('kakb->ab', miajb) * .25
    tmpa += numpy.einsum('kakb->ab', mIaJb) * .25
    tmpb  = numpy.einsum('kakb->ab', mIAJB) * .25
    tmpb += numpy.einsum('kakb->ab', miAjB) * .25
    govvv += einsum('ab,ic->iacb', tmpa, t1a)
    goVvV += einsum('ab,ic->iacb', tmpb, t1a)
    gOvVv += einsum('ab,ic->iacb', tmpa, t1b)
    gOVVV += einsum('ab,ic->iacb', tmpb, t1b)

    govvv += einsum('kaib,kc->iabc', miajb, .5 * t1a)
    goVvV += einsum('KAib,KC->iAbC', mIAjb, .5 * t1b)
    goVvV -= einsum('kAiB,kc->iAcB', miAjB, .5 * t1a)
    gOvVv += einsum('kaIB,kc->IaBc', miaJB, .5 * t1a)
    gOvVv -= einsum('KaIb,KC->IaCb', mIaJb, .5 * t1b)
    gOVVV += einsum('kaib,kc->iabc', mIAJB, .5 * t1b)
    govvv = govvv.conj()
    goVvV = goVvV.conj()
    gOvVv = gOvVv.conj()
    gOVVV = gOVVV.conj()
    govvv += einsum('ijbc,ja->iabc', l2aa, .25*t1a)
    goVvV += einsum('iJbC,JA->iAbC', l2ab, .5 *t1b)
    gOvVv += einsum('jIcB,ja->IaBc', l2ab, .5 *t1a)
    gOVVV += einsum('ijbc,ja->iabc', l2bb, .25*t1b)

    govvo = einsum('ia,jb->ibaj', l1a, t1a)
    goVvO = einsum('ia,jb->ibaj', l1a, t1b)
    gOvVo = einsum('ia,jb->ibaj', l1b, t1a)
    gOVVO = einsum('ia,jb->ibaj', l1b, t1b)

    govvo += numpy.einsum('iajb->ibaj', miajb)
    goVvO += numpy.einsum('iajb->ibaj', miaJB)
    gOvVo += numpy.einsum('iajb->ibaj', mIAjb)
    gOVVO += numpy.einsum('iajb->ibaj', mIAJB)
    goVoV = numpy.einsum('iajb->ibja', miAjB)
    gOvOv = numpy.einsum('iajb->ibja', mIaJb)

    govvo -= einsum('ikac,jc,kb->ibaj', l2aa, t1a, t1a)
    goVvO -= einsum('iKaC,JC,KB->iBaJ', l2ab, t1b, t1b)
    gOvVo -= einsum('kIcA,jc,kb->IbAj', l2ab, t1a, t1a)
    gOVVO -= einsum('ikac,jc,kb->ibaj', l2bb, t1b, t1b)
    goVoV += einsum('iKcA,jc,KB->iBjA', l2ab, t1a, t1b)
    gOvOv += einsum('kIaC,JC,kb->IbJa', l2ab, t1b, t1a)

    dovov = goovv.transpose(0,2,1,3) - goovv.transpose(0,3,1,2)
    dvvvv = gvvvv.transpose(0,2,1,3) - gvvvv.transpose(0,3,1,2)
    doooo = goooo.transpose(0,2,1,3) - goooo.transpose(0,3,1,2)
    dovvv = govvv.transpose(0,2,1,3) - govvv.transpose(0,3,1,2)
    dooov = gooov.transpose(0,2,1,3) - gooov.transpose(1,2,0,3)
    dovvo = govvo.transpose(0,2,1,3)
    dovov =(dovov + dovov.transpose(2,3,0,1)) * .5
    dvvvv = dvvvv + dvvvv.transpose(1,0,3,2).conj()
    doooo = doooo + doooo.transpose(1,0,3,2).conj()
    dovvo =(dovvo + dovvo.transpose(3,2,1,0).conj()) * .5
    doovv =-dovvo.transpose(0,3,2,1)
    dvvov = None

    dOVOV = gOOVV.transpose(0,2,1,3) - gOOVV.transpose(0,3,1,2)
    dVVVV = gVVVV.transpose(0,2,1,3) - gVVVV.transpose(0,3,1,2)
    dOOOO = gOOOO.transpose(0,2,1,3) - gOOOO.transpose(0,3,1,2)
    dOVVV = gOVVV.transpose(0,2,1,3) - gOVVV.transpose(0,3,1,2)
    dOOOV = gOOOV.transpose(0,2,1,3) - gOOOV.transpose(1,2,0,3)
    dOVVO = gOVVO.transpose(0,2,1,3)
    dOVOV =(dOVOV + dOVOV.transpose(2,3,0,1)) * .5
    dVVVV = dVVVV + dVVVV.transpose(1,0,3,2).conj()
    dOOOO = dOOOO + dOOOO.transpose(1,0,3,2).conj()
    dOVVO =(dOVVO + dOVVO.transpose(3,2,1,0).conj()) * .5
    dOOVV =-dOVVO.transpose(0,3,2,1)
    dVVOV = None

    dovOV = goOvV.transpose(0,2,1,3)
    dvvVV = gvVvV.transpose(0,2,1,3) * 2
    dooOO = goOoO.transpose(0,2,1,3) * 2
    dovVV = goVvV.transpose(0,2,1,3)
    dooOV = goOoV.transpose(0,2,1,3)
    dovVO = goVvO.transpose(0,2,1,3)

    dOVvv = gOvVv.transpose(0,2,1,3)
    dOOov = gOoOv.transpose(0,2,1,3)
    dOVvo = gOvVo.transpose(0,2,1,3)
    dooVV = goVoV.transpose(0,2,1,3)
    dOOvv = gOvOv.transpose(0,2,1,3)

    dvvVV = dvvVV + dvvVV.transpose(1,0,3,2).conj()
    dooOO = dooOO + dooOO.transpose(1,0,3,2).conj()
    dovVO = (dovVO + dOVvo.transpose(3,2,1,0).conj()) * .5
    dooVV =-(dooVV + dooVV.transpose(1,0,3,2).conj()) * .5
    dvvOV = None

    dOVov = None
    dVVvv = None
    dOOoo = None
    dOVvo =  dovVO.transpose(3,2,1,0).conj()
    dOOvv =-(dOOvv + dOOvv.transpose(1,0,3,2).conj()) * .5
    dVVov = None

    if compress_vvvv:
        nocca, noccb, nvira, nvirb = t2ab.shape
        idxa = numpy.tril_indices(nvira)
        idxa = idxa[0] * nvira + idxa[1]
        idxb = numpy.tril_indices(nvirb)
        idxb = idxb[0] * nvirb + idxb[1]
        dvvvv = dvvvv + dvvvv.transpose(1,0,2,3)
        dvvvv = lib.take_2d(dvvvv.reshape(nvira**2,nvira**2), idxa, idxa)
        dvvvv *= .5
        dvvVV = dvvVV + dvvVV.transpose(1,0,2,3)
        dvvVV = lib.take_2d(dvvVV.reshape(nvira**2,nvirb**2), idxa, idxb)
        dvvVV *= .5
        dVVVV = dVVVV + dVVVV.transpose(1,0,2,3)
        dVVVV = lib.take_2d(dVVVV.reshape(nvirb**2,nvirb**2), idxb, idxb)
        dVVVV *= .5

    return ((dovov, dovOV, dOVov, dOVOV),
            (dvvvv, dvvVV, dVVvv, dVVVV),
            (doooo, dooOO, dOOoo, dOOOO),
            (doovv, dooVV, dOOvv, dOOVV),
            (dovvo, dovVO, dOVvo, dOVVO),
            (dvvov, dvvOV, dVVov, dVVOV),
            (dovvv, dovVV, dOVvv, dOVVV),
            (dooov, dooOV, dOOov, dOOOV))

def _gamma2_intermediates(cc, t1, t2, l1, l2, compress_vvvv=False):
    #TODO: h5fobj = lib.H5TmpFile()
    h5fobj = None
    d2 = _gamma2_outcore(cc, t1, t2, l1, l2, h5fobj, compress_vvvv)
    return d2

def make_rdm1(mycc, t1, t2, l1, l2, ao_repr=False, with_frozen=True, with_mf=True):
    r'''
    One-particle spin density matrices dm1a, dm1b in MO basis (the
    occupied-virtual blocks due to the orbital response contribution are not
    included).

    dm1a[p,q] = <q_alpha^\dagger p_alpha>
    dm1b[p,q] = <q_beta^\dagger p_beta>

    The convention of 1-pdm is based on McWeeney's book, Eq (5.4.20).
    '''
    d1 = _gamma1_intermediates(mycc, t1, t2, l1, l2)
    return _make_rdm1(mycc, d1, with_frozen=with_frozen, ao_repr=ao_repr, with_mf=with_mf)

# spin-orbital rdm2 in Chemist's notation
def make_rdm2(mycc, t1, t2, l1, l2, ao_repr=False, with_frozen=True, with_dm1=True):
    r'''
    Two-particle spin density matrices dm2aa, dm2ab, dm2bb in MO basis

    dm2aa[p,q,r,s] = <q_alpha^\dagger s_alpha^\dagger r_alpha p_alpha>
    dm2ab[p,q,r,s] = <q_alpha^\dagger s_beta^\dagger r_beta p_alpha>
    dm2bb[p,q,r,s] = <q_beta^\dagger s_beta^\dagger r_beta p_beta>

    (p,q correspond to one particle and r,s correspond to another particle)
    Two-particle density matrix should be contracted to integrals with the
    pattern below to compute energy

    E = numpy.einsum('pqrs,pqrs', eri_aa, dm2_aa)
    E+= numpy.einsum('pqrs,pqrs', eri_ab, dm2_ab)
    E+= numpy.einsum('pqrs,rspq', eri_ba, dm2_ab)
    E+= numpy.einsum('pqrs,pqrs', eri_bb, dm2_bb)

    where eri_aa[p,q,r,s] = (p_alpha q_alpha | r_alpha s_alpha )
    eri_ab[p,q,r,s] = ( p_alpha q_alpha | r_beta s_beta )
    eri_ba[p,q,r,s] = ( p_beta q_beta | r_alpha s_alpha )
    eri_bb[p,q,r,s] = ( p_beta q_beta | r_beta s_beta )
    '''
    d1 = _gamma1_intermediates(mycc, t1, t2, l1, l2)
    d2 = _gamma2_intermediates(mycc, t1, t2, l1, l2)
    return _make_rdm2(mycc, d1, d2, with_dm1=with_dm1, with_frozen=with_frozen,
                      ao_repr=ao_repr)

def _make_rdm1(mycc, d1, with_frozen=True, ao_repr=False, with_mf=True):
    doo, dOO = d1[0]
    dov, dOV = d1[1]
    dvo, dVO = d1[2]
    dvv, dVV = d1[3]
    nocca, nvira = dov.shape
    noccb, nvirb = dOV.shape
    nmoa = nocca + nvira
    nmob = noccb + nvirb

    dm1a = numpy.empty((nmoa,nmoa), dtype=doo.dtype)
    dm1a[:nocca,:nocca] = doo + doo.conj().T
    dm1a[:nocca,nocca:] = dov + dvo.conj().T
    dm1a[nocca:,:nocca] = dm1a[:nocca,nocca:].conj().T
    dm1a[nocca:,nocca:] = dvv + dvv.conj().T
    dm1a *= .5

    dm1b = numpy.empty((nmob,nmob), dtype=dOO.dtype)
    dm1b[:noccb,:noccb] = dOO + dOO.conj().T
    dm1b[:noccb,noccb:] = dOV + dVO.conj().T
    dm1b[noccb:,:noccb] = dm1b[:noccb,noccb:].conj().T
    dm1b[noccb:,noccb:] = dVV + dVV.conj().T
    dm1b *= .5
    if with_mf:
        dm1a[numpy.diag_indices(nocca)] += 1
        dm1b[numpy.diag_indices(noccb)] += 1

    if with_frozen and mycc.frozen is not None:
        nmoa = mycc.mo_occ[0].size
        nmob = mycc.mo_occ[1].size
        nocca = numpy.count_nonzero(mycc.mo_occ[0] > 0)
        noccb = numpy.count_nonzero(mycc.mo_occ[1] > 0)
        rdm1a = numpy.zeros((nmoa,nmoa), dtype=dm1a.dtype)
        rdm1b = numpy.zeros((nmob,nmob), dtype=dm1b.dtype)
        if with_mf:
            rdm1a[numpy.diag_indices(nocca)] = 1
            rdm1b[numpy.diag_indices(noccb)] = 1
        moidx = mycc.get_frozen_mask()
        moidxa = numpy.where(moidx[0])[0]
        moidxb = numpy.where(moidx[1])[0]
        rdm1a[moidxa[:,None],moidxa] = dm1a
        rdm1b[moidxb[:,None],moidxb] = dm1b
        dm1a = rdm1a
        dm1b = rdm1b

    if ao_repr:
        mo_a, mo_b = mycc.mo_coeff
        dm1a = lib.einsum('pi,ij,qj->pq', mo_a, dm1a, mo_a)
        dm1b = lib.einsum('pi,ij,qj->pq', mo_b, dm1b, mo_b)
    return dm1a, dm1b

def _make_rdm2(mycc, d1, d2, with_dm1=True, with_frozen=True, ao_repr=False):
    dovov, dovOV, dOVov, dOVOV = d2[0]
    dvvvv, dvvVV, dVVvv, dVVVV = d2[1]
    doooo, dooOO, dOOoo, dOOOO = d2[2]
    doovv, dooVV, dOOvv, dOOVV = d2[3]
    dovvo, dovVO, dOVvo, dOVVO = d2[4]
    dvvov, dvvOV, dVVov, dVVOV = d2[5]
    dovvv, dovVV, dOVvv, dOVVV = d2[6]
    dooov, dooOV, dOOov, dOOOV = d2[7]
    nocca, nvira, noccb, nvirb = dovOV.shape
    nmoa = nocca + nvira
    nmob = noccb + nvirb

    dm2aa = numpy.empty((nmoa,nmoa,nmoa,nmoa), dtype=doovv.dtype)
    dm2ab = numpy.empty((nmoa,nmoa,nmob,nmob), dtype=doovv.dtype)
    dm2bb = numpy.empty((nmob,nmob,nmob,nmob), dtype=doovv.dtype)

# dm2aa
    dovov = numpy.asarray(dovov)
    dm2aa[:nocca,nocca:,:nocca,nocca:] = dovov
    dm2aa[nocca:,:nocca,nocca:,:nocca] = dm2aa[:nocca,nocca:,:nocca,nocca:].transpose(1,0,3,2).conj()
    dovov = None

    #assert(abs(doovv+dovvo.transpose(0,3,2,1)).max() == 0)
    dovvo = numpy.asarray(dovvo)
    dm2aa[:nocca,:nocca,nocca:,nocca:] =-dovvo.transpose(0,3,2,1)
    dm2aa[nocca:,nocca:,:nocca,:nocca] = dm2aa[:nocca,:nocca,nocca:,nocca:].transpose(2,3,0,1)
    dm2aa[:nocca,nocca:,nocca:,:nocca] = dovvo
    dm2aa[nocca:,:nocca,:nocca,nocca:] = dm2aa[:nocca,nocca:,nocca:,:nocca].transpose(1,0,3,2).conj()
    dovvo = None

    if len(dvvvv.shape) == 2:
        dvvvv = ao2mo.restore(1, dvvvv, nvira)
    dm2aa[nocca:,nocca:,nocca:,nocca:] = dvvvv
    dm2aa[:nocca,:nocca,:nocca,:nocca] = doooo

    dovvv = numpy.asarray(dovvv)
    dm2aa[:nocca,nocca:,nocca:,nocca:] = dovvv
    dm2aa[nocca:,nocca:,:nocca,nocca:] = dovvv.transpose(2,3,0,1)
    dm2aa[nocca:,nocca:,nocca:,:nocca] = dovvv.transpose(3,2,1,0).conj()
    dm2aa[nocca:,:nocca,nocca:,nocca:] = dovvv.transpose(1,0,3,2).conj()
    dovvv = None

    dooov = numpy.asarray(dooov)
    dm2aa[:nocca,:nocca,:nocca,nocca:] = dooov
    dm2aa[:nocca,nocca:,:nocca,:nocca] = dooov.transpose(2,3,0,1)
    dm2aa[:nocca,:nocca,nocca:,:nocca] = dooov.transpose(1,0,3,2).conj()
    dm2aa[nocca:,:nocca,:nocca,:nocca] = dooov.transpose(3,2,1,0).conj()
    dooov = None

# dm2bb
    dOVOV = numpy.asarray(dOVOV)
    dm2bb[:noccb,noccb:,:noccb,noccb:] = dOVOV
    dm2bb[noccb:,:noccb,noccb:,:noccb] = dm2bb[:noccb,noccb:,:noccb,noccb:].transpose(1,0,3,2).conj()
    dOVOV = None

    dOVVO = numpy.asarray(dOVVO)
    dm2bb[:noccb,:noccb,noccb:,noccb:] =-dOVVO.transpose(0,3,2,1)
    dm2bb[noccb:,noccb:,:noccb,:noccb] = dm2bb[:noccb,:noccb,noccb:,noccb:].transpose(2,3,0,1)
    dm2bb[:noccb,noccb:,noccb:,:noccb] = dOVVO
    dm2bb[noccb:,:noccb,:noccb,noccb:] = dm2bb[:noccb,noccb:,noccb:,:noccb].transpose(1,0,3,2).conj()
    dOVVO = None

    if len(dVVVV.shape) == 2:
        dVVVV = ao2mo.restore(1, dVVVV, nvirb)
    dm2bb[noccb:,noccb:,noccb:,noccb:] = dVVVV
    dm2bb[:noccb,:noccb,:noccb,:noccb] = dOOOO

    dOVVV = numpy.asarray(dOVVV)
    dm2bb[:noccb,noccb:,noccb:,noccb:] = dOVVV
    dm2bb[noccb:,noccb:,:noccb,noccb:] = dOVVV.transpose(2,3,0,1)
    dm2bb[noccb:,noccb:,noccb:,:noccb] = dOVVV.transpose(3,2,1,0).conj()
    dm2bb[noccb:,:noccb,noccb:,noccb:] = dOVVV.transpose(1,0,3,2).conj()
    dOVVV = None

    dOOOV = numpy.asarray(dOOOV)
    dm2bb[:noccb,:noccb,:noccb,noccb:] = dOOOV
    dm2bb[:noccb,noccb:,:noccb,:noccb] = dOOOV.transpose(2,3,0,1)
    dm2bb[:noccb,:noccb,noccb:,:noccb] = dOOOV.transpose(1,0,3,2).conj()
    dm2bb[noccb:,:noccb,:noccb,:noccb] = dOOOV.transpose(3,2,1,0).conj()
    dOOOV = None

# dm2ab
    dovOV = numpy.asarray(dovOV)
    dm2ab[:nocca,nocca:,:noccb,noccb:] = dovOV
    dm2ab[nocca:,:nocca,noccb:,:noccb] = dm2ab[:nocca,nocca:,:noccb,noccb:].transpose(1,0,3,2).conj()
    dovOV = None

    dovVO = numpy.asarray(dovVO)
    dm2ab[:nocca,:nocca,noccb:,noccb:] = dooVV
    dm2ab[nocca:,nocca:,:noccb,:noccb] = dOOvv.transpose(2,3,0,1)
    dm2ab[:nocca,nocca:,noccb:,:noccb] = dovVO
    dm2ab[nocca:,:nocca,:noccb,noccb:] = dovVO.transpose(1,0,3,2).conj()
    dovVO = None

    if len(dvvVV.shape) == 2:
        idxa = numpy.tril_indices(nvira)
        dvvVV1 = lib.unpack_tril(dvvVV)
        dvvVV = numpy.empty((nvira,nvira,nvirb,nvirb))
        dvvVV[idxa] = dvvVV1
        dvvVV[idxa[1],idxa[0]] = dvvVV1
        dvvVV1 = None
    dm2ab[nocca:,nocca:,noccb:,noccb:] = dvvVV
    dm2ab[:nocca,:nocca,:noccb,:noccb] = dooOO

    dovVV = numpy.asarray(dovVV)
    dm2ab[:nocca,nocca:,noccb:,noccb:] = dovVV
    dm2ab[nocca:,nocca:,:noccb,noccb:] = dOVvv.transpose(2,3,0,1)
    dm2ab[nocca:,nocca:,noccb:,:noccb] = dOVvv.transpose(3,2,1,0).conj()
    dm2ab[nocca:,:nocca,noccb:,noccb:] = dovVV.transpose(1,0,3,2).conj()
    dovVV = None

    dooOV = numpy.asarray(dooOV)
    dm2ab[:nocca,:nocca,:noccb,noccb:] = dooOV
    dm2ab[:nocca,nocca:,:noccb,:noccb] = dOOov.transpose(2,3,0,1)
    dm2ab[:nocca,:nocca,noccb:,:noccb] = dooOV.transpose(1,0,3,2).conj()
    dm2ab[nocca:,:nocca,:noccb,:noccb] = dOOov.transpose(3,2,1,0).conj()
    dooOV = None

    if with_frozen and mycc.frozen is not None:
        nmoa0 = dm2aa.shape[0]
        nmob0 = dm2bb.shape[0]
        nmoa = mycc.mo_occ[0].size
        nmob = mycc.mo_occ[1].size
        nocca = numpy.count_nonzero(mycc.mo_occ[0] > 0)
        noccb = numpy.count_nonzero(mycc.mo_occ[1] > 0)

        rdm2aa = numpy.zeros((nmoa,nmoa,nmoa,nmoa), dtype=dm2aa.dtype)
        rdm2ab = numpy.zeros((nmoa,nmoa,nmob,nmob), dtype=dm2ab.dtype)
        rdm2bb = numpy.zeros((nmob,nmob,nmob,nmob), dtype=dm2bb.dtype)
        moidxa, moidxb = mycc.get_frozen_mask()
        moidxa = numpy.where(moidxa)[0]
        moidxb = numpy.where(moidxb)[0]
        idxa = (moidxa.reshape(-1,1) * nmoa + moidxa).ravel()
        idxb = (moidxb.reshape(-1,1) * nmob + moidxb).ravel()
        lib.takebak_2d(rdm2aa.reshape(nmoa**2,nmoa**2),
                       dm2aa.reshape(nmoa0**2,nmoa0**2), idxa, idxa)
        lib.takebak_2d(rdm2bb.reshape(nmob**2,nmob**2),
                       dm2bb.reshape(nmob0**2,nmob0**2), idxb, idxb)
        lib.takebak_2d(rdm2ab.reshape(nmoa**2,nmob**2),
                       dm2ab.reshape(nmoa0**2,nmob0**2), idxa, idxb)
        dm2aa, dm2ab, dm2bb = rdm2aa, rdm2ab, rdm2bb

    if with_dm1:
        dm1a, dm1b = _make_rdm1(mycc, d1, with_frozen=with_frozen)
        dm1a[numpy.diag_indices(nocca)] -= 1
        dm1b[numpy.diag_indices(noccb)] -= 1

        for i in range(nocca):
            dm2aa[i,i,:,:] += dm1a
            dm2aa[:,:,i,i] += dm1a
            dm2aa[:,i,i,:] -= dm1a
            dm2aa[i,:,:,i] -= dm1a.T
            dm2ab[i,i,:,:] += dm1b
        for i in range(noccb):
            dm2bb[i,i,:,:] += dm1b
            dm2bb[:,:,i,i] += dm1b
            dm2bb[:,i,i,:] -= dm1b
            dm2bb[i,:,:,i] -= dm1b.T
            dm2ab[:,:,i,i] += dm1a

        for i in range(nocca):
            for j in range(nocca):
                dm2aa[i,i,j,j] += 1
                dm2aa[i,j,j,i] -= 1
        for i in range(noccb):
            for j in range(noccb):
                dm2bb[i,i,j,j] += 1
                dm2bb[i,j,j,i] -= 1
        for i in range(nocca):
            for j in range(noccb):
                dm2ab[i,i,j,j] += 1

    dm2aa = dm2aa.transpose(1,0,3,2)
    dm2ab = dm2ab.transpose(1,0,3,2)
    dm2bb = dm2bb.transpose(1,0,3,2)

    if ao_repr:
        from pyscf.cc import ccsd_rdm
        dm2aa = ccsd_rdm._rdm2_mo2ao(dm2aa, mycc.mo_coeff[0])
        dm2bb = ccsd_rdm._rdm2_mo2ao(dm2bb, mycc.mo_coeff[1])
        dm2ab = _dm2ab_mo2ao(dm2ab, mycc.mo_coeff[0], mycc.mo_coeff[1])
    return dm2aa, dm2ab, dm2bb


def _dm2ab_mo2ao(dm2, mo_a, mo_b):
    return lib.einsum('ijkl,pi,qj,rk,sl->pqrs', dm2, mo_a, mo_a.conj(),
                      mo_b, mo_b.conj())


if __name__ == '__main__':
    from functools import reduce
    from pyscf import gto
    from pyscf import scf
    from pyscf.cc import uccsd

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = '631g'
    mol.spin = 2
    mol.build()
    mf = scf.UHF(mol).run()

    mycc = uccsd.UCCSD(mf)
    mycc.frozen = 2
    ecc, t1, t2 = mycc.kernel()
    l1, l2 = mycc.solve_lambda()
    dm1a,dm1b = make_rdm1(mycc, t1, t2, l1, l2)
    dm2aa,dm2ab,dm2bb = make_rdm2(mycc, t1, t2, l1, l2)
    mo_a = mf.mo_coeff[0]
    mo_b = mf.mo_coeff[1]
    nmoa = mo_a.shape[1]
    nmob = mo_b.shape[1]
    eriaa = ao2mo.kernel(mf._eri, mo_a, compact=False).reshape([nmoa]*4)
    eribb = ao2mo.kernel(mf._eri, mo_b, compact=False).reshape([nmob]*4)
    eriab = ao2mo.kernel(mf._eri, (mo_a,mo_a,mo_b,mo_b), compact=False)
    eriab = eriab.reshape([nmoa,nmoa,nmob,nmob])
    hcore = mf.get_hcore()
    h1a = reduce(numpy.dot, (mo_a.T.conj(), hcore, mo_a))
    h1b = reduce(numpy.dot, (mo_b.T.conj(), hcore, mo_b))
    e1 = numpy.einsum('ij,ji', h1a, dm1a)
    e1+= numpy.einsum('ij,ji', h1b, dm1b)
    e1+= numpy.einsum('ijkl,ijkl', eriaa, dm2aa) * .5
    e1+= numpy.einsum('ijkl,ijkl', eriab, dm2ab)
    e1+= numpy.einsum('ijkl,ijkl', eribb, dm2bb) * .5
    e1+= mol.energy_nuc()
    print(e1 - mycc.e_tot)

    from pyscf.fci import direct_uhf
    mol = gto.Mole()
    mol.verbose = 0
    mol.atom = [
        ['H', ( 1.,-1.    , 0.   )],
        ['H', ( 0.,-1.    ,-1.   )],
        ['H', ( 1.,-0.5   , 0.   )],
        ['H', ( 0.,-1.    , 1.   )],
    ]
    mol.charge = 2
    mol.spin = 2
    mol.basis = '6-31g'
    mol.build()
    mf = scf.UHF(mol).run(init_guess='hcore', conv_tol=1.)
    ehf0 = mf.e_tot - mol.energy_nuc()
    mycc = uccsd.UCCSD(mf).run()
    mycc.solve_lambda()
    eri_aa = ao2mo.kernel(mf._eri, mf.mo_coeff[0])
    eri_bb = ao2mo.kernel(mf._eri, mf.mo_coeff[1])
    eri_ab = ao2mo.kernel(mf._eri, [mf.mo_coeff[0], mf.mo_coeff[0],
                                    mf.mo_coeff[1], mf.mo_coeff[1]])
    h1a = reduce(numpy.dot, (mf.mo_coeff[0].T, mf.get_hcore(), mf.mo_coeff[0]))
    h1b = reduce(numpy.dot, (mf.mo_coeff[1].T, mf.get_hcore(), mf.mo_coeff[1]))
    efci, fcivec = direct_uhf.kernel((h1a,h1b), (eri_aa,eri_ab,eri_bb),
                                     h1a.shape[0], mol.nelec)
    dm1ref, dm2ref = direct_uhf.make_rdm12s(fcivec, h1a.shape[0], mol.nelec)
    t1, t2 = mycc.t1, mycc.t2
    l1, l2 = mycc.l1, mycc.l2
    rdm1 = make_rdm1(mycc, t1, t2, l1, l2)
    rdm2 = make_rdm2(mycc, t1, t2, l1, l2)
    print('dm1a', abs(dm1ref[0] - rdm1[0]).max())
    print('dm1b', abs(dm1ref[1] - rdm1[1]).max())
    print('dm2aa', abs(dm2ref[0] - rdm2[0]).max())
    print('dm2ab', abs(dm2ref[1] - rdm2[1]).max())
    print('dm2bb', abs(dm2ref[2] - rdm2[2]).max())


# pyscf/cc/uccsd_slow.py

#!/usr/bin/env python
# Copyright 2014-2021 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.

from functools import reduce

import numpy
import numpy as np

from pyscf import lib
from pyscf import ao2mo
from pyscf.lib import logger
from pyscf.cc import ccsd
from pyscf.cc import uccsd
from pyscf.mp import ump2
from pyscf.cc import gintermediates as imd

#einsum = np.einsum
einsum = lib.einsum

# This is unrestricted (U)CCSD, i.e. spin-orbital form.


def update_amps(cc, t1, t2, eris):
    nocc, nvir = t1.shape
    fock = eris.fock

    fov = fock[:nocc,nocc:]
    foo = fock[:nocc,:nocc]
    fvv = fock[nocc:,nocc:]

    tau = imd.make_tau(t2, t1, t1)

    Fvv = imd.cc_Fvv(t1, t2, eris)
    Foo = imd.cc_Foo(t1, t2, eris)
    Fov = imd.cc_Fov(t1, t2, eris)
    Woooo = imd.cc_Woooo(t1, t2, eris)
    Wvvvv = imd.cc_Wvvvv(t1, t2, eris)
    Wovvo = imd.cc_Wovvo(t1, t2, eris)

    # Move energy terms to the other side
    Fvv -= np.diag(np.diag(fvv))
    Foo -= np.diag(np.diag(foo))

    # T1 equation
    t1new = np.array(fov).conj()
    t1new +=  einsum('ie,ae->ia', t1, Fvv)
    t1new += -einsum('ma,mi->ia', t1, Foo)
    t1new +=  einsum('imae,me->ia', t2, Fov)
    t1new += -einsum('nf,naif->ia', t1, eris.ovov)
    t1new += -0.5*einsum('imef,maef->ia', t2, eris.ovvv)
    t1new += -0.5*einsum('mnae,mnie->ia', t2, eris.ooov)

    # T2 equation
    t2new = np.array(eris.oovv).conj()
    Ftmp = Fvv - 0.5*einsum('mb,me->be', t1, Fov)
    tmp = einsum('ijae,be->ijab', t2, Ftmp)
    t2new += (tmp - tmp.transpose(0,1,3,2))
    Ftmp = Foo + 0.5*einsum('je,me->mj', t1, Fov)
    tmp = einsum('imab,mj->ijab', t2, Ftmp)
    t2new -= (tmp - tmp.transpose(1,0,2,3))
    t2new += 0.5*einsum('mnab,mnij->ijab', tau, Woooo)
    t2new += 0.5*einsum('ijef,abef->ijab', tau, Wvvvv)
    tmp = einsum('imae,mbej->ijab', t2, Wovvo)
    tmp -= -einsum('ie,ma,mbje->ijab', t1, t1, eris.ovov)
    t2new += (tmp - tmp.transpose(0,1,3,2)
              - tmp.transpose(1,0,2,3) + tmp.transpose(1,0,3,2) )
    tmp = einsum('ie,jeba->ijab', t1, np.array(eris.ovvv).conj())
    t2new += (tmp - tmp.transpose(1,0,2,3))
    tmp = einsum('ma,mbij->ijab', t1, eris.ovoo)
    t2new -= (tmp - tmp.transpose(0,1,3,2))

    mo_e = eris.fock.diagonal()
    eia = mo_e[:nocc,None] - mo_e[None,nocc:] - cc.level_shift
    eijab = lib.direct_sum('ia,jb->ijab', eia, eia)
    t1new /= eia
    t2new /= eijab

    return t1new, t2new


def energy(cc, t1, t2, eris):
    nocc, nvir = t1.shape
    fock = eris.fock
    e = einsum('ia,ia', fock[:nocc,nocc:], t1)
    e += 0.25*np.einsum('ijab,ijab', t2, eris.oovv)
    e += 0.5 *np.einsum('ia,jb,ijab', t1, t1, eris.oovv)
    return e.real


get_frozen_mask = ump2.get_frozen_mask


class UCCSD(ccsd.CCSD):

    def __init__(self, mf, frozen=None, mo_coeff=None, mo_occ=None):
        ccsd.CCSD.__init__(self, mf, frozen, mo_coeff, mo_occ)
        # Spin-orbital CCSD needs a stricter tolerance than spatial-orbital
        self.conv_tol_normt = 1e-6

    @property
    def nocc(self):
        nocca, noccb = self.get_nocc()
        return nocca + noccb

    @property
    def nmo(self):
        nmoa, nmob = self.get_nmo()
        return nmoa + nmob

    get_nocc = uccsd.get_nocc
    get_nmo = uccsd.get_nmo
    get_frozen_mask = get_frozen_mask

    def init_amps(self, eris):
        time0 = logger.process_clock(), logger.perf_counter()
        mo_e = eris.fock.diagonal()
        nocc = self.nocc
        eia = mo_e[:nocc,None] - mo_e[None,nocc:]
        eijab = lib.direct_sum('ia,jb->ijab',eia,eia)
        t1 = eris.fock[:nocc,nocc:] / eia
        eris_oovv = np.array(eris.oovv)
        t2 = eris_oovv/eijab
        self.emp2 = 0.25*einsum('ijab,ijab',t2,eris_oovv.conj()).real
        logger.info(self, 'Init t2, MP2 energy = %.15g', self.emp2)
        logger.timer(self, 'init mp2', *time0)
        return self.emp2, t1, t2

    energy = energy

    def kernel(self, t1=None, t2=None, eris=None, mbpt2=False):
        return self.ccsd(t1, t2, eris, mbpt2)
    def ccsd(self, t1=None, t2=None, eris=None, mbpt2=False):
        '''Ground-state unrestricted (U)CCSD.

        Kwargs:
            mbpt2 : bool
                Use one-shot MBPT2 approximation to CCSD.
        '''
        if mbpt2:
            #cctyp = 'MBPT2'
            #self.e_corr, self.t1, self.t2 = self.init_amps(eris)
            raise NotImplementedError

        if eris is None:
            eris = self.ao2mo(self.mo_coeff)
        self.eris = eris
        return ccsd.CCSD.ccsd(self, t1, t2, eris)

    def ao2mo(self, mo_coeff=None):
        return _PhysicistsERIs(self, mo_coeff)

    def update_amps(self, t1, t2, eris):
        return update_amps(self, t1, t2, eris)

    def nip(self):
        nocc = self.nocc
        nvir = self.nmo - nocc
        self._nip = nocc + nocc*(nocc-1)//2*nvir
        return self._nip

    def nea(self):
        nocc = self.nocc
        nvir = self.nmo - nocc
        self._nea = nvir + nocc*nvir*(nvir-1)//2
        return self._nea

    def nee(self):
        nocc = self.nocc
        nvir = self.nmo - nocc
        self._nee = nocc*nvir + nocc*(nocc-1)//2*nvir*(nvir-1)//2
        return self._nee

    def ipccsd_matvec(self, vector):
        # Ref: Tu, Wang, and Li, J. Chem. Phys. 136, 174102 (2012) Eqs.(8)-(9)
        if not getattr(self, 'imds', None):
            self.imds = _IMDS(self)
        if not self.imds.made_ip_imds:
            self.imds.make_ip()
        imds = self.imds

        r1,r2 = self.vector_to_amplitudes_ip(vector)

        # Eq. (8)
        Hr1 = -einsum('mi,m->i',imds.Foo,r1)
        Hr1 += einsum('me,mie->i',imds.Fov,r2)
        Hr1 += -0.5*einsum('nmie,mne->i',imds.Wooov,r2)
        # Eq. (9)
        Hr2 =  einsum('ae,ije->ija',imds.Fvv,r2)
        tmp1 = einsum('mi,mja->ija',imds.Foo,r2)
        Hr2 += (-tmp1 + tmp1.transpose(1,0,2))
        Hr2 += -einsum('maji,m->ija',imds.Wovoo,r1)
        Hr2 += 0.5*einsum('mnij,mna->ija',imds.Woooo,r2)
        tmp2 = einsum('maei,mje->ija',imds.Wovvo,r2)
        Hr2 += (tmp2 - tmp2.transpose(1,0,2))
        Hr2 += 0.5*einsum('mnef,ijae,mnf->ija',imds.Woovv,self.t2,r2)

        vector = self.amplitudes_to_vector_ip(Hr1,Hr2)
        return vector

    def ipccsd_diag(self):
        if not getattr(self, 'imds', None):
            self.imds = _IMDS(self)
        if not self.imds.made_ip_imds:
            self.imds.make_ip()
        imds = self.imds

        t1, t2 = self.t1, self.t2
        nocc, nvir = t1.shape

        Hr1 = -np.diag(imds.Foo)
        Hr2 = np.zeros((nocc,nocc,nvir),dtype=t1.dtype)
        for i in range(nocc):
            for j in range(nocc):
                for a in range(nvir):
                    Hr2[i,j,a] += imds.Fvv[a,a]
                    Hr2[i,j,a] += -imds.Foo[i,i]
                    Hr2[i,j,a] += -imds.Foo[j,j]
                    Hr2[i,j,a] += 0.5*(imds.Woooo[i,j,i,j]-imds.Woooo[j,i,i,j])
                    Hr2[i,j,a] += imds.Wovvo[i,a,a,i]
                    Hr2[i,j,a] += imds.Wovvo[j,a,a,j]
                    Hr2[i,j,a] += 0.5*(np.dot(imds.Woovv[i,j,:,a],t2[i,j,a,:])
                                       -np.dot(imds.Woovv[j,i,:,a],t2[i,j,a,:]))

        vector = self.amplitudes_to_vector_ip(Hr1,Hr2)
        return vector

    def vector_to_amplitudes_ip(self,vector):
        nocc = self.nocc
        nvir = self.nmo - nocc
        r1 = vector[:nocc].copy()
        r2 = np.zeros((nocc,nocc,nvir), vector.dtype)
        index = nocc
        for i in range(nocc):
            for j in range(i):
                for a in range(nvir):
                    r2[i,j,a] =  vector[index]
                    r2[j,i,a] = -vector[index]
                    index += 1
        return [r1,r2]

    def amplitudes_to_vector_ip(self,r1,r2):
        nocc = self.nocc
        nvir = self.nmo - nocc
        size = nocc + nocc*(nocc-1)//2*nvir
        vector = np.zeros(size, r1.dtype)
        vector[:nocc] = r1.copy()
        index = nocc
        for i in range(nocc):
            for j in range(i):
                for a in range(nvir):
                    vector[index] = r2[i,j,a]
                    index += 1
        return vector

    def eaccsd_matvec(self,vector):
        # Ref: Nooijen and Bartlett, J. Chem. Phys. 102, 3629 (1994) Eqs.(30)-(31)
        if not getattr(self, 'imds', None):
            self.imds = _IMDS(self)
        if not self.imds.made_ea_imds:
            self.imds.make_ea()
        imds = self.imds

        r1,r2 = self.vector_to_amplitudes_ea(vector)

        # Eq. (30)
        Hr1 = einsum('ac,c->a',imds.Fvv,r1)
        Hr1 += einsum('ld,lad->a',imds.Fov,r2)
        Hr1 += 0.5*einsum('alcd,lcd->a',imds.Wvovv,r2)
        # Eq. (31)
        Hr2 = einsum('abcj,c->jab',imds.Wvvvo,r1)
        tmp1 = einsum('ac,jcb->jab',imds.Fvv,r2)
        Hr2 += (tmp1 - tmp1.transpose(0,2,1))
        Hr2 += -einsum('lj,lab->jab',imds.Foo,r2)
        tmp2 = einsum('lbdj,lad->jab',imds.Wovvo,r2)
        Hr2 += (tmp2 - tmp2.transpose(0,2,1))
        nvir = self.nmo-self.nocc
        for a in range(nvir):
            Hr2[:,a,:] += 0.5*einsum('bcd,jcd->jb',imds.Wvvvv[a],r2)
        Hr2 += -0.5*einsum('klcd,lcd,kjab->jab',imds.Woovv,r2,self.t2)

        vector = self.amplitudes_to_vector_ea(Hr1,Hr2)
        return vector

    def eaccsd_diag(self):
        if not getattr(self, 'imds', None):
            self.imds = _IMDS(self)
        if not self.imds.made_ea_imds:
            self.imds.make_ea()
        imds = self.imds

        t1, t2 = self.t1, self.t2
        nocc, nvir = t1.shape

        Hr1 = np.diag(imds.Fvv)
        Hr2 = np.zeros((nocc,nvir,nvir),dtype=t1.dtype)
        for a in range(nvir):
            _Wvvvva = np.array(imds.Wvvvv[a])
            for b in range(a):
                for j in range(nocc):
                    Hr2[j,a,b] += imds.Fvv[a,a]
                    Hr2[j,a,b] += imds.Fvv[b,b]
                    Hr2[j,a,b] += -imds.Foo[j,j]
                    Hr2[j,a,b] += imds.Wovvo[j,b,b,j]
                    Hr2[j,a,b] += imds.Wovvo[j,a,a,j]
                    Hr2[j,a,b] += 0.5*(_Wvvvva[b,a,b]-_Wvvvva[b,b,a])
                    Hr2[j,a,b] -= 0.5*(np.dot(imds.Woovv[:,j,a,b],t2[:,j,a,b]) -
                                       np.dot(imds.Woovv[:,j,b,a],t2[:,j,a,b]))

        vector = self.amplitudes_to_vector_ea(Hr1,Hr2)
        return vector

    def vector_to_amplitudes_ea(self,vector):
        nocc = self.nocc
        nvir = self.nmo - nocc
        r1 = vector[:nvir].copy()
        r2 = np.zeros((nocc,nvir,nvir), vector.dtype)
        index = nvir
        for i in range(nocc):
            for a in range(nvir):
                for b in range(a):
                    r2[i,a,b] =  vector[index]
                    r2[i,b,a] = -vector[index]
                    index += 1
        return [r1,r2]

    def amplitudes_to_vector_ea(self,r1,r2):
        nocc = self.nocc
        nvir = self.nmo - nocc
        size = nvir + nvir*(nvir-1)//2*nocc
        vector = np.zeros(size, r1.dtype)
        vector[:nvir] = r1.copy()
        index = nvir
        for i in range(nocc):
            for a in range(nvir):
                for b in range(a):
                    vector[index] = r2[i,a,b]
                    index += 1
        return vector

    def eeccsd_matvec(self,vector):
        # Ref: Wang, Tu, and Wang, J. Chem. Theory Comput. 10, 5567 (2014) Eqs.(9)-(10)
        # Note: Last line in Eq. (10) is superfluous.
        # See, e.g. Gwaltney, Nooijen, and Barlett, Chem. Phys. Lett. 248, 189 (1996)
        if not getattr(self, 'imds', None):
            self.imds = _IMDS(self)
        if not self.imds.made_ee_imds:
            self.imds.make_ee()
        imds = self.imds

        r1,r2 = self.vector_to_amplitudes_ee(vector)

        # Eq. (9)
        Hr1 = einsum('ae,ie->ia',imds.Fvv,r1)
        Hr1 += -einsum('mi,ma->ia',imds.Foo,r1)
        Hr1 += einsum('me,imae->ia',imds.Fov,r2)
        Hr1 += einsum('maei,me->ia',imds.Wovvo,r1)
        Hr1 += -0.5*einsum('mnie,mnae->ia',imds.Wooov,r2)
        Hr1 += 0.5*einsum('amef,imef->ia',imds.Wvovv,r2)
        # Eq. (10)
        tmpab = einsum('be,ijae->ijab',imds.Fvv,r2)
        tmpab += -0.5*einsum('mnef,ijae,mnbf->ijab',imds.Woovv,self.t2,r2)
        tmpab += -einsum('mbij,ma->ijab',imds.Wovoo,r1)
        tmpab += -einsum('amef,ijfb,me->ijab',imds.Wvovv,self.t2,r1)
        tmpij = -einsum('mj,imab->ijab',imds.Foo,r2)
        tmpij += -0.5*einsum('mnef,imab,jnef->ijab',imds.Woovv,self.t2,r2)
        tmpij += einsum('abej,ie->ijab',imds.Wvvvo,r1)
        tmpij += einsum('mnie,njab,me->ijab',imds.Wooov,self.t2,r1)

        tmpabij = einsum('mbej,imae->ijab',imds.Wovvo,r2)

        Hr2 = (tmpab - tmpab.transpose(0,1,3,2)
               + tmpij - tmpij.transpose(1,0,2,3)
               + 0.5*einsum('mnij,mnab->ijab',imds.Woooo,r2)
               + 0.5*einsum('abef,ijef->ijab',imds.Wvvvv,r2)
               + tmpabij - tmpabij.transpose(0,1,3,2)
               - tmpabij.transpose(1,0,2,3) + tmpabij.transpose(1,0,3,2) )

        vector = self.amplitudes_to_vector_ee(Hr1,Hr2)
        return vector

    def eeccsd_diag(self):
        if not getattr(self, 'imds', None):
            self.imds = _IMDS(self)
        if not self.imds.made_ee_imds:
            self.imds.make_ee()
        imds = self.imds

        t1, t2 = self.t1, self.t2
        nocc, nvir = t1.shape

        Hr1 = np.zeros((nocc,nvir), dtype=t1.dtype)
        Hr2 = np.zeros((nocc,nocc,nvir,nvir), dtype=t1.dtype)
        for i in range(nocc):
            for a in range(nvir):
                Hr1[i,a] = imds.Fvv[a,a] - imds.Foo[i,i] + imds.Wovvo[i,a,a,i]
        for a in range(nvir):
            tmp = 0.5*(np.einsum('ijeb,ijbe->ijb', imds.Woovv, t2) -
                       np.einsum('jieb,ijbe->ijb', imds.Woovv, t2))
            Hr2[:,:,:,a] += imds.Fvv[a,a] + tmp
            Hr2[:,:,a,:] += imds.Fvv[a,a] + tmp
            _Wvvvva = np.array(imds.Wvvvv[a])
            for b in range(a):
                Hr2[:,:,a,b] += 0.5*(_Wvvvva[b,a,b]-_Wvvvva[b,b,a])
            for i in range(nocc):
                tmp = imds.Wovvo[i,a,a,i]
                Hr2[:,i,:,a] += tmp
                Hr2[i,:,:,a] += tmp
                Hr2[:,i,a,:] += tmp
                Hr2[i,:,a,:] += tmp
        for i in range(nocc):
            tmp = 0.5*(np.einsum('kjab,jkab->jab', imds.Woovv, t2) -
                       np.einsum('kjba,jkab->jab', imds.Woovv, t2))
            Hr2[:,i,:,:] += -imds.Foo[i,i] + tmp
            Hr2[i,:,:,:] += -imds.Foo[i,i] + tmp
            for j in range(i):
                Hr2[i,j,:,:] += 0.5*(imds.Woooo[i,j,i,j]-imds.Woooo[j,i,i,j])

        vector = self.amplitudes_to_vector_ee(Hr1,Hr2)
        return vector

    def vector_to_amplitudes_ee(self,vector):
        nocc = self.nocc
        nvir = self.nmo - nocc
        r1 = vector[:nocc*nvir].copy().reshape((nocc,nvir))
        r2 = np.zeros((nocc,nocc,nvir,nvir), vector.dtype)
        index = nocc*nvir
        for i in range(nocc):
            for j in range(i):
                for a in range(nvir):
                    for b in range(a):
                        r2[i,j,a,b] =  vector[index]
                        r2[j,i,a,b] = -vector[index]
                        r2[i,j,b,a] = -vector[index]
                        r2[j,i,b,a] =  vector[index]
                        index += 1
        return [r1,r2]

    def amplitudes_to_vector_ee(self,r1,r2):
        nocc = self.nocc
        nvir = self.nmo - nocc
        size = nocc*nvir + nocc*(nocc-1)//2*nvir*(nvir-1)//2
        vector = np.zeros(size, r1.dtype)
        vector[:nocc*nvir] = r1.copy().reshape(nocc*nvir)
        index = nocc*nvir
        for i in range(nocc):
            for j in range(i):
                for a in range(nvir):
                    for b in range(a):
                        vector[index] = r2[i,j,a,b]
                        index += 1
        return vector

    def amplitudes_from_rccsd(self, t1, t2):
        '''Convert spatial orbital T1,T2 to spin-orbital T1,T2'''
        nocc, nvir = t1.shape
        nocc2 = nocc * 2
        nvir2 = nvir * 2
        t1s = np.zeros((nocc2,nvir2))
        t1s[:nocc,:nvir] = t1
        t1s[nocc:,nvir:] = t1

        t2s = np.zeros((nocc2,nocc2,nvir2,nvir2))
        t2s[:nocc,nocc:,:nvir,nvir:] = t2
        t2s[nocc:,:nocc,nvir:,:nvir] = t2
        t2s[:nocc,nocc:,nvir:,:nvir] =-t2.transpose(0,1,3,2)
        t2s[nocc:,:nocc,:nvir,nvir:] =-t2.transpose(0,1,3,2)
        t2s[:nocc,:nocc,:nvir,:nvir] = t2 - t2.transpose(0,1,3,2)
        t2s[nocc:,nocc:,nvir:,nvir:] = t2 - t2.transpose(0,1,3,2)
        return t1s, t2s


class _PhysicistsERIs:
    def __init__(self, cc, mo_coeff=None, method='incore',
                 ao2mofn=ao2mo.outcore.general_iofree):
        cput0 = (logger.process_clock(), logger.perf_counter())
        moidx = cc.get_frozen_mask()
        if mo_coeff is None:
            self.mo_coeff = mo_coeff = [cc.mo_coeff[0][:,moidx[0]],
                                        cc.mo_coeff[1][:,moidx[1]]]
        else:
            self.mo_coeff = mo_coeff = [mo_coeff[0][:,moidx[0]],
                                        mo_coeff[1][:,moidx[1]]]

        nocc = cc.nocc
        nmo = cc.nmo
        nvir = nmo - nocc
        mem_incore = nmo**4*2 * 8/1e6
        mem_now = lib.current_memory()[0]

        self.fock, so_coeff, spin = uspatial2spin(cc, moidx, mo_coeff)

        log = logger.Logger(cc.stdout, cc.verbose)
        if (method == 'incore' and (mem_incore+mem_now < cc.max_memory)
            or cc.mol.incore_anyway):
            eri = ao2mofn(cc._scf.mol, (so_coeff,so_coeff,so_coeff,so_coeff), compact=0)
            if mo_coeff[0].dtype == np.float: eri = eri.real
            eri = eri.reshape((nmo,)*4)
            for i in range(nmo):
                for j in range(i):
                    if spin[i] != spin[j]:
                        eri[i,j,:,:] = eri[j,i,:,:] = 0.
                        eri[:,:,i,j] = eri[:,:,j,i] = 0.
            eri1 = eri - eri.transpose(0,3,2,1)
            eri1 = eri1.transpose(0,2,1,3)

            self.dtype = eri1.dtype
            self.oooo = eri1[:nocc,:nocc,:nocc,:nocc].copy()
            self.ooov = eri1[:nocc,:nocc,:nocc,nocc:].copy()
            self.ovoo = eri1[:nocc,nocc:,:nocc,:nocc].copy()
            self.oovv = eri1[:nocc,:nocc,nocc:,nocc:].copy()
            self.ovov = eri1[:nocc,nocc:,:nocc,nocc:].copy()
            self.ovvv = eri1[:nocc,nocc:,nocc:,nocc:].copy()
            self.vvvv = eri1[nocc:,nocc:,nocc:,nocc:].copy()
        else:
            self.feri1 = lib.H5TmpFile()
            orbo = so_coeff[:,:nocc]
            orbv = so_coeff[:,nocc:]
            if mo_coeff[0].dtype == np.complex128: ds_type = 'c16'
            else: ds_type = 'f8'
            self.oooo = self.feri1.create_dataset('oooo', (nocc,nocc,nocc,nocc), ds_type)
            self.ooov = self.feri1.create_dataset('ooov', (nocc,nocc,nocc,nvir), ds_type)
            self.ovoo = self.feri1.create_dataset('ovoo', (nocc,nvir,nocc,nocc), ds_type)
            self.oovv = self.feri1.create_dataset('oovv', (nocc,nocc,nvir,nvir), ds_type)
            self.ovov = self.feri1.create_dataset('ovov', (nocc,nvir,nocc,nvir), ds_type)
            self.ovvv = self.feri1.create_dataset('ovvv', (nocc,nvir,nvir,nvir), ds_type)
            self.vvvv = self.feri1.create_dataset('vvvv', (nvir,nvir,nvir,nvir), ds_type)

            cput1 = logger.process_clock(), logger.perf_counter()
            # <ij||pq> = <ij|pq> - <ij|qp> = (ip|jq) - (iq|jp)
            buf = ao2mofn(cc._scf.mol, (orbo,so_coeff,orbo,so_coeff), compact=0)
            if mo_coeff[0].dtype == np.float: buf = buf.real
            buf = buf.reshape((nocc,nmo,nocc,nmo))
            for i in range(nocc):
                for p in range(nmo):
                    if spin[i] != spin[p]:
                        buf[i,p,:,:] = 0.
                        buf[:,:,i,p] = 0.
            buf1 = buf - buf.transpose(0,3,2,1)
            buf1 = buf1.transpose(0,2,1,3)
            cput1 = log.timer_debug1('transforming oopq', *cput1)
            self.dtype = buf1.dtype
            self.oooo[:,:,:,:] = buf1[:,:,:nocc,:nocc]
            self.ooov[:,:,:,:] = buf1[:,:,:nocc,nocc:]
            self.oovv[:,:,:,:] = buf1[:,:,nocc:,nocc:]

            cput1 = logger.process_clock(), logger.perf_counter()
            # <ia||pq> = <ia|pq> - <ia|qp> = (ip|aq) - (iq|ap)
            buf = ao2mofn(cc._scf.mol, (orbo,so_coeff,orbv,so_coeff), compact=0)
            if mo_coeff[0].dtype == np.float: buf = buf.real
            buf = buf.reshape((nocc,nmo,nvir,nmo))
            for p in range(nmo):
                for i in range(nocc):
                    if spin[i] != spin[p]:
                        buf[i,p,:,:] = 0.
                for a in range(nvir):
                    if spin[nocc+a] != spin[p]:
                        buf[:,:,a,p] = 0.
            buf1 = buf - buf.transpose(0,3,2,1)
            buf1 = buf1.transpose(0,2,1,3)
            cput1 = log.timer_debug1('transforming ovpq', *cput1)
            self.ovoo[:,:,:,:] = buf1[:,:,:nocc,:nocc]
            self.ovov[:,:,:,:] = buf1[:,:,:nocc,nocc:]
            self.ovvv[:,:,:,:] = buf1[:,:,nocc:,nocc:]

            for a in range(nvir):
                orbva = orbv[:,a].reshape(-1,1)
                buf = ao2mofn(cc._scf.mol, (orbva,orbv,orbv,orbv), compact=0)
                if mo_coeff[0].dtype == np.float: buf = buf.real
                buf = buf.reshape((1,nvir,nvir,nvir))
                for b in range(nvir):
                    if spin[nocc+a] != spin[nocc+b]:
                        buf[0,b,:,:] = 0.
                    for c in range(nvir):
                        if spin[nocc+b] != spin[nocc+c]:
                            buf[:,:,b,c] = buf[:,:,c,b] = 0.
                buf1 = buf - buf.transpose(0,3,2,1)
                buf1 = buf1.transpose(0,2,1,3)
                self.vvvv[a] = buf1[:]

            cput1 = log.timer_debug1('transforming vvvv', *cput1)

        log.timer('CCSD integral transformation', *cput0)


def uspatial2spin(cc, moidx, mo_coeff):
    '''Convert the results of an unrestricted mean-field calculation to spin-orbital form.

    Spin-orbital ordering is determined by orbital energy without regard for spin.

    Returns:
        fock : (nso,nso) ndarray
            The Fock matrix in the basis of spin-orbitals
        so_coeff : (nao, nso) ndarray
            The matrix of spin-orbital coefficients in the AO basis
        spin : (nso,) ndarary
            The spin (0 or 1) of each spin-orbital
    '''

    dm = cc._scf.make_rdm1(cc.mo_coeff, cc.mo_occ)
    fockao = cc._scf.get_hcore() + cc._scf.get_veff(cc.mol, dm)
    fockab = list()
    for a in range(2):
        fockab.append( reduce(numpy.dot, (mo_coeff[a].T, fockao[a], mo_coeff[a])) )

    nocc = cc.nocc
    nao = cc.mo_coeff[0].shape[0]
    nmo = cc.nmo
    so_coeff = np.zeros((nao, nmo), dtype=mo_coeff[0].dtype)
    nocc_a = int(sum(cc.mo_occ[0]*moidx[0]))
    nocc_b = int(sum(cc.mo_occ[1]*moidx[1]))
    nmo_a = fockab[0].shape[0]
    nmo_b = fockab[1].shape[0]
    nvir_a = nmo_a - nocc_a
    #nvir_b = nmo_b - nocc_b
    oa = range(0,nocc_a)
    ob = range(nocc_a,nocc)
    va = range(nocc,nocc+nvir_a)
    vb = range(nocc+nvir_a,nmo)
    spin = np.zeros(nmo, dtype=int)
    spin[oa] = 0
    spin[ob] = 1
    spin[va] = 0
    spin[vb] = 1
    so_coeff[:,oa] = mo_coeff[0][:,:nocc_a]
    so_coeff[:,ob] = mo_coeff[1][:,:nocc_b]
    so_coeff[:,va] = mo_coeff[0][:,nocc_a:nmo_a]
    so_coeff[:,vb] = mo_coeff[1][:,nocc_b:nmo_b]

    fock = np.zeros((nmo, nmo), dtype=fockab[0].dtype)
    fock[np.ix_(oa,oa)] = fockab[0][:nocc_a,:nocc_a]
    fock[np.ix_(oa,va)] = fockab[0][:nocc_a,nocc_a:]
    fock[np.ix_(va,oa)] = fockab[0][nocc_a:,:nocc_a]
    fock[np.ix_(va,va)] = fockab[0][nocc_a:,nocc_a:]
    fock[np.ix_(ob,ob)] = fockab[1][:nocc_b,:nocc_b]
    fock[np.ix_(ob,vb)] = fockab[1][:nocc_b,nocc_b:]
    fock[np.ix_(vb,ob)] = fockab[1][nocc_b:,:nocc_b]
    fock[np.ix_(vb,vb)] = fockab[1][nocc_b:,nocc_b:]

# Do not sort because it's different to the orbital ordering generated by
# get_frozen_mask function in AO-direct vvvv contraction
#    idxo = np.diagonal(fock[:nocc,:nocc]).argsort()
#    idxv = nocc + np.diagonal(fock[nocc:,nocc:]).argsort()
#    idx = np.concatenate((idxo,idxv))
#    spin = spin[idx]
#    so_coeff = so_coeff[:,idx]
#    fock = fock[:, idx][idx]

    return fock, so_coeff, spin


class _IMDS:
    # Exactly the same as RCCSD IMDS except
    # -- rintermediates --> uintermediates
    # -- Loo, Lvv, cc_Fov --> Foo, Fvv, Fov
    # -- One less 2-virtual intermediate
    def __init__(self, cc):
        self.cc = cc
        self._made_shared = False
        self.made_ip_imds = False
        self.made_ea_imds = False
        self.made_ee_imds = False

    def _make_shared(self):
        cput0 = (logger.process_clock(), logger.perf_counter())
        log = logger.Logger(self.cc.stdout, self.cc.verbose)

        t1,t2,eris = self.cc.t1, self.cc.t2, self.cc.eris
        self.Foo = imd.Foo(t1,t2,eris)
        self.Fvv = imd.Fvv(t1,t2,eris)
        self.Fov = imd.Fov(t1,t2,eris)

        # 2 virtuals
        self.Wovvo = imd.Wovvo(t1,t2,eris)
        self.Woovv = eris.oovv

        log.timer('EOM-CCSD shared intermediates', *cput0)

    def make_ip(self):
        if self._made_shared is False:
            self._make_shared()
            self._made_shared = True

        cput0 = (logger.process_clock(), logger.perf_counter())
        log = logger.Logger(self.cc.stdout, self.cc.verbose)

        t1,t2,eris = self.cc.t1, self.cc.t2, self.cc.eris

        # 0 or 1 virtuals
        self.Woooo = imd.Woooo(t1,t2,eris)
        self.Wooov = imd.Wooov(t1,t2,eris)
        self.Wovoo = imd.Wovoo(t1,t2,eris)

        self.made_ip_imds = True
        log.timer('EOM-CCSD IP intermediates', *cput0)

    def make_ea(self):
        if self._made_shared is False:
            self._make_shared()
            self._made_shared = True

        cput0 = (logger.process_clock(), logger.perf_counter())
        log = logger.Logger(self.cc.stdout, self.cc.verbose)

        t1,t2,eris = self.cc.t1, self.cc.t2, self.cc.eris

        # 3 or 4 virtuals
        self.Wvovv = imd.Wvovv(t1,t2,eris)
        self.Wvvvv = imd.Wvvvv(t1,t2,eris)
        self.Wvvvo = imd.Wvvvo(t1,t2,eris,self.Wvvvv)

        self.made_ea_imds = True
        log.timer('EOM-CCSD EA intermediates', *cput0)

    def make_ee(self):
        if self._made_shared is False:
            self._make_shared()
            self._made_shared = True

        cput0 = (logger.process_clock(), logger.perf_counter())
        log = logger.Logger(self.cc.stdout, self.cc.verbose)

        t1,t2,eris = self.cc.t1, self.cc.t2, self.cc.eris

        if self.made_ip_imds is False:
            # 0 or 1 virtuals
            self.Woooo = imd.Woooo(t1,t2,eris)
            self.Wooov = imd.Wooov(t1,t2,eris)
            self.Wovoo = imd.Wovoo(t1,t2,eris)
        if self.made_ea_imds is False:
            # 3 or 4 virtuals
            self.Wvovv = imd.Wvovv(t1,t2,eris)
            self.Wvvvv = imd.Wvvvv(t1,t2,eris)
            self.Wvvvo = imd.Wvvvo(t1,t2,eris,self.Wvvvv)

        self.made_ee_imds = True
        log.timer('EOM-CCSD EE intermediates', *cput0)

if __name__ == '__main__':
    from pyscf import scf
    from pyscf import gto
    mol = gto.Mole()
    mol.atom = [['O', (0.,   0., 0.)],
                ['O', (1.21, 0., 0.)]]
    mol.basis = 'cc-pvdz'
    mol.spin = 2
    mol.build()
    mf = scf.UHF(mol)
    print(mf.scf())

    # Freeze 1s electrons
    frozen = [[0,1], [0,1]]
    # also acceptable
    #frozen = 4
    ucc = UCCSD(mf, frozen=frozen)
    ecc, t1, t2 = ucc.kernel()
    print(ecc - -0.3486987472235819)
    exit()

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = 'cc-pvdz'
    mol.spin = 0
    mol.build()
    mf = scf.UHF(mol)
    print(mf.scf())

    mycc = UCCSD(mf)
    ecc, t1, t2 = mycc.kernel()
    print(ecc - -0.2133432712431435)
    e,v = mycc.ipccsd(nroots=8)
    print(e[0] - 0.4335604332073799)
    print(e[2] - 0.5187659896045407)
    print(e[4] - 0.6782876002229172)

    mycc.verbose = 5
    e,v = mycc.eaccsd(nroots=8)
    print(e[0] - 0.16737886338859731)
    print(e[2] - 0.24027613852009164)
    print(e[4] - 0.51006797826488071)
    print("e=", e)

    e,v = mycc.eeccsd(nroots=4)
    print(e[0] - 0.2757159395886167)
    print(e[1] - 0.2757159395886167)
    print(e[2] - 0.2757159395886167)
    print(e[3] - 0.3005716731825082)


# pyscf/cc/uccsd_t.py

#!/usr/bin/env python
# Copyright 2014-2020 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: Qiming Sun <osirpt.sun@gmail.com>
#

'''
UCCSD(T)
'''


import ctypes
import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf.cc import _ccsd

def kernel(mycc, eris, t1=None, t2=None, verbose=logger.NOTE):
    cpu1 = cpu0 = (logger.process_clock(), logger.perf_counter())
    log = logger.new_logger(mycc, verbose)
    if t1 is None: t1 = mycc.t1
    if t2 is None: t2 = mycc.t2
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2

    nocca, noccb = mycc.nocc
    nmoa = eris.focka.shape[0]
    nmob = eris.fockb.shape[0]
    nvira = nmoa - nocca
    nvirb = nmob - noccb

    if mycc.incore_complete:
        ftmp = None
    else:
        ftmp = lib.H5TmpFile()
    t1aT = t1a.T.copy()
    t1bT = t1b.T.copy()
    t2aaT = t2aa.transpose(2,3,0,1).copy()
    t2bbT = t2bb.transpose(2,3,0,1).copy()

    eris_vooo = numpy.asarray(eris.ovoo).transpose(1,3,0,2).conj().copy()
    eris_VOOO = numpy.asarray(eris.OVOO).transpose(1,3,0,2).conj().copy()
    eris_vOoO = numpy.asarray(eris.ovOO).transpose(1,3,0,2).conj().copy()
    eris_VoOo = numpy.asarray(eris.OVoo).transpose(1,3,0,2).conj().copy()

    eris_vvop, eris_VVOP, eris_vVoP, eris_VvOp = _sort_eri(mycc, eris, ftmp, log)
    cpu1 = log.timer_debug1('UCCSD(T) sort_eri', *cpu1)

    dtype = numpy.result_type(t1a.dtype, t2aa.dtype, eris_vooo.dtype)
    et_sum = numpy.zeros(1, dtype=dtype)
    mem_now = lib.current_memory()[0]
    max_memory = max(0, mycc.max_memory - mem_now)
    # aaa
    bufsize = max(8, int((max_memory*.5e6/8-nocca**3*3*lib.num_threads())*.4/(nocca*nmoa)))
    log.debug('max_memory %d MB (%d MB in use)', max_memory, mem_now)
    orbsym = numpy.zeros(nocca, dtype=int)
    contract = _gen_contract_aaa(t1aT, t2aaT, eris_vooo, eris.focka,
                                 eris.mo_energy[0], orbsym, log)
    with lib.call_in_background(contract, sync=not mycc.async_io) as ctr:
        for a0, a1 in reversed(list(lib.prange_tril(0, nvira, bufsize))):
            cache_row_a = numpy.asarray(eris_vvop[a0:a1,:a1], order='C')
            if a0 == 0:
                cache_col_a = cache_row_a
            else:
                cache_col_a = numpy.asarray(eris_vvop[:a0,a0:a1], order='C')
            ctr(et_sum, a0, a1, a0, a1, (cache_row_a,cache_col_a,
                                         cache_row_a,cache_col_a))

            for b0, b1 in lib.prange_tril(0, a0, bufsize/8):
                cache_row_b = numpy.asarray(eris_vvop[b0:b1,:b1], order='C')
                if b0 == 0:
                    cache_col_b = cache_row_b
                else:
                    cache_col_b = numpy.asarray(eris_vvop[:b0,b0:b1], order='C')
                ctr(et_sum, a0, a1, b0, b1, (cache_row_a,cache_col_a,
                                             cache_row_b,cache_col_b))
    cpu1 = log.timer_debug1('contract_aaa', *cpu1)

    # bbb
    bufsize = max(8, int((max_memory*.5e6/8-noccb**3*3*lib.num_threads())*.4/(noccb*nmob)))
    log.debug('max_memory %d MB (%d MB in use)', max_memory, mem_now)
    orbsym = numpy.zeros(noccb, dtype=int)
    contract = _gen_contract_aaa(t1bT, t2bbT, eris_VOOO, eris.fockb,
                                 eris.mo_energy[1], orbsym, log)
    with lib.call_in_background(contract, sync=not mycc.async_io) as ctr:
        for a0, a1 in reversed(list(lib.prange_tril(0, nvirb, bufsize))):
            cache_row_a = numpy.asarray(eris_VVOP[a0:a1,:a1], order='C')
            if a0 == 0:
                cache_col_a = cache_row_a
            else:
                cache_col_a = numpy.asarray(eris_VVOP[:a0,a0:a1], order='C')
            ctr(et_sum, a0, a1, a0, a1, (cache_row_a,cache_col_a,
                                         cache_row_a,cache_col_a))

            for b0, b1 in lib.prange_tril(0, a0, bufsize/8):
                cache_row_b = numpy.asarray(eris_VVOP[b0:b1,:b1], order='C')
                if b0 == 0:
                    cache_col_b = cache_row_b
                else:
                    cache_col_b = numpy.asarray(eris_VVOP[:b0,b0:b1], order='C')
                ctr(et_sum, a0, a1, b0, b1, (cache_row_a,cache_col_a,
                                             cache_row_b,cache_col_b))
    cpu1 = log.timer_debug1('contract_bbb', *cpu1)

    # Cache t2abT in t2ab to reduce memory footprint
    assert(t2ab.flags.c_contiguous)
    t2abT = lib.transpose(t2ab.copy().reshape(nocca*noccb,nvira*nvirb), out=t2ab)
    t2abT = t2abT.reshape(nvira,nvirb,nocca,noccb)
    # baa
    bufsize = int(max(12, (max_memory*.5e6/8-noccb*nocca**2*5)*.7/(nocca*nmob)))
    ts = t1aT, t1bT, t2aaT, t2abT
    fock = (eris.focka, eris.fockb)
    vooo = (eris_vooo, eris_vOoO, eris_VoOo)
    contract = _gen_contract_baa(ts, vooo, fock, eris.mo_energy, orbsym, log)
    with lib.call_in_background(contract, sync=not mycc.async_io) as ctr:
        for a0, a1 in lib.prange(0, nvirb, int(bufsize/nvira+1)):
            cache_row_a = numpy.asarray(eris_VvOp[a0:a1,:], order='C')
            cache_col_a = numpy.asarray(eris_vVoP[:,a0:a1], order='C')
            for b0, b1 in lib.prange_tril(0, nvira, bufsize/6/2):
                cache_row_b = numpy.asarray(eris_vvop[b0:b1,:b1], order='C')
                cache_col_b = numpy.asarray(eris_vvop[:b0,b0:b1], order='C')
                ctr(et_sum, a0, a1, b0, b1, (cache_row_a,cache_col_a,
                                             cache_row_b,cache_col_b))
    cpu1 = log.timer_debug1('contract_baa', *cpu1)

    t2baT = numpy.ndarray((nvirb,nvira,noccb,nocca), buffer=t2abT,
                          dtype=t2abT.dtype)
    t2baT[:] = t2abT.copy().transpose(1,0,3,2)
    # abb
    ts = t1bT, t1aT, t2bbT, t2baT
    fock = (eris.fockb, eris.focka)
    mo_energy = (eris.mo_energy[1], eris.mo_energy[0])
    vooo = (eris_VOOO, eris_VoOo, eris_vOoO)
    contract = _gen_contract_baa(ts, vooo, fock, mo_energy, orbsym, log)
    for a0, a1 in lib.prange(0, nvira, int(bufsize/nvirb+1)):
        with lib.call_in_background(contract, sync=not mycc.async_io) as ctr:
            cache_row_a = numpy.asarray(eris_vVoP[a0:a1,:], order='C')
            cache_col_a = numpy.asarray(eris_VvOp[:,a0:a1], order='C')
            for b0, b1 in lib.prange_tril(0, nvirb, bufsize/6/2):
                cache_row_b = numpy.asarray(eris_VVOP[b0:b1,:b1], order='C')
                cache_col_b = numpy.asarray(eris_VVOP[:b0,b0:b1], order='C')
                ctr(et_sum, a0, a1, b0, b1, (cache_row_a,cache_col_a,
                                             cache_row_b,cache_col_b))
    cpu1 = log.timer_debug1('contract_abb', *cpu1)

    # Restore t2ab
    lib.transpose(t2baT.transpose(1,0,3,2).copy().reshape(nvira*nvirb,nocca*noccb),
                  out=t2ab)
    et_sum *= .25
    if abs(et_sum[0].imag) > 1e-4:
        logger.warn(mycc, 'Non-zero imaginary part of UCCSD(T) energy was found %s',
                    et_sum[0])
    et = et_sum[0].real
    log.timer('UCCSD(T)', *cpu0)
    log.note('UCCSD(T) correction = %.15g', et)
    return et

def _gen_contract_aaa(t1T, t2T, vooo, fock, mo_energy, orbsym, log):
    nvir, nocc = t1T.shape
    mo_energy = numpy.asarray(mo_energy, order='C')
    fvo = fock[nocc:,:nocc].copy()

    cpu2 = [logger.process_clock(), logger.perf_counter()]
    orbsym = numpy.hstack((numpy.sort(orbsym[:nocc]),numpy.sort(orbsym[nocc:])))
    o_ir_loc = numpy.append(0, numpy.cumsum(numpy.bincount(orbsym[:nocc], minlength=8)))
    v_ir_loc = numpy.append(0, numpy.cumsum(numpy.bincount(orbsym[nocc:], minlength=8)))
    o_sym = orbsym[:nocc]
    oo_sym = (o_sym[:,None] ^ o_sym).ravel()
    oo_ir_loc = numpy.append(0, numpy.cumsum(numpy.bincount(oo_sym, minlength=8)))
    nirrep = max(oo_sym) + 1

    orbsym   = orbsym.astype(numpy.int32)
    o_ir_loc = o_ir_loc.astype(numpy.int32)
    v_ir_loc = v_ir_loc.astype(numpy.int32)
    oo_ir_loc = oo_ir_loc.astype(numpy.int32)
    dtype = numpy.result_type(t2T.dtype, vooo.dtype, fock.dtype)
    if dtype == numpy.complex128:
        drv = _ccsd.libcc.CCuccsd_t_zaaa
    else:
        drv = _ccsd.libcc.CCuccsd_t_aaa
    def contract(et_sum, a0, a1, b0, b1, cache):
        cache_row_a, cache_col_a, cache_row_b, cache_col_b = cache
        drv(et_sum.ctypes.data_as(ctypes.c_void_p),
            mo_energy.ctypes.data_as(ctypes.c_void_p),
            t1T.ctypes.data_as(ctypes.c_void_p),
            t2T.ctypes.data_as(ctypes.c_void_p),
            vooo.ctypes.data_as(ctypes.c_void_p),
            fvo.ctypes.data_as(ctypes.c_void_p),
            ctypes.c_int(nocc), ctypes.c_int(nvir),
            ctypes.c_int(a0), ctypes.c_int(a1),
            ctypes.c_int(b0), ctypes.c_int(b1),
            ctypes.c_int(nirrep),
            o_ir_loc.ctypes.data_as(ctypes.c_void_p),
            v_ir_loc.ctypes.data_as(ctypes.c_void_p),
            oo_ir_loc.ctypes.data_as(ctypes.c_void_p),
            orbsym.ctypes.data_as(ctypes.c_void_p),
            cache_row_a.ctypes.data_as(ctypes.c_void_p),
            cache_col_a.ctypes.data_as(ctypes.c_void_p),
            cache_row_b.ctypes.data_as(ctypes.c_void_p),
            cache_col_b.ctypes.data_as(ctypes.c_void_p))
        cpu2[:] = log.timer_debug1('contract %d:%d,%d:%d'%(a0,a1,b0,b1), *cpu2)
    return contract

def _gen_contract_baa(ts, vooo, fock, mo_energy, orbsym, log):
    t1aT, t1bT, t2aaT, t2abT = ts
    focka, fockb = fock
    vooo, vOoO, VoOo = vooo
    nvira, nocca = t1aT.shape
    nvirb, noccb = t1bT.shape
    mo_ea = numpy.asarray(mo_energy[0], order='C')
    mo_eb = numpy.asarray(mo_energy[1], order='C')
    fvo = focka[nocca:,:nocca].copy()
    fVO = fockb[noccb:,:noccb].copy()

    cpu2 = [logger.process_clock(), logger.perf_counter()]
    dtype = numpy.result_type(t2aaT.dtype, vooo.dtype)
    if dtype == numpy.complex128:
        drv = _ccsd.libcc.CCuccsd_t_zbaa
    else:
        drv = _ccsd.libcc.CCuccsd_t_baa
    def contract(et_sum, a0, a1, b0, b1, cache):
        cache_row_a, cache_col_a, cache_row_b, cache_col_b = cache
        drv(et_sum.ctypes.data_as(ctypes.c_void_p),
            mo_ea.ctypes.data_as(ctypes.c_void_p),
            mo_eb.ctypes.data_as(ctypes.c_void_p),
            t1aT.ctypes.data_as(ctypes.c_void_p),
            t1bT.ctypes.data_as(ctypes.c_void_p),
            t2aaT.ctypes.data_as(ctypes.c_void_p),
            t2abT.ctypes.data_as(ctypes.c_void_p),
            vooo.ctypes.data_as(ctypes.c_void_p),
            vOoO.ctypes.data_as(ctypes.c_void_p),
            VoOo.ctypes.data_as(ctypes.c_void_p),
            fvo.ctypes.data_as(ctypes.c_void_p),
            fVO.ctypes.data_as(ctypes.c_void_p),
            ctypes.c_int(nocca), ctypes.c_int(noccb),
            ctypes.c_int(nvira), ctypes.c_int(nvirb),
            ctypes.c_int(a0), ctypes.c_int(a1),
            ctypes.c_int(b0), ctypes.c_int(b1),
            cache_row_a.ctypes.data_as(ctypes.c_void_p),
            cache_col_a.ctypes.data_as(ctypes.c_void_p),
            cache_row_b.ctypes.data_as(ctypes.c_void_p),
            cache_col_b.ctypes.data_as(ctypes.c_void_p))
        cpu2[:] = log.timer_debug1('contract %d:%d,%d:%d'%(a0,a1,b0,b1), *cpu2)
    return contract

def _sort_eri(mycc, eris, h5tmp, log):
    cpu1 = (logger.process_clock(), logger.perf_counter())
    nocca, noccb = mycc.nocc
    nmoa = eris.focka.shape[0]
    nmob = eris.fockb.shape[0]
    nvira = nmoa - nocca
    nvirb = nmob - noccb

    if mycc.t2 is None:
        dtype = eris.ovov.dtype
    else:
        dtype = numpy.result_type(mycc.t2[0], eris.ovov.dtype)

    if mycc.incore_complete or h5tmp is None:
        eris_vvop = numpy.empty((nvira,nvira,nocca,nmoa), dtype)
        eris_VVOP = numpy.empty((nvirb,nvirb,noccb,nmob), dtype)
        eris_vVoP = numpy.empty((nvira,nvirb,nocca,nmob), dtype)
        eris_VvOp = numpy.empty((nvirb,nvira,noccb,nmoa), dtype)
    else:
        eris_vvop = h5tmp.create_dataset('vvop', (nvira,nvira,nocca,nmoa), dtype)
        eris_VVOP = h5tmp.create_dataset('VVOP', (nvirb,nvirb,noccb,nmob), dtype)
        eris_vVoP = h5tmp.create_dataset('vVoP', (nvira,nvirb,nocca,nmob), dtype)
        eris_VvOp = h5tmp.create_dataset('VvOp', (nvirb,nvira,noccb,nmoa), dtype)

    max_memory = max(2000, mycc.max_memory - lib.current_memory()[0])
    max_memory = min(8000, max_memory*.9)

    blksize = min(nvira, max(16, int(max_memory*1e6/8/(nvira*nocca*nmoa))))
    with lib.call_in_background(eris_vvop.__setitem__, sync=not mycc.async_io) as save:
        bufopv = numpy.empty((nocca,nmoa,nvira), dtype=dtype)
        buf1 = numpy.empty_like(bufopv)
        for j0, j1 in lib.prange(0, nvira, blksize):
            ovov = numpy.asarray(eris.ovov[:,j0:j1])
            ovvv = eris.get_ovvv(slice(None), slice(j0,j1))
            for j in range(j0,j1):
                bufopv[:,:nocca,:] = ovov[:,j-j0].conj()
                bufopv[:,nocca:,:] = ovvv[:,j-j0].conj()
                save(j, bufopv.transpose(2,0,1))
                bufopv, buf1 = buf1, bufopv
            ovov = ovvv = None
            cpu1 = log.timer_debug1('transpose %d:%d'%(j0,j1), *cpu1)

    blksize = min(nvirb, max(16, int(max_memory*1e6/8/(nvirb*noccb*nmob))))
    with lib.call_in_background(eris_VVOP.__setitem__, sync=not mycc.async_io) as save:
        bufopv = numpy.empty((noccb,nmob,nvirb), dtype=dtype)
        buf1 = numpy.empty_like(bufopv)
        for j0, j1 in lib.prange(0, nvirb, blksize):
            ovov = numpy.asarray(eris.OVOV[:,j0:j1])
            ovvv = eris.get_OVVV(slice(None), slice(j0,j1))
            for j in range(j0,j1):
                bufopv[:,:noccb,:] = ovov[:,j-j0].conj()
                bufopv[:,noccb:,:] = ovvv[:,j-j0].conj()
                save(j, bufopv.transpose(2,0,1))
                bufopv, buf1 = buf1, bufopv
            ovov = ovvv = None
            cpu1 = log.timer_debug1('transpose %d:%d'%(j0,j1), *cpu1)

    blksize = min(nvira, max(16, int(max_memory*1e6/8/(nvirb*nocca*nmob))))
    with lib.call_in_background(eris_vVoP.__setitem__, sync=not mycc.async_io) as save:
        bufopv = numpy.empty((nocca,nmob,nvirb), dtype=dtype)
        buf1 = numpy.empty_like(bufopv)
        for j0, j1 in lib.prange(0, nvira, blksize):
            ovov = numpy.asarray(eris.ovOV[:,j0:j1])
            ovvv = eris.get_ovVV(slice(None), slice(j0,j1))
            for j in range(j0,j1):
                bufopv[:,:noccb,:] = ovov[:,j-j0].conj()
                bufopv[:,noccb:,:] = ovvv[:,j-j0].conj()
                save(j, bufopv.transpose(2,0,1))
                bufopv, buf1 = buf1, bufopv
            ovov = ovvv = None
            cpu1 = log.timer_debug1('transpose %d:%d'%(j0,j1), *cpu1)

    blksize = min(nvirb, max(16, int(max_memory*1e6/8/(nvira*noccb*nmoa))))
    OVov = numpy.asarray(eris.ovOV).transpose(2,3,0,1)
    with lib.call_in_background(eris_VvOp.__setitem__, sync=not mycc.async_io) as save:
        bufopv = numpy.empty((noccb,nmoa,nvira), dtype=dtype)
        buf1 = numpy.empty_like(bufopv)
        for j0, j1 in lib.prange(0, nvirb, blksize):
            ovov = OVov[:,j0:j1]
            ovvv = eris.get_OVvv(slice(None), slice(j0,j1))
            for j in range(j0,j1):
                bufopv[:,:nocca,:] = ovov[:,j-j0].conj()
                bufopv[:,nocca:,:] = ovvv[:,j-j0].conj()
                save(j, bufopv.transpose(2,0,1))
                bufopv, buf1 = buf1, bufopv
            ovov = ovvv = None
            cpu1 = log.timer_debug1('transpose %d:%d'%(j0,j1), *cpu1)
    return eris_vvop, eris_VVOP, eris_vVoP, eris_VvOp


if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf
    from pyscf import cc

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -.757 , .587)],
        [1 , (0. ,  .757 , .587)]]

    mol.basis = '631g'
    mol.build()
    rhf = scf.RHF(mol)
    rhf.conv_tol = 1e-14
    rhf.scf()
    mcc = cc.CCSD(rhf)
    mcc.conv_tol = 1e-12
    mcc.ccsd()
    t1a = t1b = mcc.t1
    t2ab = mcc.t2
    t2aa = t2bb = t2ab - t2ab.transpose(1,0,2,3)
    mycc = cc.UCCSD(scf.addons.convert_to_uhf(rhf))
    eris = mycc.ao2mo()
    e3a = kernel(mycc, eris, (t1a,t1b), (t2aa,t2ab,t2bb))
    print(e3a - -0.00099642337843278096)

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -.757 , .587)],
        [1 , (0. ,  .757 , .587)]]
    mol.spin = 2
    mol.basis = '3-21g'
    mol.build()
    mf = scf.UHF(mol).run(conv_tol=1e-14)
    nao, nmo = mf.mo_coeff[0].shape
    numpy.random.seed(10)
    mf.mo_coeff = numpy.random.random((2,nao,nmo))

    numpy.random.seed(12)
    nocca, noccb = mol.nelec
    nmo = mf.mo_occ[0].size
    nvira = nmo - nocca
    nvirb = nmo - noccb
    t1a  = .1 * numpy.random.random((nocca,nvira))
    t1b  = .1 * numpy.random.random((noccb,nvirb))
    t2aa = .1 * numpy.random.random((nocca,nocca,nvira,nvira))
    t2aa = t2aa - t2aa.transpose(0,1,3,2)
    t2aa = t2aa - t2aa.transpose(1,0,2,3)
    t2bb = .1 * numpy.random.random((noccb,noccb,nvirb,nvirb))
    t2bb = t2bb - t2bb.transpose(0,1,3,2)
    t2bb = t2bb - t2bb.transpose(1,0,2,3)
    t2ab = .1 * numpy.random.random((nocca,noccb,nvira,nvirb))
    t1 = t1a, t1b
    t2 = t2aa, t2ab, t2bb
    mycc = cc.UCCSD(mf)
    eris = mycc.ao2mo(mf.mo_coeff)
    e3a = kernel(mycc, eris, [t1a,t1b], [t2aa, t2ab, t2bb])
    print(e3a - 9877.2780859693339)

    mycc = cc.GCCSD(scf.addons.convert_to_ghf(mf))
    eris = mycc.ao2mo()
    t1 = mycc.spatial2spin(t1, eris.orbspin)
    t2 = mycc.spatial2spin(t2, eris.orbspin)
    from pyscf.cc import gccsd_t_slow
    et = gccsd_t_slow.kernel(mycc, eris, t1, t2)
    print(et - 9877.2780859693339)



# pyscf/cc/uccsd_t_lambda.py

#!/usr/bin/env python
# Copyright 2014-2020 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: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Spin-free lambda equation of UHF-CCSD(T)
'''

import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf.cc import ccsd_lambda
from pyscf.cc import uccsd_lambda


def kernel(mycc, eris=None, t1=None, t2=None, l1=None, l2=None,
           max_cycle=50, tol=1e-8, verbose=logger.INFO):
    return ccsd_lambda.kernel(mycc, eris, t1, t2, l1, l2, max_cycle, tol,
                              verbose, make_intermediates, update_lambda)

def make_intermediates(mycc, t1, t2, eris):
    def p6(t):
        return (t + t.transpose(1,2,0,4,5,3) +
                t.transpose(2,0,1,5,3,4) + t.transpose(0,2,1,3,5,4) +
                t.transpose(2,1,0,5,4,3) + t.transpose(1,0,2,4,3,5))
    def r6(w):
        return (w + w.transpose(2,0,1,3,4,5) + w.transpose(1,2,0,3,4,5)
                - w.transpose(2,1,0,3,4,5) - w.transpose(0,2,1,3,4,5)
                - w.transpose(1,0,2,3,4,5))

    imds = uccsd_lambda.make_intermediates(mycc, t1, t2, eris)

    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nocca, noccb = t2ab.shape[:2]
    mo_ea, mo_eb = eris.mo_energy
    eia = mo_ea[:nocca,None] - mo_ea[nocca:]
    eIA = mo_eb[:noccb,None] - mo_eb[noccb:]
    fvo = eris.focka[nocca:,:nocca]
    fVO = eris.fockb[noccb:,:noccb]

    # aaa
    d3 = lib.direct_sum('ia+jb+kc->ijkabc', eia, eia, eia)
    w = numpy.einsum('ijae,kceb->ijkabc', t2aa, numpy.asarray(eris.get_ovvv()).conj())
    w-= numpy.einsum('mkbc,iajm->ijkabc', t2aa, numpy.asarray(eris.ovoo).conj())
    v = numpy.einsum('jbkc,ia->ijkabc', numpy.asarray(eris.ovov).conj(), t1a)
    v+= numpy.einsum('jkbc,ai->ijkabc', t2aa, fvo) * .5

    rw = r6(p6(w)) / d3
    imds.l1a_t = numpy.einsum('ijkabc,jbkc->ia', rw.conj(),
                              numpy.asarray(eris.ovov)) / eia * .25
    wvd = r6(p6(w * 2 + v)) / d3
    l2_t  = numpy.einsum('ijkabc,kceb->ijae', wvd, numpy.asarray(eris.get_ovvv()).conj())
    l2_t -= numpy.einsum('ijkabc,iajm->mkbc', wvd, numpy.asarray(eris.ovoo).conj())
    l2_t = l2_t + l2_t.transpose(1,0,3,2)
    l2_t += numpy.einsum('ijkabc,ai->jkbc', rw, fvo)
    imds.l2aa_t = l2_t.conj() / lib.direct_sum('ia+jb->ijab', eia, eia) * .5

    # bbb
    d3 = lib.direct_sum('ia+jb+kc->ijkabc', eIA, eIA, eIA)
    w = numpy.einsum('ijae,kceb->ijkabc', t2bb, numpy.asarray(eris.get_OVVV()).conj())
    w-= numpy.einsum('imab,kcjm->ijkabc', t2bb, numpy.asarray(eris.OVOO).conj())
    v = numpy.einsum('jbkc,ia->ijkabc', numpy.asarray(eris.OVOV).conj(), t1b)
    v+= numpy.einsum('jkbc,ai->ijkabc', t2bb, fVO) * .5

    rw = r6(p6(w)) / d3
    imds.l1b_t = numpy.einsum('ijkabc,jbkc->ia', rw.conj(),
                              numpy.asarray(eris.OVOV)) / eIA * .25
    wvd = r6(p6(w * 2 + v)) / d3
    l2_t  = numpy.einsum('ijkabc,kceb->ijae', wvd, numpy.asarray(eris.get_OVVV()).conj())
    l2_t -= numpy.einsum('ijkabc,iajm->mkbc', wvd, numpy.asarray(eris.OVOO).conj())
    l2_t = l2_t + l2_t.transpose(1,0,3,2)
    l2_t += numpy.einsum('ijkabc,ai->jkbc', rw, fVO)
    imds.l2bb_t = l2_t.conj() / lib.direct_sum('ia+jb->ijab', eIA, eIA) * .5

    # baa
    def r4(w):
        w = w - w.transpose(0,2,1,3,4,5)
        w = w + w.transpose(0,2,1,3,5,4)
        return w
    d3 = lib.direct_sum('ia+jb+kc->ijkabc', eIA, eia, eia)
    w  = numpy.einsum('jIeA,kceb->IjkAbc', t2ab, numpy.asarray(eris.get_ovvv()).conj()) * 2
    w += numpy.einsum('jIbE,kcEA->IjkAbc', t2ab, numpy.asarray(eris.get_ovVV()).conj()) * 2
    w += numpy.einsum('jkbe,IAec->IjkAbc', t2aa, numpy.asarray(eris.get_OVvv()).conj())
    w -= numpy.einsum('mIbA,kcjm->IjkAbc', t2ab, numpy.asarray(eris.ovoo).conj()) * 2
    w -= numpy.einsum('jMbA,kcIM->IjkAbc', t2ab, numpy.asarray(eris.ovOO).conj()) * 2
    w -= numpy.einsum('jmbc,IAkm->IjkAbc', t2aa, numpy.asarray(eris.OVoo).conj())
    v  = numpy.einsum('jbkc,IA->IjkAbc', numpy.asarray(eris.ovov).conj(), t1b)
    v += numpy.einsum('kcIA,jb->IjkAbc', numpy.asarray(eris.ovOV).conj(), t1a)
    v += numpy.einsum('kcIA,jb->IjkAbc', numpy.asarray(eris.ovOV).conj(), t1a)
    v += numpy.einsum('jkbc,AI->IjkAbc', t2aa, fVO) * .5
    v += numpy.einsum('kIcA,bj->IjkAbc', t2ab, fvo) * 2

    rw = r4(w) / d3
    imds.l1a_t += numpy.einsum('ijkabc,kcia->jb', rw.conj(),
                               numpy.asarray(eris.ovOV)) / eia * .5
    imds.l1b_t += numpy.einsum('ijkabc,jbkc->ia', rw.conj(),
                               numpy.asarray(eris.ovov)) / eIA * .25
    wvd = r4(w * 2 + v) / d3
    l2_t  = numpy.einsum('ijkabc,iaec->jkbe', wvd, numpy.asarray(eris.get_OVvv()).conj())
    l2_t -= numpy.einsum('ijkabc,iakm->jmbc', wvd, numpy.asarray(eris.OVoo).conj())
    l2_t = l2_t + l2_t.transpose(1,0,3,2)
    l2_t += numpy.einsum('ijkabc,ai->jkbc', rw, fVO)
    imds.l2aa_t += l2_t.conj() / lib.direct_sum('ia+jb->ijab', eia, eia) * .5
    l2_t  = numpy.einsum('ijkabc,kceb->jiea', wvd, numpy.asarray(eris.get_ovvv()).conj())
    l2_t += numpy.einsum('ijkabc,kcea->jibe', wvd, numpy.asarray(eris.get_ovVV()).conj())
    l2_t -= numpy.einsum('ijkabc,kcjm->miba', wvd, numpy.asarray(eris.ovoo).conj())
    l2_t -= numpy.einsum('ijkabc,kcim->jmba', wvd, numpy.asarray(eris.ovOO).conj())
    l2_t += numpy.einsum('ijkabc,bj->kica', rw, fvo)
    imds.l2ab_t = l2_t.conj() / lib.direct_sum('ia+jb->ijab', eia, eIA) * .5

    # bba
    d3 = lib.direct_sum('ia+jb+kc->ijkabc', eia, eIA, eIA)
    w  = numpy.einsum('ijae,kceb->ijkabc', t2ab, numpy.asarray(eris.get_OVVV()).conj()) * 2
    w += numpy.einsum('ijeb,kcea->ijkabc', t2ab, numpy.asarray(eris.get_OVvv()).conj()) * 2
    w += numpy.einsum('jkbe,iaec->ijkabc', t2bb, numpy.asarray(eris.get_ovVV()).conj())
    w -= numpy.einsum('imab,kcjm->ijkabc', t2ab, numpy.asarray(eris.OVOO).conj()) * 2
    w -= numpy.einsum('mjab,kcim->ijkabc', t2ab, numpy.asarray(eris.OVoo).conj()) * 2
    w -= numpy.einsum('jmbc,iakm->ijkabc', t2bb, numpy.asarray(eris.ovOO).conj())
    v  = numpy.einsum('jbkc,ia->ijkabc', numpy.asarray(eris.OVOV).conj(), t1a)
    v += numpy.einsum('iakc,jb->ijkabc', numpy.asarray(eris.ovOV).conj(), t1b)
    v += numpy.einsum('iakc,jb->ijkabc', numpy.asarray(eris.ovOV).conj(), t1b)
    v += numpy.einsum('JKBC,ai->iJKaBC', t2bb, fvo) * .5
    v += numpy.einsum('iKaC,BJ->iJKaBC', t2ab, fVO) * 2

    rw = r4(w) / d3
    imds.l1a_t += numpy.einsum('ijkabc,jbkc->ia', rw.conj(),
                               numpy.asarray(eris.OVOV)) / eia * .25
    imds.l1b_t += numpy.einsum('ijkabc,iakc->jb', rw.conj(),
                               numpy.asarray(eris.ovOV)) / eIA * .5
    wvd = r4(w * 2 + v) / d3
    l2_t  = numpy.einsum('ijkabc,iaec->jkbe', wvd, numpy.asarray(eris.get_ovVV()).conj())
    l2_t -= numpy.einsum('ijkabc,iakm->jmbc', wvd, numpy.asarray(eris.ovOO).conj())
    l2_t = l2_t + l2_t.transpose(1,0,3,2)
    l2_t += numpy.einsum('ijkabc,ai->jkbc', rw, fvo)
    imds.l2bb_t += l2_t.conj() / lib.direct_sum('ia+jb->ijab', eIA, eIA) * .5
    l2_t  = numpy.einsum('ijkabc,kceb->ijae', wvd, numpy.asarray(eris.get_OVVV()).conj())
    l2_t += numpy.einsum('ijkabc,kcea->ijeb', wvd, numpy.asarray(eris.get_OVvv()).conj())
    l2_t -= numpy.einsum('ijkabc,kcjm->imab', wvd, numpy.asarray(eris.OVOO).conj())
    l2_t -= numpy.einsum('ijkabc,kcim->mjab', wvd, numpy.asarray(eris.OVoo).conj())
    l2_t += numpy.einsum('ijkabc,bj->ikac', rw, fVO)
    imds.l2ab_t += l2_t.conj() / lib.direct_sum('ia+jb->ijab', eia, eIA) * .5

    return imds


def update_lambda(mycc, t1, t2, l1, l2, eris=None, imds=None):
    if eris is None: eris = mycc.ao2mo()
    if imds is None: imds = make_intermediates(mycc, t1, t2, eris)
    l1, l2 = uccsd_lambda.update_lambda(mycc, t1, t2, l1, l2, eris, imds)
    l1a, l1b = l1
    l2aa, l2ab, l2bb = l2
    l1a += imds.l1a_t
    l1b += imds.l1b_t
    l2aa += imds.l2aa_t
    l2ab += imds.l2ab_t
    l2bb += imds.l2bb_t
    return (l1a, l1b), (l2aa, l2ab, l2bb)


if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf
    from pyscf import cc

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -0.757 , 0.587)],
        [1 , (0. , 0.757  , 0.587)]]
    mol.basis = '631g'
    mol.spin = 0
    mol.build()
    mf0 = mf = scf.RHF(mol).run(conv_tol=1)
    mf = scf.addons.convert_to_uhf(mf)
    mycc = cc.UCCSD(mf)
    eris = mycc.ao2mo()

    from pyscf.cc import ccsd_t_lambda_slow as ccsd_t_lambda
    mycc0 = cc.CCSD(mf0)
    eris0 = mycc0.ao2mo()
    mycc0.kernel(eris=eris0)
    t1 = mycc0.t1
    t2 = mycc0.t2
    imds = ccsd_t_lambda.make_intermediates(mycc0, t1, t2, eris0)
    l1, l2 = ccsd_t_lambda.update_lambda(mycc0, t1, t2, t1, t2, eris0, imds)
    l1ref, l2ref = ccsd_t_lambda.update_lambda(mycc0, t1, t2, l1, l2, eris0, imds)

    t1 = (t1, t1)
    t2aa = t2 - t2.transpose(1,0,2,3)
    t2 = (t2aa, t2, t2aa)
    l1 = (l1, l1)
    l2aa = l2 - l2.transpose(1,0,2,3)
    l2 = (l2aa, l2, l2aa)
    imds = make_intermediates(mycc, t1, t2, eris)
    l1, l2 = update_lambda(mycc, t1, t2, l1, l2, eris, imds)
    print(abs(l2[1]-l2[1].transpose(1,0,2,3)-l2[0]).max())
    print(abs(l2[1]-l2[1].transpose(0,1,3,2)-l2[2]).max())
    print(abs(l1[0]-l1ref).max())
    print(abs(l2[1]-l2ref).max())


# pyscf/cc/uccsd_t_rdm.py

#!/usr/bin/env python
# Copyright 2014-2020 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: Qiming Sun <osirpt.sun@gmail.com>
#

import numpy
from pyscf import lib
from pyscf.cc import uccsd_rdm

def _gamma1_intermediates(mycc, t1, t2, l1, l2, eris=None, for_grad=False):
    d1 = uccsd_rdm._gamma1_intermediates(mycc, t1, t2, l1, l2)

    if eris is None: eris = mycc.ao2mo()

    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nocca, noccb, nvira, nvirb = t2ab.shape
    mo_ea, mo_eb = eris.mo_energy
    eia = mo_ea[:nocca,None] - mo_ea[nocca:]
    eIA = mo_eb[:noccb,None] - mo_eb[noccb:]
    fvo = eris.focka[nocca:,:nocca]
    fVO = eris.fockb[noccb:,:noccb]

    # aaa
    d3 = lib.direct_sum('ia+jb+kc->ijkabc', eia, eia, eia)
    w = numpy.einsum('ijae,kceb->ijkabc', t2aa, numpy.asarray(eris.get_ovvv()).conj())
    w-= numpy.einsum('mkbc,iajm->ijkabc', t2aa, numpy.asarray(eris.ovoo).conj())
    v = numpy.einsum('jbkc,ia->ijkabc', numpy.asarray(eris.ovov).conj(), t1a)
    v+= numpy.einsum('jkbc,ai->ijkabc', t2aa, fvo) * .5

    rw = p6(r6(w)) / d3
    wvd = p6(w + v) / d3
    goo = numpy.einsum('iklabc,jklabc->ij', wvd.conj(), rw) * .125
    gvv = numpy.einsum('ijkacd,ijkbcd->ab', wvd, rw.conj()) * .125
    gvo = numpy.einsum('jkbc,ijkabc->ai', t2aa.conj(), rw) * .125

    # bbb
    d3 = lib.direct_sum('ia+jb+kc->ijkabc', eIA, eIA, eIA)
    w = numpy.einsum('ijae,kceb->ijkabc', t2bb, numpy.asarray(eris.get_OVVV()).conj())
    w-= numpy.einsum('imab,kcjm->ijkabc', t2bb, numpy.asarray(eris.OVOO).conj())
    v = numpy.einsum('jbkc,ia->ijkabc', numpy.asarray(eris.OVOV).conj(), t1b)
    v+= numpy.einsum('jkbc,ai->ijkabc', t2bb, fVO) * .5

    rw = p6(r6(w)) / d3
    wvd = p6(w + v) / d3
    gOO = numpy.einsum('iklabc,jklabc->ij', wvd.conj(), rw) * .125
    gVV = numpy.einsum('ijkacd,ijkbcd->ab', wvd, rw.conj()) * .125
    gVO = numpy.einsum('jkbc,ijkabc->ai', t2bb.conj(), rw) * .125

    # baa
    d3 = lib.direct_sum('ia+jb+kc->ijkabc', eIA, eia, eia)
    w  = numpy.einsum('jIeA,kceb->IjkAbc', t2ab, numpy.asarray(eris.get_ovvv()).conj()) * 2
    w += numpy.einsum('jIbE,kcEA->IjkAbc', t2ab, numpy.asarray(eris.get_ovVV()).conj()) * 2
    w += numpy.einsum('jkbe,IAec->IjkAbc', t2aa, numpy.asarray(eris.get_OVvv()).conj())
    w -= numpy.einsum('mIbA,kcjm->IjkAbc', t2ab, numpy.asarray(eris.ovoo).conj()) * 2
    w -= numpy.einsum('jMbA,kcIM->IjkAbc', t2ab, numpy.asarray(eris.ovOO).conj()) * 2
    w -= numpy.einsum('jmbc,IAkm->IjkAbc', t2aa, numpy.asarray(eris.OVoo).conj())
    v  = numpy.einsum('jbkc,IA->IjkAbc', numpy.asarray(eris.ovov).conj(), t1b)
    v += numpy.einsum('kcIA,jb->IjkAbc', numpy.asarray(eris.ovOV).conj(), t1a)
    v += numpy.einsum('kcIA,jb->IjkAbc', numpy.asarray(eris.ovOV).conj(), t1a)
    v += numpy.einsum('jkbc,AI->IjkAbc', t2aa, fVO) * .5
    v += numpy.einsum('kIcA,bj->IjkAbc', t2ab, fvo) * 2

    rw = r4(w) / d3
    wvd = (w + v) / d3
    goo += numpy.einsum('kilabc,kjlabc->ij', wvd.conj(), rw) * .25
    goo += numpy.einsum('kliabc,kljabc->ij', wvd.conj(), rw) * .25
    gOO += numpy.einsum('iklabc,jklabc->ij', wvd.conj(), rw) * .25
    gvv += numpy.einsum('ijkcad,ijkcbd->ab', wvd, rw.conj()) * .25
    gvv += numpy.einsum('ijkcda,ijkcdb->ab', wvd, rw.conj()) * .25
    gVV += numpy.einsum('ijkacd,ijkbcd->ab', wvd, rw.conj()) * .25
    gvo += numpy.einsum('kica,ijkabc->bj', t2ab.conj(), rw) * .5
    gVO += numpy.einsum('jkbc,ijkabc->ai', t2aa.conj(), rw) * .125

    # bba
    d3 = lib.direct_sum('ia+jb+kc->ijkabc', eia, eIA, eIA)
    w  = numpy.einsum('ijae,kceb->ijkabc', t2ab, numpy.asarray(eris.get_OVVV()).conj()) * 2
    w += numpy.einsum('ijeb,kcea->ijkabc', t2ab, numpy.asarray(eris.get_OVvv()).conj()) * 2
    w += numpy.einsum('jkbe,iaec->ijkabc', t2bb, numpy.asarray(eris.get_ovVV()).conj())
    w -= numpy.einsum('imab,kcjm->ijkabc', t2ab, numpy.asarray(eris.OVOO).conj()) * 2
    w -= numpy.einsum('mjab,kcim->ijkabc', t2ab, numpy.asarray(eris.OVoo).conj()) * 2
    w -= numpy.einsum('jmbc,iakm->ijkabc', t2bb, numpy.asarray(eris.ovOO).conj())
    v  = numpy.einsum('jbkc,ia->ijkabc', numpy.asarray(eris.OVOV).conj(), t1a)
    v += numpy.einsum('iakc,jb->ijkabc', numpy.asarray(eris.ovOV).conj(), t1b)
    v += numpy.einsum('iakc,jb->ijkabc', numpy.asarray(eris.ovOV).conj(), t1b)
    v += numpy.einsum('JKBC,ai->iJKaBC', t2bb, fvo) * .5
    v += numpy.einsum('iKaC,BJ->iJKaBC', t2ab, fVO) * 2

    rw = r4(w) / d3
    wvd = (w + v) / d3
    goo += numpy.einsum('iklabc,jklabc->ij', wvd.conj(), rw) * .25
    gOO += numpy.einsum('kilabc,kjlabc->ij', wvd.conj(), rw) * .25
    gOO += numpy.einsum('kliabc,kljabc->ij', wvd.conj(), rw) * .25
    gvv += numpy.einsum('ijkacd,ijkbcd->ab', wvd, rw.conj()) * .25
    gVV += numpy.einsum('ijkcad,ijkcbd->ab', wvd, rw.conj()) * .25
    gVV += numpy.einsum('ijkcda,ijkcdb->ab', wvd, rw.conj()) * .25
    gVO += numpy.einsum('ikac,ijkabc->bj', t2ab.conj(), rw) * .5
    gvo += numpy.einsum('jkbc,ijkabc->ai', t2bb.conj(), rw) * .125

    doo, dOO = d1[0]
    dov, dOV = d1[1]
    dvo, dVO = d1[2]
    dvv, dVV = d1[3]

    if for_grad:
        doo -= goo
        dOO -= gOO
        dvv += gvv
        dVV += gVV
    else:
        doo[numpy.diag_indices(nocca)] -= goo.diagonal()
        dOO[numpy.diag_indices(noccb)] -= gOO.diagonal()
        dvv[numpy.diag_indices(nvira)] += gvv.diagonal()
        dVV[numpy.diag_indices(nvirb)] += gVV.diagonal()

    dvo += gvo
    dVO += gVO

    return d1

# gamma2 intermediates in Chemist's notation
def _gamma2_intermediates(mycc, t1, t2, l1, l2, eris=None,
                          compress_vvvv=False):
    d2 = uccsd_rdm._gamma2_intermediates(mycc, t1, t2, l1, l2)

    if eris is None: eris = mycc.ao2mo()

    dovov, dovOV, dOVov, dOVOV = d2[0]
    dvvvv, dvvVV, dVVvv, dVVVV = d2[1]
    doooo, dooOO, dOOoo, dOOOO = d2[2]
    doovv, dooVV, dOOvv, dOOVV = d2[3]
    dovvo, dovVO, dOVvo, dOVVO = d2[4]
    dvvov, dvvOV, dVVov, dVVOV = d2[5]
    dovvv, dovVV, dOVvv, dOVVV = d2[6]
    dooov, dooOV, dOOov, dOOOV = d2[7]

    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nocca, noccb, nvira, nvirb = t2ab.shape
    nmoa = eris.focka.shape[0]
    nmob = eris.fockb.shape[0]
    mo_ea, mo_eb = eris.mo_energy
    eia = mo_ea[:nocca,None] - mo_ea[nocca:]
    eIA = mo_eb[:noccb,None] - mo_eb[noccb:]
    fvo = eris.focka[nocca:,:nocca]
    fVO = eris.fockb[noccb:,:noccb]

    # aaa
    d3 = lib.direct_sum('ia+jb+kc->ijkabc', eia, eia, eia)
    w = numpy.einsum('ijae,kceb->ijkabc', t2aa, numpy.asarray(eris.get_ovvv()).conj())
    w-= numpy.einsum('mkbc,iajm->ijkabc', t2aa, numpy.asarray(eris.ovoo).conj())
    v = numpy.einsum('jbkc,ia->ijkabc', numpy.asarray(eris.ovov).conj(), t1a)
    v+= numpy.einsum('jkbc,ai->ijkabc', t2aa, fvo) * .5

    rw = r6(p6(w)) / d3
    wvd = r6(p6(w * 2 + v)) / d3
    dovov += numpy.einsum('ia,ijkabc->jbkc', t1a, rw.conj()) * 0.25
# *(1/8) instead of (1/4) because ooov appears 4 times in the 2pdm tensor due
# to symmetrization, and its contribution is scaled by 1/2 in Tr(H,2pdm)
    dooov -= numpy.einsum('mkbc,ijkabc->jmia', t2aa, wvd.conj()) * .125
    dovvv += numpy.einsum('kjcf,ijkabc->iafb', t2aa, wvd.conj()) * .125

    # bbb
    d3 = lib.direct_sum('ia+jb+kc->ijkabc', eIA, eIA, eIA)
    w = numpy.einsum('ijae,kceb->ijkabc', t2bb, numpy.asarray(eris.get_OVVV()).conj())
    w-= numpy.einsum('imab,kcjm->ijkabc', t2bb, numpy.asarray(eris.OVOO).conj())
    v = numpy.einsum('jbkc,ia->ijkabc', numpy.asarray(eris.OVOV).conj(), t1b)
    v+= numpy.einsum('jkbc,ai->ijkabc', t2bb, fVO) * .5

    rw = r6(p6(w)) / d3
    wvd = r6(p6(w * 2 + v)) / d3
    dOVOV += numpy.einsum('ia,ijkabc->jbkc', t1b, rw.conj()) * .25
    dOOOV -= numpy.einsum('mkbc,ijkabc->jmia', t2bb, wvd.conj()) * .125
    dOVVV += numpy.einsum('kjcf,ijkabc->iafb', t2bb, wvd.conj()) * .125

    # baa
    d3 = lib.direct_sum('ia+jb+kc->ijkabc', eIA, eia, eia)
    w  = numpy.einsum('jIeA,kceb->IjkAbc', t2ab, numpy.asarray(eris.get_ovvv()).conj()) * 2
    w += numpy.einsum('jIbE,kcEA->IjkAbc', t2ab, numpy.asarray(eris.get_ovVV()).conj()) * 2
    w += numpy.einsum('jkbe,IAec->IjkAbc', t2aa, numpy.asarray(eris.get_OVvv()).conj())
    w -= numpy.einsum('mIbA,kcjm->IjkAbc', t2ab, numpy.asarray(eris.ovoo).conj()) * 2
    w -= numpy.einsum('jMbA,kcIM->IjkAbc', t2ab, numpy.asarray(eris.ovOO).conj()) * 2
    w -= numpy.einsum('jmbc,IAkm->IjkAbc', t2aa, numpy.asarray(eris.OVoo).conj())
    v  = numpy.einsum('jbkc,IA->IjkAbc', numpy.asarray(eris.ovov).conj(), t1b)
    v += numpy.einsum('kcIA,jb->IjkAbc', numpy.asarray(eris.ovOV).conj(), t1a)
    v += numpy.einsum('kcIA,jb->IjkAbc', numpy.asarray(eris.ovOV).conj(), t1a)
    v += numpy.einsum('jkbc,AI->IjkAbc', t2aa, fVO) * .5
    v += numpy.einsum('kIcA,bj->IjkAbc', t2ab, fvo) * 2

    rw = r4(w) / d3
    wvd = r4(w * 2 + v) / d3
    dovvv += numpy.einsum('jiea,ijkabc->kceb', t2ab, wvd.conj()) * .25
    dovVV += numpy.einsum('jibe,ijkabc->kcea', t2ab, wvd.conj()) * .25
    dOVvv += numpy.einsum('jkbe,ijkabc->iaec', t2aa, wvd.conj()) * .125
    dooov -= numpy.einsum('miba,ijkabc->jmkc', t2ab, wvd.conj()) * .25
    dOOov -= numpy.einsum('jmba,ijkabc->imkc', t2ab, wvd.conj()) * .25
    dooOV -= numpy.einsum('jmbc,ijkabc->kmia', t2aa, wvd.conj()) * .125
    dovov += numpy.einsum('ia,ijkabc->jbkc', t1b, rw.conj()) * .25
    #dOVov += numpy.einsum('jb,ijkabc->iakc', t1a, rw.conj()) * .25
    dovOV += numpy.einsum('jb,ijkabc->kcia', t1a, rw.conj()) * .25

    # bba
    d3 = lib.direct_sum('ia+jb+kc->ijkabc', eia, eIA, eIA)
    w  = numpy.einsum('ijae,kceb->ijkabc', t2ab, numpy.asarray(eris.get_OVVV()).conj()) * 2
    w += numpy.einsum('ijeb,kcea->ijkabc', t2ab, numpy.asarray(eris.get_OVvv()).conj()) * 2
    w += numpy.einsum('jkbe,iaec->ijkabc', t2bb, numpy.asarray(eris.get_ovVV()).conj())
    w -= numpy.einsum('imab,kcjm->ijkabc', t2ab, numpy.asarray(eris.OVOO).conj()) * 2
    w -= numpy.einsum('mjab,kcim->ijkabc', t2ab, numpy.asarray(eris.OVoo).conj()) * 2
    w -= numpy.einsum('jmbc,iakm->ijkabc', t2bb, numpy.asarray(eris.ovOO).conj())
    v  = numpy.einsum('jbkc,ia->ijkabc', numpy.asarray(eris.OVOV).conj(), t1a)
    v += numpy.einsum('iakc,jb->ijkabc', numpy.asarray(eris.ovOV).conj(), t1b)
    v += numpy.einsum('iakc,jb->ijkabc', numpy.asarray(eris.ovOV).conj(), t1b)
    v += numpy.einsum('JKBC,ai->iJKaBC', t2bb, fvo) * .5
    v += numpy.einsum('iKaC,BJ->iJKaBC', t2ab, fVO) * 2

    rw = r4(w) / d3
    wvd = r4(w * 2 + v) / d3
    dOVVV += numpy.einsum('ijae,ijkabc->kceb', t2ab, wvd.conj()) * .25
    dOVvv += numpy.einsum('ijeb,ijkabc->kcea', t2ab, wvd.conj()) * .25
    dovVV += numpy.einsum('jkbe,ijkabc->iaec', t2bb, wvd.conj()) * .125
    dOOOV -= numpy.einsum('imab,ijkabc->jmkc', t2ab, wvd.conj()) * .25
    dooOV -= numpy.einsum('mjab,ijkabc->imkc', t2ab, wvd.conj()) * .25
    dOOov -= numpy.einsum('jmbc,ijkabc->kmia', t2bb, wvd.conj()) * .125
    dOVOV += numpy.einsum('ia,ijkabc->jbkc', t1a, rw.conj()) * .25
    dovOV += numpy.einsum('jb,ijkabc->iakc', t1b, rw.conj()) * .25
    #dOVov += numpy.einsum('jb,ijkabc->kcia', t1b, rw.conj()) * .25

    if compress_vvvv:
        nmoa, nmob = mycc.nmo
        nocca, noccb, nvira, nvirb = t2ab.shape
        idxa = numpy.tril_indices(nvira)
        idxa = idxa[0] * nvira + idxa[1]
        idxb = numpy.tril_indices(nvirb)
        idxb = idxb[0] * nvirb + idxb[1]
        dvvvv = dvvvv + dvvvv.transpose(1,0,2,3)
        dvvvv = lib.take_2d(dvvvv.reshape(nvira**2,nvira**2), idxa, idxa)
        dvvvv *= .5
        dvvVV = dvvVV + dvvVV.transpose(1,0,2,3)
        dvvVV = lib.take_2d(dvvVV.reshape(nvira**2,nvirb**2), idxa, idxb)
        dVVVV = dVVVV + dVVVV.transpose(1,0,2,3)
        dVVVV = lib.take_2d(dVVVV.reshape(nvirb**2,nvirb**2), idxb, idxb)
        dVVVV *= .5

        d2 = ((dovov, dovOV, dOVov, dOVOV),
              (dvvvv, dvvVV, dVVvv, dVVVV),
              (doooo, dooOO, dOOoo, dOOOO),
              (doovv, dooVV, dOOvv, dOOVV),
              (dovvo, dovVO, dOVvo, dOVVO),
              (dvvov, dvvOV, dVVov, dVVOV),
              (dovvv, dovVV, dOVvv, dOVVV),
              (dooov, dooOV, dOOov, dOOOV))

    return d2

def _gamma2_outcore(mycc, t1, t2, l1, l2, eris, h5fobj, compress_vvvv=False):
    return _gamma2_intermediates(mycc, t1, t2, l1, l2, eris, compress_vvvv)

def make_rdm1(mycc, t1, t2, l1, l2, eris=None, ao_repr=False):
    d1 = _gamma1_intermediates(mycc, t1, t2, l1, l2, eris)
    return uccsd_rdm._make_rdm1(mycc, d1, True, ao_repr=ao_repr)

# rdm2 in Chemist's notation
def make_rdm2(mycc, t1, t2, l1, l2, eris=None):
    d1 = _gamma1_intermediates(mycc, t1, t2, l1, l2, eris)
    d2 = _gamma2_intermediates(mycc, t1, t2, l1, l2, eris)
    return uccsd_rdm._make_rdm2(mycc, d1, d2, True, True)

def p6(t):
    return (t + t.transpose(1,2,0,4,5,3) +
            t.transpose(2,0,1,5,3,4) + t.transpose(0,2,1,3,5,4) +
            t.transpose(2,1,0,5,4,3) + t.transpose(1,0,2,4,3,5))

def r6(w):
    return (w + w.transpose(2,0,1,3,4,5) + w.transpose(1,2,0,3,4,5)
            - w.transpose(2,1,0,3,4,5) - w.transpose(0,2,1,3,4,5)
            - w.transpose(1,0,2,3,4,5))

def r4(w):
    w = w - w.transpose(0,2,1,3,4,5)
    w = w + w.transpose(0,2,1,3,5,4)
    return w

if __name__ == '__main__':
    from functools import reduce
    from pyscf import gto
    from pyscf import scf
    from pyscf import ao2mo
    from pyscf import cc
    from pyscf.cc import uccsd_t_slow
    from pyscf.cc import uccsd_t_lambda

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -.957 , .587)],
        [1 , (0.2,  .757 , .487)]]
    mol.basis = '631g'
    mol.build()
    mf0 = mf = scf.RHF(mol).run(conv_tol=1.)
    mf = scf.addons.convert_to_uhf(mf)

    from pyscf.cc import ccsd_t_lambda_slow as ccsd_t_lambda
    from pyscf.cc import ccsd_t_rdm_slow as ccsd_t_rdm
    mycc0 = cc.CCSD(mf0)
    eris0 = mycc0.ao2mo()
    mycc0.kernel(eris=eris0)
    t1 = mycc0.t1
    t2 = mycc0.t2
    imds = ccsd_t_lambda.make_intermediates(mycc0, t1, t2, eris0)
    l1, l2 = ccsd_t_lambda.update_lambda(mycc0, t1, t2, t1, t2, eris0, imds)
    dm1ref = ccsd_t_rdm.make_rdm1(mycc0, t1, t2, l1, l2, eris0)
    dm2ref = ccsd_t_rdm.make_rdm2(mycc0, t1, t2, l1, l2, eris0)

    t1 = (t1, t1)
    t2aa = t2 - t2.transpose(1,0,2,3)
    t2 = (t2aa, t2, t2aa)
    l1 = (l1, l1)
    l2aa = l2 - l2.transpose(1,0,2,3)
    l2 = (l2aa, l2, l2aa)
    mycc = cc.UCCSD(mf)
    eris = mycc.ao2mo()
    dm1 = make_rdm1(mycc, t1, t2, l1, l2, eris)
    dm2 = make_rdm2(mycc, t1, t2, l1, l2, eris)
    trdm1 = dm1[0] + dm1[1]
    trdm2 = dm2[0] + dm2[1] + dm2[1].transpose(2,3,0,1) + dm2[2]
    print(abs(trdm1 - dm1ref).max())
    print(abs(trdm2 - dm2ref).max())

    ecc = mycc.kernel(eris=eris)[0]
    l1, l2 = mycc.solve_lambda(eris=eris)
    e3ref = mycc.e_tot + mycc.ccsd_t()

    nmoa, nmob = mycc.nmo
    eri_aa = ao2mo.kernel(mf._eri, mf.mo_coeff[0], compact=False).reshape([nmoa]*4)
    eri_bb = ao2mo.kernel(mf._eri, mf.mo_coeff[1], compact=False).reshape([nmob]*4)
    eri_ab = ao2mo.kernel(mf._eri, [mf.mo_coeff[k] for k in (0,0,1,1)],
                          compact=False).reshape(nmoa,nmoa,nmob,nmob)
    dm1 = make_rdm1(mycc, t1, t2, l1, l2, eris=eris)
    dm2 = make_rdm2(mycc, t1, t2, l1, l2, eris=eris)
    h1a = reduce(numpy.dot, (mf.mo_coeff[0].T, mf.get_hcore(), mf.mo_coeff[0]))
    h1b = reduce(numpy.dot, (mf.mo_coeff[1].T, mf.get_hcore(), mf.mo_coeff[1]))
    e3 =(numpy.einsum('ij,ji->', h1a, dm1[0]) +
         numpy.einsum('ij,ji->', h1b, dm1[1]) +
         numpy.einsum('ijkl,ijkl->', eri_aa, dm2[0])*.5 +
         numpy.einsum('ijkl,ijkl->', eri_bb, dm2[2])*.5 +
         numpy.einsum('ijkl,ijkl->', eri_ab, dm2[1])    +
         mf.mol.energy_nuc())
    print(e3 - e3ref)



# pyscf/cc/uccsd_t_slow.py

#!/usr/bin/env python
# Copyright 2014-2020 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: Qiming Sun <osirpt.sun@gmail.com>
#

import numpy
from pyscf import lib
from pyscf.cc import uccsd

'''
UCCSD(T)
'''

def kernel(mcc, eris, t1=None, t2=None):
    if t1 is None or t2 is None:
        t1, t2 = mcc.t1, mcc.t2

    def p6(t):
        return (t + t.transpose(1,2,0,4,5,3) +
                t.transpose(2,0,1,5,3,4) + t.transpose(0,2,1,3,5,4) +
                t.transpose(2,1,0,5,4,3) + t.transpose(1,0,2,4,3,5))
    def r6(w):
        return (w + w.transpose(2,0,1,3,4,5) + w.transpose(1,2,0,3,4,5)
                - w.transpose(2,1,0,3,4,5) - w.transpose(0,2,1,3,4,5)
                - w.transpose(1,0,2,3,4,5))

    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nocca, noccb = t2ab.shape[:2]
    mo_ea, mo_eb = eris.mo_energy
    eia = mo_ea[:nocca,None] - mo_ea[nocca:]
    eIA = mo_eb[:noccb,None] - mo_eb[noccb:]
    fvo = eris.focka[nocca:,:nocca]
    fVO = eris.fockb[noccb:,:noccb]

    # aaa
    d3 = lib.direct_sum('ia+jb+kc->ijkabc', eia, eia, eia)
    w = numpy.einsum('ijae,kceb->ijkabc', t2aa, numpy.asarray(eris.get_ovvv()).conj())
    w-= numpy.einsum('mkbc,iajm->ijkabc', t2aa, numpy.asarray(eris.ovoo).conj())
    r = r6(w)
    v = numpy.einsum('jbkc,ia->ijkabc', numpy.asarray(eris.ovov).conj(), t1a)
    v+= numpy.einsum('jkbc,ai->ijkabc', t2aa, fvo) * .5
    wvd = p6(w + v) / d3
    et = numpy.einsum('ijkabc,ijkabc', wvd.conj(), r)

    # bbb
    d3 = lib.direct_sum('ia+jb+kc->ijkabc', eIA, eIA, eIA)
    w = numpy.einsum('ijae,kceb->ijkabc', t2bb, numpy.asarray(eris.get_OVVV()).conj())
    w-= numpy.einsum('imab,kcjm->ijkabc', t2bb, numpy.asarray(eris.OVOO).conj())
    r = r6(w)
    v = numpy.einsum('jbkc,ia->ijkabc', numpy.asarray(eris.OVOV).conj(), t1b)
    v+= numpy.einsum('jkbc,ai->ijkabc', t2bb, fVO) * .5
    wvd = p6(w + v) / d3
    et += numpy.einsum('ijkabc,ijkabc', wvd.conj(), r)

    # baa
    w  = numpy.einsum('jIeA,kceb->IjkAbc', t2ab, numpy.asarray(eris.get_ovvv()).conj()) * 2
    w += numpy.einsum('jIbE,kcEA->IjkAbc', t2ab, numpy.asarray(eris.get_ovVV()).conj()) * 2
    w += numpy.einsum('jkbe,IAec->IjkAbc', t2aa, numpy.asarray(eris.get_OVvv()).conj())
    w -= numpy.einsum('mIbA,kcjm->IjkAbc', t2ab, numpy.asarray(eris.ovoo).conj()) * 2
    w -= numpy.einsum('jMbA,kcIM->IjkAbc', t2ab, numpy.asarray(eris.ovOO).conj()) * 2
    w -= numpy.einsum('jmbc,IAkm->IjkAbc', t2aa, numpy.asarray(eris.OVoo).conj())
    r = w - w.transpose(0,2,1,3,4,5)
    r = r + r.transpose(0,2,1,3,5,4)
    v  = numpy.einsum('jbkc,IA->IjkAbc', numpy.asarray(eris.ovov).conj(), t1b)
    v += numpy.einsum('kcIA,jb->IjkAbc', numpy.asarray(eris.ovOV).conj(), t1a)
    v += numpy.einsum('kcIA,jb->IjkAbc', numpy.asarray(eris.ovOV).conj(), t1a)
    v += numpy.einsum('jkbc,AI->IjkAbc', t2aa, fVO) * .5
    v += numpy.einsum('kIcA,bj->IjkAbc', t2ab, fvo) * 2
    w += v
    d3 = lib.direct_sum('ia+jb+kc->ijkabc', eIA, eia, eia)
    r /= d3
    et += numpy.einsum('ijkabc,ijkabc', w.conj(), r)

    # bba
    w  = numpy.einsum('ijae,kceb->ijkabc', t2ab, numpy.asarray(eris.get_OVVV()).conj()) * 2
    w += numpy.einsum('ijeb,kcea->ijkabc', t2ab, numpy.asarray(eris.get_OVvv()).conj()) * 2
    w += numpy.einsum('jkbe,iaec->ijkabc', t2bb, numpy.asarray(eris.get_ovVV()).conj())
    w -= numpy.einsum('imab,kcjm->ijkabc', t2ab, numpy.asarray(eris.OVOO).conj()) * 2
    w -= numpy.einsum('mjab,kcim->ijkabc', t2ab, numpy.asarray(eris.OVoo).conj()) * 2
    w -= numpy.einsum('jmbc,iakm->ijkabc', t2bb, numpy.asarray(eris.ovOO).conj())
    r = w - w.transpose(0,2,1,3,4,5)
    r = r + r.transpose(0,2,1,3,5,4)
    v  = numpy.einsum('jbkc,ia->ijkabc', numpy.asarray(eris.OVOV).conj(), t1a)
    v += numpy.einsum('iakc,jb->ijkabc', numpy.asarray(eris.ovOV).conj(), t1b)
    v += numpy.einsum('iakc,jb->ijkabc', numpy.asarray(eris.ovOV).conj(), t1b)
    v += numpy.einsum('JKBC,ai->iJKaBC', t2bb, fvo) * .5
    v += numpy.einsum('iKaC,BJ->iJKaBC', t2ab, fVO) * 2
    w += v
    d3 = lib.direct_sum('ia+jb+kc->ijkabc', eia, eIA, eIA)
    r /= d3
    et += numpy.einsum('ijkabc,ijkabc', w.conj(), r)

    et *= .25
    return et


if __name__ == '__main__':
    from pyscf import gto
    from pyscf import scf
    from pyscf import cc

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -.757 , .587)],
        [1 , (0. ,  .757 , .587)]]

    mol.basis = '631g'
    mol.build()
    rhf = scf.RHF(mol)
    rhf.conv_tol = 1e-14
    rhf.scf()
    mcc = cc.CCSD(rhf)
    mcc.conv_tol = 1e-14
    mcc.ccsd()
    t1a = t1b = mcc.t1
    t2ab = mcc.t2
    t2aa = t2bb = t2ab - t2ab.transpose(1,0,2,3)
    mcc = uccsd.UCCSD(scf.addons.convert_to_uhf(rhf))
    e3a = kernel(mcc, mcc.ao2mo(), (t1a,t1b), (t2aa,t2ab,t2bb))
    print(e3a - -0.00099642337843278096)

    mol = gto.Mole()
    mol.atom = [
        [8 , (0. , 0.     , 0.)],
        [1 , (0. , -.757 , .587)],
        [1 , (0. ,  .757 , .587)]]
    mol.spin = 2
    mol.basis = '3-21g'
    mol.build()
    mf = scf.UHF(mol).run(conv_tol=1e-14)
    nao, nmo = mf.mo_coeff[0].shape
    numpy.random.seed(10)
    mf.mo_coeff = numpy.random.random((2,nao,nmo))

    numpy.random.seed(12)
    nocca, noccb = mol.nelec
    nmo = mf.mo_occ[0].size
    nvira = nmo - nocca
    nvirb = nmo - noccb
    t1a  = .1 * numpy.random.random((nocca,nvira))
    t1b  = .1 * numpy.random.random((noccb,nvirb))
    t2aa = .1 * numpy.random.random((nocca,nocca,nvira,nvira))
    t2aa = t2aa - t2aa.transpose(0,1,3,2)
    t2aa = t2aa - t2aa.transpose(1,0,2,3)
    t2bb = .1 * numpy.random.random((noccb,noccb,nvirb,nvirb))
    t2bb = t2bb - t2bb.transpose(0,1,3,2)
    t2bb = t2bb - t2bb.transpose(1,0,2,3)
    t2ab = .1 * numpy.random.random((nocca,noccb,nvira,nvirb))
    t1 = t1a, t1b
    t2 = t2aa, t2ab, t2bb
    mcc = uccsd.UCCSD(scf.addons.convert_to_uhf(mf))
    e3a = kernel(mcc, mcc.ao2mo(), [t1a,t1b], [t2aa, t2ab, t2bb])
    print(e3a - 9877.2780859693339)

    mycc = cc.GCCSD(scf.addons.convert_to_ghf(mf))
    eris = mycc.ao2mo()
    t1 = mycc.spatial2spin(t1, eris.orbspin)
    t2 = mycc.spatial2spin(t2, eris.orbspin)
    from pyscf.cc import gccsd_t
    et = gccsd_t.kernel(mycc, eris, t1, t2)
    print(et - 9877.2780859693339)


    mol = gto.M()
    numpy.random.seed(12)
    nocca, noccb, nvira, nvirb = 3, 2, 4, 5
    nmo = nocca + nvira
    eris = cc.uccsd._ChemistsERIs()
    eri1 = (numpy.random.random((3,nmo,nmo,nmo,nmo)) +
            numpy.random.random((3,nmo,nmo,nmo,nmo)) * .8j - .5-.4j)
    eri1 = eri1 + eri1.transpose(0,2,1,4,3).conj()
    eri1[0] = eri1[0] + eri1[0].transpose(2,3,0,1)
    eri1[2] = eri1[2] + eri1[2].transpose(2,3,0,1)
    eri1 *= .1
    eris.ovov = eri1[0,nocca:,:nocca,nocca:,:nocca].transpose(1,0,3,2).conj()
    eris.ovvv = eri1[0,nocca:,:nocca,nocca:,nocca:].transpose(1,0,3,2).conj()
    eris.ovoo = eri1[0,nocca:,:nocca,:nocca,:nocca].transpose(1,0,3,2).conj()
    eris.OVOV = eri1[2,noccb:,:noccb,noccb:,:noccb].transpose(1,0,3,2).conj()
    eris.OVVV = eri1[2,noccb:,:noccb,noccb:,noccb:].transpose(1,0,3,2).conj()
    eris.OVOO = eri1[2,noccb:,:noccb,:noccb,:noccb].transpose(1,0,3,2).conj()
    eris.ovOV = eri1[1,nocca:,:nocca,noccb:,:noccb].transpose(1,0,3,2).conj()
    eris.ovVV = eri1[1,nocca:,:nocca,noccb:,noccb:].transpose(1,0,3,2).conj()
    eris.ovOO = eri1[1,nocca:,:nocca,:noccb,:noccb].transpose(1,0,3,2).conj()
    eris.OVov = eri1[1,nocca:,:nocca,noccb:,:noccb].transpose(3,2,1,0).conj()
    eris.OVvv = eri1[1,nocca:,nocca:,noccb:,:noccb].transpose(3,2,1,0).conj()
    eris.OVoo = eri1[1,:nocca,:nocca,noccb:,:noccb].transpose(3,2,1,0).conj()
    t1a  = .1 * numpy.random.random((nocca,nvira)) + numpy.random.random((nocca,nvira))*.1j
    t1b  = .1 * numpy.random.random((noccb,nvirb)) + numpy.random.random((noccb,nvirb))*.1j
    t2aa = .1 * numpy.random.random((nocca,nocca,nvira,nvira)) + numpy.random.random((nocca,nocca,nvira,nvira))*.1j
    t2aa = t2aa - t2aa.transpose(0,1,3,2)
    t2aa = t2aa - t2aa.transpose(1,0,2,3)
    t2bb = .1 * numpy.random.random((noccb,noccb,nvirb,nvirb)) + numpy.random.random((noccb,noccb,nvirb,nvirb))*.1j
    t2bb = t2bb - t2bb.transpose(0,1,3,2)
    t2bb = t2bb - t2bb.transpose(1,0,2,3)
    t2ab = .1 * numpy.random.random((nocca,noccb,nvira,nvirb)) + numpy.random.random((nocca,noccb,nvira,nvirb))*.1j
    f = (numpy.random.random((2,nmo,nmo)) * .4 +
         numpy.random.random((2,nmo,nmo)) * .4j)
    eris.focka = f[0]+f[0].T.conj() + numpy.diag(numpy.arange(nmo))
    eris.fockb = f[1]+f[1].T.conj() + numpy.diag(numpy.arange(nmo))
    t1 = t1a, t1b
    t2 = t2aa, t2ab, t2bb
    mcc = cc.UCCSD(scf.UHF(mol))
    print(kernel(mcc, eris, t1, t2) - (-0.056092415718338388-0.011390417704868244j))


# pyscf/cc/uintermediates.py

#!/usr/bin/env python
# Copyright 2014-2020 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.

import numpy as np
from pyscf import lib
from pyscf import ao2mo
from pyscf.cc.rintermediates import _get_vvvv  # noqa
from pyscf.cc.ccsd import BLKMIN

# Ref: Gauss and Stanton, J. Chem. Phys. 103, 3561 (1995) Table III


def make_tau(t2, t1, r1, fac=1, out=None):
    t1a, t1b = t1
    r1a, r1b = r1
    tau1aa = make_tau_aa(t2[0], t1a, r1a, fac, out)
    tau1bb = make_tau_aa(t2[2], t1b, r1b, fac, out)
    tau1ab = make_tau_ab(t2[1], t1, r1, fac, out)
    return tau1aa, tau1ab, tau1bb

def make_tau_aa(t2aa, t1a, r1a, fac=1, out=None):
    tau1aa = np.einsum('ia,jb->ijab', t1a, r1a)
    tau1aa-= np.einsum('ia,jb->jiab', t1a, r1a)
    tau1aa = tau1aa - tau1aa.transpose(0,1,3,2)
    tau1aa *= fac * .5
    tau1aa += t2aa
    return tau1aa

def make_tau_ab(t2ab, t1, r1, fac=1, out=None):
    t1a, t1b = t1
    r1a, r1b = r1
    tau1ab = np.einsum('ia,jb->ijab', t1a, r1b)
    tau1ab+= np.einsum('ia,jb->ijab', r1a, t1b)
    tau1ab *= fac * .5
    tau1ab += t2ab
    return tau1ab

def Foo(t1, t2, eris):
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nocca, noccb, nvira, nvirb = t2ab.shape

    eris_ovov = np.asarray(eris.ovov)
    eris_OVOV = np.asarray(eris.OVOV)
    eris_ovOV = np.asarray(eris.ovOV)
    tilaa, tilab, tilbb = make_tau(t2, t1, t1, fac=0.5)
    Fooa  = lib.einsum('inef,menf->mi', tilaa, eris_ovov)
    Fooa += lib.einsum('iNeF,meNF->mi', tilab, eris_ovOV)
    Foob  = lib.einsum('inef,menf->mi', tilbb, eris_OVOV)
    Foob += lib.einsum('nIfE,nfME->MI', tilab, eris_ovOV)

    eris_ovoo = np.asarray(eris.ovoo)
    eris_OVOO = np.asarray(eris.OVOO)
    eris_OVoo = np.asarray(eris.OVoo)
    eris_ovOO = np.asarray(eris.ovOO)
    Fooa += np.einsum('ne,nemi->mi', t1a, eris_ovoo)
    Fooa -= np.einsum('ne,meni->mi', t1a, eris_ovoo)
    Fooa += np.einsum('NE,NEmi->mi', t1b, eris_OVoo)
    Foob += np.einsum('ne,nemi->mi', t1b, eris_OVOO)
    Foob -= np.einsum('ne,meni->mi', t1b, eris_OVOO)
    Foob += np.einsum('ne,neMI->MI', t1a, eris_ovOO)

    fooa = eris.focka[:nocca,:nocca]
    foob = eris.fockb[:noccb,:noccb]
    fova = eris.focka[:nocca,nocca:]
    fovb = eris.fockb[:noccb,noccb:]
    Fova, Fovb = Fov(t1, t2, eris)

    Fooa += fooa + 0.5*lib.einsum('me,ie->mi', Fova+fova, t1a)
    Foob += foob + 0.5*lib.einsum('me,ie->mi', Fovb+fovb, t1b)
    return Fooa, Foob

def Fvv(t1, t2, eris):
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nocca, noccb, nvira, nvirb = t2ab.shape

    Fvva = 0
    Fvvb = 0

    tauaa, tauab, taubb = make_tau(t2, t1, t1)
    #:eris_ovvv = lib.unpack_tril(np.asarray(eris.ovvv).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvira,nvira)
    #:ovvv = eris_ovvv - eris_ovvv.transpose(0,3,2,1)
    #:self.Fvva  = np.einsum('mf,mfae->ae', t1a, ovvv)
    mem_now = lib.current_memory()[0]
    max_memory = max(0, lib.param.MAX_MEMORY - mem_now)
    blksize = min(nocca, max(BLKMIN, int(max_memory*1e6/8/(nvira**3*3))))
    for p0,p1 in lib.prange(0, nocca, blksize):
        ovvv = eris.get_ovvv(slice(p0,p1))  # ovvv = eris.ovvv[p0:p1]
        ovvv = ovvv - ovvv.transpose(0,3,2,1)
        Fvva += np.einsum('mf,mfae->ae', t1a[p0:p1], ovvv)
        ovvv = None

    #:eris_OVVV = lib.unpack_tril(np.asarray(eris.OVVV).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvirb,nvirb)
    #:OVVV = eris_OVVV - eris_OVVV.transpose(0,3,2,1)
    #:self.Fvvb  = np.einsum('mf,mfae->ae', t1b, OVVV)
    #:self.wOVVO = lib.einsum('jf,mebf->mbej', t1b, OVVV)
    #:self.wOVOO  = 0.5 * lib.einsum('mebf,ijef->mbij', OVVV, taubb)
    blksize = min(noccb, max(BLKMIN, int(max_memory*1e6/8/(nvirb**3*3))))
    for p0, p1 in lib.prange(0, noccb, blksize):
        OVVV = eris.get_OVVV(slice(p0,p1))  # OVVV = eris.OVVV[p0:p1]
        OVVV = OVVV - OVVV.transpose(0,3,2,1)
        Fvvb += np.einsum('mf,mfae->ae', t1b[p0:p1], OVVV)
        OVVV = None

    #:eris_ovVV = lib.unpack_tril(np.asarray(eris.ovVV).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvirb,nvirb)
    #:self.Fvvb += np.einsum('mf,mfAE->AE', t1a, eris_ovVV)
    #:self.woVvO = lib.einsum('JF,meBF->mBeJ', t1b, eris_ovVV)
    #:self.woVVo = lib.einsum('jf,mfBE->mBEj',-t1a, eris_ovVV)
    #:self.woVoO  = 0.5 * lib.einsum('meBF,iJeF->mBiJ', eris_ovVV, tauab)
    #:self.woVoO += 0.5 * lib.einsum('mfBE,iJfE->mBiJ', eris_ovVV, tauab)
    blksize = min(nocca, max(BLKMIN, int(max_memory*1e6/8/(nvira*nvirb**2*3))))
    for p0,p1 in lib.prange(0, nocca, blksize):
        ovVV = eris.get_ovVV(slice(p0,p1))  # ovVV = eris.ovVV[p0:p1]
        Fvvb += np.einsum('mf,mfAE->AE', t1a[p0:p1], ovVV)
        ovVV = None

    #:eris_OVvv = lib.unpack_tril(np.asarray(eris.OVvv).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvira,nvira)
    #:self.Fvva += np.einsum('MF,MFae->ae', t1b, eris_OVvv)
    #:self.wOvVo = lib.einsum('jf,MEbf->MbEj', t1a, eris_OVvv)
    #:self.wOvvO = lib.einsum('JF,MFbe->MbeJ',-t1b, eris_OVvv)
    #:self.wOvOo  = 0.5 * lib.einsum('MEbf,jIfE->MbIj', eris_OVvv, tauab)
    #:self.wOvOo += 0.5 * lib.einsum('MFbe,jIeF->MbIj', eris_OVvv, tauab)
    blksize = min(noccb, max(BLKMIN, int(max_memory*1e6/8/(nvirb*nvira**2*3))))
    for p0, p1 in lib.prange(0, noccb, blksize):
        OVvv = eris.get_OVvv(slice(p0,p1))  # OVvv = eris.OVvv[p0:p1]
        Fvva += np.einsum('MF,MFae->ae', t1b[p0:p1], OVvv)
        OVvv = None

    eris_ovov = np.asarray(eris.ovov)
    eris_OVOV = np.asarray(eris.OVOV)
    eris_ovOV = np.asarray(eris.ovOV)
    tilaa, tilab, tilbb = make_tau(t2, t1, t1, fac=0.5)
    Fvva -= lib.einsum('mnaf,menf->ae', tilaa, eris_ovov)
    Fvva -= lib.einsum('mNaF,meNF->ae', tilab, eris_ovOV)
    Fvvb -= lib.einsum('mnaf,menf->ae', tilbb, eris_OVOV)
    Fvvb -= lib.einsum('nMfA,nfME->AE', tilab, eris_ovOV)

    fvva = eris.focka[nocca:,nocca:]
    fvvb = eris.fockb[noccb:,noccb:]
    fova = eris.focka[:nocca,nocca:]
    fovb = eris.fockb[:noccb,noccb:]
    Fova, Fovb = Fov(t1, t2, eris)

    Fvva += fvva - 0.5*lib.einsum('me,ma->ae', Fova+fova, t1a)
    Fvvb += fvvb - 0.5*lib.einsum('me,ma->ae', Fovb+fovb, t1b)
    return Fvva, Fvvb

def Fov(t1, t2, eris):
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nocca, noccb, nvira, nvirb = t2ab.shape

    fova = eris.focka[:nocca,nocca:]
    fovb = eris.fockb[:noccb,noccb:]

    eris_ovov = np.asarray(eris.ovov)
    eris_OVOV = np.asarray(eris.OVOV)
    eris_ovOV = np.asarray(eris.ovOV)
    ovov = eris_ovov - eris_ovov.transpose(0,3,2,1)
    OVOV = eris_OVOV - eris_OVOV.transpose(0,3,2,1)

    Fova = np.einsum('nf,menf->me', t1a,      ovov)
    Fova+= np.einsum('NF,meNF->me', t1b, eris_ovOV)
    Fova += fova
    Fovb = np.einsum('nf,menf->me', t1b,      OVOV)
    Fovb+= np.einsum('nf,nfME->ME', t1a, eris_ovOV)
    Fovb += fovb
    return Fova, Fovb

def Woooo(t1, t2, eris):
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    eris_ovoo = np.asarray(eris.ovoo)
    eris_OVOO = np.asarray(eris.OVOO)
    eris_OVoo = np.asarray(eris.OVoo)
    eris_ovOO = np.asarray(eris.ovOO)
    ovoo = eris_ovoo - eris_ovoo.transpose(2,1,0,3)
    OVOO = eris_OVOO - eris_OVOO.transpose(2,1,0,3)
    woooo = lib.einsum('je,nemi->minj', t1a,      ovoo)
    wOOOO = lib.einsum('je,nemi->minj', t1b,      OVOO)
    wooOO = lib.einsum('JE,NEmi->miNJ', t1b, eris_OVoo)
    woOOo = lib.einsum('je,meNI->mINj',-t1a, eris_ovOO)

    woooo += np.asarray(eris.oooo)
    wOOOO += np.asarray(eris.OOOO)
    wooOO += np.asarray(eris.ooOO)
    woooo = woooo - woooo.transpose(0,3,2,1)
    wOOOO = wOOOO - wOOOO.transpose(0,3,2,1)
    wooOO = wooOO - woOOo.transpose(0,3,2,1)

    eris_ovov = np.asarray(eris.ovov)
    eris_OVOV = np.asarray(eris.OVOV)
    eris_ovOV = np.asarray(eris.ovOV)
    ovov = eris_ovov - eris_ovov.transpose(0,3,2,1)
    OVOV = eris_OVOV - eris_OVOV.transpose(0,3,2,1)
    tauaa, tauab, taubb = make_tau(t2, t1, t1)
    woooo += 0.5*lib.einsum('ijef,menf->minj', tauaa,      ovov)
    wOOOO += 0.5*lib.einsum('ijef,menf->minj', taubb,      OVOV)
    wooOO +=     lib.einsum('iJeF,meNF->miNJ', tauab, eris_ovOV)
    wOOoo = None
    return woooo, wooOO, wOOoo, wOOOO

def Wooov(t1, t2, eris):
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    dtype = np.result_type(t1a, t1b, t2aa, t2ab, t2bb)

    eris_ovoo = np.asarray(eris.ovoo)
    eris_OVOO = np.asarray(eris.OVOO)
    eris_OVoo = np.asarray(eris.OVoo)
    eris_ovOO = np.asarray(eris.ovOO)
    ovoo = eris_ovoo - eris_ovoo.transpose(2,1,0,3)
    OVOO = eris_OVOO - eris_OVOO.transpose(2,1,0,3)
    wooov = np.array(     ovoo.transpose(2,3,0,1), dtype=dtype)
    wOOOV = np.array(     OVOO.transpose(2,3,0,1), dtype=dtype)
    wooOV = np.array(eris_OVoo.transpose(2,3,0,1), dtype=dtype)
    wOOov = np.array(eris_ovOO.transpose(2,3,0,1), dtype=dtype)
    eris_ovoo = eris_OVOO = eris_ovOO = eris_OVoo = None

    eris_ovov = np.asarray(eris.ovov)
    eris_OVOV = np.asarray(eris.OVOV)
    eris_ovOV = np.asarray(eris.ovOV)
    ovov = eris_ovov - eris_ovov.transpose(0,3,2,1)
    OVOV = eris_OVOV - eris_OVOV.transpose(0,3,2,1)

    wooov += lib.einsum('if,mfne->mine', t1a,      ovov)
    wOOOV += lib.einsum('if,mfne->mine', t1b,      OVOV)
    wooOV += lib.einsum('if,mfNE->miNE', t1a, eris_ovOV)
    wOOov += lib.einsum('IF,neMF->MIne', t1b, eris_ovOV)
    return wooov, wooOV, wOOov, wOOOV

def Woovo(t1, t2, eris):
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nocca, noccb, nvira, nvirb = t2ab.shape
    dtype = np.result_type(t1a, t1b, t2aa, t2ab, t2bb)

    wovoo = np.zeros((nocca,nvira,nocca,nocca), dtype=dtype)
    wOVOO = np.zeros((noccb,nvirb,noccb,noccb), dtype=dtype)
    woVoO = np.zeros((nocca,nvirb,nocca,noccb), dtype=dtype)
    wOvOo = np.zeros((noccb,nvira,noccb,nocca), dtype=dtype)

    tauaa, tauab, taubb = make_tau(t2, t1, t1)
    #:eris_ovvv = lib.unpack_tril(np.asarray(eris.ovvv).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvira,nvira)
    #:ovvv = eris_ovvv - eris_ovvv.transpose(0,3,2,1)
    #:self.wovoo  = 0.5 * lib.einsum('mebf,ijef->mbij', eris_ovvv, tauaa)
    #:self.wovoo -= 0.5 * lib.einsum('mfbe,ijef->mbij', eris_ovvv, tauaa)
    mem_now = lib.current_memory()[0]
    max_memory = max(0, lib.param.MAX_MEMORY - mem_now)
    blksize = min(nocca, max(BLKMIN, int(max_memory*1e6/8/(nvira**3*3))))
    for p0,p1 in lib.prange(0, nocca, blksize):
        ovvv = eris.get_ovvv(slice(p0,p1))  # ovvv = eris.ovvv[p0:p1]
        ovvv = ovvv - ovvv.transpose(0,3,2,1)
        wovoo[p0:p1] = 0.5 * lib.einsum('mebf,ijef->mbij', ovvv, tauaa)
        ovvv = None

    #:eris_OVVV = lib.unpack_tril(np.asarray(eris.OVVV).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvirb,nvirb)
    #:OVVV = eris_OVVV - eris_OVVV.transpose(0,3,2,1)
    #:self.wOVOO  = 0.5 * lib.einsum('mebf,ijef->mbij', OVVV, taubb)
    blksize = min(noccb, max(BLKMIN, int(max_memory*1e6/8/(nvirb**3*3))))
    for p0, p1 in lib.prange(0, noccb, blksize):
        OVVV = eris.get_OVVV(slice(p0,p1))  # OVVV = eris.OVVV[p0:p1]
        OVVV = OVVV - OVVV.transpose(0,3,2,1)
        wOVOO[p0:p1] = 0.5 * lib.einsum('mebf,ijef->mbij', OVVV, taubb)
        OVVV = None

    #:eris_ovVV = lib.unpack_tril(np.asarray(eris.ovVV).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvirb,nvirb)
    #:self.Fvvb += np.einsum('mf,mfAE->AE', t1a, eris_ovVV)
    #:self.woVvO = lib.einsum('JF,meBF->mBeJ', t1b, eris_ovVV)
    #:self.woVVo = lib.einsum('jf,mfBE->mBEj',-t1a, eris_ovVV)
    #:self.woVoO  = 0.5 * lib.einsum('meBF,iJeF->mBiJ', eris_ovVV, tauab)
    #:self.woVoO += 0.5 * lib.einsum('mfBE,iJfE->mBiJ', eris_ovVV, tauab)
    blksize = min(nocca, max(BLKMIN, int(max_memory*1e6/8/(nvira*nvirb**2*3))))
    for p0,p1 in lib.prange(0, nocca, blksize):
        ovVV = eris.get_ovVV(slice(p0,p1))  # ovVV = eris.ovVV[p0:p1]
        woVoO[p0:p1] = 0.5 * lib.einsum('meBF,iJeF->mBiJ', ovVV, tauab)
        woVoO[p0:p1]+= 0.5 * lib.einsum('mfBE,iJfE->mBiJ', ovVV, tauab)
        ovVV = None

    #:eris_OVvv = lib.unpack_tril(np.asarray(eris.OVvv).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvira,nvira)
    #:self.Fvva += np.einsum('MF,MFae->ae', t1b, eris_OVvv)
    #:self.wOvVo = lib.einsum('jf,MEbf->MbEj', t1a, eris_OVvv)
    #:self.wOvvO = lib.einsum('JF,MFbe->MbeJ',-t1b, eris_OVvv)
    #:self.wOvOo  = 0.5 * lib.einsum('MEbf,jIfE->MbIj', eris_OVvv, tauab)
    #:self.wOvOo += 0.5 * lib.einsum('MFbe,jIeF->MbIj', eris_OVvv, tauab)
    blksize = min(noccb, max(BLKMIN, int(max_memory*1e6/8/(nvirb*nvira**2*3))))
    for p0, p1 in lib.prange(0, noccb, blksize):
        OVvv = eris.get_OVvv(slice(p0,p1))  # OVvv = eris.OVvv[p0:p1]
        wOvOo[p0:p1] = 0.5 * lib.einsum('MEbf,jIfE->MbIj', OVvv, tauab)
        wOvOo[p0:p1]+= 0.5 * lib.einsum('MFbe,jIeF->MbIj', OVvv, tauab)
        OVvv = None

    eris_ovov = np.asarray(eris.ovov)
    eris_OVOV = np.asarray(eris.OVOV)
    eris_ovOV = np.asarray(eris.ovOV)
    ovov = eris_ovov - eris_ovov.transpose(0,3,2,1)
    OVOV = eris_OVOV - eris_OVOV.transpose(0,3,2,1)
    tilaa, tilab, tilbb = make_tau(t2, t1, t1, fac=0.5)

    eris_ovoo = np.asarray(eris.ovoo)
    eris_OVOO = np.asarray(eris.OVOO)
    eris_OVoo = np.asarray(eris.OVoo)
    eris_ovOO = np.asarray(eris.ovOO)
    ovoo = eris_ovoo - eris_ovoo.transpose(2,1,0,3)
    OVOO = eris_OVOO - eris_OVOO.transpose(2,1,0,3)
    tmpaa = lib.einsum('nemi,jnbe->mbij',      ovoo, t2aa)
    tmpaa+= lib.einsum('NEmi,jNbE->mbij', eris_OVoo, t2ab)
    tmpbb = lib.einsum('nemi,jnbe->mbij',      OVOO, t2bb)
    tmpbb+= lib.einsum('neMI,nJeB->MBIJ', eris_ovOO, t2ab)
    woVoO += lib.einsum('nemi,nJeB->mBiJ',      ovoo, t2ab)
    woVoO += lib.einsum('NEmi,JNBE->mBiJ', eris_OVoo, t2bb)
    woVoO -= lib.einsum('meNI,jNeB->mBjI', eris_ovOO, t2ab)
    wOvOo += lib.einsum('NEMI,jNbE->MbIj',      OVOO, t2ab)
    wOvOo += lib.einsum('neMI,jnbe->MbIj', eris_ovOO, t2aa)
    wOvOo -= lib.einsum('MEni,nJbE->MbJi', eris_OVoo, t2ab)
    wovoo += tmpaa - tmpaa.transpose(0,1,3,2)
    wOVOO += tmpbb - tmpbb.transpose(0,1,3,2)

    tmp1aa = lib.einsum('njbf,menf->mbej', t2aa,      ovov)
    tmp1aa-= lib.einsum('jNbF,meNF->mbej', t2ab, eris_ovOV)
    tmp1bb = lib.einsum('njbf,menf->mbej', t2bb,      OVOV)
    tmp1bb-= lib.einsum('nJfB,nfME->MBEJ', t2ab, eris_ovOV)
    tmp1ab = lib.einsum('NJBF,meNF->mBeJ', t2bb, eris_ovOV)
    tmp1ab-= lib.einsum('nJfB,menf->mBeJ', t2ab,      ovov)
    tmp1ba = lib.einsum('njbf,nfME->MbEj', t2aa, eris_ovOV)
    tmp1ba-= lib.einsum('jNbF,MENF->MbEj', t2ab,      OVOV)
    tmp1abba =-lib.einsum('jNfB,mfNE->mBEj', t2ab, eris_ovOV)
    tmp1baab =-lib.einsum('nJbF,neMF->MbeJ', t2ab, eris_ovOV)
    tmpaa = lib.einsum('ie,mbej->mbij', t1a, tmp1aa)
    tmpbb = lib.einsum('ie,mbej->mbij', t1b, tmp1bb)
    tmpab = lib.einsum('ie,mBeJ->mBiJ', t1a, tmp1ab)
    tmpab-= lib.einsum('IE,mBEj->mBjI', t1b, tmp1abba)
    tmpba = lib.einsum('IE,MbEj->MbIj', t1b, tmp1ba)
    tmpba-= lib.einsum('ie,MbeJ->MbJi', t1a, tmp1baab)
    wovoo -= tmpaa - tmpaa.transpose(0,1,3,2)
    wOVOO -= tmpbb - tmpbb.transpose(0,1,3,2)
    woVoO -= tmpab
    wOvOo -= tmpba
    eris_ovov = eris_OVOV = eris_ovOV = None
    eris_ovoo = np.asarray(eris.ovoo)
    eris_OVOO = np.asarray(eris.OVOO)
    eris_ovOO = np.asarray(eris.ovOO)
    eris_OVoo = np.asarray(eris.OVoo)
    wovoo += eris_ovoo.transpose(3,1,2,0) - eris_ovoo.transpose(2,1,0,3)
    wOVOO += eris_OVOO.transpose(3,1,2,0) - eris_OVOO.transpose(2,1,0,3)
    woVoO += eris_OVoo.transpose(3,1,2,0)
    wOvOo += eris_ovOO.transpose(3,1,2,0)
    eris_ovoo = eris_OVOO = eris_ovOO = eris_OVoo = None

    eris_ovvo = np.asarray(eris.ovvo)
    eris_OVVO = np.asarray(eris.OVVO)
    eris_OVvo = np.asarray(eris.OVvo)
    eris_ovVO = np.asarray(eris.ovVO)
    eris_oovv = np.asarray(eris.oovv)
    eris_OOVV = np.asarray(eris.OOVV)
    eris_OOvv = np.asarray(eris.OOvv)
    eris_ooVV = np.asarray(eris.ooVV)
    tmpaa = lib.einsum('ie,mebj->mbij', t1a, eris_ovvo)
    tmpbb = lib.einsum('ie,mebj->mbij', t1b, eris_OVVO)
    tmpaa-= lib.einsum('ie,mjbe->mbij', t1a, eris_oovv)
    tmpbb-= lib.einsum('ie,mjbe->mbij', t1b, eris_OOVV)
    woVoO += lib.einsum('ie,meBJ->mBiJ', t1a, eris_ovVO)
    woVoO -= lib.einsum('IE,mjBE->mBjI',-t1b, eris_ooVV)
    wOvOo += lib.einsum('IE,MEbj->MbIj', t1b, eris_OVvo)
    wOvOo -= lib.einsum('ie,MJbe->MbJi',-t1a, eris_OOvv)
    wovoo += tmpaa - tmpaa.transpose(0,1,3,2)
    wOVOO += tmpbb - tmpbb.transpose(0,1,3,2)

    woooo, wooOO, wOOoo, wOOOO = Woooo(t1, t2, eris)
    Fova, Fovb = Fov(t1, t2, eris)
    wovoo -= lib.einsum('me,ijbe->mbij', Fova, t2aa)
    wOVOO -= lib.einsum('me,ijbe->mbij', Fovb, t2bb)
    woVoO += lib.einsum('me,iJeB->mBiJ', Fova, t2ab)
    wOvOo += lib.einsum('ME,jIbE->MbIj', Fovb, t2ab)
    wovoo -= lib.einsum('nb,minj->mbij', t1a, woooo)
    wOVOO -= lib.einsum('nb,minj->mbij', t1b, wOOOO)
    woVoO -= lib.einsum('NB,miNJ->mBiJ', t1b, wooOO)
    wOvOo -= lib.einsum('nb,njMI->MbIj', t1a, wooOO)
    eris_ovvo = eris_OVVO = eris_OVvo = eris_ovVO = None
    eris_oovv = eris_OOVV = eris_OOvv = eris_ooVV = None

    woovo = wovoo.transpose(0,2,1,3)
    wOOVO = wOVOO.transpose(0,2,1,3)
    wooVO = woVoO.transpose(0,2,1,3)
    wOOvo = wOvOo.transpose(0,2,1,3)
    return woovo, wooVO, wOOvo, wOOVO

def Wovvo(t1, t2, eris):
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nocca, noccb, nvira, nvirb = t2ab.shape
    dtype = np.result_type(t1a, t1b, t2aa, t2ab, t2bb)

    wovvo = np.zeros((nocca,nvira,nvira,nocca), dtype=dtype)
    wOVVO = np.zeros((noccb,nvirb,nvirb,noccb), dtype=dtype)
    woVvO = np.zeros((nocca,nvirb,nvira,noccb), dtype=dtype)
    woVVo = np.zeros((nocca,nvirb,nvirb,nocca), dtype=dtype)
    wOvVo = np.zeros((noccb,nvira,nvirb,nocca), dtype=dtype)
    wOvvO = np.zeros((noccb,nvira,nvira,noccb), dtype=dtype)

    tauaa, tauab, taubb = make_tau(t2, t1, t1)
    #:eris_ovvv = lib.unpack_tril(np.asarray(eris.ovvv).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvira,nvira)
    #:ovvv = eris_ovvv - eris_ovvv.transpose(0,3,2,1)
    #:self.Fvva  = np.einsum('mf,mfae->ae', t1a, ovvv)
    #:self.wovvo = lib.einsum('jf,mebf->mbej', t1a, ovvv)
    #:self.wovoo  = 0.5 * lib.einsum('mebf,ijef->mbij', eris_ovvv, tauaa)
    #:self.wovoo -= 0.5 * lib.einsum('mfbe,ijef->mbij', eris_ovvv, tauaa)
    mem_now = lib.current_memory()[0]
    max_memory = max(0, lib.param.MAX_MEMORY - mem_now)
    blksize = min(nocca, max(BLKMIN, int(max_memory*1e6/8/(nvira**3*3))))
    for p0,p1 in lib.prange(0, nocca, blksize):
        ovvv = eris.get_ovvv(slice(p0,p1))  # ovvv = eris.ovvv[p0:p1]
        ovvv = ovvv - ovvv.transpose(0,3,2,1)
        wovvo[p0:p1] = lib.einsum('jf,mebf->mbej', t1a, ovvv)
        ovvv = None

    #:eris_OVVV = lib.unpack_tril(np.asarray(eris.OVVV).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvirb,nvirb)
    #:OVVV = eris_OVVV - eris_OVVV.transpose(0,3,2,1)
    #:self.Fvvb  = np.einsum('mf,mfae->ae', t1b, OVVV)
    #:self.wOVVO = lib.einsum('jf,mebf->mbej', t1b, OVVV)
    #:self.wOVOO  = 0.5 * lib.einsum('mebf,ijef->mbij', OVVV, taubb)
    blksize = min(noccb, max(BLKMIN, int(max_memory*1e6/8/(nvirb**3*3))))
    for p0, p1 in lib.prange(0, noccb, blksize):
        OVVV = eris.get_OVVV(slice(p0,p1))  # OVVV = eris.OVVV[p0:p1]
        OVVV = OVVV - OVVV.transpose(0,3,2,1)
        wOVVO[p0:p1] = lib.einsum('jf,mebf->mbej', t1b, OVVV)
        OVVV = None

    #:eris_ovVV = lib.unpack_tril(np.asarray(eris.ovVV).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvirb,nvirb)
    #:self.Fvvb += np.einsum('mf,mfAE->AE', t1a, eris_ovVV)
    #:self.woVvO = lib.einsum('JF,meBF->mBeJ', t1b, eris_ovVV)
    #:self.woVVo = lib.einsum('jf,mfBE->mBEj',-t1a, eris_ovVV)
    #:self.woVoO  = 0.5 * lib.einsum('meBF,iJeF->mBiJ', eris_ovVV, tauab)
    #:self.woVoO += 0.5 * lib.einsum('mfBE,iJfE->mBiJ', eris_ovVV, tauab)
    blksize = min(nocca, max(BLKMIN, int(max_memory*1e6/8/(nvira*nvirb**2*3))))
    for p0,p1 in lib.prange(0, nocca, blksize):
        ovVV = eris.get_ovVV(slice(p0,p1))  # ovVV = eris.ovVV[p0:p1]
        woVvO[p0:p1] = lib.einsum('JF,meBF->mBeJ', t1b, ovVV)
        woVVo[p0:p1] = lib.einsum('jf,mfBE->mBEj',-t1a, ovVV)
        ovVV = None

    #:eris_OVvv = lib.unpack_tril(np.asarray(eris.OVvv).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvira,nvira)
    #:self.Fvva += np.einsum('MF,MFae->ae', t1b, eris_OVvv)
    #:self.wOvVo = lib.einsum('jf,MEbf->MbEj', t1a, eris_OVvv)
    #:self.wOvvO = lib.einsum('JF,MFbe->MbeJ',-t1b, eris_OVvv)
    #:self.wOvOo  = 0.5 * lib.einsum('MEbf,jIfE->MbIj', eris_OVvv, tauab)
    #:self.wOvOo += 0.5 * lib.einsum('MFbe,jIeF->MbIj', eris_OVvv, tauab)
    blksize = min(noccb, max(BLKMIN, int(max_memory*1e6/8/(nvirb*nvira**2*3))))
    for p0, p1 in lib.prange(0, noccb, blksize):
        OVvv = eris.get_OVvv(slice(p0,p1))  # OVvv = eris.OVvv[p0:p1]
        wOvVo[p0:p1] = lib.einsum('jf,MEbf->MbEj', t1a, OVvv)
        wOvvO[p0:p1] = lib.einsum('JF,MFbe->MbeJ',-t1b, OVvv)
        OVvv = None

    eris_ovov = np.asarray(eris.ovov)
    eris_OVOV = np.asarray(eris.OVOV)
    eris_ovOV = np.asarray(eris.ovOV)
    ovov = eris_ovov - eris_ovov.transpose(0,3,2,1)
    OVOV = eris_OVOV - eris_OVOV.transpose(0,3,2,1)
    wovvo -= lib.einsum('jnfb,menf->mbej', t2aa,      ovov)
    wovvo += lib.einsum('jNbF,meNF->mbej', t2ab, eris_ovOV)
    wOVVO -= lib.einsum('jnfb,menf->mbej', t2bb,      OVOV)
    wOVVO += lib.einsum('nJfB,nfME->MBEJ', t2ab, eris_ovOV)
    woVvO += lib.einsum('nJfB,menf->mBeJ', t2ab,      ovov)
    woVvO -= lib.einsum('JNFB,meNF->mBeJ', t2bb, eris_ovOV)
    wOvVo -= lib.einsum('jnfb,nfME->MbEj', t2aa, eris_ovOV)
    wOvVo += lib.einsum('jNbF,MENF->MbEj', t2ab,      OVOV)
    woVVo += lib.einsum('jNfB,mfNE->mBEj', t2ab, eris_ovOV)
    wOvvO += lib.einsum('nJbF,neMF->MbeJ', t2ab, eris_ovOV)

    eris_ovoo = np.asarray(eris.ovoo)
    eris_OVOO = np.asarray(eris.OVOO)
    eris_OVoo = np.asarray(eris.OVoo)
    eris_ovOO = np.asarray(eris.ovOO)
    eris_ovoo = eris_ovoo + np.einsum('nfme,jf->menj', eris_ovov, t1a)
    eris_OVOO = eris_OVOO + np.einsum('nfme,jf->menj', eris_OVOV, t1b)
    eris_OVoo = eris_OVoo + np.einsum('nfme,jf->menj', eris_ovOV, t1a)
    eris_ovOO = eris_ovOO + np.einsum('menf,jf->menj', eris_ovOV, t1b)
    ovoo = eris_ovoo - eris_ovoo.transpose(2,1,0,3)
    OVOO = eris_OVOO - eris_OVOO.transpose(2,1,0,3)
    wovvo += lib.einsum('nb,nemj->mbej', t1a,      ovoo)
    wOVVO += lib.einsum('nb,nemj->mbej', t1b,      OVOO)
    woVvO -= lib.einsum('NB,meNJ->mBeJ', t1b, eris_ovOO)
    wOvVo -= lib.einsum('nb,MEnj->MbEj', t1a, eris_OVoo)
    woVVo += lib.einsum('NB,NEmj->mBEj', t1b, eris_OVoo)
    wOvvO += lib.einsum('nb,neMJ->MbeJ', t1a, eris_ovOO)
    eris_ovoo = eris_OVOO = eris_ovOO = eris_OVoo = None

    eris_ovvo = np.asarray(eris.ovvo)
    eris_OVVO = np.asarray(eris.OVVO)
    eris_OVvo = np.asarray(eris.OVvo)
    eris_ovVO = np.asarray(eris.ovVO)
    eris_oovv = np.asarray(eris.oovv)
    eris_OOVV = np.asarray(eris.OOVV)
    eris_OOvv = np.asarray(eris.OOvv)
    eris_ooVV = np.asarray(eris.ooVV)

    wovvo = wovvo.transpose(0,2,1,3)
    wOVVO = wOVVO.transpose(0,2,1,3)
    wovVO = woVvO.transpose(0,2,1,3)
    wOVvo = wOvVo.transpose(0,2,1,3)
    woVVo = woVVo.transpose(0,2,1,3)
    wOvvO = wOvvO.transpose(0,2,1,3)

    wovvo += eris_ovvo
    wOVVO += eris_OVVO
    wovVO += eris_ovVO
    wOVvo += eris_OVvo
    wovvo -= eris_oovv.transpose(0,3,2,1)
    wOVVO -= eris_OOVV.transpose(0,3,2,1)
    woVVo -= eris_ooVV.transpose(0,3,2,1)
    wOvvO -= eris_OOvv.transpose(0,3,2,1)
    return wovvo, wovVO, wOVvo, wOVVO, woVVo, wOvvO

def Wvvov(t1, t2, eris):
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nocca, noccb, nvira, nvirb = t2ab.shape

    #:Wamef = einsum('na,nmef->amef', -t1, eris.oovv)
    #:Wamef -= np.asarray(eris.ovvv).transpose(1,0,2,3)
    eris_ovov = np.asarray(eris.ovov)
    eris_OVOV = np.asarray(eris.OVOV)
    eris_ovOV = np.asarray(eris.ovOV)
    Waemf = lib.einsum('na,nemf->aemf',-t1a, eris_ovov)
    Waemf+= np.asarray(eris.get_ovvv()).transpose(2,3,0,1)
    Waemf = Waemf - Waemf.transpose(0,3,2,1)

    WaeMF = lib.einsum('na,nemf->aemf',-t1a, eris_ovOV)
    WaeMF+= np.asarray(eris.get_OVvv()).transpose(2,3,0,1)

    WAEmf = lib.einsum('na,mfne->aemf',-t1b, eris_ovOV)
    WAEmf+= np.asarray(eris.get_ovVV()).transpose(2,3,0,1)

    WAEMF = lib.einsum('na,nemf->aemf',-t1b, eris_OVOV)
    WAEMF+= np.asarray(eris.get_OVVV()).transpose(2,3,0,1)
    WAEMF = WAEMF - WAEMF.transpose(0,3,2,1)
    return Waemf, WaeMF, WAEmf, WAEMF

def Wvvvo(t1, t2, eris):
    t1a, t1b = t1
    t2aa, t2ab, t2bb = t2
    nocca, noccb, nvira, nvirb = t2ab.shape
    dtype = np.result_type(t1a, t1b, t2aa, t2ab, t2bb)

    _wvovv = np.empty((nvira,nocca,nvira,nvira), dtype=dtype)
    _wVOVV = np.empty((nvirb,noccb,nvirb,nvirb), dtype=dtype)
    _wvOvV = np.empty((nvira,noccb,nvira,nvirb), dtype=dtype)
    _wVoVv = np.empty((nvirb,nocca,nvirb,nvira), dtype=dtype)

    Fova, Fovb = Fov(t1, t2, eris)

    # 3 or 4 virtuals
    eris_ovoo = np.asarray(eris.ovoo)
    eris_ovov = np.asarray(eris.ovov)
    eris_ovOV = np.asarray(eris.ovOV)
    ovov = eris_ovov - eris_ovov.transpose(0,3,2,1)
    eris_oovv = np.asarray(eris.oovv)
    eris_ovvo = np.asarray(eris.ovvo)
    oovv = eris_oovv - eris_ovvo.transpose(0,3,2,1)
    eris_oovv = eris_ovvo = None
    tauaa, tauab, taubb = make_tau(t2, t1, t1)
    #:wvovv  = .5 * lib.einsum('meni,mnab->eiab', eris_ovoo, tauaa)
    #:wvovv -= .5 * lib.einsum('me,miab->eiab', self.Fova, t2aa)
    #:tmp1aa = lib.einsum('nibf,menf->mbei', t2aa,      ovov)
    #:tmp1aa-= lib.einsum('iNbF,meNF->mbei', t2ab, eris_ovOV)
    #:wvovv+= lib.einsum('ma,mbei->eiab', t1a, tmp1aa)
    #:wvovv+= lib.einsum('ma,mibe->eiab', t1a,      oovv)
    for p0, p1 in lib.prange(0, nvira, nocca):
        wvovv  = .5*lib.einsum('meni,mnab->eiab', eris_ovoo[:,p0:p1], tauaa)
        wvovv -= .5*lib.einsum('me,miab->eiab', Fova[:,p0:p1], t2aa)

        tmp1aa = lib.einsum('nibf,menf->mbei', t2aa, ovov[:,p0:p1])
        tmp1aa-= lib.einsum('iNbF,meNF->mbei', t2ab, eris_ovOV[:,p0:p1])
        wvovv += lib.einsum('ma,mbei->eiab', t1a, tmp1aa)
        wvovv += lib.einsum('ma,mibe->eiab', t1a, oovv[:,:,:,p0:p1])
        _wvovv[p0:p1] = wvovv
        tmp1aa = None
    eris_ovov = eris_ovoo = eris_ovOV = None

    #:eris_ovvv = lib.unpack_tril(np.asarray(eris.ovvv).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvira,nvira)
    #:ovvv = eris_ovvv - eris_ovvv.transpose(0,3,2,1)
    #:wvovv += lib.einsum('mebf,miaf->eiab',      ovvv, t2aa)
    #:eris_OVvv = lib.unpack_tril(np.asarray(eris.OVvv).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvira,nvira)
    #:wvovv += lib.einsum('MFbe,iMaF->eiab', eris_OVvv, t2ab)
    #:wvovv += eris_ovvv.transpose(2,0,3,1).conj()
    #:self.wvovv -= wvovv - wvovv.transpose(0,1,3,2)
    mem_now = lib.current_memory()[0]
    max_memory = max(0, lib.param.MAX_MEMORY - mem_now)
    blksize = min(nocca, max(BLKMIN, int(max_memory*1e6/8/(nvira**3*6))))
    for i0,i1 in lib.prange(0, nocca, blksize):
        wvovv = _wvovv[:,i0:i1]
        for p0,p1 in lib.prange(0, noccb, blksize):
            OVvv = eris.get_OVvv(slice(p0,p1))  # OVvv = eris.OVvv[p0:p1]
            wvovv -= lib.einsum('MFbe,iMaF->eiab', OVvv, t2ab[i0:i1,p0:p1])
            OVvv = None
        for p0,p1 in lib.prange(0, nocca, blksize):
            ovvv = eris.get_ovvv(slice(p0,p1))  # ovvv = eris.ovvv[p0:p1]
            if p0 == i0:
                wvovv += ovvv.transpose(2,0,3,1).conj()
            ovvv = ovvv - ovvv.transpose(0,3,2,1)
            wvovv -= lib.einsum('mebf,miaf->eiab', ovvv, t2aa[p0:p1,i0:i1])
            ovvv = None
        wvovv = wvovv - wvovv.transpose(0,1,3,2)
        _wvovv[:,i0:i1] = wvovv

    eris_OVOO = np.asarray(eris.OVOO)
    eris_OVOV = np.asarray(eris.OVOV)
    eris_ovOV = np.asarray(eris.ovOV)
    OVOV = eris_OVOV - eris_OVOV.transpose(0,3,2,1)
    eris_OOVV = np.asarray(eris.OOVV)
    eris_OVVO = np.asarray(eris.OVVO)
    OOVV = eris_OOVV - eris_OVVO.transpose(0,3,2,1)
    eris_OOVV = eris_OVVO = None
    #:wVOVV  = .5*lib.einsum('meni,mnab->eiab', eris_OVOO, taubb)
    #:wVOVV -= .5*lib.einsum('me,miab->eiab', self.Fovb, t2bb)
    #:tmp1bb = lib.einsum('nibf,menf->mbei', t2bb,      OVOV)
    #:tmp1bb-= lib.einsum('nIfB,nfME->MBEI', t2ab, eris_ovOV)
    #:wVOVV += lib.einsum('ma,mbei->eiab', t1b, tmp1bb)
    #:wVOVV += lib.einsum('ma,mibe->eiab', t1b,      OOVV)
    for p0, p1 in lib.prange(0, nvirb, noccb):
        wVOVV  = .5*lib.einsum('meni,mnab->eiab', eris_OVOO[:,p0:p1], taubb)
        wVOVV -= .5*lib.einsum('me,miab->eiab', Fovb[:,p0:p1], t2bb)

        tmp1bb = lib.einsum('nibf,menf->mbei', t2bb, OVOV[:,p0:p1])
        tmp1bb-= lib.einsum('nIfB,nfME->MBEI', t2ab, eris_ovOV[:,:,:,p0:p1])
        wVOVV += lib.einsum('ma,mbei->eiab', t1b, tmp1bb)
        wVOVV += lib.einsum('ma,mibe->eiab', t1b, OOVV[:,:,:,p0:p1])
        _wVOVV[p0:p1] = wVOVV
        tmp1bb = None
    eris_OVOV = eris_OVOO = eris_ovOV = None

    #:eris_OVVV = lib.unpack_tril(np.asarray(eris.OVVV).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvirb,nvirb)
    #:OVVV = eris_OVVV - eris_OVVV.transpose(0,3,2,1)
    #:wVOVV -= lib.einsum('MEBF,MIAF->EIAB',      OVVV, t2bb)
    #:eris_ovVV = lib.unpack_tril(np.asarray(eris.ovVV).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvirb,nvirb)
    #:wVOVV -= lib.einsum('mfBE,mIfA->EIAB', eris_ovVV, t2ab)
    #:wVOVV += eris_OVVV.transpose(2,0,3,1).conj()
    #:self.wVOVV += wVOVV - wVOVV.transpose(0,1,3,2)
    blksize = min(noccb, max(BLKMIN, int(max_memory*1e6/8/(nvirb**3*6))))
    for i0,i1 in lib.prange(0, noccb, blksize):
        wVOVV = _wVOVV[:,i0:i1]
        for p0,p1 in lib.prange(0, nocca, blksize):
            ovVV = eris.get_ovVV(slice(p0,p1))  # ovVV = eris.ovVV[p0:p1]
            wVOVV -= lib.einsum('mfBE,mIfA->EIAB', ovVV, t2ab[p0:p1,i0:i1])
            ovVV = None
        for p0,p1 in lib.prange(0, noccb, blksize):
            OVVV = eris.get_OVVV(slice(p0,p1))  # OVVV = eris.OVVV[p0:p1]
            if p0 == i0:
                wVOVV += OVVV.transpose(2,0,3,1).conj()
            OVVV = OVVV - OVVV.transpose(0,3,2,1)
            wVOVV -= lib.einsum('mebf,miaf->eiab', OVVV, t2bb[p0:p1,i0:i1])
            OVVV = None
        wVOVV = wVOVV - wVOVV.transpose(0,1,3,2)
        _wVOVV[:,i0:i1] = wVOVV

    eris_ovOV = np.asarray(eris.ovOV)
    eris_ovOO = np.asarray(eris.ovOO)
    eris_OOvv = np.asarray(eris.OOvv)
    eris_ovVO = np.asarray(eris.ovVO)
    #:self.wvOvV = lib.einsum('meNI,mNaB->eIaB', eris_ovOO, tauab)
    #:self.wvOvV -= lib.einsum('me,mIaB->eIaB', self.Fova, t2ab)
    #:tmp1ab = lib.einsum('NIBF,meNF->mBeI', t2bb, eris_ovOV)
    #:tmp1ab-= lib.einsum('nIfB,menf->mBeI', t2ab,      ovov)
    #:tmp1baab = lib.einsum('nIbF,neMF->MbeI', t2ab, eris_ovOV)
    #:tmpab = lib.einsum('ma,mBeI->eIaB', t1a, tmp1ab)
    #:tmpab+= lib.einsum('MA,MbeI->eIbA', t1b, tmp1baab)
    #:tmpab-= lib.einsum('MA,MIbe->eIbA', t1b, eris_OOvv)
    #:tmpab-= lib.einsum('ma,meBI->eIaB', t1a, eris_ovVO)
    #:self.wvOvV += tmpab
    for p0, p1 in lib.prange(0, nvira, nocca):
        wvOvV  = lib.einsum('meNI,mNaB->eIaB', eris_ovOO[:,p0:p1], tauab)
        wvOvV -= lib.einsum('me,mIaB->eIaB', Fova[:,p0:p1], t2ab)
        tmp1ab = lib.einsum('NIBF,meNF->mBeI', t2bb, eris_ovOV[:,p0:p1])
        tmp1ab-= lib.einsum('nIfB,menf->mBeI', t2ab, ovov[:,p0:p1])
        wvOvV+= lib.einsum('ma,mBeI->eIaB', t1a, tmp1ab)
        tmp1ab = None
        tmp1baab = lib.einsum('nIbF,neMF->MbeI', t2ab, eris_ovOV[:,p0:p1])
        wvOvV+= lib.einsum('MA,MbeI->eIbA', t1b, tmp1baab)
        tmp1baab = None
        wvOvV-= lib.einsum('MA,MIbe->eIbA', t1b, eris_OOvv[:,:,:,p0:p1])
        wvOvV-= lib.einsum('ma,meBI->eIaB', t1a, eris_ovVO[:,p0:p1])
        _wvOvV[p0:p1] = wvOvV
    eris_ovOV = eris_ovOO = eris_OOvv = eris_ovVO = None

    #:eris_ovvv = lib.unpack_tril(np.asarray(eris.ovvv).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvira,nvira)
    #:ovvv = eris_ovvv - eris_ovvv.transpose(0,3,2,1)
    #:self.wvOvV -= lib.einsum('mebf,mIfA->eIbA',      ovvv, t2ab)
    #:eris_ovVV = lib.unpack_tril(np.asarray(eris.ovVV).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvirb,nvirb)
    #:self.wvOvV -= lib.einsum('meBF,mIaF->eIaB', eris_ovVV, t2ab)
    #:eris_OVvv = lib.unpack_tril(np.asarray(eris.OVvv).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvira,nvira)
    #:self.wvOvV -= lib.einsum('MFbe,MIAF->eIbA', eris_OVvv, t2bb)
    #:self.wvOvV += eris_OVvv.transpose(2,0,3,1).conj()
    blksize = min(noccb, max(BLKMIN, int(max_memory*1e6/8/(nvira**3*6))))
    for i0,i1 in lib.prange(0, noccb, blksize):
        wvOvV = _wvOvV[:,i0:i1]
        for p0,p1 in lib.prange(0, nocca, blksize):
            ovVV = eris.get_ovVV(slice(p0,p1))  # ovVV = eris.ovVV[p0:p1]
            wvOvV -= lib.einsum('meBF,mIaF->eIaB', ovVV, t2ab[p0:p1,i0:i1])
            ovVV = None
        for p0,p1 in lib.prange(0, nocca, blksize):
            ovvv = eris.get_ovvv(slice(p0,p1))  # ovvv = eris.ovvv[p0:p1]
            ovvv = ovvv - ovvv.transpose(0,3,2,1)
            wvOvV -= lib.einsum('mebf,mIfA->eIbA',ovvv, t2ab[p0:p1,i0:i1])
            ovvv = None
        _wvOvV[:,i0:i1] = wvOvV

    blksize = min(noccb, max(BLKMIN, int(max_memory*1e6/8/(nvirb*nvira**2*3))))
    for i0,i1 in lib.prange(0, noccb, blksize):
        wvOvV = _wvOvV[:,i0:i1]
        for p0,p1 in lib.prange(0, noccb, blksize):
            OVvv = eris.get_OVvv(slice(p0,p1))  # OVvv = eris.OVvv[p0:p1]
            if p0 == i0:
                wvOvV += OVvv.transpose(2,0,3,1).conj()
            wvOvV -= lib.einsum('MFbe,MIAF->eIbA', OVvv, t2bb[p0:p1,i0:i1])
            OVvv = None
        _wvOvV[:,i0:i1] = wvOvV

    eris_ovOV = np.asarray(eris.ovOV)
    eris_OVoo = np.asarray(eris.OVoo)
    eris_ooVV = np.asarray(eris.ooVV)
    eris_OVvo = np.asarray(eris.OVvo)
    #:self.wVoVv = lib.einsum('MEni,nMbA->EiAb', eris_OVoo, tauab)
    #:self.wVoVv -= lib.einsum('ME,iMbA->EiAb', self.Fovb, t2ab)
    #:tmp1ba = lib.einsum('nibf,nfME->MbEi', t2aa, eris_ovOV)
    #:tmp1ba-= lib.einsum('iNbF,MENF->MbEi', t2ab,      OVOV)
    #:tmp1abba = lib.einsum('iNfB,mfNE->mBEi', t2ab, eris_ovOV)
    #:tmpba = lib.einsum('MA,MbEi->EiAb', t1b, tmp1ba)
    #:tmpba+= lib.einsum('ma,mBEi->EiBa', t1a, tmp1abba)
    #:tmpba-= lib.einsum('ma,miBE->EiBa', t1a, eris_ooVV)
    #:tmpba-= lib.einsum('MA,MEbi->EiAb', t1b, eris_OVvo)
    #:self.wVoVv += tmpba
    for p0, p1 in lib.prange(0, nvirb, noccb):
        wVoVv  = lib.einsum('MEni,nMbA->EiAb', eris_OVoo[:,p0:p1], tauab)
        wVoVv -= lib.einsum('ME,iMbA->EiAb', Fovb[:,p0:p1], t2ab)
        tmp1ba = lib.einsum('nibf,nfME->MbEi', t2aa, eris_ovOV[:,:,:,p0:p1])
        tmp1ba-= lib.einsum('iNbF,MENF->MbEi', t2ab, OVOV[:,p0:p1])
        wVoVv += lib.einsum('MA,MbEi->EiAb', t1b, tmp1ba)
        tmp1ba = None
        tmp1abba = lib.einsum('iNfB,mfNE->mBEi', t2ab, eris_ovOV[:,:,:,p0:p1])
        wVoVv += lib.einsum('ma,mBEi->EiBa', t1a, tmp1abba)
        tmp1abba = None
        wVoVv -= lib.einsum('ma,miBE->EiBa', t1a, eris_ooVV[:,:,:,p0:p1])
        wVoVv -= lib.einsum('MA,MEbi->EiAb', t1b, eris_OVvo[:,p0:p1])
        _wVoVv[p0:p1] = wVoVv
    eris_ovOV = eris_OVoo = eris_ooVV = eris_OVvo = None

    #:eris_OVVV = lib.unpack_tril(np.asarray(eris.OVVV).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvirb,nvirb)
    #:OVVV = eris_OVVV - eris_OVVV.transpose(0,3,2,1)
    #:self.wVoVv -= lib.einsum('MEBF,iMaF->EiBa',      OVVV, t2ab)
    #:eris_OVvv = lib.unpack_tril(np.asarray(eris.OVvv).reshape(noccb*nvirb,-1)).reshape(noccb,nvirb,nvira,nvira)
    #:self.wVoVv -= lib.einsum('MEbf,iMfA->EiAb', eris_OVvv, t2ab)
    #:eris_ovVV = lib.unpack_tril(np.asarray(eris.ovVV).reshape(nocca*nvira,-1)).reshape(nocca,nvira,nvirb,nvirb)
    #:self.wVoVv -= lib.einsum('mfBE,miaf->EiBa', eris_ovVV, t2aa)
    #:self.wVoVv += eris_ovVV.transpose(2,0,3,1).conj()
    blksize = min(noccb, max(BLKMIN, int(max_memory*1e6/8/(nvirb**3*6))))
    for i0,i1 in lib.prange(0, nocca, blksize):
        wVoVv = _wVoVv[:,i0:i1]
        for p0,p1 in lib.prange(0, noccb, blksize):
            OVvv = eris.get_OVvv(slice(p0,p1))  # OVvv = eris.OVvv[p0:p1]
            wVoVv -= lib.einsum('MEbf,iMfA->EiAb', OVvv, t2ab[i0:i1,p0:p1])
            OVvv = None
        for p0,p1 in lib.prange(0, noccb, blksize):
            OVVV = eris.get_OVVV(slice(p0,p1))  # OVVV = eris.OVVV[p0:p1]
            OVVV = OVVV - OVVV.transpose(0,3,2,1)
            wVoVv -= lib.einsum('MEBF,iMaF->EiBa', OVVV, t2ab[i0:i1,p0:p1])
            OVVV = None
        _wVoVv[:,i0:i1] = wVoVv

    blksize = min(nocca, max(BLKMIN, int(max_memory*1e6/8/(nvira*nvirb**2*3))))
    for i0,i1 in lib.prange(0, nocca, blksize):
        wVoVv = _wVoVv[:,i0:i1]
        for p0,p1 in lib.prange(0, nocca, blksize):
            ovVV = eris.get_ovVV(slice(p0,p1))  # ovVV = eris.ovVV[p0:p1]
            if p0 == i0:
                wVoVv += ovVV.transpose(2,0,3,1).conj()
            wVoVv -= lib.einsum('mfBE,miaf->EiBa', ovVV, t2aa[p0:p1,i0:i1])
            ovVV = None
        _wVoVv[:,i0:i1] = wVoVv
    wvvvo = _wvovv.transpose(2,0,3,1)
    wVVvo = _wVoVv.transpose(2,0,3,1)
    wvvVO = _wvOvV.transpose(2,0,3,1)
    wVVVO = _wVOVV.transpose(2,0,3,1)
    return wvvvo, wvvVO, wVVvo, wVVVO

def _get_vvVV(eris):
    if eris.vvVV is None and getattr(eris, 'VVL', None) is not None:  # DF eris
        vvL = np.asarray(eris.vvL)
        VVL = np.asarray(eris.VVL)
        vvVV = lib.dot(vvL, VVL.T)
    elif len(eris.vvVV.shape) == 2:
        vvVV = np.asarray(eris.vvVV)
    else:
        return eris.vvVV

    nvira = int(np.sqrt(vvVV.shape[0]*2))
    nvirb = int(np.sqrt(vvVV.shape[1]*2))
    vvVV1 = np.zeros((nvira**2,nvirb**2))
    vtrila = np.tril_indices(nvira)
    vtrilb = np.tril_indices(nvirb)
    lib.takebak_2d(vvVV1, vvVV, vtrila[0]*nvira+vtrila[1], vtrilb[0]*nvirb+vtrilb[1])
    lib.takebak_2d(vvVV1, vvVV, vtrila[1]*nvira+vtrila[0], vtrilb[1]*nvirb+vtrilb[0])
    lib.takebak_2d(vvVV1, vvVV, vtrila[0]*nvira+vtrila[1], vtrilb[1]*nvirb+vtrilb[0])
    lib.takebak_2d(vvVV1, vvVV, vtrila[1]*nvira+vtrila[0], vtrilb[0]*nvirb+vtrilb[1])
    return vvVV1.reshape(nvira,nvira,nvirb,nvirb)

def _get_VVVV(eris):
    if eris.VVVV is None and getattr(eris, 'VVL', None) is not None:  # DF eris
        VVL = np.asarray(eris.VVL)
        nvir = int(np.sqrt(eris.VVL.shape[0]*2))
        return ao2mo.restore(1, lib.dot(VVL, VVL.T), nvir)
    elif len(eris.VVVV.shape) == 2:
        nvir = int(np.sqrt(eris.VVVV.shape[0]*2))
        return ao2mo.restore(1, np.asarray(eris.VVVV), nvir)
    else:
        return eris.VVVV


# pyscf/cc/uintermediates_slow.py

#!/usr/bin/env python
# Copyright 2014-2021 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.

import tempfile
import h5py
import numpy as np
from pyscf import lib

#einsum = np.einsum
einsum = lib.einsum

# Ref: Gauss and Stanton, J. Chem. Phys. 103, 3561 (1995) Table III

# Section (a)

def make_tau(t2, t1a, t1b, fac=1, out=None):
    tmp = einsum('ia,jb->ijab',t1a,t1b)
    t1t1 = tmp - tmp.transpose(1,0,2,3) - tmp.transpose(0,1,3,2) + tmp.transpose(1,0,3,2)
    tau1 = t2 + fac*0.50*t1t1
    return tau1

def cc_Fvv(t1,t2,eris):
    nocc, nvir = t1.shape
    fov = eris.fock[:nocc,nocc:]
    fvv = eris.fock[nocc:,nocc:]
    eris_vovv = np.array(eris.ovvv).transpose(1,0,3,2)
    tau_tilde = make_tau(t2,t1,t1,fac=0.5)
    Fae = ( fvv - 0.5*einsum('me,ma->ae',fov,t1)
            + einsum('mf,amef->ae',t1,eris_vovv)
            - 0.5*einsum('mnaf,mnef->ae',tau_tilde,eris.oovv) )
    return Fae

def cc_Foo(t1,t2,eris):
    nocc, nvir = t1.shape
    fov = eris.fock[:nocc,nocc:]
    foo = eris.fock[:nocc,:nocc]
    tau_tilde = make_tau(t2,t1,t1,fac=0.5)
    Fmi = (foo + 0.5*einsum('me,ie->mi',fov,t1)
           + einsum('ne,mnie->mi',t1,eris.ooov)
           + 0.5*einsum('inef,mnef->mi',tau_tilde,eris.oovv) )
    return Fmi

def cc_Fov(t1,t2,eris):
    nocc, nvir = t1.shape
    fov = eris.fock[:nocc,nocc:]
    Fme = fov + einsum('nf,mnef->me',t1,eris.oovv)
    return Fme

def cc_Woooo(t1,t2,eris):
    tau = make_tau(t2,t1,t1)
    tmp = einsum('je,mnie->mnij',t1,eris.ooov)
    Wmnij = eris.oooo + tmp - tmp.transpose(0,1,3,2)
    Wmnij += 0.25*einsum('ijef,mnef->mnij',tau,eris.oovv)
    return Wmnij

def cc_Wvvvv(t1,t2,eris):
    tau = make_tau(t2,t1,t1)
    #eris_vovv = np.array(eris.ovvv).transpose(1,0,3,2)
    #tmp = einsum('mb,amef->abef',t1,eris_vovv)
    #Wabef = eris.vvvv - tmp + tmp.transpose(1,0,2,3)
    #Wabef += 0.25*einsum('mnab,mnef->abef',tau,eris.oovv)
    if t1.dtype == np.complex128: ds_type = 'c16'
    else: ds_type = 'f8'
    _tmpfile1 = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
    fimd = h5py.File(_tmpfile1.name)
    nocc, nvir = t1.shape
    Wabef = fimd.create_dataset('vvvv', (nvir,nvir,nvir,nvir), ds_type)
    for a in range(nvir):
        Wabef[a] = eris.vvvv[a]
        Wabef[a] -= einsum('mb,mfe->bef',t1,eris.ovvv[:,a,:,:])
        Wabef[a] += einsum('m,mbfe->bef',t1[:,a],eris.ovvv)
        Wabef[a] += 0.25*einsum('mnb,mnef->bef',tau[:,:,a,:],eris.oovv)
    return Wabef

def cc_Wovvo(t1,t2,eris):
    eris_ovvo = -np.array(eris.ovov).transpose(0,1,3,2)
    eris_oovo = -np.array(eris.ooov).transpose(0,1,3,2)
    Wmbej = eris_ovvo.copy()
    Wmbej +=  einsum('jf,mbef->mbej',t1,eris.ovvv)
    Wmbej += -einsum('nb,mnej->mbej',t1,eris_oovo)
    Wmbej += -0.5*einsum('jnfb,mnef->mbej',t2,eris.oovv)
    Wmbej += -einsum('jf,nb,mnef->mbej',t1,t1,eris.oovv)
    return Wmbej

### Section (b)

def Fvv(t1,t2,eris):
    ccFov = cc_Fov(t1,t2,eris)
    Fae = cc_Fvv(t1,t2,eris) - 0.5*einsum('ma,me->ae',t1,ccFov)
    return Fae

def Foo(t1,t2,eris):
    ccFov = cc_Fov(t1,t2,eris)
    Fmi = cc_Foo(t1,t2,eris) + 0.5*einsum('ie,me->mi',t1,ccFov)
    return Fmi

def Fov(t1,t2,eris):
    Fme = cc_Fov(t1,t2,eris)
    return Fme

def Woooo(t1,t2,eris):
    tau = make_tau(t2,t1,t1)
    Wmnij = cc_Woooo(t1,t2,eris) + 0.25*einsum('ijef,mnef->mnij',tau,eris.oovv)
    return Wmnij

def Wvvvv(t1,t2,eris):
    tau = make_tau(t2,t1,t1)
    #Wabef = cc_Wvvvv(t1,t2,eris) + 0.25*einsum('mnab,mnef->abef',tau,eris.oovv)
    if t1.dtype == np.complex128: ds_type = 'c16'
    else: ds_type = 'f8'
    _tmpfile1 = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
    fimd = h5py.File(_tmpfile1.name)
    nocc, nvir = t1.shape
    Wabef = fimd.create_dataset('vvvv', (nvir,nvir,nvir,nvir), ds_type)
    #_cc_Wvvvv = cc_Wvvvv(t1,t2,eris)
    for a in range(nvir):
        #Wabef[a] = _cc_Wvvvv[a]
        Wabef[a] = eris.vvvv[a]
        Wabef[a] -= einsum('mb,mfe->bef',t1,eris.ovvv[:,a,:,:])
        Wabef[a] += einsum('m,mbfe->bef',t1[:,a],eris.ovvv)
        #Wabef[a] += 0.25*einsum('mnb,mnef->bef',tau[:,:,a,:],eris.oovv)

        #Wabef[a] += 0.25*einsum('mnb,mnef->bef',tau[:,:,a,:],eris.oovv)
        Wabef[a] += 0.5*einsum('mnb,mnef->bef',tau[:,:,a,:],eris.oovv)
    return Wabef

def Wovvo(t1,t2,eris):
    Wmbej = cc_Wovvo(t1,t2,eris) - 0.5*einsum('jnfb,mnef->mbej',t2,eris.oovv)
    return Wmbej

def Wooov(t1,t2,eris):
    Wmnie = eris.ooov + einsum('if,mnfe->mnie',t1,eris.oovv)
    return Wmnie

def Wvovv(t1,t2,eris):
    eris_vovv = -np.array(eris.ovvv).transpose(1,0,2,3)
    Wamef = eris_vovv - einsum('na,nmef->amef',t1,eris.oovv)
    return Wamef

def Wovoo(t1,t2,eris):
    eris_ovvo = -np.array(eris.ovov).transpose(0,1,3,2)
    tmp1 = einsum('mnie,jnbe->mbij',eris.ooov,t2)
    tmp2 = (einsum('ie,mbej->mbij',t1,eris_ovvo)
            - einsum('ie,njbf,mnef->mbij',t1,t2,eris.oovv))
    FFov = Fov(t1,t2,eris)
    WWoooo = Woooo(t1,t2,eris)
    tau = make_tau(t2,t1,t1)
    Wmbij = (eris.ovoo - einsum('me,ijbe->mbij',FFov,t2)
             - einsum('nb,mnij->mbij',t1,WWoooo)
             + 0.5 * einsum('mbef,ijef->mbij',eris.ovvv,tau)
             + tmp1 - tmp1.transpose(0,1,3,2)
             + tmp2 - tmp2.transpose(0,1,3,2) )
    return Wmbij

def Wvvvo(t1,t2,eris,_Wvvvv=None):
    eris_ovvo = -np.array(eris.ovov).transpose(0,1,3,2)
    eris_vvvo = -np.array(eris.ovvv).transpose(2,3,1,0).conj()
    eris_oovo = -np.array(eris.ooov).transpose(0,1,3,2)
    tmp1 = einsum('mbef,miaf->abei',eris.ovvv,t2)
    tmp2 = (einsum('ma,mbei->abei',t1,eris_ovvo)
            - einsum('ma,nibf,mnef->abei',t1,t2,eris.oovv))
    FFov = Fov(t1,t2,eris)
    tau = make_tau(t2,t1,t1)
    Wabei = eris_vvvo
    Wabei += -einsum('me,miab->abei',FFov,t2)
    Wabei += 0.5 * einsum('mnei,mnab->abei',eris_oovo,tau)
    Wabei += -tmp1 + tmp1.transpose(1,0,2,3)
    Wabei += -tmp2 + tmp2.transpose(1,0,2,3)
    nocc,nvir = t1.shape
    if _Wvvvv is None:
        _Wvvvv = Wvvvv(t1,t2,eris)
    for a in range(nvir):
        Wabei[a] += einsum('if,bef->bei',t1,_Wvvvv[a])
    return Wabei

def Wvvvo_incore(t1,t2,eris):
    eris_ovvo = -np.array(eris.ovov).transpose(0,1,3,2)
    eris_vvvo = -np.array(eris.ovvv).transpose(2,3,1,0).conj()
    eris_oovo = -np.array(eris.ooov).transpose(0,1,3,2)
    tmp1 = einsum('mbef,miaf->abei',eris.ovvv,t2)
    tmp2 = (einsum('ma,mbei->abei',t1,eris_ovvo)
            - einsum('ma,nibf,mnef->abei',t1,t2,eris.oovv))
    FFov = Fov(t1,t2,eris)
    WWvvvv = Wvvvv(t1,t2,eris)
    tau = make_tau(t2,t1,t1)
    Wabei = (eris_vvvo - einsum('me,miab->abei',FFov,t2)
             + einsum('if,abef->abei',t1,WWvvvv)
             + 0.5 * einsum('mnei,mnab->abei',eris_oovo,tau)
             - tmp1 + tmp1.transpose(1,0,2,3)
             - tmp2 + tmp2.transpose(1,0,2,3) )
    return Wabei


# pyscf/cc/ccsd_renorm_t_slow.py
#!/usr/bin/env python
# Copyright 2014-2020 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: Qiming Sun <osirpt.sun@gmail.com>
#

import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf.cc.ccsd_t_slow import r3 

'''
Renormalized perturbative triples correction on CCSD (R-CCSD(T))
'''

# t3 as ijkabc

# Comput. Phys. Commun. 149, 71 (2002); DOI:10.1016/S0010-4655(02)00598-2
def kernel(mycc, eris, t1=None, t2=None, verbose=logger.NOTE):
    if isinstance(verbose, logger.Logger):
        log = verbose
    else:
        log = logger.Logger(mycc.stdout, verbose)

    if t1 is None: t1 = mycc.t1
    if t2 is None: t2 = mycc.t2

    t1T = t1.T
    t2T = t2.transpose(2,3,0,1)

    nocc, nvir = t1.shape
    mo_e = eris.mo_energy
    e_occ, e_vir = mo_e[:nocc], mo_e[nocc:]
    eijk = lib.direct_sum('i,j,k->ijk', e_occ, e_occ, e_occ)

    eris_vvov = eris.get_ovvv().conj().transpose(1,3,0,2)
    eris_vooo = numpy.asarray(eris.ovoo).conj().transpose(1,0,3,2)
    eris_vvoo = numpy.asarray(eris.ovov).conj().transpose(1,3,0,2)
    fvo = eris.fock[nocc:,:nocc]
    def get_w(a, b, c):
        w = numpy.einsum('if,fkj->ijk', eris_vvov[a,b], t2T[c,:])
        w-= numpy.einsum('ijm,mk->ijk', eris_vooo[a,:], t2T[b,c])
        return w
    def get_v(a, b, c):
        v = numpy.einsum('ij,k->ijk', eris_vvoo[a,b], t1T[c])
        v+= numpy.einsum('ij,k->ijk', t2T[a,b], fvo[c])
        return v
    def get_y(a, b, c):
        y = numpy.einsum('i,j,k->ijk', t1T[a], t1T[b], t1T[c]) / 3.
        y+= numpy.einsum('i,jk->ijk', t1T[a], t2T[b,c])
        return y 

    et = 0
    dn = .5
    dn += numpy.einsum('ia,ia', t1, t1)
    tmpt = t2 - .5 * t2.transpose(0,1,3,2)
    tmpc = t2 + numpy.einsum('ia,jb->ijab', t1, t1) 
    dn += numpy.einsum('ijab,ijab', tmpt, tmpc)
    for a in range(nvir):
        for b in range(a+1):
            for c in range(b+1):
                d3 = eijk - e_vir[a] - e_vir[b] - e_vir[c]
                if a == c:  # a == b == c
                    d3 *= 6
                elif a == b or b == c:
                    d3 *= 2

                wabc = get_w(a, b, c)
                wacb = get_w(a, c, b)
                wbac = get_w(b, a, c)
                wbca = get_w(b, c, a)
                wcab = get_w(c, a, b)
                wcba = get_w(c, b, a)
                vabc = get_v(a, b, c)
                vacb = get_v(a, c, b)
                vbac = get_v(b, a, c)
                vbca = get_v(b, c, a)
                vcab = get_v(c, a, b)
                vcba = get_v(c, b, a)
                zabc = r3(wabc + .5 * vabc) / d3
                zacb = r3(wacb + .5 * vacb) / d3
                zbac = r3(wbac + .5 * vbac) / d3
                zbca = r3(wbca + .5 * vbca) / d3
                zcab = r3(wcab + .5 * vcab) / d3
                zcba = r3(wcba + .5 * vcba) / d3
                yabc = .5 * get_y(a, b, c)
                yacb = .5 * get_y(a, c, b)
                ybac = .5 * get_y(b, a, c)
                ybca = .5 * get_y(b, c, a)
                ycab = .5 * get_y(c, a, b)
                ycba = .5 * get_y(c, b, a)

                et+= numpy.einsum('ijk,ijk', wabc, zabc.conj())
                et+= numpy.einsum('ikj,ijk', wacb, zabc.conj())
                et+= numpy.einsum('jik,ijk', wbac, zabc.conj())
                et+= numpy.einsum('jki,ijk', wbca, zabc.conj())
                et+= numpy.einsum('kij,ijk', wcab, zabc.conj())
                et+= numpy.einsum('kji,ijk', wcba, zabc.conj())

                et+= numpy.einsum('ijk,ijk', wacb, zacb.conj())
                et+= numpy.einsum('ikj,ijk', wabc, zacb.conj())
                et+= numpy.einsum('jik,ijk', wcab, zacb.conj())
                et+= numpy.einsum('jki,ijk', wcba, zacb.conj())
                et+= numpy.einsum('kij,ijk', wbac, zacb.conj())
                et+= numpy.einsum('kji,ijk', wbca, zacb.conj())

                et+= numpy.einsum('ijk,ijk', wbac, zbac.conj())
                et+= numpy.einsum('ikj,ijk', wbca, zbac.conj())
                et+= numpy.einsum('jik,ijk', wabc, zbac.conj())
                et+= numpy.einsum('jki,ijk', wacb, zbac.conj())
                et+= numpy.einsum('kij,ijk', wcba, zbac.conj())
                et+= numpy.einsum('kji,ijk', wcab, zbac.conj())

                et+= numpy.einsum('ijk,ijk', wbca, zbca.conj())
                et+= numpy.einsum('ikj,ijk', wbac, zbca.conj())
                et+= numpy.einsum('jik,ijk', wcba, zbca.conj())
                et+= numpy.einsum('jki,ijk', wcab, zbca.conj())
                et+= numpy.einsum('kij,ijk', wabc, zbca.conj())
                et+= numpy.einsum('kji,ijk', wacb, zbca.conj())

                et+= numpy.einsum('ijk,ijk', wcab, zcab.conj())
                et+= numpy.einsum('ikj,ijk', wcba, zcab.conj())
                et+= numpy.einsum('jik,ijk', wacb, zcab.conj())
                et+= numpy.einsum('jki,ijk', wabc, zcab.conj())
                et+= numpy.einsum('kij,ijk', wbca, zcab.conj())
                et+= numpy.einsum('kji,ijk', wbac, zcab.conj())

                et+= numpy.einsum('ijk,ijk', wcba, zcba.conj())
                et+= numpy.einsum('ikj,ijk', wcab, zcba.conj())
                et+= numpy.einsum('jik,ijk', wbca, zcba.conj())
                et+= numpy.einsum('jki,ijk', wbac, zcba.conj())
                et+= numpy.einsum('kij,ijk', wacb, zcba.conj())
                et+= numpy.einsum('kji,ijk', wabc, zcba.conj())

                dn+= numpy.einsum('ijk,ijk', yabc, zabc.conj())
                dn+= numpy.einsum('ikj,ijk', yacb, zabc.conj())
                dn+= numpy.einsum('jik,ijk', ybac, zabc.conj())
                dn+= numpy.einsum('jki,ijk', ybca, zabc.conj())
                dn+= numpy.einsum('kij,ijk', ycab, zabc.conj())
                dn+= numpy.einsum('kji,ijk', ycba, zabc.conj())

                dn+= numpy.einsum('ijk,ijk', yacb, zacb.conj())
                dn+= numpy.einsum('ikj,ijk', yabc, zacb.conj())
                dn+= numpy.einsum('jik,ijk', ycab, zacb.conj())
                dn+= numpy.einsum('jki,ijk', ycba, zacb.conj())
                dn+= numpy.einsum('kij,ijk', ybac, zacb.conj())
                dn+= numpy.einsum('kji,ijk', ybca, zacb.conj())

                dn+= numpy.einsum('ijk,ijk', ybac, zbac.conj())
                dn+= numpy.einsum('ikj,ijk', ybca, zbac.conj())
                dn+= numpy.einsum('jik,ijk', yabc, zbac.conj())
                dn+= numpy.einsum('jki,ijk', yacb, zbac.conj())
                dn+= numpy.einsum('kij,ijk', ycba, zbac.conj())
                dn+= numpy.einsum('kji,ijk', ycab, zbac.conj())

                dn+= numpy.einsum('ijk,ijk', ybca, zbca.conj())
                dn+= numpy.einsum('ikj,ijk', ybac, zbca.conj())
                dn+= numpy.einsum('jik,ijk', ycba, zbca.conj())
                dn+= numpy.einsum('jki,ijk', ycab, zbca.conj())
                dn+= numpy.einsum('kij,ijk', yabc, zbca.conj())
                dn+= numpy.einsum('kji,ijk', yacb, zbca.conj())

                dn+= numpy.einsum('ijk,ijk', ycab, zcab.conj())
                dn+= numpy.einsum('ikj,ijk', ycba, zcab.conj())
                dn+= numpy.einsum('jik,ijk', yacb, zcab.conj())
                dn+= numpy.einsum('jki,ijk', yabc, zcab.conj())
                dn+= numpy.einsum('kij,ijk', ybca, zcab.conj())
                dn+= numpy.einsum('kji,ijk', ybac, zcab.conj())

                dn+= numpy.einsum('ijk,ijk', ycba, zcba.conj())
                dn+= numpy.einsum('ikj,ijk', ycab, zcba.conj())
                dn+= numpy.einsum('jik,ijk', ybca, zcba.conj())
                dn+= numpy.einsum('jki,ijk', ybac, zcba.conj())
                dn+= numpy.einsum('kij,ijk', yacb, zcba.conj())
                dn+= numpy.einsum('kji,ijk', yabc, zcba.conj())
    et *= 2
    log.info('  CCSD(T) correction = %.15g', et)
    dn *= 2
    log.info('R-CCSD(T) correction = %.15g', et / dn)
    return et / dn

if __name__ == '__main__':
    import pyscf
    from pyscf.data import nist
    re = 1.7328 * nist.BOHR
    r = re * 3 
    
    mol = pyscf.M(
        atom = 'H 0 0 0; F 0 0 %f' % r,
        basis = 'dz')
    
    mf = mol.RHF().run(max_cycle=100, verbose=0)

    print('> Comput. Phys. Commun. 149, 71 (2002); DOI:10.1016/S0010-4655(02)00598-2')    
    cisolver = pyscf.fci.FCI(mf)
    e_fci = cisolver.kernel()[0]
    print(' E(FCI)          /  Eh = %.6f' % e_fci)
    
    mycc = mf.CCSD().run(verbose=0)
    e_ccsd = mf.e_tot + mycc.e_corr
    print(' CCSD error      / mEh = %.3f' % ((e_ccsd - e_fci) * 1e3))
    
    et = mycc.ccsd_t()
    e_ccsd_t = e_ccsd + et
    print(' CCSD(T) error   / mEh = %.3f' % ((e_ccsd_t - e_fci) * 1e3))
      
    eris = mycc.ao2mo(mycc.mo_coeff)
    et = kernel(mycc, eris, mycc.t1, mycc.t2) 
    e_r_ccsd_t = e_ccsd + et
    print(' R-CCSD(T) error / mEh = %.3f' % ((e_r_ccsd_t - e_fci) * 1e3))
    
