Python numpy.tril_indices() Examples

def vector_to_amplitudes_ea(vector, kshift, nkpts, nmo, nocc, kconserv):
    nvir = nmo - nocc

    r1 = vector[:nvir].copy()
    r2_tril = vector[nvir:].copy().reshape(nocc*nkpts*nvir*(nkpts*nvir-1)//2)
    r2 = np.zeros((nkpts,nkpts,nocc,nvir,nvir), dtype=vector.dtype)

    index = 0
    for kj, ka in itertools.product(range(nkpts), repeat=2):
        kb = kconserv[kshift,ka,kj]
        if ka < kb:
            idx, idy = np.tril_indices(nvir, 0)
        else:
            idx, idy = np.tril_indices(nvir, -1)
        tmp = r2_tril[index:index + nocc*len(idy)].reshape(-1,nocc)
        r2[kj,ka,:,idx,idy] = tmp
        r2[kj,kb,:,idy,idx] = -tmp
        index = index + nocc*len(idy)

    return [r1,r2] 

def reorder4irrep(eri, norb, link_index, orbsym, offdiag=0):
    if orbsym is None:
        return eri, link_index, numpy.array(norb, dtype=numpy.int32)
    orbsym = numpy.asarray(orbsym)
# map irrep IDs of Dooh or Coov to D2h, C2v
# see symm.basis.linearmole_symm_descent
    orbsym = orbsym % 10
# irrep of (ij| pair
    trilirrep = (orbsym[:,None]^orbsym)[numpy.tril_indices(norb, offdiag)]
# and the number of occurence for each irrep
    dimirrep = numpy.array(numpy.bincount(trilirrep), dtype=numpy.int32)
# we sort the irreps of (ij| pair, to group the pairs which have same irreps
# "order" is irrep-id-sorted index. The (ij| paired is ordered that the
# pair-id given by order[0] comes first in the sorted pair
# "rank" is a sorted "order". Given nth (ij| pair, it returns the place(rank)
# of the sorted pair
    order = numpy.asarray(trilirrep.argsort(), dtype=numpy.int32)
    rank = numpy.asarray(order.argsort(), dtype=numpy.int32)
    eri = lib.take_2d(eri, order, order)
    link_index_irrep = link_index.copy()
    link_index_irrep[:,:,0] = rank[link_index[:,:,0]]
    return numpy.asarray(eri, order='C'), link_index_irrep, dimirrep 

def reorder_eri(eri, norb, orbsym):
    if orbsym is None:
        return [eri], numpy.arange(norb), numpy.zeros(norb,dtype=numpy.int32)
# map irrep IDs of Dooh or Coov to D2h, C2v
# see symm.basis.linearmole_symm_descent
    orbsym = numpy.asarray(orbsym) % 10
# irrep of (ij| pair
    trilirrep = (orbsym[:,None]^orbsym)[numpy.tril_indices(norb)]
# and the number of occurence for each irrep
    dimirrep = numpy.asarray(numpy.bincount(trilirrep), dtype=numpy.int32)
# we sort the irreps of (ij| pair, to group the pairs which have same irreps
# "order" is irrep-id-sorted index. The (ij| paired is ordered that the
# pair-id given by order[0] comes first in the sorted pair
# "rank" is a sorted "order". Given nth (ij| pair, it returns the place(rank)
# of the sorted pair
    old_eri_irrep = numpy.asarray(trilirrep, dtype=numpy.int32)
    rank_in_irrep = numpy.empty_like(old_eri_irrep)
    p0 = 0
    eri_irs = [numpy.zeros((0,0))] * TOTIRREPS
    for ir, nnorb in enumerate(dimirrep):
        idx = numpy.asarray(numpy.where(trilirrep == ir)[0], dtype=numpy.int32)
        rank_in_irrep[idx] = numpy.arange(nnorb, dtype=numpy.int32)
        eri_irs[ir] = lib.take_2d(eri, idx, idx)
        p0 += nnorb
    return eri_irs, rank_in_irrep, old_eri_irrep 

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

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

def unpack_eri_2s(eri, norb):

    n_oo = norb * (norb + 1) // 2
    ind_oo = np.tril_indices(norb)

    eri_ = None

    if len(eri.shape) == 2:
        if (eri.shape[0] != n_oo or eri.shape[1] != n_oo):
            raise TypeError("ERI dimensions don't match")

        temp = np.zeros((n_oo, norb, norb))
        temp[:, ind_oo[0], ind_oo[1]] = eri
        temp[:, ind_oo[1], ind_oo[0]] = eri
        eri_ = np.zeros((norb, norb, norb, norb))
        eri_[ind_oo[0], ind_oo[1]] = temp
        eri_[ind_oo[1], ind_oo[0]] = temp
    else:
            raise RuntimeError("ERI does not have a correct dimension")

    return eri_ 

def unpack_eri_2(eri, norb1, norb2):

    n_oo1 = norb1 * (norb1 + 1) // 2
    ind_oo1 = np.tril_indices(norb1)
    n_oo2 = norb2 * (norb2 + 1) // 2
    ind_oo2 = np.tril_indices(norb2)

    eri_ = None

    if len(eri.shape) == 2:
        if (eri.shape[0] != n_oo1 or eri.shape[1] != n_oo2):
            raise TypeError("ERI dimensions don't match")

        temp = np.zeros((n_oo1, norb2, norb2))
        temp[:, ind_oo2[0], ind_oo2[1]] = eri
        temp[:, ind_oo2[1], ind_oo2[0]] = eri
        eri_ = np.zeros((norb1, norb1, norb2, norb2))
        eri_[ind_oo1[0], ind_oo1[1]] = temp
        eri_[ind_oo1[1], ind_oo1[0]] = temp
    else:
            raise RuntimeError("ERI does not have a correct dimension")

    return eri_ 

def test_fill_k(self):
        fill = 'PBCnr3c_fill_ks2'
        out = run3c(fill, kband)
        self.assertAlmostEqual(finger(out), 6.6227481557912071-0.80306690266835279j, 9)

        out0 = run3c('PBCnr3c_fill_kks2', kband)
        idx = numpy.tril_indices(cell.nao_nr())
        self.assertTrue(numpy.allclose(out0[0][idx], out[0]))
        self.assertTrue(numpy.allclose(out0[2][idx], out[1]))

        fill = 'PBCnr3c_fill_ks1'
        out = run3c(fill, kband)
        self.assertAlmostEqual(finger(out), -14.883220617173009-3.9101277614911112j, 9)
        self.assertTrue(numpy.allclose(out0[0], out[0]))
        self.assertTrue(numpy.allclose(out0[2], out[1]))

        fill = 'PBCnr3c_fill_ks2'
        out = run3c(fill, numpy.zeros((1,3)))
        self.assertAlmostEqual(finger(out), -1.3258082818877739, 9) 

def amplitudes_to_cisdvec(c0, c1, c2):
    c1a, c1b = c1
    c2aa, c2ab, c2bb = c2
    nocca, nvira = c1a.shape
    noccb, nvirb = c1b.shape
    def trilidx(n):
        idx = numpy.tril_indices(n, -1)
        return idx[0] * n + idx[1]
    ooidxa = trilidx(nocca)
    vvidxa = trilidx(nvira)
    ooidxb = trilidx(noccb)
    vvidxb = trilidx(nvirb)
    size = (1, nocca*nvira, noccb*nvirb, nocca*noccb*nvira*nvirb,
            len(ooidxa)*len(vvidxa), len(ooidxb)*len(vvidxb))
    loc = numpy.cumsum(size)
    civec = numpy.empty(loc[-1], dtype=c2ab.dtype)
    civec[0] = c0
    civec[loc[0]:loc[1]] = c1a.ravel()
    civec[loc[1]:loc[2]] = c1b.ravel()
    civec[loc[2]:loc[3]] = c2ab.ravel()
    lib.take_2d(c2aa.reshape(nocca**2,nvira**2), ooidxa, vvidxa, out=civec[loc[3]:loc[4]])
    lib.take_2d(c2bb.reshape(noccb**2,nvirb**2), ooidxb, vvidxb, out=civec[loc[4]:loc[5]])
    return civec 

def _exact_paaa(self, mo, u, out=None):
        nmo = mo.shape[1]
        ncore = self.ncore
        ncas = self.ncas
        nocc = ncore + ncas
        mo1 = numpy.dot(mo, u)
        mo1_cas = mo1[:,ncore:nocc]
        mos = (mo1_cas, mo1_cas, mo1, mo1_cas)
        if self._scf._eri is None:
            aapa = ao2mo.general(self.mol, mos)
        else:
            aapa = ao2mo.general(self._scf._eri, mos)
        paaa = numpy.empty((nmo*ncas,ncas*ncas))
        buf = numpy.empty((ncas,ncas,nmo*ncas))
        for ij, (i, j) in enumerate(zip(*numpy.tril_indices(ncas))):
            buf[i,j] = buf[j,i] = aapa[ij]
        paaa = lib.transpose(buf.reshape(ncas*ncas,-1), out=out)
        return paaa.reshape(nmo,ncas,ncas,ncas) 

def test_lower_triangular(self, index_class):

        # make an index for each dataframe with a new index name
        index_a = pd.Index(self.a.index, name='index')
        df_a = pd.DataFrame(self.a, index=index_a)
        pairs = index_class.index(df_a)

        # expected
        levels = [df_a.index.values, df_a.index.values]
        codes = np.tril_indices(len(df_a.index), k=-1)

        full_pairs = pd.MultiIndex(levels=levels,
                                   codes=codes,
                                   verify_integrity=False)

        # all pairs are in the lower triangle of the matrix.
        assert len(pairs.difference(full_pairs)) == 0 

def __init__(self, features, using_cache=False, identity_init=True, eps=1e-3):
        super().__init__(features, using_cache)

        self.eps = eps

        self.lower_indices = np.tril_indices(features, k=-1)
        self.upper_indices = np.triu_indices(features, k=1)
        self.diag_indices = np.diag_indices(features)

        n_triangular_entries = ((features - 1) * features) // 2

        self.lower_entries = nn.Parameter(torch.zeros(n_triangular_entries))
        self.upper_entries = nn.Parameter(torch.zeros(n_triangular_entries))
        self.unconstrained_upper_diag = nn.Parameter(torch.zeros(features))

        self._initialize(identity_init) 

def _initialize_error_cov_unstructured(self):
        # Initialize the parameters
        k_endog = self.k_endog
        self.parameters['error_cov'] = int(k_endog * (k_endog + 1) / 2)

        # Setup fixed components of state space matrices

        # Setup indices of state space matrices
        self._idx_lower_error_cov = np.tril_indices(self.k_endog)
        if self.error_order > 0:
            start = self.k_factors
            end = self.k_factors + self.k_endog
            self._idx_error_cov = (
                np.s_['state_cov', start:end, start:end])
        else:
            self._idx_error_cov = np.s_['obs_cov', :, :] 

def _helmert_contrast(self, levels):
        n = len(levels)
        #http://www.ats.ucla.edu/stat/sas/webbooks/reg/chapter5/sasreg5.htm#HELMERT
        #contr = np.eye(n - 1)
        #int_range = np.arange(n - 1., 1, -1)
        #denom = np.repeat(int_range, np.arange(n - 2, 0, -1))
        #contr[np.tril_indices(n - 1, -1)] = -1. / denom

        #http://www.ats.ucla.edu/stat/r/library/contrast_coding.htm#HELMERT
        #contr = np.zeros((n - 1., n - 1))
        #int_range = np.arange(n, 1, -1)
        #denom = np.repeat(int_range[:-1], np.arange(n - 2, 0, -1))
        #contr[np.diag_indices(n - 1)] = (int_range - 1.) / int_range
        #contr[np.tril_indices(n - 1, -1)] = -1. / denom
        #contr = np.vstack((contr, -1./int_range))

        #r-like
        contr = np.zeros((n, n - 1))
        contr[1:][np.diag_indices(n - 1)] = np.arange(1, n)
        contr[np.triu_indices(n - 1)] = -1
        return contr 

def _diff_contrast(self, levels):
        nlevels = len(levels)
        contr = np.zeros((nlevels, nlevels-1))
        int_range = np.arange(1, nlevels)
        upper_int = np.repeat(int_range, int_range)
        row_i, col_i = np.triu_indices(nlevels-1)
        # we want to iterate down the columns not across the rows
        # it would be nice if the index functions had a row/col order arg
        col_order = np.argsort(col_i)
        contr[row_i[col_order],
              col_i[col_order]] = (upper_int-nlevels)/float(nlevels)
        lower_int = np.repeat(int_range, int_range[::-1])
        row_i, col_i = np.tril_indices(nlevels-1)
        # we want to iterate down the columns not across the rows
        col_order = np.argsort(col_i)
        contr[row_i[col_order]+1, col_i[col_order]] = lower_int/float(nlevels)
        return contr 

def vector_to_amplitudes_ip(vector, kshift, nkpts, nmo, nocc, kconserv):
    nvir = nmo - nocc

    r1 = vector[:nocc].copy()
    r2_tril = vector[nocc:].copy().reshape(nkpts*nocc*(nkpts*nocc-1)//2,nvir)
    idx, idy = np.tril_indices(nkpts*nocc, -1)
    r2 = np.zeros((nkpts*nocc,nkpts*nocc,nvir), dtype=vector.dtype)
    r2[idx, idy] = r2_tril
    r2[idy, idx] = -r2_tril
    r2 = r2.reshape(nkpts,nocc,nkpts,nocc,nvir).transpose(0,2,1,3,4)
    return [r1,r2] 

def amplitudes_to_vector_ip(r1, r2, kshift, kconserv):
    nkpts, nocc, nvir = np.asarray(r2.shape)[[0,2,4]]
    # From symmetry for aaa and bbb terms, only store lower
    # triangular part (ki,i) < (kj,j)
    idx, idy = np.tril_indices(nkpts*nocc, -1)
    r2 = r2.transpose(0,2,1,3,4).reshape(nkpts*nocc,nkpts*nocc,nvir)
    return np.hstack((r1, r2[idx,idy].ravel())) 

def amplitudes_to_vector_ea(r1, r2, kshift, kconserv):
    nkpts, nocc, nvir = np.asarray(r2.shape)[[0,2,3]]
    r2_tril = np.zeros((nocc*nkpts*nvir*(nkpts*nvir-1)//2), dtype=r2.dtype)
    index = 0
    for kj, ka in itertools.product(range(nkpts), repeat=2):
        kb = kconserv[kshift,ka,kj]
        if ka < kb:
            idx, idy = np.tril_indices(nvir, 0)
        else:
            idx, idy = np.tril_indices(nvir, -1)
        r2_tril[index:index + nocc*len(idy)] = r2[kj,ka,:,idx,idy].reshape(-1)
        index = index + nocc*len(idy)
    vector = np.hstack((r1, r2_tril))
    return vector 

def amplitudes_to_vector_singlet(r1, r2, kconserv):
    '''Transform 3- and 7-dimensional arrays, r1 and r2, to a 1-dimensional
    array with unique indices.

    For example:
        r1: t_{i k_i}^{a k_a}
        r2: t_{i k_i, j k_j}^{a k_a, b k_b}
        return: a vector with all r1 elements, and r2 elements whose indices
    satisfy (i k_i a k_a) >= (j k_j b k_b)
    '''
    cput0 = (time.clock(), time.time())
    log = logger.Logger(sys.stdout, logger.DEBUG)
    # r1 indices: k_i, i, a
    nkpts, nocc, nvir = np.asarray(r1.shape)[[0, 1, 2]]
    nov = nocc * nvir

    # r2 indices (old): k_i, k_J, k_a, i, J, a, B
    # r2 indices (new): (k_i, k_a), (k_J), (i, a), (J, B)
    r2 = r2.transpose(0,2,1,3,5,4,6).reshape(nkpts**2, nkpts, nov, nov)

    idx, idy = np.tril_indices(nov)
    nov2_tril = nov * (nov + 1) // 2
    nov2 = nov * nov

    vector = np.empty(r2.size, dtype=r2.dtype)
    offset = 0
    for ki, ka, kj in kpts_helper.loop_kkk(nkpts):
        kb = kconserv[ki, ka, kj]
        kika = ki * nkpts + ka
        kjkb = kj * nkpts + kb
        r2ovov = r2[kika, kj]
        if kika == kjkb:
            vector[offset:offset+nov2_tril] = r2ovov[idx, idy]
            offset += nov2_tril
        elif kika > kjkb:
            vector[offset:offset+nov2] = r2ovov.ravel()
            offset += nov2

    vector = np.hstack((r1.ravel(), vector[:offset]))
    log.timer("amplitudes_to_vector_singlet", *cput0)
    return vector 

def amplitudes_to_vector_ip(r1, r2, kshift, kconserv):
    r1a, r1b = r1
    r2aaa, r2baa, r2abb, r2bbb = r2
    nkpts = r2aaa.shape[0]
    nocca, noccb = r1a.shape[0], r1b.shape[0]
    nvira, nvirb = r2aaa.shape[-1], r2bbb.shape[-1]
    # From symmetry for aaa and bbb terms, only store lower
    # triangular part (ki,i) < (kj,j)
    idxa, idya = np.tril_indices(nkpts*nocca, -1)
    idxb, idyb = np.tril_indices(nkpts*noccb, -1)
    r2aaa = r2aaa.transpose(0,2,1,3,4).reshape(nkpts*nocca,nkpts*nocca,nvira)
    r2bbb = r2bbb.transpose(0,2,1,3,4).reshape(nkpts*noccb,nkpts*noccb,nvirb)
    return np.hstack((r1a, r1b, r2aaa[idxa,idya].ravel(),
                      r2baa.ravel(), r2abb.ravel(),
                      r2bbb[idxb,idyb].ravel())) 

def vector_to_amplitudes_ip(vector, kshift, nkpts, nmo, nocc, kconserv):
    nocca, noccb = nocc
    nmoa, nmob = nmo
    nvira, nvirb = nmoa-nocca, nmob-noccb

    sizes = (nocca, noccb, (nkpts*nocca)*(nkpts*nocca-1)*nvira//2,
             nkpts**2*noccb*nocca*nvira, nkpts**2*nocca*noccb*nvirb,
             nkpts*noccb*(nkpts*noccb-1)*nvirb//2)
    sections = np.cumsum(sizes[:-1])
    r1a, r1b, r2a, r2baa, r2abb, r2b = np.split(vector, sections)

    r2a = r2a.reshape(nkpts*nocca*(nkpts*nocca-1)//2,nvira)
    r2b = r2b.reshape(nkpts*noccb*(nkpts*noccb-1)//2,nvirb)

    idxa, idya = np.tril_indices(nkpts*nocca, -1)
    idxb, idyb = np.tril_indices(nkpts*noccb, -1)

    r2aaa = np.zeros((nkpts*nocca,nkpts*nocca,nvira), dtype=r2a.dtype)
    r2aaa[idxa,idya] = r2a.copy()
    r2aaa[idya,idxa] = -r2a.copy()  # Fill in value :  kj, j < ki, i
    r2aaa = r2aaa.reshape(nkpts,nocca,nkpts,nocca,nvira)
    r2aaa = r2aaa.transpose(0,2,1,3,4)
    r2baa = r2baa.reshape(nkpts,nkpts,noccb,nocca,nvira).copy()
    r2abb = r2abb.reshape(nkpts,nkpts,nocca,noccb,nvirb).copy()
    r2bbb = np.zeros((nkpts*noccb,nkpts*noccb,nvirb), dtype=r2b.dtype)
    r2bbb[idxb,idyb] = r2b.copy()
    r2bbb[idyb,idxb] = -r2b.copy()  # Fill in value :  kj, j < ki, i
    r2bbb = r2bbb.reshape(nkpts,noccb,nkpts,noccb,nvirb)
    r2bbb = r2bbb.transpose(0,2,1,3,4)

    r1 = (r1a.copy(), r1b.copy())
    r2 = (r2aaa, r2baa, r2abb, r2bbb)
    return r1, r2 

def amplitudes_to_vector_ea(r1, r2, kshift, kconserv):
    r1a, r1b = r1
    r2a, r2aba, r2bab, r2b = r2
    nkpts = r2a.shape[0]
    nocca, noccb = r2a.shape[2], r2b.shape[2]
    nvira, nvirb = r2a.shape[3], r2b.shape[3]
    # From symmetry for aaa and bbb terms, only store lower
    # triangular part (ka,a) < (kb,b)
    r2aaa = np.zeros((nocca*nkpts*nvira*(nkpts*nvira-1))//2, dtype=r2a.dtype)
    r2bbb = np.zeros((noccb*nkpts*nvirb*(nkpts*nvirb-1))//2, dtype=r2b.dtype)

    index = 0
    for kj, ka in itertools.product(range(nkpts), repeat=2):
        kb = kconserv[kshift,ka,kj]
        if ka < kb:  # Take diagonal part
            idxa, idya = np.tril_indices(nvira, 0)
        else:  # Don't take diagonal (equal to zero)
            idxa, idya = np.tril_indices(nvira, -1)
        r2aaa[index:index + nocca*len(idya)] = r2a[kj,ka,:,idxa,idya].reshape(-1)
        index = index + nocca*len(idya)

    index = 0
    for kj, ka in itertools.product(range(nkpts), repeat=2):
        kb = kconserv[kshift,ka,kj]
        if ka < kb:  # Take diagonal part
            idxb, idyb = np.tril_indices(nvirb, 0)
        else:
            idxb, idyb = np.tril_indices(nvirb, -1)
        r2bbb[index:index + noccb*len(idyb)] = r2b[kj,ka,:,idxb,idyb].reshape(-1)
        index = index + noccb*len(idyb)

    return np.hstack((r1a, r1b, r2aaa.ravel(),
                      r2aba.ravel(), r2bab.ravel(),
                      r2bbb.ravel())) 

def test_aux_e1(self):
        tmpfile = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
        numpy.random.seed(1)
        kptij_lst = numpy.random.random((3,2,3))
        kptij_lst[0] = 0
        outcore.aux_e1(cell, cell, tmpfile.name, aosym='s2', comp=1,
                       kptij_lst=kptij_lst, verbose=0)
        refk = incore.aux_e2(cell, cell, aosym='s2', kptij_lst=kptij_lst)
        with h5py.File(tmpfile.name, 'r') as f:
            nao = cell.nao_nr()
            idx = numpy.tril_indices(nao)
            idx = idx[0] * nao + idx[1]
            self.assertTrue(numpy.allclose(refk[0,idx], f['eri_mo/0'].value.T))
            self.assertTrue(numpy.allclose(refk[1], f['eri_mo/1'].value.T))
            self.assertTrue(numpy.allclose(refk[2], f['eri_mo/2'].value.T)) 

def writeIntegralFile(DMRGCI, h1eff, eri_cas, ncas, nelec, ecore=0):
    if isinstance(nelec, (int, numpy.integer)):
        neleca = nelec//2 + nelec%2
        nelecb = nelec - neleca
    else :
        neleca, nelecb = nelec

    # The name of the FCIDUMP file, default is "FCIDUMP".
    integralFile = os.path.join(DMRGCI.runtimeDir, DMRGCI.integralFile)
    if DMRGCI.groupname is not None and DMRGCI.orbsym is not []:
# First removing the symmetry forbidden integrals. This has been done using
# the pyscf internal irrep-IDs (stored in DMRGCI.orbsym)
        orbsym = numpy.asarray(DMRGCI.orbsym) % 10
        pair_irrep = (orbsym.reshape(-1,1) ^ orbsym)[numpy.tril_indices(ncas)]
        sym_forbid = pair_irrep.reshape(-1,1) != pair_irrep.ravel()
        eri_cas = ao2mo.restore(4, eri_cas, ncas)
        eri_cas[sym_forbid] = 0
        eri_cas = ao2mo.restore(8, eri_cas, ncas)
       #orbsym = numpy.asarray(dmrg_sym.convert_orbsym(DMRGCI.groupname, DMRGCI.orbsym))
       #eri_cas = pyscf.ao2mo.restore(8, eri_cas, ncas)
# Then convert the pyscf internal irrep-ID to molpro irrep-ID
        orbsym = numpy.asarray(dmrg_sym.convert_orbsym(DMRGCI.groupname, orbsym))
    else:
        orbsym = []
        eri_cas = ao2mo.restore(8, eri_cas, ncas)
    if not os.path.exists(DMRGCI.scratchDirectory):
        os.makedirs(DMRGCI.scratchDirectory)
    if not os.path.exists(DMRGCI.runtimeDir):
        os.makedirs(DMRGCI.runtimeDir)

    tools.fcidump.from_integrals(integralFile, h1eff, eri_cas, ncas,
                                 neleca+nelecb, ecore, ms=abs(neleca-nelecb),
                                 orbsym=orbsym)
    return integralFile 

def pack_uniq_var(self, mat):
        nmo = mat.shape[0]
        idx = numpy.tril_indices(nmo, -1)
        return mat[idx] 

def unpack_uniq_var(self, v):
        nmo = int(numpy.sqrt(v.size*2)) + 1
        idx = numpy.tril_indices(nmo, -1)
        mat = numpy.zeros((nmo,nmo))
        mat[idx] = v
        return mat - mat.conj().T 

def _gamma2_intermediates(mycc, t1, t2, l1, l2, eris=None,
                          compress_vvvv=False):
    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 _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