def test_spin_squre(self):
        ss = fci.spin_op.spin_square(ci0, norb, nelec)
        self.assertAlmostEqual(ss[0], 6, 9)
        ss = fci.spin_op.spin_square0(ci0, norb, nelec)
        self.assertAlmostEqual(ss[0], 6, 9)

        numpy.random.seed(1)
        u,w,v = numpy.linalg.svd(numpy.random.random((norb,6)))
        u = u[:,:6]
        h1a = h1[:6,:6]
        h1b = reduce(numpy.dot, (v.T, h1a, v))
        h2aa = ao2mo.restore(1, h2, norb)[:6,:6,:6,:6]
        h2ab = lib.einsum('klpq,pi,qj->klij', h2aa, v, v)
        h2bb = lib.einsum('pqkl,pi,qj->ijkl', h2ab, v, v)
        e1, ci1 = fci.direct_uhf.kernel((h1a,h1b), (h2aa,h2ab,h2bb), 6, (3,2))
        ss = fci.spin_op.spin_square(ci1, 6, (3,2), mo_coeff=(numpy.eye(6),v))[0]
        self.assertAlmostEqual(ss, 3.75, 8)

        numpy.random.seed(1)
        n = fci.cistring.num_strings(6,3)
        ci1 = numpy.random.random((n,n))
        ss1 = numpy.einsum('ij,ij->', ci1, fci.spin_op.contract_ss(ci1, 6, 6))
        self.assertAlmostEqual(ss1, fci.spin_op.spin_square(ci1, 6, 6)[0], 12)

        na = fci.cistring.num_strings(6,4)
        nb = fci.cistring.num_strings(6,2)
        ci1 = numpy.random.random((na,nb))
        ss1 = numpy.einsum('ij,ij->', ci1, fci.spin_op.contract_ss(ci1, 6, (4,2)))
        self.assertAlmostEqual(ss1, fci.spin_op.spin_square(ci1, 6, (4,2))[0], 12)

        numpy.random.seed(1)
        n = fci.cistring.num_strings(10,5)
        ci1 = numpy.random.random((n,n))
        ss1 = numpy.einsum('ij,ij->', ci1, fci.spin_op.contract_ss(ci1, 10, 10))
        self.assertAlmostEqual(ss1, fci.spin_op.spin_square(ci1, 10, 10)[0], 8)

def Wvvov(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)

    #: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 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.finger(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.finger(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_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 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 vind(xys):
        nz = len(xys)
        dm1a = numpy.empty((nz,nao,nao))
        dm1b = numpy.empty((nz,nao,nao))
        for i in range(nz):
            xa, xb, ya, yb = numpy.split(xys[i], offsets)
            dmx = reduce(numpy.dot, (orbva, xa.reshape(nvira,nocca)  , orboa.T))
            dmy = reduce(numpy.dot, (orboa, ya.reshape(nvira,nocca).T, orbva.T))
            dm1a[i] = dmx + dmy  # AX + BY
            dmx = reduce(numpy.dot, (orbvb, xb.reshape(nvirb,noccb)  , orbob.T))
            dmy = reduce(numpy.dot, (orbob, yb.reshape(nvirb,noccb).T, orbvb.T))
            dm1b[i] = dmx + dmy  # AX + BY

        v1ao = vresp(numpy.stack((dm1a,dm1b)))
        v1voa = lib.einsum('xpq,pi,qj->xij', v1ao[0], orbva, orboa).reshape(nz,-1)
        v1vob = lib.einsum('xpq,pi,qj->xij', v1ao[1], orbvb, orbob).reshape(nz,-1)
        v1ova = lib.einsum('xpq,pi,qj->xji', v1ao[0], orboa, orbva).reshape(nz,-1)
        v1ovb = lib.einsum('xpq,pi,qj->xji', v1ao[1], orbob, orbvb).reshape(nz,-1)

        for i in range(nz):
            xa, xb, ya, yb = numpy.split(xys[i], offsets)
            v1voa[i] += (e_ai_a - freq - diag[0]) * xa
            v1voa[i] /= diag[0]
            v1vob[i] += (e_ai_b - freq - diag[1]) * xb
            v1vob[i] /= diag[1]
            v1ova[i] += (e_ai_a + freq - diag[2]) * ya
            v1ova[i] /= diag[2]
            v1ovb[i] += (e_ai_b + freq - diag[3]) * yb
            v1ovb[i] /= diag[3]
        v = numpy.hstack((v1voa, v1vob, v1ova, v1ovb))
        return v

def _add_vvvv_full(mycc, t1T, t2T, 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 = time.clock(), time.time()
    log = logger.Logger(mycc.stdout, mycc.verbose)

    nvir_seg, nvir, nocc = t2T.shape[:3]
    vloc0, vloc1 = _task_location(nvir, rank)
    nocc2 = nocc*(nocc+1)//2
    if t1T is None:
        tau = lib.pack_tril(t2T.reshape(nvir_seg*nvir,nocc2))
    else:
        tau = t2T + numpy.einsum('ai,bj->abij', t1T[vloc0:vloc1], t1T)
        tau = lib.pack_tril(tau.reshape(nvir_seg*nvir,nocc2))
    tau = tau.reshape(nvir_seg,nvir,nocc2)

    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)

        ao_loc = mycc.mol.ao_loc_nr()
        nao, nmo = mo.shape
        ntasks = mpi.pool.size
        task_sh_locs = lib.misc._balanced_partition(ao_loc, ntasks)
        ao_loc0 = ao_loc[task_sh_locs[rank  ]]
        ao_loc1 = ao_loc[task_sh_locs[rank+1]]

        orbv = mo[:,nocc:]
        tau = lib.einsum('abij,pb->apij', tau, orbv)
        tau_priv = numpy.zeros((ao_loc1-ao_loc0,nao,nocc,nocc))
        for task_id, tau in _rotate_tensor_block(tau):
            loc0, loc1 = _task_location(nvir, task_id)
            tau_priv += lib.einsum('pa,abij->pbij', orbv[ao_loc0:ao_loc1,loc0:loc1], tau)
        tau = None
        time1 = log.timer_debug1('vvvv-tau mo2ao', *time0)

        buf = _contract_vvvv_t2(mycc, None, tau_priv, task_sh_locs, None, log)
        buf = buf.reshape(tau_priv.shape)
        tau_priv = None
        time1 = log.timer_debug1('vvvv-tau contraction', *time1)

        buf = lib.einsum('apij,pb->abij', buf, orbv)
        Ht2 = numpy.ndarray(t2T.shape, buffer=out)
        Ht2[:] = 0
        for task_id, buf in _rotate_tensor_block(buf):
            ao_loc0 = ao_loc[task_sh_locs[task_id  ]]
            ao_loc1 = ao_loc[task_sh_locs[task_id+1]]
            Ht2 += lib.einsum('pa,pbij->abij', orbv[ao_loc0:ao_loc1,vloc0:vloc1], buf)

        time1 = log.timer_debug1('vvvv-tau ao2mo', *time1)
    else:
        raise NotImplementedError
    return Ht2.reshape(t2T.shape)

def para(magobj, gauge_orig=None, h1=None, s1=None, with_cphf=None):
    '''Paramagnetic susceptibility tensor

    Kwargs:
        h1: A list of arrays. Shapes are [(3,nmo_a,nocc_a), (3,nmo_b,nocc_b)]
            First order Fock matrices in MO basis.
        s1: A list of arrays. Shapes are [(3,nmo_a,nocc_a), (3,nmo_b,nocc_b)]
            First order overlap matrices in MO basis.
        with_cphf : boolean or  function(dm_mo) => v1_mo
            If a boolean value is given, the value determines whether CPHF
            equation will be solved or not. The induced potential will be
            generated by the function gen_vind.
            If a function is given, CPHF equation will be solved, and the
            given function is used to compute induced potential
    '''
    log = logger.Logger(magobj.stdout, magobj.verbose)
    cput1 = (time.clock(), time.time())

    mol = magobj.mol
    mf = magobj._scf
    mo_energy = magobj._scf.mo_energy
    mo_coeff = magobj._scf.mo_coeff
    mo_occ = magobj._scf.mo_occ
    orboa = mo_coeff[0][:,mo_occ[0] > 0]
    orbob = mo_coeff[1][:,mo_occ[1] > 0]

    if h1 is None:
        # Imaginary part of F10
        dm0 = (numpy.dot(orboa, orboa.T), numpy.dot(orbob, orbob.T))
        h1 = magobj.get_fock(dm0, gauge_orig)
        h1 = (lib.einsum('xpq,pi,qj->xij', h1[0], mo_coeff[0].conj(), orboa),
              lib.einsum('xpq,pi,qj->xij', h1[1], mo_coeff[1].conj(), orbob))
        cput1 = log.timer('first order Fock matrix', *cput1)
    if s1 is None:
        # Imaginary part of S10
        s1 = magobj.get_ovlp(mol, gauge_orig)
        s1 = (lib.einsum('xpq,pi,qj->xij', s1, mo_coeff[0].conj(), orboa),
              lib.einsum('xpq,pi,qj->xij', s1, mo_coeff[1].conj(), orbob))

    with_cphf = magobj.cphf
    mo1, mo_e1 = uhf_nmr.solve_mo1(magobj, mo_energy, mo_coeff, mo_occ,
                                   h1, s1, with_cphf)
    cput1 = logger.timer(magobj, 'solving mo1 eqn', *cput1)

    occidxa = mo_occ[0] > 0
    occidxb = mo_occ[1] > 0
    mag_para = numpy.einsum('yji,xji->xy', mo1[0], h1[0])
    mag_para+= numpy.einsum('yji,xji->xy', mo1[1], h1[1])
    mag_para-= numpy.einsum('yji,xji,i->xy', mo1[0], s1[0], mo_energy[0][occidxa])
    mag_para-= numpy.einsum('yji,xji,i->xy', mo1[1], s1[1], mo_energy[1][occidxb])
    # + c.c.
    mag_para = mag_para + mag_para.conj()

    mag_para-= numpy.einsum('xij,yij->xy', s1[0][:,occidxa], mo_e1[0])
    mag_para-= numpy.einsum('xij,yij->xy', s1[1][:,occidxb], mo_e1[1])
    return -mag_para

def para(magobj, gauge_orig=None, h1=None, s1=None, with_cphf=None):
    '''Paramagnetic susceptibility tensor

    Kwargs:
        h1: (3,nmo,nocc) array
            First order Fock matrix in MO basis.
        s1: (3,nmo,nocc) array
            First order overlap matrix in MO basis.
        with_cphf : boolean or  function(dm_mo) => v1_mo
            If a boolean value is given, the value determines whether CPHF
            equation will be solved or not. The induced potential will be
            generated by the function gen_vind.
            If a function is given, CPHF equation will be solved, and the
            given function is used to compute induced potential
    '''
    log = logger.Logger(magobj.stdout, magobj.verbose)
    cput1 = (time.clock(), time.time())

    mol = magobj.mol
    mf = magobj._scf
    mo_energy = mf.mo_energy
    mo_coeff = mf.mo_coeff
    mo_occ = mf.mo_occ
    occidx = mo_occ > 0
    orbo = mo_coeff[:,occidx]

    if h1 is None:
        # Imaginary part of F10
        dm0 = mf.make_rdm1(mo_coeff, mo_occ)
        h1 = lib.einsum('xpq,pi,qj->xij', magobj.get_fock(dm0, gauge_orig),
                        mo_coeff.conj(), orbo)
    if s1 is None:
        # Imaginary part of S10
        s1 = lib.einsum('xpq,pi,qj->xij', magobj.get_ovlp(mol, gauge_orig),
                        mo_coeff.conj(), orbo)
    cput1 = log.timer('first order Fock matrix', *cput1)

    with_cphf = magobj.cphf
    mo1, mo_e1 = rhf_nmr.solve_mo1(magobj, mo_energy, mo_coeff, mo_occ,
                                   h1, s1, with_cphf)
    cput1 = logger.timer(magobj, 'solving mo1 eqn', *cput1)

    mag_para = numpy.einsum('yji,xji->xy', mo1, h1)
    mag_para-= numpy.einsum('yji,xji,i->xy', mo1, s1, mo_energy[occidx])
    # + c.c.
    mag_para = mag_para + mag_para.conj()

    mag_para-= numpy.einsum('xij,yij->xy', s1[:,occidx], mo_e1)

    # *2 for double occupancy.
    mag_para *= 2
    return -mag_para

def hyper_polarizability(polobj, with_cphf=True):
    from pyscf.prop.nmr import uhf as uhf_nmr
    log = logger.new_logger(polobj)
    mf = polobj._scf
    mol = mf.mol
    mo_energy = mf.mo_energy
    mo_coeff = mf.mo_coeff
    mo_occ = mf.mo_occ
    occidxa = mo_occ[0] > 0
    occidxb = mo_occ[1] > 0
    mo0a, mo0b = mo_coeff
    orboa = mo0a[:, occidxa]
    orbva = mo0a[:,~occidxa]
    orbob = mo0b[:, occidxb]
    orbvb = mo0b[:,~occidxb]

    charges = mol.atom_charges()
    coords  = mol.atom_coords()
    charge_center = numpy.einsum('i,ix->x', charges, coords) / charges.sum()
    with mol.with_common_orig(charge_center):
        int_r = mol.intor_symmetric('int1e_r', comp=3)

    h1a = lib.einsum('xpq,pi,qj->xij', int_r, mo0a.conj(), orboa)
    h1b = lib.einsum('xpq,pi,qj->xij', int_r, mo0b.conj(), orbob)
    s1a = numpy.zeros_like(h1a)
    s1b = numpy.zeros_like(h1b)
    vind = polobj.gen_vind(mf, mo_coeff, mo_occ)
    if with_cphf:
        mo1, e1 = ucphf.solve(vind, mo_energy, mo_occ, (h1a,h1b), (s1a,s1b),
                              polobj.max_cycle_cphf, polobj.conv_tol, verbose=log)
    else:
        mo1, e1 = uhf_nmr._solve_mo1_uncoupled(mo_energy, mo_occ, (h1a,h1b),
                                               (s1a,s1b))
    mo1a = lib.einsum('xqi,pq->xpi', mo1[0], mo0a)
    mo1b = lib.einsum('xqi,pq->xpi', mo1[1], mo0b)

    dm1a = lib.einsum('xpi,qi->xpq', mo1a, orboa)
    dm1b = lib.einsum('xpi,qi->xpq', mo1b, orbob)
    dm1a = dm1a + dm1a.transpose(0,2,1)
    dm1b = dm1b + dm1b.transpose(0,2,1)
    vresp = _gen_uhf_response(mf, hermi=1)
    h1ao = int_r + vresp(numpy.stack((dm1a, dm1b)))
    s0 = mf.get_ovlp()
    e3  = lib.einsum('xpq,ypi,zqi->xyz', h1ao[0], mo1a, mo1a)
    e3 += lib.einsum('xpq,ypi,zqi->xyz', h1ao[1], mo1b, mo1b)
    e3 -= lib.einsum('pq,xpi,yqj,zij->xyz', s0, mo1a, mo1a, e1[0])
    e3 -= lib.einsum('pq,xpi,yqj,zij->xyz', s0, mo1b, mo1b, e1[1])
    e3 = (e3 + e3.transpose(1,2,0) + e3.transpose(2,0,1) +
          e3.transpose(0,2,1) + e3.transpose(1,0,2) + e3.transpose(2,1,0))
    e3 = -e3
    log.debug('Static hyper polarizability tensor\n%s', e3)
    return e3

def _make_rdm1(mycc, d1, with_frozen=True, ao_repr=False):
    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
    dm1a[numpy.diag_indices(nocca)] += 1

    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
    dm1b[numpy.diag_indices(noccb)] += 1

    if with_frozen and not (mycc.frozen is 0 or mycc.frozen is 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)
        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 test_uccsd_rdm2_mo2ao(self):
        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.spin = 2
        mol.basis = '631g'
        mol.build(0, 0)
        mf = scf.UHF(mol)
        mf.conv_tol_grad = 1e-8
        mf.kernel()
        mycc = cc.UCCSD(mf)
        mycc.diis_start_cycle = 1
        mycc.conv_tol = 1e-10
        eris = mycc.ao2mo()
        ecc, t1, t2 = mycc.kernel(eris=eris)
        l1, l2 = mycc.solve_lambda(eris=eris)
        fdm2 = lib.H5TmpFile()
        d2 = cc.uccsd_rdm._gamma2_outcore(mycc, t1, t2, l1, l2, fdm2, True)

        nao = mycc.mo_coeff[0].shape[0]
        ref = cc.uccsd_rdm._make_rdm2(mycc, None, d2, with_dm1=False)
        aa = lib.einsum('ijkl,pi,qj,rk,sl->pqrs', ref[0], mycc.mo_coeff[0],
                        mycc.mo_coeff[0], mycc.mo_coeff[0], mycc.mo_coeff[0])
        ab = lib.einsum('ijkl,pi,qj,rk,sl->pqrs', ref[1], mycc.mo_coeff[0],
                        mycc.mo_coeff[0], mycc.mo_coeff[1], mycc.mo_coeff[1])
        bb = lib.einsum('ijkl,pi,qj,rk,sl->pqrs', ref[2], mycc.mo_coeff[1],
                        mycc.mo_coeff[1], mycc.mo_coeff[1], mycc.mo_coeff[1])
        aa = aa + aa.transpose(0,1,3,2)
        aa = aa + aa.transpose(1,0,2,3)
        aa = ao2mo.restore(4, aa, nao) * .5
        bb = bb + bb.transpose(0,1,3,2)
        bb = bb + bb.transpose(1,0,2,3)
        bb = ao2mo.restore(4, bb, nao) * .5
        ab = ab + ab.transpose(0,1,3,2)
        ab = ab + ab.transpose(1,0,2,3)
        ab = ao2mo.restore(4, ab, nao) * .5
        ref = (aa+ab, bb+ab.T)
        rdm2 = uccsd_grad._rdm2_mo2ao(mycc, d2, mycc.mo_coeff)
        self.assertAlmostEqual(abs(ref[0]-rdm2[0]).max(), 0, 10)
        self.assertAlmostEqual(abs(ref[1]-rdm2[1]).max(), 0, 10)
        uccsd_grad._rdm2_mo2ao(mycc, d2, mycc.mo_coeff, fdm2)
        self.assertAlmostEqual(abs(ref[0]-fdm2['dm2aa+ab'].value).max(), 0, 10)
        self.assertAlmostEqual(abs(ref[1]-fdm2['dm2bb+ab'].value).max(), 0, 10)
        self.assertAlmostEqual(lib.finger(rdm2[0]), -1.6247203743431637, 7)
        self.assertAlmostEqual(lib.finger(rdm2[1]), -0.44062825991527471, 7)

def kernel(method, efg_nuc=None):
    log = lib.logger.Logger(method.stdout, method.verbose)
    cell = method.cell
    if efg_nuc is None:
        efg_nuc = range(cell.natm)

    dm = method.make_rdm1()
    if isinstance(method, scf.khf.KSCF):
        if isinstance(dm[0][0], numpy.ndarray) and dm[0][0].ndim == 2:
            dm = dm[0] + dm[1]  # KUHF density matrix
    elif isinstance(method, scf.hf.SCF):
        if isinstance(dm[0], numpy.ndarray) and dm[0].ndim == 2:
            dm = dm[0] + dm[1]  # UHF density matrix
    else:
        mo = method.mo_coeff
        if isinstance(dm[0][0], numpy.ndarray) and dm[0][0].ndim == 2:
            dm_a = [lib.einsum('pi,ij,qj->pq', c, dm[0][k], c.conj())
                    for k, c in enumerate(mo)]
            dm_b = [lib.einsum('pi,ij,qj->pq', c, dm[1][k], c.conj())
                    for k, c in enumerate(mo)]
            dm = lib.asarray(dm_a) + lib.asarray(dm_b)
        else:
            dm = lib.asarray([lib.einsum('pi,ij,qj->pq', c, dm[k], c.conj())
                              for k, c in enumerate(mo)])

    if isinstance(method, scf.hf.SCF):
        with_df = getattr(method, 'with_df', None)
        with_x2c = getattr(method, 'with_x2c', None)
    else:
        with_df = getattr(method._scf, 'with_df', None)
        with_x2c = getattr(method._scf, 'with_x2c', None)
    if with_x2c:
        raise NotImplementedError

    log.info('\nElectric Field Gradient Tensor Results')
    if isinstance(with_df, df.fft.FFTDF):
        efg_e = _fft_quad_integrals(with_df, dm, efg_nuc)
    else:
        efg_e = _aft_quad_integrals(with_df, dm, efg_nuc)
    efg = []
    for i, atm_id in enumerate(efg_nuc):
        efg_nuc = _get_quad_nuc(cell, atm_id)
        v = efg_nuc - efg_e[i]
        efg.append(v)

        rhf_efg._analyze(cell, atm_id, v, log)

    return numpy.asarray(efg)

    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 _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

def _make_rdm1(mycc, d1, with_frozen=True, ao_repr=False):
    '''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
    dm1[numpy.diag_indices(nocc)] += 2

    if with_frozen and not (mycc.frozen is 0 or mycc.frozen is None):
        nmo = mycc.mo_occ.size
        nocc = numpy.count_nonzero(mycc.mo_occ > 0)
        rdm1 = numpy.zeros((nmo,nmo), dtype=dm1.dtype)
        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

def polarizability_with_freq(polobj, freq=None):
    from pyscf.prop.nmr import rhf as rhf_nmr
    log = logger.new_logger(polobj)
    mf = polobj._scf
    mol = mf.mol
    mo_energy = mf.mo_energy
    mo_coeff = mf.mo_coeff
    mo_occ = mf.mo_occ
    occidx = mo_occ > 0
    orbo = mo_coeff[:, occidx]
    orbv = mo_coeff[:,~occidx]

    charges = mol.atom_charges()
    coords  = mol.atom_coords()
    charge_center = numpy.einsum('i,ix->x', charges, coords) / charges.sum()
    with mol.with_common_orig(charge_center):
        int_r = mol.intor_symmetric('int1e_r', comp=3)

    h1 = lib.einsum('xpq,pi,qj->xij', int_r, orbv.conj(), orbo)
    mo1 = cphf_with_freq(mf, mo_energy, mo_occ, h1, freq,
                         polobj.max_cycle_cphf, polobj.conv_tol, verbose=log)[0]

    e2 =  numpy.einsum('xpi,ypi->xy', h1, mo1[0])
    e2 += numpy.einsum('xpi,ypi->xy', h1, mo1[1])

    # *-1 from the definition of dipole moment. *2 for double occupancy
    e2 *= -2
    log.debug('Polarizability tensor with freq %s', freq)
    log.debug('%s', e2)
    return e2

def get_mo_pairs_G(mydf, mo_coeffs, kpts=numpy.zeros((2,3)), q=None,
                   compact=getattr(__config__, 'pbc_df_mo_pairs_compact', False)):
    '''Calculate forward fourier transform (G|ij) of all MO pairs.

    Args:
        mo_coeff: length-2 list of (nao,nmo) ndarrays
            The two sets of MO coefficients to use in calculating the
            product |ij).

    Returns:
        mo_pairs_G : (ngrids, nmoi*nmoj) ndarray
            The FFT of the real-space MO pairs.
    '''
    if kpts is None: kpts = numpy.zeros((2,3))
    cell = mydf.cell
    kpts = numpy.asarray(kpts)
    q = kpts[1] - kpts[0]
    coords = cell.gen_uniform_grids(mydf.mesh)
    nmoi = mo_coeffs[0].shape[1]
    nmoj = mo_coeffs[1].shape[1]
    ngrids = len(coords)

    mo_pairs_G = numpy.empty((ngrids,nmoi,nmoj), dtype=numpy.complex)
    for pqkR, pqkI, p0, p1 \
            in mydf.pw_loop(mydf.mesh, kptijkl[:2], q,
                            max_memory=max_memory, aosym='s2'):
        pqk = (pqkR + pqkI*1j).reshape(nao,nao,-1)
        mo_pairs_G[p0:p1] = lib.einsum('pqk,pi,qj->kij', pqk, *mo_coeffs[:2])
    return mo_pairs_G.reshape(ngrids,nmoi*nmoj)

def project_dm_nr2nr(mol1, dm1, mol2):
    r''' Project density matrix representation from basis set 1 (mol1) to basis
    set 2 (mol2).

    .. math::

        |AO2\rangle DM_AO2 \langle AO2|

        = |AO2\rangle P DM_AO1 P \langle AO2|

        DM_AO2 = P DM_AO1 P

        P = S_{AO2}^{-1}\langle AO2|AO1\rangle

    There are three relevant functions:
    :func:`project_dm_nr2nr` is the projection for non-relativistic (scalar) basis.
    :func:`project_dm_nr2r` projects from non-relativistic to relativistic basis.
    :func:`project_dm_r2r`  is the projection between relativistic (spinor) basis.
    '''
    s22 = mol2.intor_symmetric('int1e_ovlp')
    s21 = mole.intor_cross('int1e_ovlp', mol2, mol1)
    p21 = lib.cho_solve(s22, s21)
    if isinstance(dm1, numpy.ndarray) and dm1.ndim == 2:
        return reduce(numpy.dot, (p21, dm1, p21.conj().T))
    else:
        return lib.einsum('pi,nij,qj->npq', p21, dm1, p21.conj())

def contract_2e(eri, fcivec, norb, nelec, opt=None):
    if isinstance(nelec, (int, numpy.integer)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec
    link_indexa = cistring.gen_linkstr_index(range(norb), neleca)
    link_indexb = cistring.gen_linkstr_index(range(norb), nelecb)
    na = cistring.num_strings(norb, neleca)
    nb = cistring.num_strings(norb, nelecb)
    ci0 = fcivec.reshape(na,nb)
    t1 = numpy.zeros((norb,norb,na,nb))
    for str0, tab in enumerate(link_indexa):
        for a, i, str1, sign in tab:
            t1[a,i,str1] += sign * ci0[str0]
    for str0, tab in enumerate(link_indexb):
        for a, i, str1, sign in tab:
            t1[a,i,:,str1] += sign * ci0[:,str0]

    t1 = lib.einsum('bjai,aiAB->bjAB', eri.reshape([norb]*4), t1)

    fcinew = numpy.zeros_like(ci0)
    for str0, tab in enumerate(link_indexa):
        for a, i, str1, sign in tab:
            fcinew[str1] += sign * t1[a,i,str0]
    for str0, tab in enumerate(link_indexb):
        for a, i, str1, sign in tab:
            fcinew[:,str1] += sign * t1[a,i,:,str0]
    return fcinew.reshape(fcivec.shape)