Python numpy.einsum方法代碼示例

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import einsum [as 別名]
def validate_and_fill_geometry(geom=None, tooclose=0.1, copy=True):
    """Check `geom` for overlapping atoms. Return flattened"""

    npgeom = np.array(geom, copy=copy, dtype=np.float).reshape((-1, 3))

    # Upper triangular
    metric = tooclose ** 2
    tooclose_inds = []
    for x in range(npgeom.shape[0]):
        diffs = npgeom[x] - npgeom[x + 1 :]
        dists = np.einsum("ij,ij->i", diffs, diffs)

        # Record issues
        if np.any(dists < metric):
            indices = np.where(dists < metric)[0]
            tooclose_inds.extend([(x, y, dist) for y, dist in zip(indices + x + 1, dists[indices] ** 0.5)])

    if tooclose_inds:
        raise ValidationError(
            """Following atoms are too close: {}""".format([(i, j, dist) for i, j, dist in tooclose_inds])
        )

    return {"geom": npgeom.reshape((-1))} 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import einsum [as 別名]
def generate_inflate(cloth):
    """Blow it up baby!"""    
    
    tri_nor = cloth.normals #* cloth.ob.modeling_cloth_inflate # non-unit calculated by tri_normals_in_place() per each triangle
    #tri_nor /= np.einsum("ij, ij->i", tri_nor, tri_nor)[:, nax]
    
    # reshape for add.at
    shape = cloth.inflate.shape
    
    cloth.inflate += tri_nor[:, nax] * cloth.ob.modeling_cloth_inflate# * cloth.tri_mix
    
    cloth.inflate.shape = (shape[0] * 3, 3)
    cloth.inflate *= cloth.tri_mix


    np.add.at(cloth.vel, cloth.tridex.ravel(), cloth.inflate)
    cloth.inflate.shape = shape
    cloth.inflate *= 0 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import einsum [as 別名]
def tda_denisty_matrix(td, state_id):
    '''
    Taking the TDA amplitudes as the CIS coefficients, calculate the density
    matrix (in AO basis) of the excited states
    '''
    cis_t1 = td.xy[state_id][0]
    dm_oo =-np.einsum('ia,ka->ik', cis_t1.conj(), cis_t1)
    dm_vv = np.einsum('ia,ic->ac', cis_t1, cis_t1.conj())

    # The ground state density matrix in mo_basis
    mf = td._scf
    dm = np.diag(mf.mo_occ)

    # Add CIS contribution
    nocc = cis_t1.shape[0]
    dm[:nocc,:nocc] += dm_oo * 2
    dm[nocc:,nocc:] += dm_vv * 2

    # Transform density matrix to AO basis
    mo = mf.mo_coeff
    dm = np.einsum('pi,ij,qj->pq', mo, dm, mo.conj())
    return dm

# Density matrix for the 3rd excited state 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import einsum [as 別名]
def ss(mol):
    n = mol.nao_nr()
    mat1 = mol.intor('int2e_ip1ip2_sph').reshape(3,3,n,n,n,n) # <nabla1 nabla2 | 1 2>
    mat2 =-mat1.transpose(0,1,2,3,5,4) # <nabla1 2 | 1 nabla2>
    mat3 =-mat2.transpose(1,0,3,2,4,5) # <1 nabla2 | nabla1 2>
    mat4 = mat1.transpose(0,1,3,2,5,4) # <1 2 | nabla1 nabla2>
    mat = mat1 - mat2 - mat3 + mat4
# Fermi contact term
    h_fc = mol.intor('int4c1e').reshape(nao,nao,nao,nao) * (4*numpy.pi/3)
    mat[0,0] -= h_fc
    mat[1,1] -= h_fc
    mat[2,2] -= h_fc

    s = lib.PauliMatrices * .5
# wxyz are the spin indices, ijkl are the AO indicies
    alpha = 137.036
    fac = alpha ** 2 / 2
    mat = numpy.einsum('swx,tyz,stijkl->wxyzijkl', s[:,0,0], s[:,0,0], mat) * fac
    return mat 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import einsum [as 別名]
def force(dm):
    # The interaction between QM atoms and MM particles
    # \sum_K d/dR (1/|r_K-R|) = \sum_K (r_K-R)/|r_K-R|^3
    qm_coords = mol.atom_coords()
    qm_charges = mol.atom_charges()
    dr = qm_coords[:,None,:] - coords
    r = numpy.linalg.norm(dr, axis=2)
    g = numpy.einsum('r,R,rRx,rR->Rx', qm_charges, charges, dr, r**-3)

    # The interaction between electron density and MM particles
    # d/dR <i| (1/|r-R|) |j> = <i| d/dR (1/|r-R|) |j> = <i| -d/dr (1/|r-R|) |j>
    #   = <d/dr i| (1/|r-R|) |j> + <i| (1/|r-R|) |d/dr j>
    for i, q in enumerate(charges):
        with mol.with_rinv_origin(coords[i]):
            v = mol.intor('int1e_iprinv')
        f =(numpy.einsum('ij,xji->x', dm, v) +
            numpy.einsum('ij,xij->x', dm, v.conj())) * -q
        g[i] += f

    # Force = -d/dR
    return -g

# The force from HF electron density
# Be careful with the unit of the MM particle coordinates. The gradients are
# computed in the atomic unit. 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import einsum [as 別名]
def myump2(mol, mo_energy, mo_coeff, mo_occ):
    o = numpy.hstack((mo_coeff[0][:,mo_occ[0]>0] ,mo_coeff[1][:,mo_occ[1]>0]))
    v = numpy.hstack((mo_coeff[0][:,mo_occ[0]==0],mo_coeff[1][:,mo_occ[1]==0]))
    eo = numpy.hstack((mo_energy[0][mo_occ[0]>0] ,mo_energy[1][mo_occ[1]>0]))
    ev = numpy.hstack((mo_energy[0][mo_occ[0]==0],mo_energy[1][mo_occ[1]==0]))
    no = o.shape[1]
    nv = v.shape[1]
    noa = sum(mo_occ[0]>0)
    nva = sum(mo_occ[0]==0)
    eri = ao2mo.general(mol, (o,v,o,v)).reshape(no,nv,no,nv)
    eri[:noa,nva:] = eri[noa:,:nva] = eri[:,:,:noa,nva:] = eri[:,:,noa:,:nva] = 0
    g = eri - eri.transpose(0,3,2,1)
    eov = eo.reshape(-1,1) - ev.reshape(-1)
    de = 1/(eov.reshape(-1,1) + eov.reshape(-1)).reshape(g.shape)
    emp2 = .25 * numpy.einsum('iajb,iajb,iajb->', g, g, de)
    return emp2 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import einsum [as 別名]
def hcore_grad_generator(x2cobj, mol=None):
    '''nuclear gradients of 1-component X2c hcore Hamiltonian  (spin-free part only)
    '''
    if mol is None: mol = x2cobj.mol
    xmol, contr_coeff = x2cobj.get_xmol(mol)

    if x2cobj.basis is not None:
        s22 = xmol.intor_symmetric('int1e_ovlp')
        s21 = gto.intor_cross('int1e_ovlp', xmol, mol)
        contr_coeff = lib.cho_solve(s22, s21)

    get_h1_xmol = gen_sf_hfw(xmol, x2cobj.approx)
    def hcore_deriv(atm_id):
        h1 = get_h1_xmol(atm_id)
        if contr_coeff is not None:
            h1 = lib.einsum('pi,xpq,qj->xij', contr_coeff, h1, contr_coeff)
        return numpy.asarray(h1)
    return hcore_deriv 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import einsum [as 別名]
def _invsqrt2(a0, a1i, a1j, a2ij):
    '''Solving first order derivative of x^2 = a^{-1}'''
    w, v = scipy.linalg.eigh(a0)
    w = 1./numpy.sqrt(w)
    a1i = reduce(numpy.dot, (v.conj().T, a1i, v))
    x1i = numpy.einsum('i,ij,j->ij', w**2, a1i, w**2) / (w[:,None] + w)

    a1j = reduce(numpy.dot, (v.conj().T, a1j, v))
    x1j = numpy.einsum('i,ij,j->ij', w**2, a1j, w**2) / (w[:,None] + w)

    a2ij = reduce(numpy.dot, (v.conj().T, a2ij, v))
    tmp = (a1i*w**2).dot(a1j)
    a2ij -= tmp + tmp.conj().T
    a2ij = -numpy.einsum('i,ij,j->ij', w**2, a2ij, w**2)
    tmp = x1i.dot(x1j)
    a2ij -= tmp + tmp.conj().T
    x2 = a2ij / (w[:,None] + w)
    x2 = reduce(numpy.dot, (v, x2, v.conj().T))
    return x2 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import einsum [as 別名]
def get_x1(mol, ia):
    h0, s0 = get_h0_s0(mol)
    h1, s1 = get_h1_s1(mol, ia)
    e0, c0 = scipy.linalg.eigh(h0, s0)
    nao = mol.nao_nr()
    cl0 = c0[:nao,nao:]
    cs0 = c0[nao:,nao:]
    x0 = scipy.linalg.solve(cl0.T, cs0.T).T
    h1 = numpy.einsum('pi,xpq,qj->xij', c0.conj(), h1, c0[:,nao:])
    s1 = numpy.einsum('pi,xpq,qj->xij', c0.conj(), s1, c0[:,nao:])
    epi = e0[:,None] - e0[nao:]
    degen_mask = abs(epi) < 1e-7
    epi[degen_mask] = 1e200
    c1 = (h1 - s1 * e0[nao:]) / -epi
    c1[:,degen_mask] = -.5 * s1[:,degen_mask]
    c1 = numpy.einsum('pq,xqi->xpi', c0, c1)
    cl1 = c1[:,:nao]
    cs1 = c1[:,nao:]
    x1 = [scipy.linalg.solve(cl0.T, (cs1[i] - x0.dot(cl1[i])).T).T
          for i in range(3)]
    return numpy.asarray(x1) 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import einsum [as 別名]
def hcore_hess_generator(x2cobj, mol=None):
    '''nuclear gradients of 1-component X2c hcore Hamiltonian  (spin-free part only)
    '''
    if mol is None: mol = x2cobj.mol
    xmol, contr_coeff = x2cobj.get_xmol(mol)

    if x2cobj.basis is not None:
        s22 = xmol.intor_symmetric('int1e_ovlp')
        s21 = gto.intor_cross('int1e_ovlp', xmol, mol)
        contr_coeff = lib.cho_solve(s22, s21)

    get_h1_xmol = gen_sf_hfw(xmol, x2cobj.approx)
    def hcore_deriv(ia, ja):
        h1 = get_h1_xmol(ia, ja)
        if contr_coeff is not None:
            h1 = lib.einsum('pi,xypq,qj->xyij', contr_coeff, h1, contr_coeff)
        return numpy.asarray(h1)
    return hcore_deriv 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import einsum [as 別名]
def normalize_dm_(mf, dm):
    '''
    Scale density matrix to make it produce the correct number of electrons.
    '''
    cell = mf.cell
    if isinstance(dm, np.ndarray) and dm.ndim == 2:
        ne = np.einsum('ij,ji->', dm, mf.get_ovlp(cell)).real
    else:
        ne = np.einsum('xij,ji->', dm, mf.get_ovlp(cell)).real
    if abs(ne - cell.nelectron).sum() > 1e-7:
        logger.debug(mf, 'Big error detected in the electron number '
                    'of initial guess density matrix (Ne/cell = %g)!\n'
                    '  This can cause huge error in Fock matrix and '
                    'lead to instability in SCF for low-dimensional '
                    'systems.\n  DM is normalized wrt the number '
                    'of electrons %s', ne, cell.nelectron)
        dm *= cell.nelectron / ne
    return dm 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import einsum [as 別名]
def spin_square(self, mo_coeff=None, s=None):
        '''Treating the k-point sampling wfn as a giant Slater determinant,
        the spin_square value is the <S^2> of the giant determinant.
        '''
        nkpts = len(self.kpts)
        if mo_coeff is None:
            mo_a = [self.mo_coeff[0][k][:,self.mo_occ[0][k]>0] for k in range(nkpts)]
            mo_b = [self.mo_coeff[1][k][:,self.mo_occ[1][k]>0] for k in range(nkpts)]
        else:
            mo_a, mo_b = mo_coeff
        if s is None:
            s = self.get_ovlp()

        nelec_a = sum([mo_a[k].shape[1] for k in range(nkpts)])
        nelec_b = sum([mo_b[k].shape[1] for k in range(nkpts)])
        ssxy = (nelec_a + nelec_b) * .5
        for k in range(nkpts):
            sij = reduce(np.dot, (mo_a[k].T.conj(), s[k], mo_b[k]))
            ssxy -= np.einsum('ij,ij->', sij.conj(), sij).real
        ssz = (nelec_b-nelec_a)**2 * .25
        ss = ssxy + ssz
        s = np.sqrt(ss+.25) - .5
        return ss, s*2+1 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import einsum [as 別名]
def energy_elec(mf, dm_kpts=None, h1e_kpts=None, vhf_kpts=None):
    '''Following pyscf.scf.hf.energy_elec()
    '''
    if dm_kpts is None: dm_kpts = mf.make_rdm1()
    if h1e_kpts is None: h1e_kpts = mf.get_hcore()
    if vhf_kpts is None: vhf_kpts = mf.get_veff(mf.cell, dm_kpts)

    nkpts = len(dm_kpts)
    e1 = 1./nkpts * np.einsum('kij,kji', dm_kpts, h1e_kpts)
    e_coul = 1./nkpts * np.einsum('kij,kji', dm_kpts, vhf_kpts) * 0.5
    mf.scf_summary['e1'] = e1.real
    mf.scf_summary['e2'] = e_coul.real
    logger.debug(mf, 'E1 = %s  E_coul = %s', e1, e_coul)
    if CHECK_COULOMB_IMAG and abs(e_coul.imag > mf.cell.precision*10):
        logger.warn(mf, "Coulomb energy has imaginary part %s. "
                    "Coulomb integrals (e-e, e-N) may not converge !",
                    e_coul.imag)
    return (e1+e_coul).real, e_coul.real 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import einsum [as 別名]
def _gamma1_intermediates(mp, t2=None):
    # Memory optimization should be here
    if t2 is None:
        t2 = mp.t2
    if t2 is None:
        raise NotImplementedError("Run kmp2.kernel with `with_t2=True`")
    nmo = mp.nmo
    nocc = mp.nocc
    nvir = nmo - nocc
    nkpts = mp.nkpts
    dtype = t2.dtype

    dm1occ = np.zeros((nkpts, nocc, nocc), dtype=dtype)
    dm1vir = np.zeros((nkpts, nvir, nvir), dtype=dtype)

    for ki in range(nkpts):
        for kj in range(nkpts):
            for ka in range(nkpts):
                kb = mp.khelper.kconserv[ki, ka, kj]

                dm1vir[kb] += einsum('ijax,ijay->yx', t2[ki][kj][ka].conj(), t2[ki][kj][ka]) * 2 -\
                              einsum('ijax,ijya->yx', t2[ki][kj][ka].conj(), t2[ki][kj][kb])
                dm1occ[kj] += einsum('ixab,iyab->xy', t2[ki][kj][ka].conj(), t2[ki][kj][ka]) * 2 -\
                              einsum('ixab,iyba->xy', t2[ki][kj][ka].conj(), t2[ki][kj][kb])
    return -dm1occ, dm1vir 

# 需要導入模塊: import numpy [as 別名]
# 或者: from numpy import einsum [as 別名]
def make_h3_2_5() -> Tuple[RestrictedHartreeFockObjective, of.MolecularData, np.
                           ndarray, np.ndarray, np.ndarray]:
    # load the molecule from moelcular data
    h3_2_5_path = os.path.join(
        hfvqe.__path__[0],
        'molecular_data/hydrogen_chains/h_3_p_sto-3g/bond_distance_2.5')

    molfile = os.path.join(h3_2_5_path,
                           'H3_plus_sto-3g_singlet_linear_r-2.5.hdf5')
    molecule = of.MolecularData(filename=molfile)
    molecule.load()

    S = np.load(os.path.join(h3_2_5_path, 'overlap.npy'))
    Hcore = np.load(os.path.join(h3_2_5_path, 'h_core.npy'))
    TEI = np.load(os.path.join(h3_2_5_path, 'tei.npy'))

    _, X = sp.linalg.eigh(Hcore, S)
    obi = of.general_basis_change(Hcore, X, (1, 0))
    tbi = np.einsum('psqr', of.general_basis_change(TEI, X, (1, 0, 1, 0)))
    molecular_hamiltonian = generate_hamiltonian(obi, tbi,
                                                 molecule.nuclear_repulsion)

    rhf_objective = RestrictedHartreeFockObjective(molecular_hamiltonian,
                                                   molecule.n_electrons)

    scipy_result = rhf_minimization(rhf_objective)
    return rhf_objective, molecule, scipy_result.x, obi, tbi