Ejemplo n.º 1
0
    def test_gga_deriv2(self):
        ng = 7
        xctype = 'GGA'
        np.random.seed(8)
        rho = np.random.rand(2, 4, ng)
        rho1 = np.random.rand(2, 4, ng)
        weight = 1

        exc, vxc, fxc, kxc = eval_xc(f'R-{xctype}', ng)
        ref = numint._rks_gga_wv1(rho[0], rho1[0], vxc, fxc, weight)
        ref[0] *= 2
        v1 = xc_deriv.transform_fxc(rho[0], vxc, fxc, xctype, spin=0)
        v1 = np.einsum('xg,xyg->yg', rho1[0], v1)
        self.assertAlmostEqual(abs(v1 - ref).max(), 0, 12)

        exc, vxc, fxc, kxc = eval_xc(f'U-{xctype}', ng)
        ref = np.array(numint._uks_gga_wv1(rho, rho1, vxc, fxc, weight))
        ref[:, 0] *= 2
        v1 = xc_deriv.transform_fxc(rho, vxc, fxc, xctype, spin=1)
        v1 = np.einsum('axg,axbyg->byg', rho1, v1)
        self.assertAlmostEqual(abs(v1 - ref).max(), 0, 12)
Ejemplo n.º 2
0
def _get_vxc_deriv1(hessobj, mo_coeff, mo_occ, max_memory):
    mol = hessobj.mol
    mf = hessobj.base
    if hessobj.grids is not None:
        grids = hessobj.grids
    else:
        grids = mf.grids
    if grids.coords is None:
        grids.build(with_non0tab=True)

    nao, nmo = mo_coeff[0].shape
    ni = mf._numint
    xctype = ni._xc_type(mf.xc)
    aoslices = mol.aoslice_by_atom()
    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_nr()
    dm0a, dm0b = mf.make_rdm1(mo_coeff, mo_occ)

    vmata = numpy.zeros((mol.natm,3,nao,nao))
    vmatb = numpy.zeros((mol.natm,3,nao,nao))
    max_memory = max(2000, max_memory-(vmata.size+vmatb.size)*8/1e6)
    if xctype == 'LDA':
        ao_deriv = 1
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rhoa = ni.eval_rho2(mol, ao[0], mo_coeff[0], mo_occ[0], mask, xctype)
            rhob = ni.eval_rho2(mol, ao[0], mo_coeff[1], mo_occ[1], mask, xctype)
            vxc, fxc = ni.eval_xc(mf.xc, (rhoa,rhob), 1, deriv=2)[1:3]
            vrho = vxc[0]
            u_u, u_d, d_d = fxc[0].T

            ao_dm0a = numint._dot_ao_dm(mol, ao[0], dm0a, mask, shls_slice, ao_loc)
            ao_dm0b = numint._dot_ao_dm(mol, ao[0], dm0b, mask, shls_slice, ao_loc)
            aow1a = numpy.einsum('xpi,p->xpi', ao[1:], weight*vrho[:,0])
            aow1b = numpy.einsum('xpi,p->xpi', ao[1:], weight*vrho[:,1])
            for ia in range(mol.natm):
                p0, p1 = aoslices[ia][2:]
# First order density = rho1 * 2.  *2 is not applied because + c.c. in the end
                rho1a = numpy.einsum('xpi,pi->xp', ao[1:,:,p0:p1], ao_dm0a[:,p0:p1])
                rho1b = numpy.einsum('xpi,pi->xp', ao[1:,:,p0:p1], ao_dm0b[:,p0:p1])

                wv = u_u * rho1a + u_d * rho1b
                wv *= weight
                aow = numpy.einsum('pi,xp->xpi', ao[0], wv)
                aow[:,:,p0:p1] += aow1a[:,:,p0:p1]
                rks_grad._d1_dot_(vmata[ia], mol, aow, ao[0], mask, ao_loc, True)

                wv = u_d * rho1a + d_d * rho1b
                wv *= weight
                aow = numpy.einsum('pi,xp->xpi', ao[0], wv)
                aow[:,:,p0:p1] += aow1b[:,:,p0:p1]
                rks_grad._d1_dot_(vmatb[ia], mol, aow, ao[0], mask, ao_loc, True)
            ao_dm0a = ao_dm0b = aow = aow1a = aow1b = None

        for ia in range(mol.natm):
            vmata[ia] = -vmata[ia] - vmata[ia].transpose(0,2,1)
            vmatb[ia] = -vmatb[ia] - vmatb[ia].transpose(0,2,1)

    elif xctype == 'GGA':
        ao_deriv = 2
        vipa = numpy.zeros((3,nao,nao))
        vipb = numpy.zeros((3,nao,nao))
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rhoa = ni.eval_rho2(mol, ao[:4], mo_coeff[0], mo_occ[0], mask, xctype)
            rhob = ni.eval_rho2(mol, ao[:4], mo_coeff[1], mo_occ[1], mask, xctype)
            vxc, fxc = ni.eval_xc(mf.xc, (rhoa,rhob), 1, deriv=2)[1:3]

            wva, wvb = numint._uks_gga_wv0((rhoa,rhob), vxc, weight)
            rks_grad._gga_grad_sum_(vipa, mol, ao, wva, mask, ao_loc)
            rks_grad._gga_grad_sum_(vipb, mol, ao, wvb, mask, ao_loc)

            ao_dm0a = [numint._dot_ao_dm(mol, ao[i], dm0a, mask, shls_slice, ao_loc)
                       for i in range(4)]
            ao_dm0b = [numint._dot_ao_dm(mol, ao[i], dm0b, mask, shls_slice, ao_loc)
                       for i in range(4)]
            for ia in range(mol.natm):
                wva = dR_rho1a = rks_hess._make_dR_rho1(ao, ao_dm0a, ia, aoslices)
                wvb = dR_rho1b = rks_hess._make_dR_rho1(ao, ao_dm0b, ia, aoslices)
                wva[0], wvb[0] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[0],dR_rho1b[0]), vxc, fxc, weight)
                wva[1], wvb[1] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[1],dR_rho1b[1]), vxc, fxc, weight)
                wva[2], wvb[2] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[2],dR_rho1b[2]), vxc, fxc, weight)

                aow = numpy.einsum('npi,Xnp->Xpi', ao[:4], wva)
                rks_grad._d1_dot_(vmata[ia], mol, aow, ao[0], mask, ao_loc, True)
                aow = numpy.einsum('npi,Xnp->Xpi', ao[:4], wvb)
                rks_grad._d1_dot_(vmatb[ia], mol, aow, ao[0], mask, ao_loc, True)
            ao_dm0a = ao_dm0b = aow = None

        for ia in range(mol.natm):
            p0, p1 = aoslices[ia][2:]
            vmata[ia,:,p0:p1] += vipa[:,p0:p1]
            vmata[ia] = -vmata[ia] - vmata[ia].transpose(0,2,1)
            vmatb[ia,:,p0:p1] += vipb[:,p0:p1]
            vmatb[ia] = -vmatb[ia] - vmatb[ia].transpose(0,2,1)

    elif xctype == 'MGGA':
        raise NotImplementedError('meta-GGA')

    return vmata, vmatb
Ejemplo n.º 3
0
def _get_vxc_deriv2(hessobj, mo_coeff, mo_occ, max_memory):
    mol = hessobj.mol
    mf = hessobj.base
    if hessobj.grids is not None:
        grids = hessobj.grids
    else:
        grids = mf.grids
    if grids.coords is None:
        grids.build(with_non0tab=True)

    nao, nmo = mo_coeff[0].shape
    ni = mf._numint
    xctype = ni._xc_type(mf.xc)
    aoslices = mol.aoslice_by_atom()
    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_nr()
    dm0a, dm0b = mf.make_rdm1(mo_coeff, mo_occ)

    vmata = numpy.zeros((mol.natm,3,3,nao,nao))
    vmatb = numpy.zeros((mol.natm,3,3,nao,nao))
    ipipa = numpy.zeros((3,3,nao,nao))
    ipipb = numpy.zeros((3,3,nao,nao))
    if xctype == 'LDA':
        ao_deriv = 1
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rhoa = ni.eval_rho2(mol, ao[0], mo_coeff[0], mo_occ[0], mask, xctype)
            rhob = ni.eval_rho2(mol, ao[0], mo_coeff[1], mo_occ[1], mask, xctype)
            vxc, fxc = ni.eval_xc(mf.xc, (rhoa,rhob), 1, deriv=2)[1:3]
            vrho = vxc[0]
            u_u, u_d, d_d = fxc[0].T

            aow = numpy.einsum('xpi,p->xpi', ao[1:4], weight*vrho[:,0])
            rks_hess._d1d2_dot_(ipipa, mol, aow, ao[1:4], mask, ao_loc, False)
            aow = numpy.einsum('xpi,p->xpi', ao[1:4], weight*vrho[:,1])
            rks_hess._d1d2_dot_(ipipb, mol, aow, ao[1:4], mask, ao_loc, False)

            ao_dm0a = numint._dot_ao_dm(mol, ao[0], dm0a, mask, shls_slice, ao_loc)
            ao_dm0b = numint._dot_ao_dm(mol, ao[0], dm0b, mask, shls_slice, ao_loc)
            for ia in range(mol.natm):
                p0, p1 = aoslices[ia][2:]
                # *2 for \nabla|ket> in rho1
                rho1a = numpy.einsum('xpi,pi->xp', ao[1:,:,p0:p1], ao_dm0a[:,p0:p1]) * 2
                rho1b = numpy.einsum('xpi,pi->xp', ao[1:,:,p0:p1], ao_dm0b[:,p0:p1]) * 2

                wv = u_u * rho1a + u_d * rho1b
                wv *= weight
                # aow ~ rho1 ~ d/dR1
                aow = numpy.einsum('pi,xp->xpi', ao[0], wv)
                rks_hess._d1d2_dot_(vmata[ia], mol, ao[1:4], aow, mask, ao_loc, False)

                wv = u_d * rho1a + d_d * rho1b
                wv *= weight
                aow = numpy.einsum('pi,xp->xpi', ao[0], wv)
                rks_hess._d1d2_dot_(vmatb[ia], mol, ao[1:4], aow, mask, ao_loc, False)
            ao_dm0a = ao_dm0b = aow = None

        for ia in range(mol.natm):
            p0, p1 = aoslices[ia][2:]
            vmata[ia,:,:,:,p0:p1] += ipipa[:,:,:,p0:p1]
            vmatb[ia,:,:,:,p0:p1] += ipipb[:,:,:,p0:p1]

    elif xctype == 'GGA':
        ao_deriv = 2
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rhoa = ni.eval_rho2(mol, ao[:4], mo_coeff[0], mo_occ[0], mask, xctype)
            rhob = ni.eval_rho2(mol, ao[:4], mo_coeff[1], mo_occ[1], mask, xctype)
            vxc, fxc = ni.eval_xc(mf.xc, (rhoa,rhob), 1, deriv=2)[1:3]

            wva, wvb = numint._uks_gga_wv0((rhoa,rhob), vxc, weight)
            aow = rks_grad._make_dR_dao_w(ao, wva)
            rks_hess._d1d2_dot_(ipipa, mol, aow, ao[1:4], mask, ao_loc, False)
            aow = rks_grad._make_dR_dao_w(ao, wvb)
            rks_hess._d1d2_dot_(ipipb, mol, aow, ao[1:4], mask, ao_loc, False)

            ao_dm0a = [numint._dot_ao_dm(mol, ao[i], dm0a, mask, shls_slice, ao_loc)
                       for i in range(4)]
            ao_dm0b = [numint._dot_ao_dm(mol, ao[i], dm0b, mask, shls_slice, ao_loc)
                       for i in range(4)]
            for ia in range(mol.natm):
                wva = dR_rho1a = rks_hess._make_dR_rho1(ao, ao_dm0a, ia, aoslices)
                wvb = dR_rho1b = rks_hess._make_dR_rho1(ao, ao_dm0b, ia, aoslices)
                wva[0], wvb[0] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[0],dR_rho1b[0]), vxc, fxc, weight)
                wva[1], wvb[1] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[1],dR_rho1b[1]), vxc, fxc, weight)
                wva[2], wvb[2] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[2],dR_rho1b[2]), vxc, fxc, weight)

                aow = rks_grad._make_dR_dao_w(ao, wva[0])
                rks_grad._d1_dot_(vmata[ia,0], mol, aow, ao[0], mask, ao_loc, True)
                aow = rks_grad._make_dR_dao_w(ao, wva[1])
                rks_grad._d1_dot_(vmata[ia,1], mol, aow, ao[0], mask, ao_loc, True)
                aow = rks_grad._make_dR_dao_w(ao, wva[2])
                rks_grad._d1_dot_(vmata[ia,2], mol, aow, ao[0], mask, ao_loc, True)
                aow = numpy.einsum('npi,Xnp->Xpi', ao[:4], wva)
                rks_hess._d1d2_dot_(vmata[ia], mol, ao[1:4], aow, mask, ao_loc, False)

                aow = rks_grad._make_dR_dao_w(ao, wvb[0])
                rks_grad._d1_dot_(vmatb[ia,0], mol, aow, ao[0], mask, ao_loc, True)
                aow = rks_grad._make_dR_dao_w(ao, wvb[1])
                rks_grad._d1_dot_(vmatb[ia,1], mol, aow, ao[0], mask, ao_loc, True)
                aow = rks_grad._make_dR_dao_w(ao, wvb[2])
                rks_grad._d1_dot_(vmatb[ia,2], mol, aow, ao[0], mask, ao_loc, True)
                aow = numpy.einsum('npi,Xnp->Xpi', ao[:4], wvb)
                rks_hess._d1d2_dot_(vmatb[ia], mol, ao[1:4], aow, mask, ao_loc, False)
            ao_dm0a = ao_dm0b = aow = None

        for ia in range(mol.natm):
            p0, p1 = aoslices[ia][2:]
            vmata[ia,:,:,:,p0:p1] += ipipa[:,:,:,p0:p1]
            vmata[ia,:,:,:,p0:p1] += ipipa[:,:,p0:p1].transpose(1,0,3,2)
            vmatb[ia,:,:,:,p0:p1] += ipipb[:,:,:,p0:p1]
            vmatb[ia,:,:,:,p0:p1] += ipipb[:,:,p0:p1].transpose(1,0,3,2)

    elif xctype == 'MGGA':
        raise NotImplementedError('meta-GGA')

    return vmata, vmatb
Ejemplo n.º 4
0
def _get_vxc_deriv1(hessobj, mo_coeff, mo_occ, max_memory):
    mol = hessobj.mol
    mf = hessobj.base
    if hessobj.grids is not None:
        grids = hessobj.grids
    else:
        grids = mf.grids
    if grids.coords is None:
        grids.build(with_non0tab=True)

    nao, nmo = mo_coeff[0].shape
    ni = mf._numint
    xctype = ni._xc_type(mf.xc)
    aoslices = mol.aoslice_by_atom()
    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_nr()
    dm0a, dm0b = mf.make_rdm1(mo_coeff, mo_occ)

    vmata = numpy.zeros((mol.natm,3,nao,nao))
    vmatb = numpy.zeros((mol.natm,3,nao,nao))
    max_memory = max(2000, max_memory-(vmata.size+vmatb.size)*8/1e6)
    if xctype == 'LDA':
        ao_deriv = 1
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rhoa = ni.eval_rho2(mol, ao[0], mo_coeff[0], mo_occ[0], mask, xctype)
            rhob = ni.eval_rho2(mol, ao[0], mo_coeff[1], mo_occ[1], mask, xctype)
            vxc, fxc = ni.eval_xc(mf.xc, (rhoa,rhob), 1, deriv=2)[1:3]
            vrho = vxc[0]
            u_u, u_d, d_d = fxc[0].T

            ao_dm0a = numint._dot_ao_dm(mol, ao[0], dm0a, mask, shls_slice, ao_loc)
            ao_dm0b = numint._dot_ao_dm(mol, ao[0], dm0b, mask, shls_slice, ao_loc)
            aow1a = numpy.einsum('xpi,p->xpi', ao[1:], weight*vrho[:,0])
            aow1b = numpy.einsum('xpi,p->xpi', ao[1:], weight*vrho[:,1])
            for ia in range(mol.natm):
                p0, p1 = aoslices[ia][2:]
# First order density = rho1 * 2.  *2 is not applied because + c.c. in the end
                rho1a = numpy.einsum('xpi,pi->xp', ao[1:,:,p0:p1], ao_dm0a[:,p0:p1])
                rho1b = numpy.einsum('xpi,pi->xp', ao[1:,:,p0:p1], ao_dm0b[:,p0:p1])

                wv = u_u * rho1a + u_d * rho1b
                wv *= weight
                aow = numpy.einsum('pi,xp->xpi', ao[0], wv)
                aow[:,:,p0:p1] += aow1a[:,:,p0:p1]
                rks_grad._d1_dot_(vmata[ia], mol, aow, ao[0], mask, ao_loc, True)

                wv = u_d * rho1a + d_d * rho1b
                wv *= weight
                aow = numpy.einsum('pi,xp->xpi', ao[0], wv)
                aow[:,:,p0:p1] += aow1b[:,:,p0:p1]
                rks_grad._d1_dot_(vmatb[ia], mol, aow, ao[0], mask, ao_loc, True)
            ao_dm0a = ao_dm0b = aow = aow1a = aow1b = None

        for ia in range(mol.natm):
            vmata[ia] = -vmata[ia] - vmata[ia].transpose(0,2,1)
            vmatb[ia] = -vmatb[ia] - vmatb[ia].transpose(0,2,1)

    elif xctype == 'GGA':
        ao_deriv = 2
        vipa = numpy.zeros((3,nao,nao))
        vipb = numpy.zeros((3,nao,nao))
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rhoa = ni.eval_rho2(mol, ao[:4], mo_coeff[0], mo_occ[0], mask, xctype)
            rhob = ni.eval_rho2(mol, ao[:4], mo_coeff[1], mo_occ[1], mask, xctype)
            vxc, fxc = ni.eval_xc(mf.xc, (rhoa,rhob), 1, deriv=2)[1:3]

            wva, wvb = numint._uks_gga_wv0((rhoa,rhob), vxc, weight)
            rks_grad._gga_grad_sum_(vipa, mol, ao, wva, mask, ao_loc)
            rks_grad._gga_grad_sum_(vipb, mol, ao, wvb, mask, ao_loc)

            ao_dm0a = [numint._dot_ao_dm(mol, ao[i], dm0a, mask, shls_slice, ao_loc)
                       for i in range(4)]
            ao_dm0b = [numint._dot_ao_dm(mol, ao[i], dm0b, mask, shls_slice, ao_loc)
                       for i in range(4)]
            for ia in range(mol.natm):
                wva = dR_rho1a = rks_hess._make_dR_rho1(ao, ao_dm0a, ia, aoslices)
                wvb = dR_rho1b = rks_hess._make_dR_rho1(ao, ao_dm0b, ia, aoslices)
                wva[0], wvb[0] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[0],dR_rho1b[0]), vxc, fxc, weight)
                wva[1], wvb[1] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[1],dR_rho1b[1]), vxc, fxc, weight)
                wva[2], wvb[2] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[2],dR_rho1b[2]), vxc, fxc, weight)

                aow = numpy.einsum('npi,Xnp->Xpi', ao[:4], wva)
                rks_grad._d1_dot_(vmata[ia], mol, aow, ao[0], mask, ao_loc, True)
                aow = numpy.einsum('npi,Xnp->Xpi', ao[:4], wvb)
                rks_grad._d1_dot_(vmatb[ia], mol, aow, ao[0], mask, ao_loc, True)
            ao_dm0a = ao_dm0b = aow = None

        for ia in range(mol.natm):
            p0, p1 = aoslices[ia][2:]
            vmata[ia,:,p0:p1] += vipa[:,p0:p1]
            vmata[ia] = -vmata[ia] - vmata[ia].transpose(0,2,1)
            vmatb[ia,:,p0:p1] += vipb[:,p0:p1]
            vmatb[ia] = -vmatb[ia] - vmatb[ia].transpose(0,2,1)

    elif xctype == 'MGGA':
        raise NotImplementedError('meta-GGA')

    return vmata, vmatb
Ejemplo n.º 5
0
def _get_vxc_deriv2(hessobj, mo_coeff, mo_occ, max_memory):
    mol = hessobj.mol
    mf = hessobj.base
    if hessobj.grids is not None:
        grids = hessobj.grids
    else:
        grids = mf.grids
    if grids.coords is None:
        grids.build(with_non0tab=True)

    nao, nmo = mo_coeff[0].shape
    ni = mf._numint
    xctype = ni._xc_type(mf.xc)
    aoslices = mol.aoslice_by_atom()
    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_nr()
    dm0a, dm0b = mf.make_rdm1(mo_coeff, mo_occ)

    vmata = numpy.zeros((mol.natm,3,3,nao,nao))
    vmatb = numpy.zeros((mol.natm,3,3,nao,nao))
    ipipa = numpy.zeros((3,3,nao,nao))
    ipipb = numpy.zeros((3,3,nao,nao))
    if xctype == 'LDA':
        ao_deriv = 1
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rhoa = ni.eval_rho2(mol, ao[0], mo_coeff[0], mo_occ[0], mask, xctype)
            rhob = ni.eval_rho2(mol, ao[0], mo_coeff[1], mo_occ[1], mask, xctype)
            vxc, fxc = ni.eval_xc(mf.xc, (rhoa,rhob), 1, deriv=2)[1:3]
            vrho = vxc[0]
            u_u, u_d, d_d = fxc[0].T

            aow = numpy.einsum('xpi,p->xpi', ao[1:4], weight*vrho[:,0])
            rks_hess._d1d2_dot_(ipipa, mol, aow, ao[1:4], mask, ao_loc, False)
            aow = numpy.einsum('xpi,p->xpi', ao[1:4], weight*vrho[:,1])
            rks_hess._d1d2_dot_(ipipb, mol, aow, ao[1:4], mask, ao_loc, False)

            ao_dm0a = numint._dot_ao_dm(mol, ao[0], dm0a, mask, shls_slice, ao_loc)
            ao_dm0b = numint._dot_ao_dm(mol, ao[0], dm0b, mask, shls_slice, ao_loc)
            for ia in range(mol.natm):
                p0, p1 = aoslices[ia][2:]
                # *2 for \nabla|ket> in rho1
                rho1a = numpy.einsum('xpi,pi->xp', ao[1:,:,p0:p1], ao_dm0a[:,p0:p1]) * 2
                rho1b = numpy.einsum('xpi,pi->xp', ao[1:,:,p0:p1], ao_dm0b[:,p0:p1]) * 2

                wv = u_u * rho1a + u_d * rho1b
                wv *= weight
                # aow ~ rho1 ~ d/dR1
                aow = numpy.einsum('pi,xp->xpi', ao[0], wv)
                rks_hess._d1d2_dot_(vmata[ia], mol, ao[1:4], aow, mask, ao_loc, False)

                wv = u_d * rho1a + d_d * rho1b
                wv *= weight
                aow = numpy.einsum('pi,xp->xpi', ao[0], wv)
                rks_hess._d1d2_dot_(vmatb[ia], mol, ao[1:4], aow, mask, ao_loc, False)
            ao_dm0a = ao_dm0b = aow = None

        for ia in range(mol.natm):
            p0, p1 = aoslices[ia][2:]
            vmata[ia,:,:,:,p0:p1] += ipipa[:,:,:,p0:p1]
            vmatb[ia,:,:,:,p0:p1] += ipipb[:,:,:,p0:p1]

    elif xctype == 'GGA':
        ao_deriv = 2
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rhoa = ni.eval_rho2(mol, ao[:4], mo_coeff[0], mo_occ[0], mask, xctype)
            rhob = ni.eval_rho2(mol, ao[:4], mo_coeff[1], mo_occ[1], mask, xctype)
            vxc, fxc = ni.eval_xc(mf.xc, (rhoa,rhob), 1, deriv=2)[1:3]

            wva, wvb = numint._uks_gga_wv0((rhoa,rhob), vxc, weight)
            aow = rks_grad._make_dR_dao_w(ao, wva)
            rks_hess._d1d2_dot_(ipipa, mol, aow, ao[1:4], mask, ao_loc, False)
            aow = rks_grad._make_dR_dao_w(ao, wvb)
            rks_hess._d1d2_dot_(ipipb, mol, aow, ao[1:4], mask, ao_loc, False)

            ao_dm0a = [numint._dot_ao_dm(mol, ao[i], dm0a, mask, shls_slice, ao_loc)
                       for i in range(4)]
            ao_dm0b = [numint._dot_ao_dm(mol, ao[i], dm0b, mask, shls_slice, ao_loc)
                       for i in range(4)]
            for ia in range(mol.natm):
                wva = dR_rho1a = rks_hess._make_dR_rho1(ao, ao_dm0a, ia, aoslices)
                wvb = dR_rho1b = rks_hess._make_dR_rho1(ao, ao_dm0b, ia, aoslices)
                wva[0], wvb[0] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[0],dR_rho1b[0]), vxc, fxc, weight)
                wva[1], wvb[1] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[1],dR_rho1b[1]), vxc, fxc, weight)
                wva[2], wvb[2] = numint._uks_gga_wv1((rhoa,rhob), (dR_rho1a[2],dR_rho1b[2]), vxc, fxc, weight)

                aow = rks_grad._make_dR_dao_w(ao, wva[0])
                rks_grad._d1_dot_(vmata[ia,0], mol, aow, ao[0], mask, ao_loc, True)
                aow = rks_grad._make_dR_dao_w(ao, wva[1])
                rks_grad._d1_dot_(vmata[ia,1], mol, aow, ao[0], mask, ao_loc, True)
                aow = rks_grad._make_dR_dao_w(ao, wva[2])
                rks_grad._d1_dot_(vmata[ia,2], mol, aow, ao[0], mask, ao_loc, True)
                aow = numpy.einsum('npi,Xnp->Xpi', ao[:4], wva)
                rks_hess._d1d2_dot_(vmata[ia], mol, ao[1:4], aow, mask, ao_loc, False)

                aow = rks_grad._make_dR_dao_w(ao, wvb[0])
                rks_grad._d1_dot_(vmatb[ia,0], mol, aow, ao[0], mask, ao_loc, True)
                aow = rks_grad._make_dR_dao_w(ao, wvb[1])
                rks_grad._d1_dot_(vmatb[ia,1], mol, aow, ao[0], mask, ao_loc, True)
                aow = rks_grad._make_dR_dao_w(ao, wvb[2])
                rks_grad._d1_dot_(vmatb[ia,2], mol, aow, ao[0], mask, ao_loc, True)
                aow = numpy.einsum('npi,Xnp->Xpi', ao[:4], wvb)
                rks_hess._d1d2_dot_(vmatb[ia], mol, ao[1:4], aow, mask, ao_loc, False)
            ao_dm0a = ao_dm0b = aow = None

        for ia in range(mol.natm):
            p0, p1 = aoslices[ia][2:]
            vmata[ia,:,:,:,p0:p1] += ipipa[:,:,:,p0:p1]
            vmata[ia,:,:,:,p0:p1] += ipipa[:,:,p0:p1].transpose(1,0,3,2)
            vmatb[ia,:,:,:,p0:p1] += ipipb[:,:,:,p0:p1]
            vmatb[ia,:,:,:,p0:p1] += ipipb[:,:,p0:p1].transpose(1,0,3,2)

    elif xctype == 'MGGA':
        raise NotImplementedError('meta-GGA')

    return vmata, vmatb
Ejemplo n.º 6
0
def _contract_xc_kernel(td_grad, xc_code, dmvo, dmoo=None, with_vxc=True,
                        with_kxc=True, max_memory=2000):
    mol = td_grad.mol
    mf = td_grad.base._scf
    grids = mf.grids

    ni = mf._numint
    xctype = ni._xc_type(xc_code)

    mo_coeff = mf.mo_coeff
    mo_occ = mf.mo_occ
    nao = mo_coeff[0].shape[0]
    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_nr()

    # dmvo ~ reduce(numpy.dot, (orbv, Xai, orbo.T))
    dmvo = [(dmvo[0] + dmvo[0].T) * .5, # because K_{ia,jb} == K_{ia,jb}
            (dmvo[1] + dmvo[1].T) * .5]

    f1vo = numpy.zeros((2,4,nao,nao))
    deriv = 2
    if dmoo is not None:
        f1oo = numpy.zeros((2,4,nao,nao))
    else:
        f1oo = None
    if with_vxc:
        v1ao = numpy.zeros((2,4,nao,nao))
    else:
        v1ao = None
    if with_kxc:
        k1ao = numpy.zeros((2,4,nao,nao))
        deriv = 3
    else:
        k1ao = None

    if xctype == 'LDA':
        ao_deriv = 1
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rho = (ni.eval_rho2(mol, ao[0], mo_coeff[0], mo_occ[0], mask, 'LDA'),
                   ni.eval_rho2(mol, ao[0], mo_coeff[1], mo_occ[1], mask, 'LDA'))
            vxc, fxc, kxc = ni.eval_xc(xc_code, rho, 1, deriv=deriv)[1:]

            u_u, u_d, d_d = fxc[0].T * weight
            rho1a = ni.eval_rho(mol, ao[0], dmvo[0], mask, 'LDA')
            rho1b = ni.eval_rho(mol, ao[0], dmvo[1], mask, 'LDA')
            aow = (numpy.einsum('pi,p,p->pi', ao[0], u_u, rho1a) +
                   numpy.einsum('pi,p,p->pi', ao[0], u_d, rho1b),
                   numpy.einsum('pi,p,p->pi', ao[0], u_d, rho1a) +
                   numpy.einsum('pi,p,p->pi', ao[0], d_d, rho1b))
            for k in range(4):
                f1vo[0,k] += numint._dot_ao_ao(mol, ao[k], aow[0], mask, shls_slice, ao_loc)
                f1vo[1,k] += numint._dot_ao_ao(mol, ao[k], aow[1], mask, shls_slice, ao_loc)
            if dmoo is not None:
                rho2a = ni.eval_rho(mol, ao[0], dmoo[0], mask, 'LDA')
                rho2b = ni.eval_rho(mol, ao[0], dmoo[1], mask, 'LDA')
                aow = (numpy.einsum('pi,p,p->pi', ao[0], u_u, rho2a) +
                       numpy.einsum('pi,p,p->pi', ao[0], u_d, rho2b),
                       numpy.einsum('pi,p,p->pi', ao[0], u_d, rho2a) +
                       numpy.einsum('pi,p,p->pi', ao[0], d_d, rho2b))
                for k in range(4):
                    f1oo[0,k] += numint._dot_ao_ao(mol, ao[k], aow[0], mask, shls_slice, ao_loc)
                    f1oo[1,k] += numint._dot_ao_ao(mol, ao[k], aow[1], mask, shls_slice, ao_loc)
            if with_vxc:
                vrho = vxc[0].T * weight
                aow = (numpy.einsum('pi,p->pi', ao[0], vrho[0]),
                       numpy.einsum('pi,p->pi', ao[0], vrho[1]))
                for k in range(4):
                    v1ao[0,k] += numint._dot_ao_ao(mol, ao[k], aow[0], mask, shls_slice, ao_loc)
                    v1ao[1,k] += numint._dot_ao_ao(mol, ao[k], aow[1], mask, shls_slice, ao_loc)
            if with_kxc:
                u_u_u, u_u_d, u_d_d, d_d_d = kxc[0].T * weight
                aow = (numpy.einsum('pi,p,p,p->pi', ao[0], u_u_u, rho1a, rho1a) +
                       numpy.einsum('pi,p,p,p->pi', ao[0], u_u_d, rho1a, rho1b)*2 +
                       numpy.einsum('pi,p,p,p->pi', ao[0], u_d_d, rho1b, rho1b),
                       numpy.einsum('pi,p,p,p->pi', ao[0], u_u_d, rho1a, rho1a) +
                       numpy.einsum('pi,p,p,p->pi', ao[0], u_d_d, rho1a, rho1b)*2 +
                       numpy.einsum('pi,p,p,p->pi', ao[0], d_d_d, rho1b, rho1b))
                for k in range(4):
                    k1ao[0,k] += numint._dot_ao_ao(mol, ao[k], aow[0], mask, shls_slice, ao_loc)
                    k1ao[1,k] += numint._dot_ao_ao(mol, ao[k], aow[1], mask, shls_slice, ao_loc)
            vxc = fxc = kxc = aow = rho = rho1 = rho2 = None

    elif xctype == 'GGA':
        def gga_sum_(vmat, ao, wv, mask):
            aow  = numpy.einsum('pi,p->pi', ao[0], wv[0])
            aow += numpy.einsum('npi,np->pi', ao[1:4], wv[1:])
            tmp = numint._dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
            vmat[0] += tmp + tmp.T
            rks_grad._gga_grad_sum_(vmat[1:], mol, ao, wv, mask, ao_loc)
        ao_deriv = 2
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rho = (ni.eval_rho2(mol, ao, mo_coeff[0], mo_occ[0], mask, 'GGA'),
                   ni.eval_rho2(mol, ao, mo_coeff[1], mo_occ[1], mask, 'GGA'))
            vxc, fxc, kxc = ni.eval_xc(xc_code, rho, 1, deriv=deriv)[1:]

            rho1 = (ni.eval_rho(mol, ao, dmvo[0], mask, 'GGA'),
                    ni.eval_rho(mol, ao, dmvo[1], mask, 'GGA'))
            wv = numint._uks_gga_wv1(rho, rho1, vxc, fxc, weight)
            gga_sum_(f1vo[0], ao, wv[0], mask)
            gga_sum_(f1vo[1], ao, wv[1], mask)

            if dmoo is not None:
                rho2 = (ni.eval_rho(mol, ao, dmoo[0], mask, 'GGA'),
                        ni.eval_rho(mol, ao, dmoo[1], mask, 'GGA'))
                wv = numint._uks_gga_wv1(rho, rho2, vxc, fxc, weight)
                gga_sum_(f1oo[0], ao, wv[0], mask)
                gga_sum_(f1oo[1], ao, wv[1], mask)
            if with_vxc:
                wv = numint._uks_gga_wv0(rho, vxc, weight)
                gga_sum_(v1ao[0], ao, wv[0], mask)
                gga_sum_(v1ao[1], ao, wv[1], mask)
            if with_kxc:
                wv = numint._uks_gga_wv2(rho, rho1, fxc, kxc, weight)
                gga_sum_(k1ao[0], ao, wv[0], mask)
                gga_sum_(k1ao[1], ao, wv[1], mask)
            vxc = fxc = kxc = rho = rho1 = None

    else:
        raise NotImplementedError('meta-GGA')

    f1vo[:,1:] *= -1
    if f1oo is not None: f1oo[:,1:] *= -1
    if v1ao is not None: v1ao[:,1:] *= -1
    if k1ao is not None: k1ao[:,1:] *= -1
    return f1vo, f1oo, v1ao, k1ao
Ejemplo n.º 7
0
Archivo: tduks.py Proyecto: tmash/pyscf
def _contract_xc_kernel(td_grad,
                        xc_code,
                        dmvo,
                        dmoo=None,
                        with_vxc=True,
                        with_kxc=True,
                        max_memory=2000):
    mol = td_grad.mol
    mf = td_grad.base._scf
    grids = mf.grids

    ni = mf._numint
    xctype = ni._xc_type(xc_code)

    mo_coeff = mf.mo_coeff
    mo_occ = mf.mo_occ
    nao = mo_coeff[0].shape[0]
    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_nr()

    # dmvo ~ reduce(numpy.dot, (orbv, Xai, orbo.T))
    dmvo = [
        (dmvo[0] + dmvo[0].T) * .5,  # because K_{ia,jb} == K_{ia,jb}
        (dmvo[1] + dmvo[1].T) * .5
    ]

    f1vo = numpy.zeros((2, 4, nao, nao))
    deriv = 2
    if dmoo is not None:
        f1oo = numpy.zeros((2, 4, nao, nao))
    else:
        f1oo = None
    if with_vxc:
        v1ao = numpy.zeros((2, 4, nao, nao))
    else:
        v1ao = None
    if with_kxc:
        k1ao = numpy.zeros((2, 4, nao, nao))
        deriv = 3
    else:
        k1ao = None

    if xctype == 'LDA':
        ao_deriv = 1
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rho = (ni.eval_rho2(mol, ao[0], mo_coeff[0], mo_occ[0], mask,
                                'LDA'),
                   ni.eval_rho2(mol, ao[0], mo_coeff[1], mo_occ[1], mask,
                                'LDA'))
            vxc, fxc, kxc = ni.eval_xc(xc_code, rho, 1, deriv=deriv)[1:]

            u_u, u_d, d_d = fxc[0].T * weight
            rho1a = ni.eval_rho(mol, ao[0], dmvo[0], mask, 'LDA')
            rho1b = ni.eval_rho(mol, ao[0], dmvo[1], mask, 'LDA')
            aow = (numpy.einsum('pi,p,p->pi', ao[0], u_u, rho1a) +
                   numpy.einsum('pi,p,p->pi', ao[0], u_d, rho1b),
                   numpy.einsum('pi,p,p->pi', ao[0], u_d, rho1a) +
                   numpy.einsum('pi,p,p->pi', ao[0], d_d, rho1b))
            for k in range(4):
                f1vo[0, k] += numint._dot_ao_ao(mol, ao[k], aow[0], mask,
                                                shls_slice, ao_loc)
                f1vo[1, k] += numint._dot_ao_ao(mol, ao[k], aow[1], mask,
                                                shls_slice, ao_loc)
            if dmoo is not None:
                rho2a = ni.eval_rho(mol, ao[0], dmoo[0], mask, 'LDA')
                rho2b = ni.eval_rho(mol, ao[0], dmoo[1], mask, 'LDA')
                aow = (numpy.einsum('pi,p,p->pi', ao[0], u_u, rho2a) +
                       numpy.einsum('pi,p,p->pi', ao[0], u_d, rho2b),
                       numpy.einsum('pi,p,p->pi', ao[0], u_d, rho2a) +
                       numpy.einsum('pi,p,p->pi', ao[0], d_d, rho2b))
                for k in range(4):
                    f1oo[0, k] += numint._dot_ao_ao(mol, ao[k], aow[0], mask,
                                                    shls_slice, ao_loc)
                    f1oo[1, k] += numint._dot_ao_ao(mol, ao[k], aow[1], mask,
                                                    shls_slice, ao_loc)
            if with_vxc:
                vrho = vxc[0].T * weight
                aow = (numpy.einsum('pi,p->pi', ao[0], vrho[0]),
                       numpy.einsum('pi,p->pi', ao[0], vrho[1]))
                for k in range(4):
                    v1ao[0, k] += numint._dot_ao_ao(mol, ao[k], aow[0], mask,
                                                    shls_slice, ao_loc)
                    v1ao[1, k] += numint._dot_ao_ao(mol, ao[k], aow[1], mask,
                                                    shls_slice, ao_loc)
            if with_kxc:
                u_u_u, u_u_d, u_d_d, d_d_d = kxc[0].T * weight
                aow = (
                    numpy.einsum('pi,p,p,p->pi', ao[0], u_u_u, rho1a, rho1a) +
                    numpy.einsum('pi,p,p,p->pi', ao[0], u_u_d, rho1a, rho1b) *
                    2 +
                    numpy.einsum('pi,p,p,p->pi', ao[0], u_d_d, rho1b, rho1b),
                    numpy.einsum('pi,p,p,p->pi', ao[0], u_u_d, rho1a, rho1a) +
                    numpy.einsum('pi,p,p,p->pi', ao[0], u_d_d, rho1a, rho1b) *
                    2 +
                    numpy.einsum('pi,p,p,p->pi', ao[0], d_d_d, rho1b, rho1b))
                for k in range(4):
                    k1ao[0, k] += numint._dot_ao_ao(mol, ao[k], aow[0], mask,
                                                    shls_slice, ao_loc)
                    k1ao[1, k] += numint._dot_ao_ao(mol, ao[k], aow[1], mask,
                                                    shls_slice, ao_loc)
            vxc = fxc = kxc = aow = rho = rho1 = rho2 = None

    elif xctype == 'GGA':

        def gga_sum_(vmat, ao, wv, mask):
            aow = numpy.einsum('pi,p->pi', ao[0], wv[0])
            aow += numpy.einsum('npi,np->pi', ao[1:4], wv[1:])
            tmp = numint._dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
            vmat[0] += tmp + tmp.T
            rks_grad._gga_grad_sum_(vmat[1:], mol, ao, wv, mask, ao_loc)

        ao_deriv = 2
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rho = (ni.eval_rho2(mol, ao, mo_coeff[0], mo_occ[0], mask, 'GGA'),
                   ni.eval_rho2(mol, ao, mo_coeff[1], mo_occ[1], mask, 'GGA'))
            vxc, fxc, kxc = ni.eval_xc(xc_code, rho, 1, deriv=deriv)[1:]

            rho1 = (ni.eval_rho(mol, ao, dmvo[0], mask, 'GGA'),
                    ni.eval_rho(mol, ao, dmvo[1], mask, 'GGA'))
            wv = numint._uks_gga_wv1(rho, rho1, vxc, fxc, weight)
            gga_sum_(f1vo[0], ao, wv[0], mask)
            gga_sum_(f1vo[1], ao, wv[1], mask)

            if dmoo is not None:
                rho2 = (ni.eval_rho(mol, ao, dmoo[0], mask, 'GGA'),
                        ni.eval_rho(mol, ao, dmoo[1], mask, 'GGA'))
                wv = numint._uks_gga_wv1(rho, rho2, vxc, fxc, weight)
                gga_sum_(f1oo[0], ao, wv[0], mask)
                gga_sum_(f1oo[1], ao, wv[1], mask)
            if with_vxc:
                wv = numint._uks_gga_wv0(rho, vxc, weight)
                gga_sum_(v1ao[0], ao, wv[0], mask)
                gga_sum_(v1ao[1], ao, wv[1], mask)
            if with_kxc:
                wv = numint._uks_gga_wv2(rho, rho1, fxc, kxc, weight)
                gga_sum_(k1ao[0], ao, wv[0], mask)
                gga_sum_(k1ao[1], ao, wv[1], mask)
            vxc = fxc = kxc = rho = rho1 = None

    else:
        raise NotImplementedError('meta-GGA')

    f1vo[:, 1:] *= -1
    if f1oo is not None: f1oo[:, 1:] *= -1
    if v1ao is not None: v1ao[:, 1:] *= -1
    if k1ao is not None: k1ao[:, 1:] *= -1
    return f1vo, f1oo, v1ao, k1ao
Ejemplo n.º 8
0
def _contract_xc_kernel(td_grad,
                        xc_code,
                        dmvo,
                        dmoo=None,
                        with_vxc=True,
                        with_kxc=True,
                        max_memory=2000):
    mol = td_grad.mol
    mf = td_grad.base._scf
    grids = mf.grids

    ni = mf._numint
    xctype = ni._xc_type(xc_code)

    mo_coeff = mf.mo_coeff
    mo_occ = mf.mo_occ
    nao = mo_coeff[0].shape[0]
    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_nr()

    # dmvo ~ reduce(numpy.dot, (orbv, Xai, orbo.T))
    dmvo = [
        (dmvo[0] + dmvo[0].T) * .5,  # because K_{ia,jb} == K_{ia,jb}
        (dmvo[1] + dmvo[1].T) * .5
    ]

    f1vo = numpy.zeros((2, 4, nao, nao))
    deriv = 2
    if dmoo is not None:
        f1oo = numpy.zeros((2, 4, nao, nao))
    else:
        f1oo = None
    if with_vxc:
        v1ao = numpy.zeros((2, 4, nao, nao))
    else:
        v1ao = None
    if with_kxc:
        k1ao = numpy.zeros((2, 4, nao, nao))
        deriv = 3
    else:
        k1ao = None

    if xctype == 'LDA':

        def lda_sum_(vmat, ao, wv, mask):
            aow = numint._scale_ao(ao[0], wv)
            for k in range(4):
                vmat[k] += numint._dot_ao_ao(mol, ao[k], aow, mask, shls_slice,
                                             ao_loc)

        ao_deriv = 1
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rho = (ni.eval_rho2(mol, ao[0], mo_coeff[0], mo_occ[0], mask,
                                xctype),
                   ni.eval_rho2(mol, ao[0], mo_coeff[1], mo_occ[1], mask,
                                xctype))
            vxc, fxc, kxc = ni.eval_xc(xc_code, rho, 1, deriv=deriv)[1:]

            u_u, u_d, d_d = fxc[0].T * weight
            rho1a = ni.eval_rho(mol, ao[0], dmvo[0], mask, xctype)
            rho1b = ni.eval_rho(mol, ao[0], dmvo[1], mask, xctype)
            lda_sum_(f1vo[0], ao, u_u * rho1a + u_d * rho1b, mask)
            lda_sum_(f1vo[1], ao, u_d * rho1a + d_d * rho1b, mask)
            if dmoo is not None:
                rho2a = ni.eval_rho(mol, ao[0], dmoo[0], mask, xctype)
                rho2b = ni.eval_rho(mol, ao[0], dmoo[1], mask, xctype)
                lda_sum_(f1oo[0], ao, u_u * rho2a + u_d * rho2b, mask)
                lda_sum_(f1oo[1], ao, u_d * rho2a + d_d * rho2b, mask)
            if with_vxc:
                vrho = vxc[0].T * weight
                lda_sum_(v1ao[0], ao, vrho[0], mask)
                lda_sum_(v1ao[1], ao, vrho[1], mask)
            if with_kxc:
                u_u_u, u_u_d, u_d_d, d_d_d = kxc[0].T * weight
                lda_sum_(
                    k1ao[0], ao, u_u_u * rho1a * rho1a +
                    u_u_d * rho1a * rho1b * 2 + u_d_d * rho1b * rho1b, mask)
                lda_sum_(
                    k1ao[1], ao, u_u_d * rho1a * rho1a +
                    u_d_d * rho1a * rho1b * 2 + d_d_d * rho1b * rho1b, mask)

    elif xctype == 'GGA':

        def gga_sum_(vmat, ao, wv, mask):
            aow = numint._scale_ao(ao[:4], wv[:4])
            tmp = numint._dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
            vmat[0] += tmp + tmp.T
            rks_grad._gga_grad_sum_(vmat[1:], mol, ao, wv, mask, ao_loc)

        ao_deriv = 2
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rho = (ni.eval_rho2(mol, ao, mo_coeff[0], mo_occ[0], mask, xctype),
                   ni.eval_rho2(mol, ao, mo_coeff[1], mo_occ[1], mask, xctype))
            vxc, fxc, kxc = ni.eval_xc(xc_code, rho, 1, deriv=deriv)[1:]

            rho1 = (ni.eval_rho(mol, ao, dmvo[0], mask, xctype),
                    ni.eval_rho(mol, ao, dmvo[1], mask, xctype))
            wv = numint._uks_gga_wv1(rho, rho1, vxc, fxc, weight)
            gga_sum_(f1vo[0], ao, wv[0], mask)
            gga_sum_(f1vo[1], ao, wv[1], mask)

            if dmoo is not None:
                rho2 = (ni.eval_rho(mol, ao, dmoo[0], mask, xctype),
                        ni.eval_rho(mol, ao, dmoo[1], mask, xctype))
                wv = numint._uks_gga_wv1(rho, rho2, vxc, fxc, weight)
                gga_sum_(f1oo[0], ao, wv[0], mask)
                gga_sum_(f1oo[1], ao, wv[1], mask)
            if with_vxc:
                wv = numint._uks_gga_wv0(rho, vxc, weight)
                gga_sum_(v1ao[0], ao, wv[0], mask)
                gga_sum_(v1ao[1], ao, wv[1], mask)
            if with_kxc:
                wv = numint._uks_gga_wv2(rho, rho1, fxc, kxc, weight)
                gga_sum_(k1ao[0], ao, wv[0], mask)
                gga_sum_(k1ao[1], ao, wv[1], mask)
            vxc = fxc = kxc = rho = rho1 = None

    elif xctype == 'MGGA':
        logger.warn(mol, 'More tests are needed for TD-MGGA')

        def mgga_sum_(vmat, ao, wv, mask):
            aow = numint._scale_ao(ao[:4], wv[:4])
            tmp = numint._dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
            aow = numint._scale_ao(ao[1], wv[5], aow)
            tmp += numint._dot_ao_ao(mol, ao[1], aow, mask, shls_slice, ao_loc)
            aow = numint._scale_ao(ao[2], wv[5], aow)
            tmp += numint._dot_ao_ao(mol, ao[2], aow, mask, shls_slice, ao_loc)
            aow = numint._scale_ao(ao[3], wv[5], aow)
            tmp += numint._dot_ao_ao(mol, ao[3], aow, mask, shls_slice, ao_loc)
            vmat[0] += tmp + tmp.T

            rks_grad._gga_grad_sum_(vmat[1:], mol, ao, wv[:4], mask, ao_loc)
            rks_grad._tau_grad_dot_(vmat[1:], mol, ao, wv[5] * 2, mask, ao_loc,
                                    True)

        ao_deriv = 2
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rho = (ni.eval_rho2(mol, ao, mo_coeff[0], mo_occ[0], mask, xctype),
                   ni.eval_rho2(mol, ao, mo_coeff[1], mo_occ[1], mask, xctype))
            vxc, fxc, kxc = ni.eval_xc(xc_code, rho, 1, deriv=deriv)[1:]

            rho1 = (ni.eval_rho(mol, ao, dmvo[0], mask, xctype),
                    ni.eval_rho(mol, ao, dmvo[1], mask, xctype))
            wv = numint._uks_mgga_wv1(rho, rho1, vxc, fxc, weight)
            mgga_sum_(f1vo[0], ao, wv[0], mask)
            mgga_sum_(f1vo[1], ao, wv[1], mask)

            if dmoo is not None:
                rho2 = (ni.eval_rho(mol, ao, dmoo[0], mask, xctype),
                        ni.eval_rho(mol, ao, dmoo[1], mask, xctype))
                wv = numint._uks_mgga_wv1(rho, rho2, vxc, fxc, weight)
                mgga_sum_(f1oo[0], ao, wv[0], mask)
                mgga_sum_(f1oo[1], ao, wv[1], mask)
            if with_vxc:
                wv = numint._uks_mgga_wv0(rho, vxc, weight)
                mgga_sum_(v1ao[0], ao, wv[0], mask)
                mgga_sum_(v1ao[1], ao, wv[1], mask)
            if with_kxc:
                wv = numint._uks_mgga_wv2(rho, rho1, fxc, kxc, weight)
                mgga_sum_(k1ao[0], ao, wv[0], mask)
                mgga_sum_(k1ao[1], ao, wv[1], mask)
            vxc = fxc = kxc = rho = rho1 = None

    elif xctype == 'HF':
        pass
    else:
        raise NotImplementedError(f'td-uks for functional {xc_code}')

    f1vo[:, 1:] *= -1
    if f1oo is not None: f1oo[:, 1:] *= -1
    if v1ao is not None: v1ao[:, 1:] *= -1
    if k1ao is not None: k1ao[:, 1:] *= -1
    return f1vo, f1oo, v1ao, k1ao
Ejemplo n.º 9
0
def _get_vxc_deriv2(hessobj, mo_coeff, mo_occ, max_memory):
    mol = hessobj.mol
    mf = hessobj.base
    if hessobj.grids is not None:
        grids = hessobj.grids
    else:
        grids = mf.grids
    if grids.coords is None:
        grids.build(with_non0tab=True)

    nao, nmo = mo_coeff[0].shape
    ni = mf._numint
    xctype = ni._xc_type(mf.xc)
    aoslices = mol.aoslice_by_atom()
    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_nr()
    dm0a, dm0b = mf.make_rdm1(mo_coeff, mo_occ)

    vmata = numpy.zeros((mol.natm, 3, 3, nao, nao))
    vmatb = numpy.zeros((mol.natm, 3, 3, nao, nao))
    ipipa = numpy.zeros((3, 3, nao, nao))
    ipipb = numpy.zeros((3, 3, nao, nao))
    if xctype == 'LDA':
        ao_deriv = 1
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rhoa = ni.eval_rho2(mol, ao[0], mo_coeff[0], mo_occ[0], mask,
                                xctype)
            rhob = ni.eval_rho2(mol, ao[0], mo_coeff[1], mo_occ[1], mask,
                                xctype)
            vxc, fxc = ni.eval_xc(mf.xc, (rhoa, rhob), 1, deriv=2)[1:3]
            vrho = vxc[0]
            u_u, u_d, d_d = fxc[0].T

            aow = numpy.einsum('xpi,p->xpi', ao[1:4], weight * vrho[:, 0])
            rks_hess._d1d2_dot_(ipipa, mol, aow, ao[1:4], mask, ao_loc, False)
            aow = numpy.einsum('xpi,p->xpi', ao[1:4], weight * vrho[:, 1])
            rks_hess._d1d2_dot_(ipipb, mol, aow, ao[1:4], mask, ao_loc, False)

            ao_dm0a = numint._dot_ao_dm(mol, ao[0], dm0a, mask, shls_slice,
                                        ao_loc)
            ao_dm0b = numint._dot_ao_dm(mol, ao[0], dm0b, mask, shls_slice,
                                        ao_loc)
            for ia in range(mol.natm):
                p0, p1 = aoslices[ia][2:]
                # *2 for \nabla|ket> in rho1
                rho1a = numpy.einsum('xpi,pi->xp', ao[1:, :, p0:p1],
                                     ao_dm0a[:, p0:p1]) * 2
                rho1b = numpy.einsum('xpi,pi->xp', ao[1:, :, p0:p1],
                                     ao_dm0b[:, p0:p1]) * 2

                wv = u_u * rho1a + u_d * rho1b
                wv *= weight
                # aow ~ rho1 ~ d/dR1
                aow = numpy.einsum('pi,xp->xpi', ao[0], wv)
                rks_hess._d1d2_dot_(vmata[ia], mol, ao[1:4], aow, mask, ao_loc,
                                    False)

                wv = u_d * rho1a + d_d * rho1b
                wv *= weight
                aow = numpy.einsum('pi,xp->xpi', ao[0], wv)
                rks_hess._d1d2_dot_(vmatb[ia], mol, ao[1:4], aow, mask, ao_loc,
                                    False)
            ao_dm0a = ao_dm0b = aow = None

        for ia in range(mol.natm):
            p0, p1 = aoslices[ia][2:]
            vmata[ia, :, :, :, p0:p1] += ipipa[:, :, :, p0:p1]
            vmatb[ia, :, :, :, p0:p1] += ipipb[:, :, :, p0:p1]

    elif xctype == 'GGA':
        ao_deriv = 2
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rhoa = ni.eval_rho2(mol, ao[:4], mo_coeff[0], mo_occ[0], mask,
                                xctype)
            rhob = ni.eval_rho2(mol, ao[:4], mo_coeff[1], mo_occ[1], mask,
                                xctype)
            vxc, fxc = ni.eval_xc(mf.xc, (rhoa, rhob), 1, deriv=2)[1:3]

            wva, wvb = numint._uks_gga_wv0((rhoa, rhob), vxc, weight)
            aow = rks_grad._make_dR_dao_w(ao, wva)
            rks_hess._d1d2_dot_(ipipa, mol, aow, ao[1:4], mask, ao_loc, False)
            aow = rks_grad._make_dR_dao_w(ao, wvb)
            rks_hess._d1d2_dot_(ipipb, mol, aow, ao[1:4], mask, ao_loc, False)

            ao_dm0a = [
                numint._dot_ao_dm(mol, ao[i], dm0a, mask, shls_slice, ao_loc)
                for i in range(4)
            ]
            ao_dm0b = [
                numint._dot_ao_dm(mol, ao[i], dm0b, mask, shls_slice, ao_loc)
                for i in range(4)
            ]
            for ia in range(mol.natm):
                wva = dR_rho1a = rks_hess._make_dR_rho1(
                    ao, ao_dm0a, ia, aoslices)
                wvb = dR_rho1b = rks_hess._make_dR_rho1(
                    ao, ao_dm0b, ia, aoslices)
                wva[0], wvb[0] = numint._uks_gga_wv1(
                    (rhoa, rhob), (dR_rho1a[0], dR_rho1b[0]), vxc, fxc, weight)
                wva[1], wvb[1] = numint._uks_gga_wv1(
                    (rhoa, rhob), (dR_rho1a[1], dR_rho1b[1]), vxc, fxc, weight)
                wva[2], wvb[2] = numint._uks_gga_wv1(
                    (rhoa, rhob), (dR_rho1a[2], dR_rho1b[2]), vxc, fxc, weight)

                aow = rks_grad._make_dR_dao_w(ao, wva[0])
                rks_grad._d1_dot_(vmata[ia, 0], mol, aow, ao[0], mask, ao_loc,
                                  True)
                aow = rks_grad._make_dR_dao_w(ao, wva[1])
                rks_grad._d1_dot_(vmata[ia, 1], mol, aow, ao[0], mask, ao_loc,
                                  True)
                aow = rks_grad._make_dR_dao_w(ao, wva[2])
                rks_grad._d1_dot_(vmata[ia, 2], mol, aow, ao[0], mask, ao_loc,
                                  True)
                aow = [numint._scale_ao(ao[:4], wva[i, :4]) for i in range(3)]
                rks_hess._d1d2_dot_(vmata[ia], mol, ao[1:4], aow, mask, ao_loc,
                                    False)

                aow = rks_grad._make_dR_dao_w(ao, wvb[0])
                rks_grad._d1_dot_(vmatb[ia, 0], mol, aow, ao[0], mask, ao_loc,
                                  True)
                aow = rks_grad._make_dR_dao_w(ao, wvb[1])
                rks_grad._d1_dot_(vmatb[ia, 1], mol, aow, ao[0], mask, ao_loc,
                                  True)
                aow = rks_grad._make_dR_dao_w(ao, wvb[2])
                rks_grad._d1_dot_(vmatb[ia, 2], mol, aow, ao[0], mask, ao_loc,
                                  True)
                aow = [numint._scale_ao(ao[:4], wvb[i, :4]) for i in range(3)]
                rks_hess._d1d2_dot_(vmatb[ia], mol, ao[1:4], aow, mask, ao_loc,
                                    False)
            ao_dm0a = ao_dm0b = aow = None

        for ia in range(mol.natm):
            p0, p1 = aoslices[ia][2:]
            vmata[ia, :, :, :, p0:p1] += ipipa[:, :, :, p0:p1]
            vmata[ia, :, :, :, p0:p1] += ipipa[:, :,
                                               p0:p1].transpose(1, 0, 3, 2)
            vmatb[ia, :, :, :, p0:p1] += ipipb[:, :, :, p0:p1]
            vmatb[ia, :, :, :, p0:p1] += ipipb[:, :,
                                               p0:p1].transpose(1, 0, 3, 2)

    elif xctype == 'MGGA':
        XX, XY, XZ = 4, 5, 6
        YX, YY, YZ = 5, 7, 8
        ZX, ZY, ZZ = 6, 8, 9
        ao_deriv = 2
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rhoa = ni.eval_rho2(mol, ao[:10], mo_coeff[0], mo_occ[0], mask,
                                xctype)
            rhob = ni.eval_rho2(mol, ao[:10], mo_coeff[1], mo_occ[1], mask,
                                xctype)
            vxc, fxc = ni.eval_xc(mf.xc, (rhoa, rhob), 1, deriv=2)[1:3]

            wva, wvb = numint._uks_mgga_wv0((rhoa, rhob), vxc, weight)
            aow = rks_grad._make_dR_dao_w(ao, wva)
            rks_hess._d1d2_dot_(ipipa, mol, aow, ao[1:4], mask, ao_loc, False)
            aow = rks_grad._make_dR_dao_w(ao, wvb)
            rks_hess._d1d2_dot_(ipipb, mol, aow, ao[1:4], mask, ao_loc, False)

            aow = [numint._scale_ao(ao[i], wva[5]) for i in range(4, 10)]
            rks_hess._d1d2_dot_(ipipa, mol, [aow[0], aow[1], aow[2]],
                                [ao[XX], ao[XY], ao[XZ]], mask, ao_loc, False)
            rks_hess._d1d2_dot_(ipipa, mol, [aow[1], aow[3], aow[4]],
                                [ao[YX], ao[YY], ao[YZ]], mask, ao_loc, False)
            rks_hess._d1d2_dot_(ipipa, mol, [aow[2], aow[4], aow[5]],
                                [ao[ZX], ao[ZY], ao[ZZ]], mask, ao_loc, False)
            aow = [numint._scale_ao(ao[i], wvb[5]) for i in range(4, 10)]
            rks_hess._d1d2_dot_(ipipb, mol, [aow[0], aow[1], aow[2]],
                                [ao[XX], ao[XY], ao[XZ]], mask, ao_loc, False)
            rks_hess._d1d2_dot_(ipipb, mol, [aow[1], aow[3], aow[4]],
                                [ao[YX], ao[YY], ao[YZ]], mask, ao_loc, False)
            rks_hess._d1d2_dot_(ipipb, mol, [aow[2], aow[4], aow[5]],
                                [ao[ZX], ao[ZY], ao[ZZ]], mask, ao_loc, False)

            ao_dm0a = [
                numint._dot_ao_dm(mol, ao[i], dm0a, mask, shls_slice, ao_loc)
                for i in range(4)
            ]
            ao_dm0b = [
                numint._dot_ao_dm(mol, ao[i], dm0b, mask, shls_slice, ao_loc)
                for i in range(4)
            ]
            for ia in range(mol.natm):
                wva = dR_rho1a = rks_hess._make_dR_rho1(
                    ao, ao_dm0a, ia, aoslices, xctype)
                wvb = dR_rho1b = rks_hess._make_dR_rho1(
                    ao, ao_dm0b, ia, aoslices, xctype)
                wva[0], wvb[0] = numint._uks_mgga_wv1(
                    (rhoa, rhob), (dR_rho1a[0], dR_rho1b[0]), vxc, fxc, weight)
                wva[1], wvb[1] = numint._uks_mgga_wv1(
                    (rhoa, rhob), (dR_rho1a[1], dR_rho1b[1]), vxc, fxc, weight)
                wva[2], wvb[2] = numint._uks_mgga_wv1(
                    (rhoa, rhob), (dR_rho1a[2], dR_rho1b[2]), vxc, fxc, weight)

                aow = rks_grad._make_dR_dao_w(ao, wva[0])
                rks_grad._d1_dot_(vmata[ia, 0], mol, aow, ao[0], mask, ao_loc,
                                  True)
                aow = rks_grad._make_dR_dao_w(ao, wva[1])
                rks_grad._d1_dot_(vmata[ia, 1], mol, aow, ao[0], mask, ao_loc,
                                  True)
                aow = rks_grad._make_dR_dao_w(ao, wva[2])
                rks_grad._d1_dot_(vmata[ia, 2], mol, aow, ao[0], mask, ao_loc,
                                  True)
                aow = [numint._scale_ao(ao[:4], wva[i, :4]) for i in range(3)]
                rks_hess._d1d2_dot_(vmata[ia], mol, ao[1:4], aow, mask, ao_loc,
                                    False)

                aow = rks_grad._make_dR_dao_w(ao, wvb[0])
                rks_grad._d1_dot_(vmatb[ia, 0], mol, aow, ao[0], mask, ao_loc,
                                  True)
                aow = rks_grad._make_dR_dao_w(ao, wvb[1])
                rks_grad._d1_dot_(vmatb[ia, 1], mol, aow, ao[0], mask, ao_loc,
                                  True)
                aow = rks_grad._make_dR_dao_w(ao, wvb[2])
                rks_grad._d1_dot_(vmatb[ia, 2], mol, aow, ao[0], mask, ao_loc,
                                  True)
                aow = [numint._scale_ao(ao[:4], wvb[i, :4]) for i in range(3)]
                rks_hess._d1d2_dot_(vmatb[ia], mol, ao[1:4], aow, mask, ao_loc,
                                    False)

                # *2 because wv[5] is scaled by 0.5 in _rks_mgga_wv1
                wv = wva[:, 5] * 2
                aow = [numint._scale_ao(ao[1], wv[i]) for i in range(3)]
                rks_hess._d1d2_dot_(vmata[ia], mol, [ao[XX], ao[XY], ao[XZ]],
                                    aow, mask, ao_loc, False)
                aow = [numint._scale_ao(ao[2], wv[i]) for i in range(3)]
                rks_hess._d1d2_dot_(vmata[ia], mol, [ao[YX], ao[YY], ao[YZ]],
                                    aow, mask, ao_loc, False)
                aow = [numint._scale_ao(ao[3], wv[i]) for i in range(3)]
                rks_hess._d1d2_dot_(vmata[ia], mol, [ao[ZX], ao[ZY], ao[ZZ]],
                                    aow, mask, ao_loc, False)
                wv = wvb[:, 5] * 2
                aow = [numint._scale_ao(ao[1], wv[i]) for i in range(3)]
                rks_hess._d1d2_dot_(vmatb[ia], mol, [ao[XX], ao[XY], ao[XZ]],
                                    aow, mask, ao_loc, False)
                aow = [numint._scale_ao(ao[2], wv[i]) for i in range(3)]
                rks_hess._d1d2_dot_(vmatb[ia], mol, [ao[YX], ao[YY], ao[YZ]],
                                    aow, mask, ao_loc, False)
                aow = [numint._scale_ao(ao[3], wv[i]) for i in range(3)]
                rks_hess._d1d2_dot_(vmatb[ia], mol, [ao[ZX], ao[ZY], ao[ZZ]],
                                    aow, mask, ao_loc, False)

        for ia in range(mol.natm):
            p0, p1 = aoslices[ia][2:]
            vmata[ia, :, :, :, p0:p1] += ipipa[:, :, :, p0:p1]
            vmata[ia, :, :, :, p0:p1] += ipipa[:, :,
                                               p0:p1].transpose(1, 0, 3, 2)
            vmatb[ia, :, :, :, p0:p1] += ipipb[:, :, :, p0:p1]
            vmatb[ia, :, :, :, p0:p1] += ipipb[:, :,
                                               p0:p1].transpose(1, 0, 3, 2)

    return vmata, vmatb