Esempio n. 1
0
        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)
Esempio n. 2
0
File: rks.py Progetto: pyscf/pyscf
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.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()
    dm0 = mf.make_rdm1(mo_coeff, mo_occ)

    v_ip = numpy.zeros((3, nao, nao))
    vmat = numpy.zeros((mol.natm, 3, nao, nao))
    max_memory = max(2000, max_memory - vmat.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):
            rho = ni.eval_rho2(mol, ao[0], mo_coeff, mo_occ, mask, xctype)
            vxc, fxc = ni.eval_xc(mf.xc, rho, 0, deriv=2)[1:3]
            vrho = vxc[0]
            frr = fxc[0]
            aow = numint._scale_ao(ao[0], weight * vrho)
            rks_grad._d1_dot_(v_ip, mol, ao[1:4], aow, mask, ao_loc, True)

            ao_dm0 = numint._dot_ao_dm(mol, ao[0], dm0, mask, shls_slice,
                                       ao_loc)
            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
                rho1 = numpy.einsum('xpi,pi->xp', ao[1:, :, p0:p1],
                                    ao_dm0[:, p0:p1])
                wv = weight * frr * rho1
                aow = [numint._scale_ao(ao[0], wv[i]) for i in range(3)]
                rks_grad._d1_dot_(vmat[ia], mol, aow, ao[0], mask, ao_loc,
                                  True)
            ao_dm0 = aow = None

    elif xctype == 'GGA':
        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[:4], mo_coeff, mo_occ, mask, xctype)
            vxc, fxc = ni.eval_xc(mf.xc, rho, 0, deriv=2)[1:3]

            wv = numint._rks_gga_wv0(rho, vxc, weight)
            rks_grad._gga_grad_sum_(v_ip, mol, ao, wv, mask, ao_loc)

            ao_dm0 = [
                numint._dot_ao_dm(mol, ao[i], dm0, mask, shls_slice, ao_loc)
                for i in range(4)
            ]
            for ia in range(mol.natm):
                wv = dR_rho1 = _make_dR_rho1(ao, ao_dm0, ia, aoslices)
                wv[0] = numint._rks_gga_wv1(rho, dR_rho1[0], vxc, fxc, weight)
                wv[1] = numint._rks_gga_wv1(rho, dR_rho1[1], vxc, fxc, weight)
                wv[2] = numint._rks_gga_wv1(rho, dR_rho1[2], vxc, fxc, weight)
                aow = [numint._scale_ao(ao[:4], wv[i, :4]) for i in range(3)]
                rks_grad._d1_dot_(vmat[ia], mol, aow, ao[0], mask, ao_loc,
                                  True)
            ao_dm0 = aow = None

    elif xctype == 'MGGA':
        if grids.level < 5:
            logger.warn(mol, 'MGGA Hessian is sensitive to dft grids.')
        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[:10], mo_coeff, mo_occ, mask, xctype)
            vxc, fxc = ni.eval_xc(mf.xc, rho, 0, deriv=2)[1:3]

            wv = numint._rks_mgga_wv0(rho, vxc, weight)
            rks_grad._gga_grad_sum_(v_ip, mol, ao, wv, mask, ao_loc)

            # *2 because wv[5] is scaled by 0.5 in _rks_mgga_wv0
            rks_grad._tau_grad_dot_(v_ip, mol, ao, wv[5] * 2, mask, ao_loc,
                                    True)

            ao_dm0 = [
                numint._dot_ao_dm(mol, ao[i], dm0, mask, shls_slice, ao_loc)
                for i in range(4)
            ]
            for ia in range(mol.natm):
                wv = dR_rho1 = _make_dR_rho1(ao, ao_dm0, ia, aoslices, xctype)
                wv[0] = numint._rks_mgga_wv1(rho, dR_rho1[0], vxc, fxc, weight)
                wv[1] = numint._rks_mgga_wv1(rho, dR_rho1[1], vxc, fxc, weight)
                wv[2] = numint._rks_mgga_wv1(rho, dR_rho1[2], vxc, fxc, weight)
                aow = [numint._scale_ao(ao[:4], wv[i, :4]) for i in range(3)]
                rks_grad._d1_dot_(vmat[ia], mol, aow, ao[0], mask, ao_loc,
                                  True)

                for j in range(1, 4):
                    aow = [numint._scale_ao(ao[j], wv[i, 5]) for i in range(3)]
                    rks_grad._d1_dot_(vmat[ia], mol, aow, ao[j], mask, ao_loc,
                                      True)
            ao_dm0 = aow = None

    for ia in range(mol.natm):
        p0, p1 = aoslices[ia][2:]
        vmat[ia, :, p0:p1] += v_ip[:, p0:p1]
        vmat[ia] = -vmat[ia] - vmat[ia].transpose(0, 2, 1)

    return vmat
Esempio n. 3
0
File: uks.py Progetto: pyscf/pyscf
def get_vxc(ni,
            mol,
            grids,
            xc_code,
            dms,
            relativity=0,
            hermi=1,
            max_memory=2000,
            verbose=None):
    xctype = ni._xc_type(xc_code)
    make_rho, nset, nao = ni._gen_rho_evaluator(mol, dms, hermi)
    ao_loc = mol.ao_loc_nr()

    vmat = numpy.zeros((2, 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):
            rho_a = make_rho(0, ao[0], mask, xctype)
            rho_b = make_rho(1, ao[0], mask, xctype)
            vxc = ni.eval_xc(xc_code, (rho_a, rho_b),
                             1,
                             relativity,
                             1,
                             verbose=verbose)[1]
            vrho = vxc[0]
            #:aow = numpy.einsum('pi,p->pi', ao[0], weight*vrho[:,0])
            aow = numint._scale_ao(ao[0], weight * vrho[:, 0])
            rks_grad._d1_dot_(vmat[0], mol, ao[1:4], aow, mask, ao_loc, True)
            #:aow = numpy.einsum('pi,p->pi', ao[0], weight*vrho[:,1])
            aow = numint._scale_ao(ao[0], weight * vrho[:, 1])
            rks_grad._d1_dot_(vmat[1], mol, ao[1:4], aow, mask, ao_loc, True)

    elif xctype == 'GGA':
        ao_deriv = 2
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rho_a = make_rho(0, ao[:4], mask, xctype)
            rho_b = make_rho(1, ao[:4], mask, xctype)
            vxc = ni.eval_xc(xc_code, (rho_a, rho_b),
                             1,
                             relativity,
                             1,
                             verbose=verbose)[1]
            wva, wvb = numint._uks_gga_wv0((rho_a, rho_b), vxc, weight)
            rks_grad._gga_grad_sum_(vmat[0], mol, ao, wva, mask, ao_loc)
            rks_grad._gga_grad_sum_(vmat[1], mol, ao, wvb, mask, ao_loc)

    elif xctype == 'NLC':
        raise NotImplementedError('NLC')

    elif xctype == 'MGGA':
        ao_deriv = 2
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rho_a = make_rho(0, ao[:10], mask, xctype)
            rho_b = make_rho(1, ao[:10], mask, xctype)
            vxc = ni.eval_xc(xc_code, (rho_a, rho_b),
                             1,
                             relativity,
                             1,
                             verbose=verbose)[1]
            wva, wvb = numint._uks_mgga_wv0((rho_a, rho_b), vxc, weight)
            rks_grad._gga_grad_sum_(vmat[0], mol, ao, wva, mask, ao_loc)
            rks_grad._gga_grad_sum_(vmat[1], mol, ao, wvb, mask, ao_loc)

            # *2 because wv[5] is scaled by 0.5 in _uks_mgga_wv0
            rks_grad._tau_grad_dot_(vmat[0], mol, ao, wva[5] * 2, mask, ao_loc,
                                    True)
            rks_grad._tau_grad_dot_(vmat[1], mol, ao, wvb[5] * 2, mask, ao_loc,
                                    True)

    exc = numpy.zeros((mol.natm, 3))
    # - sign because nabla_X = -nabla_x
    return exc, -vmat
Esempio 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 = [numint._scale_ao(ao[:4], wva[i, :4]) for i in range(3)]
                rks_grad._d1_dot_(vmata[ia], mol, aow, ao[0], mask, ao_loc,
                                  True)
                aow = [numint._scale_ao(ao[:4], wvb[i, :4]) for i in range(3)]
                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]
            vmatb[ia, :, p0:p1] += vipb[:, p0:p1]
            vmata[ia] = -vmata[ia] - vmata[ia].transpose(0, 2, 1)
            vmatb[ia] = -vmatb[ia] - vmatb[ia].transpose(0, 2, 1)

    elif xctype == 'MGGA':
        if grids.level < 5:
            logger.warn(mol, 'MGGA Hessian is sensitive to dft grids.')
        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[: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)
            rks_grad._gga_grad_sum_(vipa, mol, ao, wva, mask, ao_loc)
            rks_grad._gga_grad_sum_(vipb, mol, ao, wvb, mask, ao_loc)

            rks_grad._tau_grad_dot_(vipa, mol, ao, wva[5] * 2, mask, ao_loc,
                                    True)
            rks_grad._tau_grad_dot_(vipb, mol, ao, wvb[5] * 2, mask, ao_loc,
                                    True)

            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 = [numint._scale_ao(ao[:4], wva[i, :4]) for i in range(3)]
                rks_grad._d1_dot_(vmata[ia], mol, aow, ao[0], mask, ao_loc,
                                  True)
                aow = [numint._scale_ao(ao[:4], wvb[i, :4]) for i in range(3)]
                rks_grad._d1_dot_(vmatb[ia], mol, aow, ao[0], mask, ao_loc,
                                  True)

                for j in range(1, 4):
                    aow = [
                        numint._scale_ao(ao[j], wva[i, 5]) for i in range(3)
                    ]
                    rks_grad._d1_dot_(vmata[ia], mol, aow, ao[j], mask, ao_loc,
                                      True)
                    aow = [
                        numint._scale_ao(ao[j], wvb[i, 5]) for i in range(3)
                    ]
                    rks_grad._d1_dot_(vmatb[ia], mol, aow, ao[j], mask, ao_loc,
                                      True)

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

    return vmata, vmatb
Esempio n. 5
0
File: uks.py Progetto: pyscf/pyscf
def get_vxc_full_response(ni,
                          mol,
                          grids,
                          xc_code,
                          dms,
                          relativity=0,
                          hermi=1,
                          max_memory=2000,
                          verbose=None):
    '''Full response including the response of the grids'''
    xctype = ni._xc_type(xc_code)
    make_rho, nset, nao = ni._gen_rho_evaluator(mol, dms, hermi)
    ao_loc = mol.ao_loc_nr()
    aoslices = mol.aoslice_by_atom()

    excsum = 0
    vmat = numpy.zeros((2, 3, nao, nao))
    if xctype == 'LDA':
        ao_deriv = 1
        for atm_id, (coords, weight, weight1) \
                in enumerate(rks_grad.grids_response_cc(grids)):
            sh0, sh1 = aoslices[atm_id][:2]
            mask = gen_grid.make_mask(mol, coords)
            ao = ni.eval_ao(mol, coords, deriv=ao_deriv, non0tab=mask)
            rho_a = make_rho(0, ao[0], mask, xctype)
            rho_b = make_rho(1, ao[0], mask, xctype)
            exc, vxc = ni.eval_xc(xc_code, (rho_a, rho_b),
                                  1,
                                  relativity,
                                  1,
                                  verbose=verbose)[:2]
            vrho = vxc[0]

            vtmp = numpy.zeros((3, nao, nao))
            #:aow = numpy.einsum('pi,p->pi', ao[0], weight*vrho[:,0])
            aow = numint._scale_ao(ao[0], weight * vrho[:, 0])
            rks_grad._d1_dot_(vtmp, mol, ao[1:4], aow, mask, ao_loc, True)
            vmat[0] += vtmp
            excsum += numpy.einsum('r,r,nxr->nx', exc, rho_a + rho_b, weight1)
            excsum[atm_id] += numpy.einsum('xij,ji->x', vtmp, dms[0]) * 2

            vtmp = numpy.zeros((3, nao, nao))
            #:aow = numpy.einsum('pi,p->pi', ao[0], weight*vrho[:,1])
            aow = numint._scale_ao(ao[0], weight * vrho[:, 1])
            rks_grad._d1_dot_(vtmp, mol, ao[1:4], aow, mask, ao_loc, True)
            vmat[1] += vtmp
            excsum[atm_id] += numpy.einsum('xij,ji->x', vtmp, dms[1]) * 2

    elif xctype == 'GGA':
        ao_deriv = 2
        for atm_id, (coords, weight, weight1) \
                in enumerate(rks_grad.grids_response_cc(grids)):
            sh0, sh1 = aoslices[atm_id][:2]
            mask = gen_grid.make_mask(mol, coords)
            ao = ni.eval_ao(mol, coords, deriv=ao_deriv, non0tab=mask)
            rho_a = make_rho(0, ao[:4], mask, xctype)
            rho_b = make_rho(1, ao[:4], mask, xctype)
            exc, vxc = ni.eval_xc(xc_code, (rho_a, rho_b),
                                  1,
                                  relativity,
                                  1,
                                  verbose=verbose)[:2]
            wva, wvb = numint._uks_gga_wv0((rho_a, rho_b), vxc, weight)

            vtmp = numpy.zeros((3, nao, nao))
            rks_grad._gga_grad_sum_(vtmp, mol, ao, wva, mask, ao_loc)
            vmat[0] += vtmp
            excsum += numpy.einsum('r,r,nxr->nx', exc, rho_a[0] + rho_b[0],
                                   weight1)
            excsum[atm_id] += numpy.einsum('xij,ji->x', vtmp, dms[0]) * 2

            vtmp = numpy.zeros((3, nao, nao))
            rks_grad._gga_grad_sum_(vtmp, mol, ao, wvb, mask, ao_loc)
            vmat[1] += vtmp
            excsum[atm_id] += numpy.einsum('xij,ji->x', vtmp, dms[1]) * 2

    elif xctype == 'NLC':
        raise NotImplementedError('NLC')

    elif xctype == 'MGGA':
        ao_deriv = 2
        for atm_id, (coords, weight, weight1) \
                in enumerate(rks_grad.grids_response_cc(grids)):
            sh0, sh1 = aoslices[atm_id][:2]
            mask = gen_grid.make_mask(mol, coords)
            ao = ni.eval_ao(mol, coords, deriv=ao_deriv, non0tab=mask)
            rho_a = make_rho(0, ao[:10], mask, xctype)
            rho_b = make_rho(1, ao[:10], mask, xctype)
            exc, vxc = ni.eval_xc(xc_code, (rho_a, rho_b),
                                  1,
                                  relativity,
                                  1,
                                  verbose=verbose)[:2]
            wva, wvb = numint._uks_mgga_wv0((rho_a, rho_b), vxc, weight)

            vtmp = numpy.zeros((3, nao, nao))
            rks_grad._gga_grad_sum_(vtmp, mol, ao, wva, mask, ao_loc)
            # *2 because wv[5] is scaled by 0.5 in _uks_mgga_wv0
            rks_grad._tau_grad_dot_(vtmp, mol, ao, wva[5] * 2, mask, ao_loc,
                                    True)
            vmat[0] += vtmp
            excsum += numpy.einsum('r,r,nxr->nx', exc, rho_a[0] + rho_b[0],
                                   weight1)
            excsum[atm_id] += numpy.einsum('xij,ji->x', vtmp, dms[0]) * 2

            vtmp = numpy.zeros((3, nao, nao))
            rks_grad._gga_grad_sum_(vtmp, mol, ao, wvb, mask, ao_loc)
            rks_grad._tau_grad_dot_(vtmp, mol, ao, wvb[5] * 2, mask, ao_loc,
                                    True)
            vmat[1] += vtmp
            excsum[atm_id] += numpy.einsum('xij,ji->x', vtmp, dms[1]) * 2

    # - sign because nabla_X = -nabla_x
    return excsum, -vmat