예제 #1
0
파일: uks.py 프로젝트: chrinide/pyscf
def partial_hess_elec(hessobj, mo_energy=None, mo_coeff=None, mo_occ=None,
                      atmlst=None, max_memory=4000, verbose=None):
    log = logger.new_logger(hessobj, verbose)
    time0 = t1 = (time.clock(), time.time())

    mol = hessobj.mol
    mf = hessobj.base
    if mo_energy is None: mo_energy = mf.mo_energy
    if mo_occ is None:    mo_occ = mf.mo_occ
    if mo_coeff is None:  mo_coeff = mf.mo_coeff
    if atmlst is None: atmlst = range(mol.natm)

    nao, nmo = mo_coeff[0].shape
    mocca = mo_coeff[0][:,mo_occ[0]>0]
    moccb = mo_coeff[1][:,mo_occ[1]>0]
    dm0a = numpy.dot(mocca, mocca.T)
    dm0b = numpy.dot(moccb, moccb.T)
    dm0 = dm0a + dm0b
    # Energy weighted density matrix
    mo_ea = mo_energy[0][mo_occ[0]>0]
    mo_eb = mo_energy[1][mo_occ[1]>0]
    dme0 = numpy.einsum('pi,qi,i->pq', mocca, mocca, mo_ea)
    dme0+= numpy.einsum('pi,qi,i->pq', moccb, moccb, mo_eb)

    hcore_deriv = hessobj.hcore_generator(mol)
    s1aa, s1ab, s1a = rhf_hess.get_ovlp(mol)

    if mf.nlc != '':
        raise NotImplementedError
    #enabling range-separated hybrids
    omega, alpha, hyb = mf._numint.rsh_and_hybrid_coeff(mf.xc, spin=mol.spin)

    mem_now = lib.current_memory()[0]
    max_memory = max(2000, mf.max_memory*.9-mem_now)
    veffa_diag, veffb_diag = _get_vxc_diag(hessobj, mo_coeff, mo_occ, max_memory)
    if abs(hyb) > 1e-10:
        vj1a, vj1b, vk1a, vk1b = \
                _get_jk(mol, 'int2e_ipip1', 9, 's2kl',
                        ['lk->s1ij', dm0a, 'lk->s1ij', dm0b,
                         'jk->s1il', dm0a, 'jk->s1il', dm0b])
        vj1 = vj1a + vj1b
        veffa_diag += (vj1 - hyb * vk1a).reshape(3,3,nao,nao)
        veffb_diag += (vj1 - hyb * vk1b).reshape(3,3,nao,nao)
        if abs(omega) > 1e-10:
            with mol.with_range_coulomb(omega):
                vk1a, vk1b = \
                        _get_jk(mol, 'int2e_ipip1', 9, 's2kl',
                                ['jk->s1il', dm0a, 'jk->s1il', dm0b])
            veffa_diag -= (alpha-hyb) * vk1a.reshape(3,3,nao,nao)
            veffb_diag -= (alpha-hyb) * vk1b.reshape(3,3,nao,nao)
    else:
        vj1a, vj1b = \
                _get_jk(mol, 'int2e_ipip1', 9, 's2kl',
                        ['lk->s1ij', dm0a, 'lk->s1ij', dm0b])
        vj1 = vj1a + vj1b
        veffa_diag += vj1.reshape(3,3,nao,nao)
        veffb_diag += vj1.reshape(3,3,nao,nao)
    vj1 = vj1a = vj1b = vk1a = vk1b = None
    t1 = log.timer_debug1('contracting int2e_ipip1', *t1)

    aoslices = mol.aoslice_by_atom()
    de2 = numpy.zeros((mol.natm,mol.natm,3,3))  # (A,B,dR_A,dR_B)
    vxca, vxcb = _get_vxc_deriv2(hessobj, mo_coeff, mo_occ, max_memory)
    for i0, ia in enumerate(atmlst):
        shl0, shl1, p0, p1 = aoslices[ia]

        veffa = vxca[ia]
        veffb = vxcb[ia]
        shls_slice = (shl0, shl1) + (0, mol.nbas)*3
        if abs(hyb) > 1e-10:
            vj1a, vj1b, vk1a, vk1b, vk2a, vk2b = \
                    _get_jk(mol, 'int2e_ip1ip2', 9, 's1',
                            ['ji->s1kl', dm0a[:,p0:p1], 'ji->s1kl', dm0b[:,p0:p1],
                             'li->s1kj', dm0a[:,p0:p1], 'li->s1kj', dm0b[:,p0:p1],
                             'lj->s1ki', dm0a         , 'lj->s1ki', dm0b         ],
                            shls_slice=shls_slice)
            vj1 = vj1a + vj1b
            veffa += (vj1 * 2 - hyb * vk1a).reshape(3,3,nao,nao)
            veffb += (vj1 * 2 - hyb * vk1b).reshape(3,3,nao,nao)
            veffa[:,:,:,p0:p1] -= hyb * vk2a.reshape(3,3,nao,p1-p0)
            veffb[:,:,:,p0:p1] -= hyb * vk2b.reshape(3,3,nao,p1-p0)
            if abs(omega) > 1e-10:
                with mol.with_range_coulomb(omega):
                    vk1a, vk1b, vk2a, vk2b = \
                            _get_jk(mol, 'int2e_ip1ip2', 9, 's1',
                                    ['li->s1kj', dm0a[:,p0:p1],
                                     'li->s1kj', dm0b[:,p0:p1],
                                     'lj->s1ki', dm0a         ,
                                     'lj->s1ki', dm0b         ],
                                    shls_slice=shls_slice)
                veffa -= (alpha-hyb) * vk1a.reshape(3,3,nao,nao)
                veffb -= (alpha-hyb) * vk1b.reshape(3,3,nao,nao)
                veffa[:,:,:,p0:p1] -= (alpha-hyb) * vk2a.reshape(3,3,nao,p1-p0)
                veffb[:,:,:,p0:p1] -= (alpha-hyb) * vk2b.reshape(3,3,nao,p1-p0)
            t1 = log.timer_debug1('contracting int2e_ip1ip2 for atom %d'%ia, *t1)
            vj1a, vj1b, vk1a, vk1b = \
                    _get_jk(mol, 'int2e_ipvip1', 9, 's2kl',
                            ['lk->s1ij', dm0a         , 'lk->s1ij', dm0b         ,
                             'li->s1kj', dm0a[:,p0:p1], 'li->s1kj', dm0b[:,p0:p1]],
                            shls_slice=shls_slice)
            vj1 = vj1a + vj1b
            veffa[:,:,:,p0:p1] += vj1.transpose(0,2,1).reshape(3,3,nao,p1-p0)
            veffb[:,:,:,p0:p1] += vj1.transpose(0,2,1).reshape(3,3,nao,p1-p0)
            veffa -= hyb * vk1a.transpose(0,2,1).reshape(3,3,nao,nao)
            veffb -= hyb * vk1b.transpose(0,2,1).reshape(3,3,nao,nao)
            if abs(omega) > 1e-10:
                with mol.with_range_coulomb(omega):
                    vk1a, vk1b = _get_jk(mol, 'int2e_ipvip1', 9, 's2kl',
                                         ['li->s1kj', dm0a[:,p0:p1],
                                          'li->s1kj', dm0b[:,p0:p1]],
                                         shls_slice=shls_slice)
                veffa -= (alpha-hyb) * vk1a.transpose(0,2,1).reshape(3,3,nao,nao)
                veffb -= (alpha-hyb) * vk1b.transpose(0,2,1).reshape(3,3,nao,nao)
            t1 = log.timer_debug1('contracting int2e_ipvip1 for atom %d'%ia, *t1)
        else:
            vj1a, vj1b = \
                    _get_jk(mol, 'int2e_ip1ip2', 9, 's1',
                            ['ji->s1kl', dm0a[:,p0:p1], 'ji->s1kl', dm0b[:,p0:p1]],
                            shls_slice=shls_slice)
            vj1 = (vj1a + vj1b).reshape(3,3,nao,nao)
            veffa += vj1 * 2
            veffb += vj1 * 2
            t1 = log.timer_debug1('contracting int2e_ip1ip2 for atom %d'%ia, *t1)
            vj1a, vj1b = \
                    _get_jk(mol, 'int2e_ipvip1', 9, 's2kl',
                            ['lk->s1ij', dm0a, 'lk->s1ij', dm0b],
                            shls_slice=shls_slice)
            vj1 = vj1a + vj1b
            veffa[:,:,:,p0:p1] += vj1.transpose(0,2,1).reshape(3,3,nao,p1-p0)
            veffb[:,:,:,p0:p1] += vj1.transpose(0,2,1).reshape(3,3,nao,p1-p0)
            t1 = log.timer_debug1('contracting int2e_ipvip1 for atom %d'%ia, *t1)
        vj1 = vj1a = vj1b = vk1a = vk1b = vk2a = vk2b = None

        s1ao = numpy.zeros((3,nao,nao))
        s1ao[:,p0:p1] += s1a[:,p0:p1]
        s1ao[:,:,p0:p1] += s1a[:,p0:p1].transpose(0,2,1)
        s1ooa = numpy.einsum('xpq,pi,qj->xij', s1ao, mocca, mocca)
        s1oob = numpy.einsum('xpq,pi,qj->xij', s1ao, moccb, moccb)

        de2[i0,i0] += numpy.einsum('xypq,pq->xy', veffa_diag[:,:,p0:p1], dm0a[p0:p1])*2
        de2[i0,i0] += numpy.einsum('xypq,pq->xy', veffb_diag[:,:,p0:p1], dm0b[p0:p1])*2
        de2[i0,i0] -= numpy.einsum('xypq,pq->xy', s1aa[:,:,p0:p1], dme0[p0:p1])*2

        for j0, ja in enumerate(atmlst[:i0+1]):
            q0, q1 = aoslices[ja][2:]
            de2[i0,j0] += numpy.einsum('xypq,pq->xy', veffa[:,:,q0:q1], dm0a[q0:q1])*2
            de2[i0,j0] += numpy.einsum('xypq,pq->xy', veffb[:,:,q0:q1], dm0b[q0:q1])*2
            de2[i0,j0] -= numpy.einsum('xypq,pq->xy', s1ab[:,:,p0:p1,q0:q1], dme0[p0:p1,q0:q1])*2

            h1ao = hcore_deriv(ia, ja)
            de2[i0,j0] += numpy.einsum('xypq,pq->xy', h1ao, dm0)

        for j0 in range(i0):
            de2[j0,i0] = de2[i0,j0].T

    log.timer('UKS partial hessian', *time0)
    return de2
예제 #2
0
파일: rks.py 프로젝트: chrinide/pyscf
def partial_hess_elec(hessobj, mo_energy=None, mo_coeff=None, mo_occ=None,
                      atmlst=None, max_memory=4000, verbose=None):
    log = logger.new_logger(hessobj, verbose)
    time0 = t1 = (time.clock(), time.time())

    mol = hessobj.mol
    mf = hessobj.base
    if mo_energy is None: mo_energy = mf.mo_energy
    if mo_occ is None:    mo_occ = mf.mo_occ
    if mo_coeff is None:  mo_coeff = mf.mo_coeff
    if atmlst is None: atmlst = range(mol.natm)

    nao, nmo = mo_coeff.shape
    mocc = mo_coeff[:,mo_occ>0]
    nocc = mocc.shape[1]
    dm0 = numpy.dot(mocc, mocc.T) * 2
    # Energy weighted density matrix
    dme0 = numpy.einsum('pi,qi,i->pq', mocc, mocc, mo_energy[mo_occ>0]) * 2

    hcore_deriv = hessobj.hcore_generator(mol)
    s1aa, s1ab, s1a = rhf_hess.get_ovlp(mol)

    if mf.nlc != '':
        raise NotImplementedError
    #enabling range-separated hybrids
    omega, alpha, beta = mf._numint.rsh_coeff(mf.xc)
    if abs(omega) > 1e-10:
        hyb = alpha + beta
    else:
        hyb = mf._numint.hybrid_coeff(mf.xc, spin=mol.spin)

    mem_now = lib.current_memory()[0]
    max_memory = max(2000, mf.max_memory*.9-mem_now)
    veff_diag = _get_vxc_diag(hessobj, mo_coeff, mo_occ, max_memory)
    if abs(hyb) > 1e-10:
        vj1, vk1 = rhf_hess._get_jk(mol, 'int2e_ipip1', 9, 's2kl',
                                    ['lk->s1ij', dm0,   # vj1
                                     'jk->s1il', dm0])  # vk1
        veff_diag += (vj1 - hyb * .5 * vk1).reshape(3,3,nao,nao)
        if abs(omega) > 1e-10:
            with mol.with_range_coulomb(omega):
                vk1 = rhf_hess._get_jk(mol, 'int2e_ipip1', 9, 's2kl',
                                       ['jk->s1il', dm0])[0]
            veff_diag -= (alpha-hyb)*.5 * vk1.reshape(3,3,nao,nao)
    else:
        vj1 = rhf_hess._get_jk(mol, 'int2e_ipip1', 9, 's2kl',
                               ['lk->s1ij', dm0])[0]
        veff_diag += vj1.reshape(3,3,nao,nao)
    vj1 = vk1 = None
    t1 = log.timer_debug1('contracting int2e_ipip1', *t1)

    aoslices = mol.aoslice_by_atom()
    de2 = numpy.zeros((mol.natm,mol.natm,3,3))  # (A,B,dR_A,dR_B)
    vxc = _get_vxc_deriv2(hessobj, mo_coeff, mo_occ, max_memory)
    for i0, ia in enumerate(atmlst):
        shl0, shl1, p0, p1 = aoslices[ia]

        shls_slice = (shl0, shl1) + (0, mol.nbas)*3
        veff = vxc[ia]
        if abs(hyb) > 1e-10:
            vj1, vk1, vk2 = rhf_hess._get_jk(mol, 'int2e_ip1ip2', 9, 's1',
                                             ['ji->s1kl', dm0[:,p0:p1],  # vj1
                                              'li->s1kj', dm0[:,p0:p1],  # vk1
                                              'lj->s1ki', dm0         ], # vk2
                                             shls_slice=shls_slice)
            veff += (vj1 * 2 - hyb * .5 * vk1).reshape(3,3,nao,nao)
            veff[:,:,:,p0:p1] -= (hyb * .5 * vk2).reshape(3,3,nao,p1-p0)
            if abs(omega) > 1e-10:
                with mol.with_range_coulomb(omega):
                    vk1, vk2 = rhf_hess._get_jk(mol, 'int2e_ip1ip2', 9, 's1',
                                                ['li->s1kj', dm0[:,p0:p1],  # vk1
                                                 'lj->s1ki', dm0         ], # vk2
                                                shls_slice=shls_slice)
                veff -= (alpha-hyb)*.5 * vk1.reshape(3,3,nao,nao)
                veff[:,:,:,p0:p1] -= (alpha-hyb)*.5 * vk2.reshape(3,3,nao,p1-p0)
            t1 = log.timer_debug1('contracting int2e_ip1ip2 for atom %d'%ia, *t1)

            vj1, vk1 = rhf_hess._get_jk(mol, 'int2e_ipvip1', 9, 's2kl',
                                        ['lk->s1ij', dm0,           # vj1
                                         'li->s1kj', dm0[:,p0:p1]], # vk1
                                        shls_slice=shls_slice)
            veff[:,:,:,p0:p1] += vj1.transpose(0,2,1).reshape(3,3,nao,p1-p0)
            veff -= hyb * .5 * vk1.transpose(0,2,1).reshape(3,3,nao,nao)
            if abs(omega) > 1e-10:
                with mol.with_range_coulomb(omega):
                    vk1 = rhf_hess._get_jk(mol, 'int2e_ipvip1', 9, 's2kl',
                                           ['li->s1kj', dm0[:,p0:p1]], # vk1
                                           shls_slice=shls_slice)[0]
                veff -= (alpha-hyb)*.5 * vk1.transpose(0,2,1).reshape(3,3,nao,nao)
            t1 = log.timer_debug1('contracting int2e_ipvip1 for atom %d'%ia, *t1)
        else:
            vj1 = rhf_hess._get_jk(mol, 'int2e_ip1ip2', 9, 's1',
                                   ['ji->s1kl', dm0[:,p0:p1]],
                                   shls_slice=shls_slice)[0]
            veff += vj1.reshape(3,3,nao,nao) * 2
            t1 = log.timer_debug1('contracting int2e_ip1ip2 for atom %d'%ia, *t1)

            vj1 = rhf_hess._get_jk(mol, 'int2e_ipvip1', 9, 's2kl',
                                   ['lk->s1ij', dm0], shls_slice=shls_slice)[0]
            veff[:,:,:,p0:p1] += vj1.transpose(0,2,1).reshape(3,3,nao,p1-p0)
            t1 = log.timer_debug1('contracting int2e_ipvip1 for atom %d'%ia, *t1)
        vj1 = vk1 = vk2 = None

        s1ao = numpy.zeros((3,nao,nao))
        s1ao[:,p0:p1] += s1a[:,p0:p1]
        s1ao[:,:,p0:p1] += s1a[:,p0:p1].transpose(0,2,1)
        s1oo = numpy.einsum('xpq,pi,qj->xij', s1ao, mocc, mocc)

        de2[i0,i0] += numpy.einsum('xypq,pq->xy', veff_diag[:,:,p0:p1], dm0[p0:p1])*2
        de2[i0,i0] -= numpy.einsum('xypq,pq->xy', s1aa[:,:,p0:p1], dme0[p0:p1])*2

        for j0, ja in enumerate(atmlst[:i0+1]):
            q0, q1 = aoslices[ja][2:]
            de2[i0,j0] += numpy.einsum('xypq,pq->xy', veff[:,:,q0:q1], dm0[q0:q1])*2
            de2[i0,j0] -= numpy.einsum('xypq,pq->xy', s1ab[:,:,p0:p1,q0:q1], dme0[p0:p1,q0:q1])*2

            h1ao = hcore_deriv(ia, ja)
            de2[i0,j0] += numpy.einsum('xypq,pq->xy', h1ao, dm0)

        for j0 in range(i0):
            de2[j0,i0] = de2[i0,j0].T

    log.timer('RKS partial hessian', *time0)
    return de2
예제 #3
0
파일: rks.py 프로젝트: zwgsyntax/pyscf
def partial_hess_elec(hessobj,
                      mo_energy=None,
                      mo_coeff=None,
                      mo_occ=None,
                      atmlst=None,
                      max_memory=4000,
                      verbose=None):
    log = logger.new_logger(hessobj, verbose)
    time0 = t1 = (time.clock(), time.time())

    mol = hessobj.mol
    mf = hessobj.base
    if mo_energy is None: mo_energy = mf.mo_energy
    if mo_occ is None: mo_occ = mf.mo_occ
    if mo_coeff is None: mo_coeff = mf.mo_coeff
    if atmlst is None: atmlst = range(mol.natm)

    nao, nmo = mo_coeff.shape
    mocc = mo_coeff[:, mo_occ > 0]
    nocc = mocc.shape[1]
    dm0 = numpy.dot(mocc, mocc.T) * 2
    # Energy weighted density matrix
    dme0 = numpy.einsum('pi,qi,i->pq', mocc, mocc, mo_energy[mo_occ > 0]) * 2

    hcore_deriv = hessobj.hcore_generator(mol)
    s1aa, s1ab, s1a = rhf_hess.get_ovlp(mol)

    if mf.nlc != '':
        raise NotImplementedError
    #enabling range-separated hybrids
    omega, alpha, beta = mf._numint.rsh_coeff(mf.xc)
    if abs(omega) > 1e-10:
        hyb = alpha + beta
    else:
        hyb = mf._numint.hybrid_coeff(mf.xc, spin=mol.spin)

    mem_now = lib.current_memory()[0]
    max_memory = max(2000, mf.max_memory * .9 - mem_now)
    veff_diag = _get_vxc_diag(hessobj, mo_coeff, mo_occ, max_memory)
    if abs(hyb) > 1e-10:
        vj1, vk1 = rhf_hess._get_jk(
            mol,
            'int2e_ipip1',
            9,
            's2kl',
            [
                'lk->s1ij',
                dm0,  # vj1
                'jk->s1il',
                dm0
            ])  # vk1
        veff_diag += (vj1 - hyb * .5 * vk1).reshape(3, 3, nao, nao)
        if abs(omega) > 1e-10:
            with mol.with_range_coulomb(omega):
                vk1 = rhf_hess._get_jk(mol, 'int2e_ipip1', 9, 's2kl',
                                       ['jk->s1il', dm0])[0]
            veff_diag -= (alpha - hyb) * .5 * vk1.reshape(3, 3, nao, nao)
    else:
        vj1 = rhf_hess._get_jk(mol, 'int2e_ipip1', 9, 's2kl',
                               ['lk->s1ij', dm0])[0]
        veff_diag += vj1.reshape(3, 3, nao, nao)
    vj1 = vk1 = None
    t1 = log.timer_debug1('contracting int2e_ipip1', *t1)

    aoslices = mol.aoslice_by_atom()
    de2 = numpy.zeros((mol.natm, mol.natm, 3, 3))  # (A,B,dR_A,dR_B)
    vxc = _get_vxc_deriv2(hessobj, mo_coeff, mo_occ, max_memory)
    for i0, ia in enumerate(atmlst):
        shl0, shl1, p0, p1 = aoslices[ia]

        shls_slice = (shl0, shl1) + (0, mol.nbas) * 3
        veff = vxc[ia]
        if abs(hyb) > 1e-10:
            vj1, vk1, vk2 = rhf_hess._get_jk(
                mol,
                'int2e_ip1ip2',
                9,
                's1',
                [
                    'ji->s1kl',
                    dm0[:, p0:p1],  # vj1
                    'li->s1kj',
                    dm0[:, p0:p1],  # vk1
                    'lj->s1ki',
                    dm0
                ],  # vk2
                shls_slice=shls_slice)
            veff += (vj1 * 2 - hyb * .5 * vk1).reshape(3, 3, nao, nao)
            veff[:, :, :,
                 p0:p1] -= (hyb * .5 * vk2).reshape(3, 3, nao, p1 - p0)
            if abs(omega) > 1e-10:
                with mol.with_range_coulomb(omega):
                    vk1, vk2 = rhf_hess._get_jk(
                        mol,
                        'int2e_ip1ip2',
                        9,
                        's1',
                        [
                            'li->s1kj',
                            dm0[:, p0:p1],  # vk1
                            'lj->s1ki',
                            dm0
                        ],  # vk2
                        shls_slice=shls_slice)
                veff -= (alpha - hyb) * .5 * vk1.reshape(3, 3, nao, nao)
                veff[:, :, :, p0:p1] -= (alpha - hyb) * .5 * vk2.reshape(
                    3, 3, nao, p1 - p0)
            t1 = log.timer_debug1('contracting int2e_ip1ip2 for atom %d' % ia,
                                  *t1)

            vj1, vk1 = rhf_hess._get_jk(
                mol,
                'int2e_ipvip1',
                9,
                's2kl',
                [
                    'lk->s1ij',
                    dm0,  # vj1
                    'li->s1kj',
                    dm0[:, p0:p1]
                ],  # vk1
                shls_slice=shls_slice)
            veff[:, :, :,
                 p0:p1] += vj1.transpose(0, 2, 1).reshape(3, 3, nao, p1 - p0)
            veff -= hyb * .5 * vk1.transpose(0, 2, 1).reshape(3, 3, nao, nao)
            if abs(omega) > 1e-10:
                with mol.with_range_coulomb(omega):
                    vk1 = rhf_hess._get_jk(
                        mol,
                        'int2e_ipvip1',
                        9,
                        's2kl',
                        ['li->s1kj', dm0[:, p0:p1]],  # vk1
                        shls_slice=shls_slice)[0]
                veff -= (alpha - hyb) * .5 * vk1.transpose(0, 2, 1).reshape(
                    3, 3, nao, nao)
            t1 = log.timer_debug1('contracting int2e_ipvip1 for atom %d' % ia,
                                  *t1)
        else:
            vj1 = rhf_hess._get_jk(mol,
                                   'int2e_ip1ip2',
                                   9,
                                   's1', ['ji->s1kl', dm0[:, p0:p1]],
                                   shls_slice=shls_slice)[0]
            veff += vj1.reshape(3, 3, nao, nao) * 2
            t1 = log.timer_debug1('contracting int2e_ip1ip2 for atom %d' % ia,
                                  *t1)

            vj1 = rhf_hess._get_jk(mol,
                                   'int2e_ipvip1',
                                   9,
                                   's2kl', ['lk->s1ij', dm0],
                                   shls_slice=shls_slice)[0]
            veff[:, :, :,
                 p0:p1] += vj1.transpose(0, 2, 1).reshape(3, 3, nao, p1 - p0)
            t1 = log.timer_debug1('contracting int2e_ipvip1 for atom %d' % ia,
                                  *t1)
        vj1 = vk1 = vk2 = None

        s1ao = numpy.zeros((3, nao, nao))
        s1ao[:, p0:p1] += s1a[:, p0:p1]
        s1ao[:, :, p0:p1] += s1a[:, p0:p1].transpose(0, 2, 1)
        s1oo = numpy.einsum('xpq,pi,qj->xij', s1ao, mocc, mocc)

        de2[i0, i0] += numpy.einsum('xypq,pq->xy', veff_diag[:, :, p0:p1],
                                    dm0[p0:p1]) * 2
        de2[i0, i0] -= numpy.einsum('xypq,pq->xy', s1aa[:, :, p0:p1],
                                    dme0[p0:p1]) * 2

        for j0, ja in enumerate(atmlst[:i0 + 1]):
            q0, q1 = aoslices[ja][2:]
            de2[i0, j0] += numpy.einsum('xypq,pq->xy', veff[:, :, q0:q1],
                                        dm0[q0:q1]) * 2
            de2[i0,
                j0] -= numpy.einsum('xypq,pq->xy', s1ab[:, :, p0:p1, q0:q1],
                                    dme0[p0:p1, q0:q1]) * 2

            h1ao = hcore_deriv(ia, ja)
            de2[i0, j0] += numpy.einsum('xypq,pq->xy', h1ao, dm0)

        for j0 in range(i0):
            de2[j0, i0] = de2[i0, j0].T

    log.timer('RKS partial hessian', *time0)
    return de2
예제 #4
0
파일: uhf.py 프로젝트: chrinide/pyscf
def partial_hess_elec(hessobj, mo_energy=None, mo_coeff=None, mo_occ=None,
                      atmlst=None, max_memory=4000, verbose=None):
    log = logger.new_logger(hessobj, verbose)
    time0 = t1 = (time.clock(), time.time())

    mol = hessobj.mol
    mf = hessobj.base
    if mo_energy is None: mo_energy = mf.mo_energy
    if mo_occ is None:    mo_occ = mf.mo_occ
    if mo_coeff is None:  mo_coeff = mf.mo_coeff
    if atmlst is None: atmlst = range(mol.natm)

    nao, nmo = mo_coeff[0].shape
    mocca = mo_coeff[0][:,mo_occ[0]>0]
    moccb = mo_coeff[1][:,mo_occ[1]>0]
    dm0a = numpy.dot(mocca, mocca.T)
    dm0b = numpy.dot(moccb, moccb.T)
    dm0 = dm0a + dm0b
    # Energy weighted density matrix
    mo_ea = mo_energy[0][mo_occ[0]>0]
    mo_eb = mo_energy[1][mo_occ[1]>0]
    dme0 = numpy.einsum('pi,qi,i->pq', mocca, mocca, mo_ea)
    dme0+= numpy.einsum('pi,qi,i->pq', moccb, moccb, mo_eb)

    hcore_deriv = hessobj.hcore_generator(mol)
    s1aa, s1ab, s1a = rhf_hess.get_ovlp(mol)

    vj1a, vj1b, vk1a, vk1b = \
            _get_jk(mol, 'int2e_ipip1', 9, 's2kl',
                    ['lk->s1ij', dm0a, 'lk->s1ij', dm0b,
                     'jk->s1il', dm0a, 'jk->s1il', dm0b])
    vj1 = vj1a + vj1b
    vhfa_diag = vj1 - vk1a
    vhfb_diag = vj1 - vk1b
    vhfa_diag = vhfa_diag.reshape(3,3,nao,nao)
    vhfb_diag = vhfb_diag.reshape(3,3,nao,nao)
    vj1 = vj1a = vj1b = vk1a = vk1b = None
    t1 = log.timer_debug1('contracting int2e_ipip1', *t1)

    aoslices = mol.aoslice_by_atom()
    de2 = numpy.zeros((mol.natm,mol.natm,3,3))  # (A,B,dR_A,dR_B)
    for i0, ia in enumerate(atmlst):
        shl0, shl1, p0, p1 = aoslices[ia]
        shls_slice = (shl0, shl1) + (0, mol.nbas)*3
        vj1a, vj1b, vk1a, vk1b, vk2a, vk2b = \
                _get_jk(mol, 'int2e_ip1ip2', 9, 's1',
                        ['ji->s1kl', dm0a[:,p0:p1], 'ji->s1kl', dm0b[:,p0:p1],
                         'li->s1kj', dm0a[:,p0:p1], 'li->s1kj', dm0b[:,p0:p1],
                         'lj->s1ki', dm0a         , 'lj->s1ki', dm0b         ],
                        shls_slice=shls_slice)
        vj1 = vj1a + vj1b
        vhfa = vj1 * 2 - vk1a
        vhfb = vj1 * 2 - vk1b
        vhfa[:,:,p0:p1] -= vk2a
        vhfb[:,:,p0:p1] -= vk2b
        t1 = log.timer_debug1('contracting int2e_ip1ip2 for atom %d'%ia, *t1)
        vj1a, vj1b, vk1a, vk1b = \
                _get_jk(mol, 'int2e_ipvip1', 9, 's2kl',
                        ['lk->s1ij', dm0a         , 'lk->s1ij', dm0b         ,
                         'li->s1kj', dm0a[:,p0:p1], 'li->s1kj', dm0b[:,p0:p1]],
                        shls_slice=shls_slice)
        vj1 = vj1a + vj1b
        vhfa[:,:,p0:p1] += vj1.transpose(0,2,1)
        vhfb[:,:,p0:p1] += vj1.transpose(0,2,1)
        vhfa -= vk1a.transpose(0,2,1)
        vhfb -= vk1b.transpose(0,2,1)
        vj1 = vj1a = vj1b = vk1a = vk1b = vk2a = vk2b = None
        t1 = log.timer_debug1('contracting int2e_ipvip1 for atom %d'%ia, *t1)
        vhfa = vhfa.reshape(3,3,nao,nao)
        vhfb = vhfb.reshape(3,3,nao,nao)

        s1ao = numpy.zeros((3,nao,nao))
        s1ao[:,p0:p1] += s1a[:,p0:p1]
        s1ao[:,:,p0:p1] += s1a[:,p0:p1].transpose(0,2,1)
        s1ooa = numpy.einsum('xpq,pi,qj->xij', s1ao, mocca, mocca)
        s1oob = numpy.einsum('xpq,pi,qj->xij', s1ao, moccb, moccb)

        de2[i0,i0] += numpy.einsum('xypq,pq->xy', vhfa_diag[:,:,p0:p1], dm0a[p0:p1])*2
        de2[i0,i0] += numpy.einsum('xypq,pq->xy', vhfb_diag[:,:,p0:p1], dm0b[p0:p1])*2
        de2[i0,i0] -= numpy.einsum('xypq,pq->xy', s1aa[:,:,p0:p1], dme0[p0:p1])*2

        for j0, ja in enumerate(atmlst[:i0+1]):
            q0, q1 = aoslices[ja][2:]
            de2[i0,j0] += numpy.einsum('xypq,pq->xy', vhfa[:,:,q0:q1], dm0a[q0:q1])*2
            de2[i0,j0] += numpy.einsum('xypq,pq->xy', vhfb[:,:,q0:q1], dm0b[q0:q1])*2
            de2[i0,j0] -= numpy.einsum('xypq,pq->xy', s1ab[:,:,p0:p1,q0:q1], dme0[p0:p1,q0:q1])*2

            h1ao = hcore_deriv(ia, ja)
            de2[i0,j0] += numpy.einsum('xypq,pq->xy', h1ao, dm0)

        for j0 in range(i0):
            de2[j0,i0] = de2[i0,j0].T

    log.timer('UHF partial hessian', *time0)
    return de2
예제 #5
0
def partial_hess_elec(hessobj,
                      mo_energy=None,
                      mo_coeff=None,
                      mo_occ=None,
                      atmlst=None,
                      max_memory=4000,
                      verbose=None):
    log = logger.new_logger(hessobj, verbose)
    time0 = t1 = (time.clock(), time.time())

    mol = hessobj.mol
    mf = hessobj.base
    if mo_energy is None: mo_energy = mf.mo_energy
    if mo_occ is None: mo_occ = mf.mo_occ
    if mo_coeff is None: mo_coeff = mf.mo_coeff
    if atmlst is None: atmlst = range(mol.natm)

    nao, nmo = mo_coeff[0].shape
    mocca = mo_coeff[0][:, mo_occ[0] > 0]
    moccb = mo_coeff[1][:, mo_occ[1] > 0]
    dm0a = numpy.dot(mocca, mocca.T)
    dm0b = numpy.dot(moccb, moccb.T)
    dm0 = dm0a + dm0b
    # Energy weighted density matrix
    mo_ea = mo_energy[0][mo_occ[0] > 0]
    mo_eb = mo_energy[1][mo_occ[1] > 0]
    dme0 = numpy.einsum('pi,qi,i->pq', mocca, mocca, mo_ea)
    dme0 += numpy.einsum('pi,qi,i->pq', moccb, moccb, mo_eb)

    hcore_deriv = hessobj.hcore_generator(mol)
    s1aa, s1ab, s1a = rhf_hess.get_ovlp(mol)

    vj1a, vj1b, vk1a, vk1b = \
            _get_jk(mol, 'int2e_ipip1', 9, 's2kl',
                    ['lk->s1ij', dm0a, 'lk->s1ij', dm0b,
                     'jk->s1il', dm0a, 'jk->s1il', dm0b])
    vj1 = vj1a + vj1b
    vhfa_diag = vj1 - vk1a
    vhfb_diag = vj1 - vk1b
    vhfa_diag = vhfa_diag.reshape(3, 3, nao, nao)
    vhfb_diag = vhfb_diag.reshape(3, 3, nao, nao)
    vj1 = vj1a = vj1b = vk1a = vk1b = None
    t1 = log.timer_debug1('contracting int2e_ipip1', *t1)

    aoslices = mol.aoslice_by_atom()
    de2 = numpy.zeros((mol.natm, mol.natm, 3, 3))  # (A,B,dR_A,dR_B)
    for i0, ia in enumerate(atmlst):
        shl0, shl1, p0, p1 = aoslices[ia]
        shls_slice = (shl0, shl1) + (0, mol.nbas) * 3
        vj1a, vj1b, vk1a, vk1b, vk2a, vk2b = \
                _get_jk(mol, 'int2e_ip1ip2', 9, 's1',
                        ['ji->s1kl', dm0a[:,p0:p1], 'ji->s1kl', dm0b[:,p0:p1],
                         'li->s1kj', dm0a[:,p0:p1], 'li->s1kj', dm0b[:,p0:p1],
                         'lj->s1ki', dm0a         , 'lj->s1ki', dm0b         ],
                        shls_slice=shls_slice)
        vj1 = vj1a + vj1b
        vhfa = vj1 * 2 - vk1a
        vhfb = vj1 * 2 - vk1b
        vhfa[:, :, p0:p1] -= vk2a
        vhfb[:, :, p0:p1] -= vk2b
        t1 = log.timer_debug1('contracting int2e_ip1ip2 for atom %d' % ia, *t1)
        vj1a, vj1b, vk1a, vk1b = \
                _get_jk(mol, 'int2e_ipvip1', 9, 's2kl',
                        ['lk->s1ij', dm0a         , 'lk->s1ij', dm0b         ,
                         'li->s1kj', dm0a[:,p0:p1], 'li->s1kj', dm0b[:,p0:p1]],
                        shls_slice=shls_slice)
        vj1 = vj1a + vj1b
        vhfa[:, :, p0:p1] += vj1.transpose(0, 2, 1)
        vhfb[:, :, p0:p1] += vj1.transpose(0, 2, 1)
        vhfa -= vk1a.transpose(0, 2, 1)
        vhfb -= vk1b.transpose(0, 2, 1)
        vj1 = vj1a = vj1b = vk1a = vk1b = vk2a = vk2b = None
        t1 = log.timer_debug1('contracting int2e_ipvip1 for atom %d' % ia, *t1)
        vhfa = vhfa.reshape(3, 3, nao, nao)
        vhfb = vhfb.reshape(3, 3, nao, nao)

        s1ao = numpy.zeros((3, nao, nao))
        s1ao[:, p0:p1] += s1a[:, p0:p1]
        s1ao[:, :, p0:p1] += s1a[:, p0:p1].transpose(0, 2, 1)
        s1ooa = numpy.einsum('xpq,pi,qj->xij', s1ao, mocca, mocca)
        s1oob = numpy.einsum('xpq,pi,qj->xij', s1ao, moccb, moccb)

        de2[i0, i0] += numpy.einsum('xypq,pq->xy', vhfa_diag[:, :, p0:p1],
                                    dm0a[p0:p1]) * 2
        de2[i0, i0] += numpy.einsum('xypq,pq->xy', vhfb_diag[:, :, p0:p1],
                                    dm0b[p0:p1]) * 2
        de2[i0, i0] -= numpy.einsum('xypq,pq->xy', s1aa[:, :, p0:p1],
                                    dme0[p0:p1]) * 2

        for j0, ja in enumerate(atmlst[:i0 + 1]):
            q0, q1 = aoslices[ja][2:]
            de2[i0, j0] += numpy.einsum('xypq,pq->xy', vhfa[:, :, q0:q1],
                                        dm0a[q0:q1]) * 2
            de2[i0, j0] += numpy.einsum('xypq,pq->xy', vhfb[:, :, q0:q1],
                                        dm0b[q0:q1]) * 2
            de2[i0,
                j0] -= numpy.einsum('xypq,pq->xy', s1ab[:, :, p0:p1, q0:q1],
                                    dme0[p0:p1, q0:q1]) * 2

            h1ao = hcore_deriv(ia, ja)
            de2[i0, j0] += numpy.einsum('xypq,pq->xy', h1ao, dm0)

        for j0 in range(i0):
            de2[j0, i0] = de2[i0, j0].T

    log.timer('UHF partial hessian', *time0)
    return de2
예제 #6
0
파일: uks.py 프로젝트: zwang123/pyscf
def partial_hess_elec(hessobj,
                      mo_energy=None,
                      mo_coeff=None,
                      mo_occ=None,
                      atmlst=None,
                      max_memory=4000,
                      verbose=None):
    log = logger.new_logger(hessobj, verbose)
    time0 = t1 = (time.clock(), time.time())

    mol = hessobj.mol
    mf = hessobj.base
    if mo_energy is None: mo_energy = mf.mo_energy
    if mo_occ is None: mo_occ = mf.mo_occ
    if mo_coeff is None: mo_coeff = mf.mo_coeff
    if atmlst is None: atmlst = range(mol.natm)

    nao, nmo = mo_coeff[0].shape
    mocca = mo_coeff[0][:, mo_occ[0] > 0]
    moccb = mo_coeff[1][:, mo_occ[1] > 0]
    dm0a = numpy.dot(mocca, mocca.T)
    dm0b = numpy.dot(moccb, moccb.T)
    dm0 = dm0a + dm0b
    # Energy weighted density matrix
    mo_ea = mo_energy[0][mo_occ[0] > 0]
    mo_eb = mo_energy[1][mo_occ[1] > 0]
    dme0 = numpy.einsum('pi,qi,i->pq', mocca, mocca, mo_ea)
    dme0 += numpy.einsum('pi,qi,i->pq', moccb, moccb, mo_eb)

    hcore_deriv = hessobj.hcore_generator(mol)
    s1aa, s1ab, s1a = rhf_hess.get_ovlp(mol)

    if mf.nlc != '':
        raise NotImplementedError
    #enabling range-separated hybrids
    omega, alpha, hyb = mf._numint.rsh_and_hybrid_coeff(mf.xc, spin=mol.spin)

    mem_now = lib.current_memory()[0]
    max_memory = max(2000, mf.max_memory * .9 - mem_now)
    veffa_diag, veffb_diag = _get_vxc_diag(hessobj, mo_coeff, mo_occ,
                                           max_memory)
    if abs(hyb) > 1e-10:
        vj1a, vj1b, vk1a, vk1b = \
                _get_jk(mol, 'int2e_ipip1', 9, 's2kl',
                        ['lk->s1ij', dm0a, 'lk->s1ij', dm0b,
                         'jk->s1il', dm0a, 'jk->s1il', dm0b])
        vj1 = vj1a + vj1b
        veffa_diag += (vj1 - hyb * vk1a).reshape(3, 3, nao, nao)
        veffb_diag += (vj1 - hyb * vk1b).reshape(3, 3, nao, nao)
        if abs(omega) > 1e-10:
            with mol.with_range_coulomb(omega):
                vk1a, vk1b = \
                        _get_jk(mol, 'int2e_ipip1', 9, 's2kl',
                                ['jk->s1il', dm0a, 'jk->s1il', dm0b])
            veffa_diag -= (alpha - hyb) * vk1a.reshape(3, 3, nao, nao)
            veffb_diag -= (alpha - hyb) * vk1b.reshape(3, 3, nao, nao)
    else:
        vj1a, vj1b = \
                _get_jk(mol, 'int2e_ipip1', 9, 's2kl',
                        ['lk->s1ij', dm0a, 'lk->s1ij', dm0b])
        vj1 = vj1a + vj1b
        veffa_diag += vj1.reshape(3, 3, nao, nao)
        veffb_diag += vj1.reshape(3, 3, nao, nao)
    vj1 = vj1a = vj1b = vk1a = vk1b = None
    t1 = log.timer_debug1('contracting int2e_ipip1', *t1)

    aoslices = mol.aoslice_by_atom()
    de2 = numpy.zeros((mol.natm, mol.natm, 3, 3))  # (A,B,dR_A,dR_B)
    vxca, vxcb = _get_vxc_deriv2(hessobj, mo_coeff, mo_occ, max_memory)
    for i0, ia in enumerate(atmlst):
        shl0, shl1, p0, p1 = aoslices[ia]

        veffa = vxca[ia]
        veffb = vxcb[ia]
        shls_slice = (shl0, shl1) + (0, mol.nbas) * 3
        if abs(hyb) > 1e-10:
            vj1a, vj1b, vk1a, vk1b, vk2a, vk2b = \
                    _get_jk(mol, 'int2e_ip1ip2', 9, 's1',
                            ['ji->s1kl', dm0a[:,p0:p1], 'ji->s1kl', dm0b[:,p0:p1],
                             'li->s1kj', dm0a[:,p0:p1], 'li->s1kj', dm0b[:,p0:p1],
                             'lj->s1ki', dm0a         , 'lj->s1ki', dm0b         ],
                            shls_slice=shls_slice)
            vj1 = vj1a + vj1b
            veffa += (vj1 * 2 - hyb * vk1a).reshape(3, 3, nao, nao)
            veffb += (vj1 * 2 - hyb * vk1b).reshape(3, 3, nao, nao)
            veffa[:, :, :, p0:p1] -= hyb * vk2a.reshape(3, 3, nao, p1 - p0)
            veffb[:, :, :, p0:p1] -= hyb * vk2b.reshape(3, 3, nao, p1 - p0)
            if abs(omega) > 1e-10:
                with mol.with_range_coulomb(omega):
                    vk1a, vk1b, vk2a, vk2b = \
                            _get_jk(mol, 'int2e_ip1ip2', 9, 's1',
                                    ['li->s1kj', dm0a[:,p0:p1],
                                     'li->s1kj', dm0b[:,p0:p1],
                                     'lj->s1ki', dm0a         ,
                                     'lj->s1ki', dm0b         ],
                                    shls_slice=shls_slice)
                veffa -= (alpha - hyb) * vk1a.reshape(3, 3, nao, nao)
                veffb -= (alpha - hyb) * vk1b.reshape(3, 3, nao, nao)
                veffa[:, :, :, p0:p1] -= (alpha - hyb) * vk2a.reshape(
                    3, 3, nao, p1 - p0)
                veffb[:, :, :, p0:p1] -= (alpha - hyb) * vk2b.reshape(
                    3, 3, nao, p1 - p0)
            t1 = log.timer_debug1('contracting int2e_ip1ip2 for atom %d' % ia,
                                  *t1)
            vj1a, vj1b, vk1a, vk1b = \
                    _get_jk(mol, 'int2e_ipvip1', 9, 's2kl',
                            ['lk->s1ij', dm0a         , 'lk->s1ij', dm0b         ,
                             'li->s1kj', dm0a[:,p0:p1], 'li->s1kj', dm0b[:,p0:p1]],
                            shls_slice=shls_slice)
            vj1 = vj1a + vj1b
            veffa[:, :, :,
                  p0:p1] += vj1.transpose(0, 2, 1).reshape(3, 3, nao, p1 - p0)
            veffb[:, :, :,
                  p0:p1] += vj1.transpose(0, 2, 1).reshape(3, 3, nao, p1 - p0)
            veffa -= hyb * vk1a.transpose(0, 2, 1).reshape(3, 3, nao, nao)
            veffb -= hyb * vk1b.transpose(0, 2, 1).reshape(3, 3, nao, nao)
            if abs(omega) > 1e-10:
                with mol.with_range_coulomb(omega):
                    vk1a, vk1b = _get_jk(mol,
                                         'int2e_ipvip1',
                                         9,
                                         's2kl', [
                                             'li->s1kj', dm0a[:, p0:p1],
                                             'li->s1kj', dm0b[:, p0:p1]
                                         ],
                                         shls_slice=shls_slice)
                veffa -= (alpha - hyb) * vk1a.transpose(0, 2, 1).reshape(
                    3, 3, nao, nao)
                veffb -= (alpha - hyb) * vk1b.transpose(0, 2, 1).reshape(
                    3, 3, nao, nao)
            t1 = log.timer_debug1('contracting int2e_ipvip1 for atom %d' % ia,
                                  *t1)
        else:
            vj1a, vj1b = \
                    _get_jk(mol, 'int2e_ip1ip2', 9, 's1',
                            ['ji->s1kl', dm0a[:,p0:p1], 'ji->s1kl', dm0b[:,p0:p1]],
                            shls_slice=shls_slice)
            vj1 = (vj1a + vj1b).reshape(3, 3, nao, nao)
            veffa += vj1 * 2
            veffb += vj1 * 2
            t1 = log.timer_debug1('contracting int2e_ip1ip2 for atom %d' % ia,
                                  *t1)
            vj1a, vj1b = \
                    _get_jk(mol, 'int2e_ipvip1', 9, 's2kl',
                            ['lk->s1ij', dm0a, 'lk->s1ij', dm0b],
                            shls_slice=shls_slice)
            vj1 = vj1a + vj1b
            veffa[:, :, :,
                  p0:p1] += vj1.transpose(0, 2, 1).reshape(3, 3, nao, p1 - p0)
            veffb[:, :, :,
                  p0:p1] += vj1.transpose(0, 2, 1).reshape(3, 3, nao, p1 - p0)
            t1 = log.timer_debug1('contracting int2e_ipvip1 for atom %d' % ia,
                                  *t1)
        vj1 = vj1a = vj1b = vk1a = vk1b = vk2a = vk2b = None

        s1ao = numpy.zeros((3, nao, nao))
        s1ao[:, p0:p1] += s1a[:, p0:p1]
        s1ao[:, :, p0:p1] += s1a[:, p0:p1].transpose(0, 2, 1)
        s1ooa = numpy.einsum('xpq,pi,qj->xij', s1ao, mocca, mocca)
        s1oob = numpy.einsum('xpq,pi,qj->xij', s1ao, moccb, moccb)

        de2[i0, i0] += numpy.einsum('xypq,pq->xy', veffa_diag[:, :, p0:p1],
                                    dm0a[p0:p1]) * 2
        de2[i0, i0] += numpy.einsum('xypq,pq->xy', veffb_diag[:, :, p0:p1],
                                    dm0b[p0:p1]) * 2
        de2[i0, i0] -= numpy.einsum('xypq,pq->xy', s1aa[:, :, p0:p1],
                                    dme0[p0:p1]) * 2

        for j0, ja in enumerate(atmlst[:i0 + 1]):
            q0, q1 = aoslices[ja][2:]
            de2[i0, j0] += numpy.einsum('xypq,pq->xy', veffa[:, :, q0:q1],
                                        dm0a[q0:q1]) * 2
            de2[i0, j0] += numpy.einsum('xypq,pq->xy', veffb[:, :, q0:q1],
                                        dm0b[q0:q1]) * 2
            de2[i0,
                j0] -= numpy.einsum('xypq,pq->xy', s1ab[:, :, p0:p1, q0:q1],
                                    dme0[p0:p1, q0:q1]) * 2

            h1ao = hcore_deriv(ia, ja)
            de2[i0, j0] += numpy.einsum('xypq,pq->xy', h1ao, dm0)

        for j0 in range(i0):
            de2[j0, i0] = de2[i0, j0].T

    log.timer('UKS partial hessian', *time0)
    return de2
예제 #7
0
파일: rhf.py 프로젝트: sunqm/pyscf-test
def _partial_hess_ejk(hessobj,
                      mo_energy=None,
                      mo_coeff=None,
                      mo_occ=None,
                      atmlst=None,
                      max_memory=4000,
                      verbose=None,
                      with_k=True):
    '''Partial derivative
    '''
    log = logger.new_logger(hessobj, verbose)
    time0 = t1 = (logger.process_clock(), logger.perf_counter())

    mol = hessobj.mol
    mf = hessobj.base
    if mo_energy is None: mo_energy = mf.mo_energy
    if mo_occ is None: mo_occ = mf.mo_occ
    if mo_coeff is None: mo_coeff = mf.mo_coeff
    if atmlst is None: atmlst = range(mol.natm)

    nao, nmo = mo_coeff.shape
    mocc = mo_coeff[:, mo_occ > 0]
    mocc_2 = np.einsum('pi,i->pi', mocc, mo_occ[mo_occ > 0]**.5)
    nocc = mocc.shape[1]
    dm0 = numpy.dot(mocc, mocc.T) * 2
    # Energy weighted density matrix
    dme0 = numpy.einsum('pi,qi,i->pq', mocc, mocc, mo_energy[mo_occ > 0]) * 2

    auxmol = hessobj.base.with_df.auxmol
    naux = auxmol.nao
    nbas = mol.nbas
    auxslices = auxmol.aoslice_by_atom()
    aoslices = mol.aoslice_by_atom()
    aux_loc = auxmol.ao_loc
    blksize = min(480, hessobj.max_memory * .3e6 / 8 / nao**2)
    aux_ranges = ao2mo.outcore.balance_partition(auxmol.ao_loc, blksize)

    hcore_deriv = hessobj.hcore_generator(mol)
    s1aa, s1ab, s1a = rhf_hess.get_ovlp(mol)

    ftmp = lib.H5TmpFile()
    get_int3c = _int3c_wrapper(mol, auxmol, 'int3c2e', 's1')
    # Without RI basis response
    #    (20|0)(0|00)
    #    (11|0)(0|00)
    #    (10|0)(0|10)
    int2c = auxmol.intor('int2c2e', aosym='s1')
    int2c_low = scipy.linalg.cho_factor(int2c, lower=True)
    int2c_ip1 = auxmol.intor('int2c2e_ip1', aosym='s1')

    rhoj0_P = 0
    if with_k:
        if hessobj.max_memory * .8e6 / 8 < naux * nocc * (nocc + nao):
            raise RuntimeError(
                'Memory not enough. You need to increase mol.max_memory')
        rhok0_Pl_ = np.empty((naux, nao, nocc))
    for i, (shl0, shl1, p0, p1) in enumerate(aoslices):
        int3c = get_int3c((shl0, shl1, 0, nbas, 0, auxmol.nbas))
        rhoj0_P += np.einsum('klp,kl->p', int3c, dm0[p0:p1])
        if with_k:
            tmp = lib.einsum('ijp,jk->pik', int3c, mocc_2)
            tmp = scipy.linalg.cho_solve(int2c_low,
                                         tmp.reshape(naux, -1),
                                         overwrite_b=True)
            rhok0_Pl_[:, p0:p1] = tmp.reshape(naux, p1 - p0, nocc)
        int3c = tmp = None
    rhoj0_P = scipy.linalg.cho_solve(int2c_low, rhoj0_P)

    get_int3c_ipip1 = _int3c_wrapper(mol, auxmol, 'int3c2e_ipip1', 's1')
    vj1_diag = 0
    vk1_diag = 0
    for shl0, shl1, nL in aux_ranges:
        shls_slice = (0, nbas, 0, nbas, shl0, shl1)
        p0, p1 = aux_loc[shl0], aux_loc[shl1]
        int3c_ipip1 = get_int3c_ipip1(shls_slice)
        vj1_diag += np.einsum('xijp,p->xij', int3c_ipip1,
                              rhoj0_P[p0:p1]).reshape(3, 3, nao, nao)
        if with_k:
            tmp = lib.einsum('Plj,Jj->PlJ', rhok0_Pl_[p0:p1], mocc_2)
            vk1_diag += lib.einsum('xijp,plj->xil', int3c_ipip1,
                                   tmp).reshape(3, 3, nao, nao)
    int3c_ipip1 = get_int3c_ipip1 = tmp = None
    t1 = log.timer_debug1('contracting int2e_ipip1', *t1)

    get_int3c_ip1 = _int3c_wrapper(mol, auxmol, 'int3c2e_ip1', 's1')
    rho_ip1 = ftmp.create_dataset('rho_ip1', (nao, nao, naux, 3), 'f8')
    rhok_ip1_IkP = ftmp.create_group('rhok_ip1_IkP')
    rhok_ip1_PkI = ftmp.create_group('rhok_ip1_PkI')
    rhoj1 = np.empty((mol.natm, naux, 3))
    wj1 = np.empty((mol.natm, naux, 3))
    for i0, ia in enumerate(atmlst):
        shl0, shl1, p0, p1 = aoslices[ia]
        shls_slice = (shl0, shl1, 0, nbas, 0, auxmol.nbas)
        int3c_ip1 = get_int3c_ip1(shls_slice)
        tmp_ip1 = scipy.linalg.cho_solve(int2c_low,
                                         int3c_ip1.reshape(-1, naux).T,
                                         overwrite_b=True).reshape(
                                             naux, 3, p1 - p0, nao)
        rhoj1[i0] = np.einsum('pxij,ji->px', tmp_ip1, dm0[:, p0:p1])
        wj1[i0] = np.einsum('xijp,ji->px', int3c_ip1, dm0[:, p0:p1])
        rho_ip1[p0:p1] = tmp_ip1.transpose(2, 3, 0, 1)
        if with_k:
            tmp = lib.einsum('pykl,li->ikpy', tmp_ip1, dm0)
            rhok_ip1_IkP['%.4d' % ia] = tmp
            rhok_ip1_PkI['%.4d' % ia] = tmp.transpose(2, 1, 0, 3)
            tmp = None
    ej = lib.einsum('ipx,jpy->ijxy', rhoj1, wj1) * 4
    ek = np.zeros_like(ej)
    e1 = np.zeros_like(ej)
    rhoj1 = wj1 = None

    if with_k:
        vk2buf = 0
        for shl0, shl1, nL in aux_ranges:
            shls_slice = (0, nbas, 0, nbas, shl0, shl1)
            p0, p1 = aux_loc[shl0], aux_loc[shl1]
            int3c_ip1 = get_int3c_ip1(shls_slice)
            vk2buf += lib.einsum('xijp,pkjy->xyki', int3c_ip1,
                                 _load_dim0(rhok_ip1_PkI, p0, p1))
            int3c_ip1 = None

    get_int3c_ip2 = _int3c_wrapper(mol, auxmol, 'int3c2e_ip2', 's1')
    wj_ip2 = np.empty((naux, 3))
    wk_ip2_Ipk = ftmp.create_dataset('wk_ip2', (nao, naux, 3, nao), 'f8')
    if hessobj.auxbasis_response > 1:
        wk_ip2_P__ = np.empty((naux, 3, nocc, nocc))
    for shl0, shl1, nL in aux_ranges:
        shls_slice = (0, nbas, 0, nbas, shl0, shl1)
        p0, p1 = aux_loc[shl0], aux_loc[shl1]
        int3c_ip2 = get_int3c_ip2(shls_slice)
        wj_ip2[p0:p1] = np.einsum('yklp,lk->py', int3c_ip2, dm0)
        if with_k:
            wk_ip2_Ipk[:, p0:p1] = lib.einsum('yklp,il->ipyk', int3c_ip2, dm0)
            if hessobj.auxbasis_response > 1:
                wk_ip2_P__[p0:p1] = lib.einsum('xuvp,ui,vj->pxij', int3c_ip2,
                                               mocc_2, mocc_2)
        int3c_ip2 = None

    if hessobj.auxbasis_response > 1:
        get_int3c_ipip2 = _int3c_wrapper(mol, auxmol, 'int3c2e_ipip2', 's1')
        rhok0_P__ = lib.einsum('plj,li->pij', rhok0_Pl_, mocc_2)
        rho2c_0 = lib.einsum('pij,qji->pq', rhok0_P__, rhok0_P__)

        int2c_inv = np.linalg.inv(int2c)
        int2c_ipip1 = auxmol.intor('int2c2e_ipip1', aosym='s1')
        int2c_ip_ip = lib.einsum('xpq,qr,ysr->xyps', int2c_ip1, int2c_inv,
                                 int2c_ip1)
        int2c_ip_ip -= auxmol.intor('int2c2e_ip1ip2',
                                    aosym='s1').reshape(3, 3, naux, naux)
    int2c = int2c_low = None

    get_int3c_ipvip1 = _int3c_wrapper(mol, auxmol, 'int3c2e_ipvip1', 's1')
    get_int3c_ip1ip2 = _int3c_wrapper(mol, auxmol, 'int3c2e_ip1ip2', 's1')

    for i0, ia in enumerate(atmlst):
        shl0, shl1, p0, p1 = aoslices[ia]
        shls_slice = (shl0, shl1, 0, nbas, 0, auxmol.nbas)
        # (10|0)(0|10) without response of RI basis
        if with_k:
            int3c_ip1 = get_int3c_ip1(shls_slice)
            vk1 = lib.einsum('xijp,ikpy->xykj', int3c_ip1,
                             _load_dim0(rhok_ip1_IkP, p0, p1))
            vk1[:, :, :, p0:p1] += vk2buf[:, :, :, p0:p1]
        t1 = log.timer_debug1('contracting int2e_ip1ip2 for atom %d' % ia, *t1)
        int3c_ip1 = None

        # (11|0)(0|00) without response of RI basis
        int3c_ipvip1 = get_int3c_ipvip1(shls_slice)
        vj1 = np.einsum('xijp,p->xji', int3c_ipvip1,
                        rhoj0_P).reshape(3, 3, nao, p1 - p0)
        if with_k:
            tmp = lib.einsum('pki,ji->pkj', rhok0_Pl_, mocc_2[p0:p1])
            vk1 += lib.einsum('xijp,pki->xjk', int3c_ipvip1,
                              tmp).reshape(3, 3, nao, nao)
        t1 = log.timer_debug1('contracting int2e_ipvip1 for atom %d' % ia, *t1)
        int3c_ipvip1 = tmp = None

        e1[i0, i0] -= numpy.einsum('xypq,pq->xy', s1aa[:, :, p0:p1],
                                   dme0[p0:p1]) * 2
        ej[i0, i0] += numpy.einsum('xypq,pq->xy', vj1_diag[:, :, p0:p1],
                                   dm0[p0:p1]) * 2
        if with_k:
            ek[i0, i0] += numpy.einsum('xypq,pq->xy', vk1_diag[:, :, p0:p1],
                                       dm0[p0:p1])

        for j0, ja in enumerate(atmlst[:i0 + 1]):
            q0, q1 = aoslices[ja][2:]
            ej[i0, j0] += numpy.einsum('xypq,pq->xy', vj1[:, :, q0:q1],
                                       dm0[q0:q1, p0:p1]) * 2
            e1[i0, j0] -= numpy.einsum('xypq,pq->xy', s1ab[:, :, p0:p1, q0:q1],
                                       dme0[p0:p1, q0:q1]) * 2
            if with_k:
                ek[i0, j0] += numpy.einsum('xypq,pq->xy', vk1[:, :, q0:q1],
                                           dm0[q0:q1])

            h1ao = hcore_deriv(ia, ja)
            e1[i0, j0] += numpy.einsum('xypq,pq->xy', h1ao, dm0)

        #
        # The first order RI basis response
        #
        #    (10|1)(0|00)
        #    (10|0)(1|0)(0|00)
        #    (10|0)(0|1)(0|00)
        #    (10|0)(1|00)
        if hessobj.auxbasis_response:
            wk1_Pij = rho_ip1[p0:p1].transpose(2, 3, 0, 1)
            rhoj1_P = np.einsum('pxij,ji->px', wk1_Pij, dm0[:, p0:p1])
            # (10|1)(0|0)(0|00)
            int3c_ip1ip2 = get_int3c_ip1ip2(shls_slice)
            wj11_p = np.einsum('xijp,ji->xp', int3c_ip1ip2, dm0[:, p0:p1])
            # (10|0)(1|0)(0|00)
            wj0_01 = np.einsum('ypq,q->yp', int2c_ip1, rhoj0_P)
            if with_k:
                rhok0_P_I = lib.einsum('plj,il->pji', rhok0_Pl_, dm0[p0:p1])
                rhok0_PJI = lib.einsum('pji,Jj->pJi', rhok0_P_I, mocc_2)
                wk1_pJI = lib.einsum('ypq,qji->ypji', int2c_ip1, rhok0_PJI)
                wk1_IpJ = lib.einsum('ipyk,kj->ipyj', wk_ip2_Ipk[p0:p1], dm0)
                #rho2c_PQ = lib.einsum('qij,uj,iupx->xqp', rhok0_Pl_, mocc_2[p0:p1], rhok_ip1_IkP['%.4d'%ia])
                rho2c_PQ = lib.einsum('pxij,qji->xqp', wk1_Pij, rhok0_PJI)
            for j0, (q0, q1) in enumerate(auxslices[:, 2:]):
                # (10|1)(0|00)
                _ej = np.einsum('xp,p->x', wj11_p[:, q0:q1],
                                rhoj0_P[q0:q1]).reshape(3, 3)
                # (10|0)(0|1)(0|00)
                _ej -= lib.einsum('yqp,q,px->xy', int2c_ip1[:, q0:q1],
                                  rhoj0_P[q0:q1], rhoj1_P)
                # (10|0)(1|0)(0|00)
                _ej -= lib.einsum('px,yp->xy', rhoj1_P[q0:q1], wj0_01[:,
                                                                      q0:q1])
                # (10|0)(1|00)
                _ej += lib.einsum('px,py->xy', rhoj1_P[q0:q1], wj_ip2[q0:q1])
                if hessobj.auxbasis_response > 1:
                    ej[i0, j0] += _ej * 2
                    ej[j0, i0] += _ej.T * 2
                else:
                    ej[i0, j0] += _ej
                    ej[j0, i0] += _ej.T
                if with_k:
                    _ek = lib.einsum('xijp,pji->x', int3c_ip1ip2[:, :, :,
                                                                 q0:q1],
                                     rhok0_PJI[q0:q1]).reshape(3, 3)
                    _ek -= lib.einsum('pxij,ypji->xy', wk1_Pij[q0:q1],
                                      wk1_pJI[:, q0:q1])
                    _ek -= lib.einsum('xqp,yqp->xy', rho2c_PQ[:, q0:q1],
                                      int2c_ip1[:, q0:q1])
                    _ek += lib.einsum('pxij,ipyj->xy', wk1_Pij[q0:q1],
                                      wk1_IpJ[:, q0:q1])
                    if hessobj.auxbasis_response > 1:
                        ek[i0, j0] += _ek
                        ek[j0, i0] += _ek.T
                    else:
                        ek[i0, j0] += _ek * .5
                        ek[j0, i0] += _ek.T * .5
            int3c_ip1ip2 = rhok0_P_I = rhok0_PJI = wk1_pJI = wk1_IpJ = rho2c_PQ = None

        #
        # The second order RI basis response
        #
        if hessobj.auxbasis_response > 1:
            # (00|2)(0|00)
            # (00|0)(2|0)(0|00)
            shl0, shl1, p0, p1 = auxslices[ia]
            shls_slice = (0, nbas, 0, nbas, shl0, shl1)
            int3c_ipip2 = get_int3c_ipip2(shls_slice)
            ej[i0, i0] += np.einsum('xijp,ji,p->x', int3c_ipip2, dm0,
                                    rhoj0_P[p0:p1]).reshape(3, 3)
            ej[i0, i0] -= np.einsum('p,xpq,q->x', rhoj0_P[p0:p1],
                                    int2c_ipip1[:,
                                                p0:p1], rhoj0_P).reshape(3, 3)

            if with_k:
                rhok0_PJI = lib.einsum('Pij,Jj,Ii->PJI', rhok0_P__[p0:p1],
                                       mocc_2, mocc_2)
                ek[i0, i0] += .5 * np.einsum('xijp,pij->x', int3c_ipip2,
                                             rhok0_PJI).reshape(3, 3)
                ek[i0,
                   i0] -= .5 * np.einsum('pq,xpq->x', rho2c_0[p0:p1],
                                         int2c_ipip1[:, p0:p1]).reshape(3, 3)
                rhok0_PJI = None
            # (00|0)(1|1)(0|00)
            # (00|1)(1|0)(0|00)
            # (00|1)(0|1)(0|00)
            # (00|1)(1|00)
            rhoj1 = lib.einsum('px,pq->xq', wj_ip2[p0:p1], int2c_inv[p0:p1])
            # (00|0)(0|1)(1|0)(0|00)
            rhoj0_01 = lib.einsum('xp,pq->xq', wj0_01[:, p0:p1],
                                  int2c_inv[p0:p1])
            # (00|0)(1|0)(1|0)(0|00)
            ip1_2c_2c = lib.einsum('xpq,qr->xpr', int2c_ip1[:, p0:p1],
                                   int2c_inv)
            rhoj0_10 = lib.einsum('p,xpq->xq', rhoj0_P[p0:p1], ip1_2c_2c)
            if with_k:
                # (00|0)(0|1)(1|0)(0|00)
                ip1_rho2c = .5 * lib.einsum('xpq,qr->xpr', int2c_ip1[:, p0:p1],
                                            rho2c_0)
                rho2c_1 = lib.einsum('xrq,rp->xpq', ip1_rho2c,
                                     int2c_inv[p0:p1])
                # (00|0)(1|0)(1|0)(0|00)
                rho2c_1 += lib.einsum('xrp,rq->xpq', ip1_2c_2c, rho2c_0[p0:p1])
                # (00|1)(0|1)(0|00)
                # (00|1)(1|0)(0|00)
                int3c_ip2 = get_int3c_ip2(shls_slice)
                tmp = lib.einsum('xuvr,vj,ui->xrij', int3c_ip2, mocc_2, mocc_2)
                tmp = lib.einsum('xrij,qij,rp->xpq', tmp, rhok0_P__,
                                 int2c_inv[p0:p1])
                rho2c_1 -= tmp
                rho2c_1 -= tmp.transpose(0, 2, 1)
                int3c_ip2 = tmp = None
            for j0, (q0, q1) in enumerate(auxslices[:, 2:]):
                _ej = 0
                # (00|0)(1|1)(0|00)
                # (00|0)(1|0)(0|1)(0|00)
                _ej += .5 * np.einsum('p,xypq,q->xy', rhoj0_P[p0:p1],
                                      int2c_ip_ip[:, :, p0:p1,
                                                  q0:q1], rhoj0_P[q0:q1])
                # (00|1)(1|0)(0|00)
                _ej -= lib.einsum('xp,yp->xy', rhoj1[:, q0:q1], wj0_01[:,
                                                                       q0:q1])
                # (00|1)(1|00)
                _ej += .5 * lib.einsum('xp,py->xy', rhoj1[:, q0:q1],
                                       wj_ip2[q0:q1])
                # (00|0)(0|1)(1|0)(0|00)
                _ej += .5 * np.einsum('xp,yp->xy', rhoj0_01[:, q0:q1],
                                      wj0_01[:, q0:q1])
                # (00|1)(0|1)(0|00)
                _ej -= lib.einsum('yqp,q,xp->xy', int2c_ip1[:, q0:q1],
                                  rhoj0_P[q0:q1], rhoj1)
                # (00|0)(1|0)(1|0)(0|00)
                _ej += np.einsum('xp,yp->xy', rhoj0_10[:, q0:q1],
                                 wj0_01[:, q0:q1])
                ej[i0, j0] += _ej
                ej[j0, i0] += _ej.T
                if with_k:
                    # (00|0)(1|1)(0|00)
                    # (00|0)(1|0)(0|1)(0|00)
                    _ek = .5 * np.einsum('pq,xypq->xy', rho2c_0[p0:p1, q0:q1],
                                         int2c_ip_ip[:, :, p0:p1, q0:q1])
                    # (00|1)(0|1)(0|00)
                    # (00|1)(1|0)(0|00)
                    # (00|0)(0|1)(1|0)(0|00)
                    # (00|0)(1|0)(1|0)(0|00)
                    _ek += np.einsum('xpq,ypq->xy', rho2c_1[:, q0:q1],
                                     int2c_ip1[:, q0:q1])
                    # (00|1)(1|00)
                    _ek += .5 * lib.einsum(
                        'pxij,pq,qyij->xy', wk_ip2_P__[p0:p1],
                        int2c_inv[p0:p1, q0:q1], wk_ip2_P__[q0:q1])
                    ek[i0, j0] += _ek * .5
                    ek[j0, i0] += _ek.T * .5

    for i0, ia in enumerate(atmlst):
        for j0 in range(i0):
            e1[j0, i0] = e1[i0, j0].T
            ej[j0, i0] = ej[i0, j0].T
            ek[j0, i0] = ek[i0, j0].T

    log.timer('RHF partial hessian', *time0)
    return e1, ej, ek
예제 #8
0
def _partial_hess_ejk(hessobj,
                      mo_energy=None,
                      mo_coeff=None,
                      mo_occ=None,
                      atmlst=None,
                      max_memory=4000,
                      verbose=None,
                      with_k=True):
    log = logger.new_logger(hessobj, verbose)
    time0 = t1 = (time.clock(), time.time())

    mol = hessobj.mol
    mf = hessobj.base
    if mo_energy is None: mo_energy = mf.mo_energy
    if mo_occ is None: mo_occ = mf.mo_occ
    if mo_coeff is None: mo_coeff = mf.mo_coeff
    if atmlst is None: atmlst = range(mol.natm)

    nao, nmo = mo_coeff[0].shape
    mocca = mo_coeff[0][:, mo_occ[0] > 0]
    moccb = mo_coeff[1][:, mo_occ[1] > 0]
    dm0a = numpy.dot(mocca, mocca.T)
    dm0b = numpy.dot(moccb, moccb.T)
    dm0 = dm0a + dm0b
    # Energy weighted density matrix
    mo_ea = mo_energy[0][mo_occ[0] > 0]
    mo_eb = mo_energy[1][mo_occ[1] > 0]
    dme0 = numpy.einsum('pi,qi,i->pq', mocca, mocca, mo_ea)
    dme0 += numpy.einsum('pi,qi,i->pq', moccb, moccb, mo_eb)

    hcore_deriv = hessobj.hcore_generator(mol)
    s1aa, s1ab, s1a = rhf_hess.get_ovlp(mol)

    vj1_diag, vk1a_diag, vk1b_diag = \
            _get_jk(mol, 'int2e_ipip1', 9, 's2kl',
                    ['lk->s1ij', dm0,
                     'jk->s1il', dm0a, 'jk->s1il', dm0b],
                    vhfopt=_make_vhfopt(mol, dm0, 'ipip1', 'int2e_ipip1ipip2'))
    vj1_diag = vj1_diag.reshape(3, 3, nao, nao)
    vk1a_diag = vk1a_diag.reshape(3, 3, nao, nao)
    vk1b_diag = vk1b_diag.reshape(3, 3, nao, nao)
    t1 = log.timer_debug1('contracting int2e_ipip1', *t1)

    ip1ip2_opt = _make_vhfopt(mol, dm0, 'ip1ip2', 'int2e_ip1ip2')
    ipvip1_opt = _make_vhfopt(mol, dm0, 'ipvip1', 'int2e_ipvip1ipvip2')
    aoslices = mol.aoslice_by_atom()
    e1 = numpy.zeros((mol.natm, mol.natm, 3, 3))  # (A,B,dR_A,dR_B)
    ej = numpy.zeros((mol.natm, mol.natm, 3, 3))
    ek = numpy.zeros((mol.natm, mol.natm, 3, 3))
    for i0, ia in enumerate(atmlst):
        shl0, shl1, p0, p1 = aoslices[ia]
        shls_slice = (shl0, shl1) + (0, mol.nbas) * 3
        vj1, vk1a, vk1b, vk2a, vk2b = \
                _get_jk(mol, 'int2e_ip1ip2', 9, 's1',
                        ['ji->s1kl', dm0 [:,p0:p1],
                         'li->s1kj', dm0a[:,p0:p1], 'li->s1kj', dm0b[:,p0:p1],
                         'lj->s1ki', dm0a         , 'lj->s1ki', dm0b         ],
                        shls_slice=shls_slice, vhfopt=ip1ip2_opt)
        vk1a[:, :, p0:p1] += vk2a
        vk1b[:, :, p0:p1] += vk2b
        t1 = log.timer_debug1('contracting int2e_ip1ip2 for atom %d' % ia, *t1)
        vj2, vk2a, vk2b = \
                _get_jk(mol, 'int2e_ipvip1', 9, 's2kl',
                        ['lk->s1ij', dm0          ,
                         'li->s1kj', dm0a[:,p0:p1], 'li->s1kj', dm0b[:,p0:p1]],
                        shls_slice=shls_slice, vhfopt=ipvip1_opt)
        vj1[:, :, p0:p1] += vj2.transpose(0, 2, 1) * .5
        vk1a += vk2a.transpose(0, 2, 1)
        vk1b += vk2b.transpose(0, 2, 1)
        t1 = log.timer_debug1('contracting int2e_ipvip1 for atom %d' % ia, *t1)
        vj1 = vj1.reshape(3, 3, nao, nao)
        vk1a = vk1a.reshape(3, 3, nao, nao)
        vk1b = vk1b.reshape(3, 3, nao, nao)

        ej[i0, i0] += numpy.einsum('xypq,pq->xy', vj1_diag[:, :, p0:p1],
                                   dm0[p0:p1]) * 2
        ek[i0, i0] += numpy.einsum('xypq,pq->xy', vk1a_diag[:, :, p0:p1],
                                   dm0a[p0:p1]) * 2
        ek[i0, i0] += numpy.einsum('xypq,pq->xy', vk1b_diag[:, :, p0:p1],
                                   dm0b[p0:p1]) * 2
        e1[i0, i0] -= numpy.einsum('xypq,pq->xy', s1aa[:, :, p0:p1],
                                   dme0[p0:p1]) * 2

        for j0, ja in enumerate(atmlst[:i0 + 1]):
            q0, q1 = aoslices[ja][2:]
            ej[i0, j0] += numpy.einsum('xypq,pq->xy', vj1[:, :, q0:q1],
                                       dm0[q0:q1]) * 4
            ek[i0, j0] += numpy.einsum('xypq,pq->xy', vk1a[:, :, q0:q1],
                                       dm0a[q0:q1]) * 2
            ek[i0, j0] += numpy.einsum('xypq,pq->xy', vk1b[:, :, q0:q1],
                                       dm0b[q0:q1]) * 2
            e1[i0, j0] -= numpy.einsum('xypq,pq->xy', s1ab[:, :, p0:p1, q0:q1],
                                       dme0[p0:p1, q0:q1]) * 2

            h1ao = hcore_deriv(ia, ja)
            e1[i0, j0] += numpy.einsum('xypq,pq->xy', h1ao, dm0)

        for j0 in range(i0):
            e1[j0, i0] = e1[i0, j0].T
            ej[j0, i0] = ej[i0, j0].T
            ek[j0, i0] = ek[i0, j0].T

    log.timer('UHF partial hessian', *time0)
    return e1, ej, ek