Python numpy.tril_indices() Examples

The following are 30 code examples for showing how to use numpy.tril_indices(). These examples are extracted from open source projects. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example.

You may check out the related API usage on the sidebar.

You may also want to check out all available functions/classes of the module numpy , or try the search function .

Example 1
Project: pyscf   Author: pyscf   File: eom_kccsd_ghf.py    License: Apache License 2.0 6 votes vote down vote up
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] 
Example 2
Project: pyscf   Author: pyscf   File: selected_ci_symm.py    License: Apache License 2.0 6 votes vote down vote up
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 
Example 3
Project: pyscf   Author: pyscf   File: direct_spin1_symm.py    License: Apache License 2.0 6 votes vote down vote up
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 
Example 4
Project: pyscf   Author: pyscf   File: eom_uccsd.py    License: Apache License 2.0 6 votes vote down vote up
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 
Example 5
Project: pyscf   Author: pyscf   File: eom_uccsd.py    License: Apache License 2.0 6 votes vote down vote up
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 
Example 6
Project: pyscf   Author: pyscf   File: eom_rccsd.py    License: Apache License 2.0 6 votes vote down vote up
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) 
Example 7
Project: pyscf   Author: pyscf   File: test_gccsd_lambda.py    License: Apache License 2.0 6 votes vote down vote up
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 
Example 8
Project: pyscf   Author: pyscf   File: uadc_ao2mo.py    License: Apache License 2.0 6 votes vote down vote up
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_ 
Example 9
Project: pyscf   Author: pyscf   File: uadc_ao2mo.py    License: Apache License 2.0 6 votes vote down vote up
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_ 
Example 10
Project: pyscf   Author: pyscf   File: test_pbc_fill_int.py    License: Apache License 2.0 6 votes vote down vote up
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) 
Example 11
Project: pyscf   Author: pyscf   File: ucisd.py    License: Apache License 2.0 6 votes vote down vote up
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 
Example 12
Project: pyscf   Author: pyscf   File: mc1step.py    License: Apache License 2.0 6 votes vote down vote up
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) 
Example 13
Project: recordlinkage   Author: J535D165   File: test_indexing.py    License: BSD 3-Clause "New" or "Revised" License 6 votes vote down vote up
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 
Example 14
Project: nsf   Author: bayesiains   File: lu.py    License: MIT License 6 votes vote down vote up
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) 
Example 15
Project: vnpy_crypto   Author: birforce   File: dynamic_factor.py    License: MIT License 6 votes vote down vote up
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', :, :] 
Example 16
Project: vnpy_crypto   Author: birforce   File: contrasts.py    License: MIT License 6 votes vote down vote up
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 
Example 17
Project: vnpy_crypto   Author: birforce   File: contrasts.py    License: MIT License 6 votes vote down vote up
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 
Example 18
Project: pyscf   Author: pyscf   File: eom_kccsd_ghf.py    License: Apache License 2.0 5 votes vote down vote up
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] 
Example 19
Project: pyscf   Author: pyscf   File: eom_kccsd_ghf.py    License: Apache License 2.0 5 votes vote down vote up
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())) 
Example 20
Project: pyscf   Author: pyscf   File: eom_kccsd_ghf.py    License: Apache License 2.0 5 votes vote down vote up
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 
Example 21
Project: pyscf   Author: pyscf   File: eom_kccsd_rhf.py    License: Apache License 2.0 5 votes vote down vote up
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 
Example 22
Project: pyscf   Author: pyscf   File: eom_kccsd_uhf.py    License: Apache License 2.0 5 votes vote down vote up
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())) 
Example 23
Project: pyscf   Author: pyscf   File: eom_kccsd_uhf.py    License: Apache License 2.0 5 votes vote down vote up
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 
Example 24
Project: pyscf   Author: pyscf   File: eom_kccsd_uhf.py    License: Apache License 2.0 5 votes vote down vote up
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())) 
Example 25
Project: pyscf   Author: pyscf   File: test_outcore.py    License: Apache License 2.0 5 votes vote down vote up
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)) 
Example 26
Project: pyscf   Author: pyscf   File: dmrgci.py    License: Apache License 2.0 5 votes vote down vote up
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 
Example 27
Project: pyscf   Author: pyscf   File: ciah.py    License: Apache License 2.0 5 votes vote down vote up
def pack_uniq_var(self, mat):
        nmo = mat.shape[0]
        idx = numpy.tril_indices(nmo, -1)
        return mat[idx] 
Example 28
Project: pyscf   Author: pyscf   File: ciah.py    License: Apache License 2.0 5 votes vote down vote up
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 
Example 29
Project: pyscf   Author: pyscf   File: ccsd_t_rdm_slow.py    License: Apache License 2.0 5 votes vote down vote up
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 
Example 30
Project: pyscf   Author: pyscf   File: ccsd.py    License: Apache License 2.0 5 votes vote down vote up
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