Python numpy.einsum函数代码示例

        def myprecon(resid, eigval, eigvec):

            myprecon_cutoff = 1e-10
            local_myprecon = np.zeros([numVars + 1], dtype=float)
            for occ in range(numPairs):
                for virt in range(numVirt):
                    denominator = FOCK_mo[numPairs + virt, numPairs + virt] - FOCK_mo[occ, occ] - eigval
                    if abs(denominator) < myprecon_cutoff:
                        local_myprecon[occ + numPairs * virt] = eigvec[occ + numPairs * virt] / myprecon_cutoff
                    else:
                        # local_myprecon = eigvec / ( diag(H) - eigval ) = K^{-1} u
                        local_myprecon[occ + numPairs * virt] = eigvec[occ + numPairs * virt] / denominator
            if abs(eigval) < myprecon_cutoff:
                local_myprecon[numVars] = eigvec[numVars] / myprecon_cutoff
            else:
                local_myprecon[numVars] = -eigvec[numVars] / eigval
            # alpha_myprecon = - ( r, K^{-1} u ) / ( u, K^{-1} u )
            alpha_myprecon = -np.einsum("i,i->", local_myprecon, resid) / np.einsum("i,i->", local_myprecon, eigvec)
            # local_myprecon = r - ( r, K^{-1} u ) / ( u, K^{-1} u ) * u
            local_myprecon = resid + alpha_myprecon * eigvec
            for occ in range(numPairs):
                for virt in range(numVirt):
                    denominator = FOCK_mo[numPairs + virt, numPairs + virt] - FOCK_mo[occ, occ] - eigval
                    if abs(denominator) < myprecon_cutoff:
                        local_myprecon[occ + numPairs * virt] = -local_myprecon[occ + numPairs * virt] / myprecon_cutoff
                    else:
                        local_myprecon[occ + numPairs * virt] = -local_myprecon[occ + numPairs * virt] / denominator
            if abs(eigval) < myprecon_cutoff:
                local_myprecon[numVars] = -local_myprecon[occ + numPairs * virt] / myprecon_cutoff
            else:
                local_myprecon[numVars] = local_myprecon[occ + numPairs * virt] / eigval
            return local_myprecon

def analyze(mf, verbose=logger.DEBUG, **kwargs):
    '''Analyze the given SCF object:  print orbital energies, occupancies;
    print orbital coefficients; Mulliken population analysis
    '''
    from pyscf.lo import orth
    from pyscf.tools import dump_mat
    mo_energy = mf.mo_energy
    mo_occ = mf.mo_occ
    mo_coeff = mf.mo_coeff
    if isinstance(verbose, logger.Logger):
        log = verbose
    else:
        log = logger.Logger(mf.stdout, verbose)

    log.note('**** MO energy ****')
    if mf._focka_ao is None:
        for i,c in enumerate(mo_occ):
            log.note('MO #%-3d energy= %-18.15g occ= %g', i+1, mo_energy[i], c)
    else:
        mo_ea = numpy.einsum('ik,ik->k', mo_coeff, mf._focka_ao.dot(mo_coeff))
        mo_eb = numpy.einsum('ik,ik->k', mo_coeff, mf._fockb_ao.dot(mo_coeff))
        log.note('                Roothaan           | alpha              | beta')
        for i,c in enumerate(mo_occ):
            log.note('MO #%-3d energy= %-18.15g | %-18.15g | %-18.15g occ= %g',
                     i+1, mo_energy[i], mo_ea[i], mo_eb[i], c)
    ovlp_ao = mf.get_ovlp()
    if verbose >= logger.DEBUG:
        log.debug(' ** MO coefficients (expansion on meta-Lowdin AOs) **')
        label = mf.mol.spheric_labels(True)
        orth_coeff = orth.orth_ao(mf.mol, 'meta_lowdin', s=ovlp_ao)
        c = reduce(numpy.dot, (orth_coeff.T, ovlp_ao, mo_coeff))
        dump_mat.dump_rec(mf.stdout, c, label, start=1, **kwargs)
    dm = mf.make_rdm1(mo_coeff, mo_occ)
    return mf.mulliken_meta(mf.mol, dm, s=s, verbose=log)

def compute_pixels(orb, sgeom, times, rpy=(0.0, 0.0, 0.0)):
    """Compute cartesian coordinates of the pixels in instrument scan."""
    if isinstance(orb, (list, tuple)):
        tle1, tle2 = orb
        orb = Orbital("mysatellite", line1=tle1, line2=tle2)

    # get position and velocity for each time of each pixel
    pos, vel = orb.get_position(times, normalize=False)

    # now, get the vectors pointing to each pixel
    vectors = sgeom.vectors(pos, vel, *rpy)

    # compute intersection of lines (directed by vectors and passing through
    # (0, 0, 0)) and ellipsoid. Derived from:
    # http://en.wikipedia.org/wiki/Line%E2%80%93sphere_intersection

    # do the computation between line and ellipsoid (WGS 84)
    # NB: AAPP uses GRS 80...
    centre = -pos
    a__ = 6378.137  # km
    # b__ = 6356.75231414 # km, GRS80
    b__ = 6356.752314245  # km, WGS84
    radius = np.array([[1 / a__, 1 / a__, 1 / b__]]).T
    shape = vectors.shape

    xr_ = vectors.reshape([3, -1]) * radius
    cr_ = centre.reshape([3, -1]) * radius
    ldotc = np.einsum("ij,ij->j", xr_, cr_)
    lsq = np.einsum("ij,ij->j", xr_, xr_)
    csq = np.einsum("ij,ij->j", cr_, cr_)

    d1_ = (ldotc - np.sqrt(ldotc ** 2 - csq * lsq + lsq)) / lsq

    # return the actual pixel positions
    return vectors * d1_.reshape(shape[1:]) - centre

def contract_ep(g, fcivec, nsite, nelec, nphonon):
    if isinstance(nelec, (int, numpy.number)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec
    strsa = numpy.asarray(cistring.gen_strings4orblist(range(nsite), neleca))
    strsb = numpy.asarray(cistring.gen_strings4orblist(range(nsite), nelecb))
    cishape = make_shape(nsite, nelec, nphonon)
    na, nb = cishape[:2]
    ci0 = fcivec.reshape(cishape)
    fcinew = numpy.zeros(cishape)
    nbar = float(neleca+nelecb) / nsite

    phonon_cre = numpy.sqrt(numpy.arange(1,nphonon+1))
    for i in range(nsite):
        maska = (strsa & (1<<i)) > 0
        maskb = (strsb & (1<<i)) > 0
        e_part = numpy.zeros((na,nb))
        e_part[maska,:] += 1
        e_part[:,maskb] += 1
        e_part[:] -= float(neleca+nelecb) / nsite
        for ip in range(nphonon):
            slices1 = slices_for_cre(i, nsite, ip)
            slices0 = slices_for    (i, nsite, ip)
            fcinew[slices1] += numpy.einsum('ij...,ij...->ij...', g*phonon_cre[ip]*e_part, ci0[slices0])
            fcinew[slices0] += numpy.einsum('ij...,ij...->ij...', g*phonon_cre[ip]*e_part, ci0[slices1])
    return fcinew.reshape(fcivec.shape)

def get_pp_loc_part1(mydf, cell, kpts):
    log = logger.Logger(mydf.stdout, mydf.verbose)
    t1 = t0 = (time.clock(), time.time())
    nkpts = len(kpts)

    gs = mydf.gs
    nao = cell.nao_nr()
    Gv = cell.get_Gv(gs)
    SI = cell.get_SI(Gv)
    vpplocG = pseudo.pp_int.get_gth_vlocG_part1(cell, Gv)
    vpplocG = -1./cell.vol * numpy.einsum('ij,ij->j', SI, vpplocG)
    kpt_allow = numpy.zeros(3)
    real = gamma_point(kpts)

    if real:
        vloc = numpy.zeros((nkpts,nao**2))
    else:
        vloc = numpy.zeros((nkpts,nao**2), dtype=numpy.complex128)
    max_memory = mydf.max_memory - lib.current_memory()[0]
    for k, pqkR, pqkI, p0, p1 \
            in mydf.ft_loop(cell, mydf.gs, kpt_allow, kpts, max_memory=max_memory):
        vG = vpplocG[p0:p1]
        if not real:
            vloc[k] += numpy.einsum('k,xk->x', vG.real, pqkI) * 1j
            vloc[k] += numpy.einsum('k,xk->x', vG.imag, pqkR) *-1j
        vloc[k] += numpy.einsum('k,xk->x', vG.real, pqkR)
        vloc[k] += numpy.einsum('k,xk->x', vG.imag, pqkI)
        pqkR = pqkI = None
    t1 = log.timer_debug1('contracting vloc part1', *t1)
    return vloc.reshape(-1,nao,nao)

    def toMagnet(self, pos, vel=None):
        """ Transform the set of lab-frame coordinates into the coordinate
        frame of this array. Positions are shifted then rotated, then velocity
        vectors are rotated.

        Parameters:
            pos ((n,3) np.ndarray) (mm):
                Array of particle positions.
            vel ((n,3) np.ndarray) (mm/us):
                Array of particle velocities.
        """

        #pos = np.atleast_2d(pos)
        # Move to magnet origin
        pos[:,0] -= self.position[0]
        pos[:,1] -= self.position[1]
        pos[:,2] -= self.position[2]

        if self.angle != None:
            # Generate rotation matrix
            rot = np.dot(self._yMatrix(self.angle[1]), self._xMatrix(self.angle[0]))
            # Take the dot product of each position and velocity with this matrix.
            pos[:] = np.einsum('jk,ik->ij', rot, pos)
            if vel != None:
                vel[:] = np.einsum('jk,ik->ij', rot, vel)

        if vel==None:
            return pos
        else:
            return pos, vel

    def hop_uhf2ghf(x1):
        x1ab = []
        x1ba = []
        ip = 0
        for k in range(nkpts):
            nv = nvira[k]
            no = noccb[k]
            x1ab.append(x1[ip:ip+nv*no].reshape(nv,no))
            ip += nv * no
        for k in range(nkpts):
            nv = nvirb[k]
            no = nocca[k]
            x1ba.append(x1[ip:ip+nv*no].reshape(nv,no))
            ip += nv * no

        dm1ab = []
        dm1ba = []
        for k in range(nkpts):
            d1ab = reduce(numpy.dot, (orbva[k], x1ab[k], orbob[k].T.conj()))
            d1ba = reduce(numpy.dot, (orbvb[k], x1ba[k], orboa[k].T.conj()))
            dm1ab.append(d1ab+d1ba.T.conj())
            dm1ba.append(d1ba+d1ab.T.conj())

        v1ao = vresp1(lib.asarray([dm1ab,dm1ba]))
        x2ab = [0] * nkpts
        x2ba = [0] * nkpts
        for k in range(nkpts):
            x2ab[k] = numpy.einsum('pr,rq->pq', fvva[k], x1ab[k])
            x2ab[k]-= numpy.einsum('sq,ps->pq', foob[k], x1ab[k])
            x2ba[k] = numpy.einsum('pr,rq->pq', fvvb[k], x1ba[k])
            x2ba[k]-= numpy.einsum('qs,ps->pq', fooa[k], x1ba[k])
            x2ab[k] += reduce(numpy.dot, (orbva[k].T.conj(), v1ao[0][k], orbob[k]))
            x2ba[k] += reduce(numpy.dot, (orbvb[k].T.conj(), v1ao[1][k], orboa[k]))
        return numpy.hstack([x.real.ravel() for x in (x2ab+x2ba)])

def compute_inertia_tensor(traj):
    """Compute the inertia tensor of a trajectory.

    For each frame,
        I_{ab} = sum_{i_atoms} [m_i * (r_i^2 * d_{ab} - r_{ia} * r_{ib})]

    Parameters
    ----------
    traj : Trajectory
        Trajectory to compute inertia tensor of.

    Returns
    -------
    I_ab:  np.ndarray, shape=(traj.n_frames, 3, 3), dtype=float64
        Inertia tensors for each frame.

    """
    center_of_mass = np.expand_dims(compute_center_of_mass(traj), axis=1)
    xyz = traj.xyz - center_of_mass
    masses = np.array([atom.element.mass for atom in traj.top.atoms])

    eyes = np.empty(shape=(traj.n_frames, 3, 3), dtype=np.float64)
    eyes[:] = np.eye(3)
    A = np.einsum("i, kij->k", masses, xyz ** 2).reshape(traj.n_frames, 1, 1)
    B = np.einsum("ij..., ...jk->...ki", masses[:, np.newaxis] * xyz.T, xyz)
    return A * eyes - B

def einsum(subscripts, *tensors, **kwargs):
    '''Perform a more efficient einsum via reshaping to a matrix multiply.

    Current differences compared to numpy.einsum:
    This assumes that each repeated index is actually summed (i.e. no 'i,i->i')
    and appears only twice (i.e. no 'ij,ik,il->jkl'). The output indices must
    be explicitly specified (i.e. 'ij,j->i' and not 'ij,j').
    '''
    contract = kwargs.pop('_contract', _contract)

    subscripts = subscripts.replace(' ','')
    if len(tensors) <= 1 or '...' in subscripts:
        out = numpy.einsum(subscripts, *tensors, **kwargs)
    elif len(tensors) <= 2:
        out = _contract(subscripts, *tensors, **kwargs)
    else:
        if '->' in subscripts:
            indices_in, idx_final = subscripts.split('->')
            indices_in = indices_in.split(',')
        else:
            idx_final = ''
            indices_in = subscripts.split('->')[0].split(',')
        tensors = list(tensors)
        contraction_list = _einsum_path(subscripts, *tensors, optimize=True,
                                        einsum_call=True)[1]
        for contraction in contraction_list:
            inds, idx_rm, einsum_str, remaining = contraction[:4]
            tmp_operands = [tensors.pop(x) for x in inds]
            if len(tmp_operands) > 2:
                out = numpy.einsum(einsum_str, *tmp_operands)
            else:
                out = contract(einsum_str, *tmp_operands)
            tensors.append(out)
    return out

def Srsi(mc, dms, eris, verbose=None):
    #Subspace S_ijr^{(1)}
    mo_core, mo_cas, mo_virt = _extract_orbs(mc, mc.mo_coeff)
    dm1 = dms['1']
    dm2 = dms['2']
    ncore = mo_core.shape[1]
    ncas = mo_cas.shape[1]
    nocc = ncore + ncas
    if eris is None:
        h1e = mc.h1e_for_cas()[0]
        h2e = ao2mo.restore(1, mc.ao2mo(mo_cas), ncas).transpose(0,2,1,3)
        h2e_v = ao2mo.incore.general(mc._scf._eri,[mo_virt,mo_core,mo_virt,mo_cas],compact=False)
        h2e_v = h2e_v.reshape(mo_virt.shape[1],ncore,mo_virt.shape[1],ncas).transpose(0,2,1,3)
    else:
        h1e = eris['h1eff'][ncore:nocc,ncore:nocc]
        h2e = eris['ppaa'][ncore:nocc,ncore:nocc].transpose(0,2,1,3)
        h2e_v = eris['pacv'][nocc:].transpose(3,0,2,1)

    k27 = make_k27(h1e,h2e,dm1,dm2)
    norm = 2.0*numpy.einsum('rsip,rsia,pa->rsi',h2e_v,h2e_v,dm1)\
         - 1.0*numpy.einsum('rsip,sria,pa->rsi',h2e_v,h2e_v,dm1)
    h = 2.0*numpy.einsum('rsip,rsia,pa->rsi',h2e_v,h2e_v,k27)\
         - 1.0*numpy.einsum('rsip,sria,pa->rsi',h2e_v,h2e_v,k27)
    diff = mc.mo_energy[nocc:,None,None] + mc.mo_energy[None,nocc:,None] - mc.mo_energy[None,None,:ncore]
    return _norm_to_energy(norm, h, diff)

def Srs(mc, dms, eris=None, verbose=None):
    #Subspace S_rs^{(-2)}
    mo_core, mo_cas, mo_virt = _extract_orbs(mc, mc.mo_coeff)
    dm1 = dms['1']
    dm2 = dms['2']
    dm3 = dms['3']
    ncore = mo_core.shape[1]
    ncas = mo_cas.shape[1]
    nocc = ncore + ncas
    if mo_virt.shape[1] ==0:
        return 0, 0
    if eris is None:
        h1e = mc.h1e_for_cas()[0]
        h2e = ao2mo.restore(1, mc.ao2mo(mo_cas), ncas).transpose(0,2,1,3)
        h2e_v = ao2mo.incore.general(mc._scf._eri,[mo_virt,mo_cas,mo_virt,mo_cas],compact=False)
        h2e_v = h2e_v.reshape(mo_virt.shape[1],ncas,mo_virt.shape[1],ncas).transpose(0,2,1,3)
    else:
        h1e = eris['h1eff'][ncore:nocc,ncore:nocc]
        h2e = eris['ppaa'][ncore:nocc,ncore:nocc].transpose(0,2,1,3)
        h2e_v = eris['papa'][nocc:,:,nocc:].transpose(0,2,1,3)

# a7 is very sensitive to the accuracy of HF orbital and CI wfn
    rm2, a7 = make_a7(h1e,h2e,dm1,dm2,dm3)
    norm = 0.5*numpy.einsum('rsqp,rsba,pqba->rs',h2e_v,h2e_v,rm2)
    h = 0.5*numpy.einsum('rsqp,rsba,pqab->rs',h2e_v,h2e_v,a7)
    diff = mc.mo_energy[nocc:,None] + mc.mo_energy[None,nocc:]
    return _norm_to_energy(norm, h, diff)

def Sijr(mc, dms, eris, verbose=None):
    #Subspace S_ijr^{(1)}
    mo_core, mo_cas, mo_virt = _extract_orbs(mc, mc.mo_coeff)
    dm1 = dms['1']
    dm2 = dms['2']
    ncore = mo_core.shape[1]
    ncas = mo_cas.shape[1]
    nocc = ncore + ncas
    if eris is None:
        h1e = mc.h1e_for_cas()[0]
        h2e = ao2mo.restore(1, mc.ao2mo(mo_cas), ncas).transpose(0,2,1,3)
        h2e_v = ao2mo.incore.general(mc._scf._eri,[mo_virt,mo_core,mo_cas,mo_core],compact=False)
        h2e_v = h2e_v.reshape(mo_virt.shape[1],ncore,ncas,ncore).transpose(0,2,1,3)
    else:
        h1e = eris['h1eff'][ncore:nocc,ncore:nocc]
        h2e = eris['ppaa'][ncore:nocc,ncore:nocc].transpose(0,2,1,3)
        h2e_v = eris['pacv'][:ncore].transpose(3,1,2,0)
    if 'h1' in dms:
        hdm1 = dms['h1']
    else:
        hdm1 = make_hdm1(dm1)

    a3 = make_a3(h1e,h2e,dm1,dm2,hdm1)
    norm = 2.0*numpy.einsum('rpji,raji,pa->rji',h2e_v,h2e_v,hdm1)\
         - 1.0*numpy.einsum('rpji,raij,pa->rji',h2e_v,h2e_v,hdm1)
    h = 2.0*numpy.einsum('rpji,raji,pa->rji',h2e_v,h2e_v,a3)\
         - 1.0*numpy.einsum('rpji,raij,pa->rji',h2e_v,h2e_v,a3)

    diff = mc.mo_energy[nocc:,None,None] - mc.mo_energy[None,:ncore,None] - mc.mo_energy[None,None,:ncore]

    return _norm_to_energy(norm, h, diff)

def Sijrs(mc, eris, verbose=None):
    mo_core, mo_cas, mo_virt = _extract_orbs(mc, mc.mo_coeff)
    ncore = mo_core.shape[1]
    nvirt = mo_virt.shape[1]
    ncas = mo_cas.shape[1]
    nocc = ncore + ncas
    if eris is None:
        erifile = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
        feri = ao2mo.outcore.general(mc.mol, (mo_core,mo_virt,mo_core,mo_virt),
                                     erifile.name, verbose=mc.verbose)
    else:
        feri = eris['cvcv']

    eia = mc.mo_energy[:ncore,None] -mc.mo_energy[None,nocc:]
    norm = 0
    e = 0
    with ao2mo.load(feri) as cvcv:
        for i in range(ncore):
            djba = (eia.reshape(-1,1) + eia[i].reshape(1,-1)).ravel()
            gi = numpy.asarray(cvcv[i*nvirt:(i+1)*nvirt])
            gi = gi.reshape(nvirt,ncore,nvirt).transpose(1,2,0)
            t2i = (gi.ravel()/djba).reshape(ncore,nvirt,nvirt)
            # 2*ijab-ijba
            theta = gi*2 - gi.transpose(0,2,1)
            norm += numpy.einsum('jab,jab', gi, theta)
            e += numpy.einsum('jab,jab', t2i, theta)
    return norm, e

def make_a23(h1e,h2e,dm1,dm2,dm3):
    a23 = -numpy.einsum('ip,caib->abcp',h1e,dm2)\
          -numpy.einsum('pijk,cajbik->abcp',h2e,dm3)\
          +2.0*numpy.einsum('bp,ca->abcp',h1e,dm1)\
          +2.0*numpy.einsum('pibk,caik->abcp',h2e,dm2)

    return a23

def G_Dyson(G0, SigmaDeltaT, Sigma, map):
    Beta=map.Beta
    G=weight.Weight("SmoothT", map, "TwoSpins", "AntiSymmetric", "K","W")
    G0.FFT("K", "W")
    SigmaDeltaT.FFT("K")
    Sigma.FFT("K", "W")

    NSpin, NSub=G.NSpin, G.NSublat

    G0SigmaDeltaT=np.einsum("ijklvt,klmnv->ijmnvt",G0.Data, SigmaDeltaT.Data)
    G0Sigma=np.einsum("ijklvt,klmnvt->ijmnvt",G0.Data, Sigma.Data)

    ####correction term
    for tau in range(map.MaxTauBin):
        G0SigmaDeltaT[...,tau]*= np.cos(np.pi*map.IndexToTau(tau)/Beta)

    GS  = Beta/map.MaxTauBin*(Beta/map.MaxTauBin*G0Sigma+G0SigmaDeltaT) 
    #GS shape: NSpin,NSub,NSpin,NSub,Vol,Tau

    I=np.eye(NSpin*NSub).reshape([NSpin,NSub,NSpin,NSub])
    Denorm=I[...,np.newaxis,np.newaxis]-GS
    lu_piv,Determ=weight.LUFactor(Denorm)
    Check_Denorminator(Denorm, Determ,map)
    G.LUSolve(lu_piv, G0.Data);
    return G