,rp->rs', foo, x)
        #x2-= .5*numpy.einsum('sp,rp->rs', foo, x)
        x2-= numpy.einsum('ps,rp->rs', foo, x)

        # *2 for double occupancy
        d1 = reduce(numpy.dot, (orbv, x*2, orbo.conj().T))
        dm1 = d1 + d1.conj().T
        v1 = vind(dm1)
        x2 += reduce(numpy.dot, (orbv.conj().T, v1, orbo))
        if with_symmetry and mol.symmetry:
            x2[sym_forbid] = 0
        return x2.ravel() * 2

    return g.reshape(-1), h_op, h_diag.reshape(-1)

def gen_g_hop_rohf(mf, mo_coeff, mo_occ, fock_ao=None, h1e=None,
                   with_symmetry=True):
    if getattr(fock_ao, 'focka', None) is None:
        dm0 = mf.make_rdm1(mo_coeff, mo_occ)
        fock_ao = mf.get_fock(h1e, dm=dm0)
    fock_ao = fock_ao.focka, fock_ao.fockb
    mo_occa = occidxa = mo_occ > 0
    mo_occb = occidxb = mo_occ ==2
    ug, uh_op, uh_diag = gen_g_hop_uhf(mf, (mo_coeff,)*2, (mo_occa,mo_occb),
                                       fock_ao, None, with_symmetry)

    viridxa = ~occidxa
    viridxb = ~occidxb
    uniq_var_a = viridxa[:,None] & occidxa
    uniq_var_b = viridxb[:,None] & occidxb
    uniq_ab = uniq_var_a | uniq_var_b
    nmo = mo_coeff.shape[-1]
    nocca = numpy.count_nonzero(mo_occa)
    nvira = nmo - nocca

    def sum_ab(x):
        x1 = numpy.zeros((nmo,nmo), dtype=x.dtype)
        x1[uniq_var_a]  = x[:nvira*nocca]
        x1[uniq_var_b] += x[nvira*nocca:]
        return x1[uniq_ab]

    g = sum_ab(ug)
    h_diag = sum_ab(uh_diag)
    def h_op(x):
        x1 = numpy.zeros((nmo,nmo), dtype=x.dtype)
        # unpack ROHF rotation parameters
        x1[uniq_ab] = x
        x1 = numpy.hstack((x1[uniq_var_a],x1[uniq_var_b]))
        return sum_ab(uh_op(x1))

    return g, h_op, h_diag

def gen_g_hop_uhf(mf, mo_coeff, mo_occ, fock_ao=None, h1e=None,
                  with_symmetry=True):
    mol = mf.mol
    mo_coeff0 = mo_coeff
    if getattr(mf, '_scf', None) and mf._scf.mol != mol:
        #TODO: construct vind with dual-basis treatment, (see also JCP, 118, 9497)
        mo_coeff = (addons.project_mo_nr2nr(mf._scf.mol, mo_coeff[0], mol),
                    addons.project_mo_nr2nr(mf._scf.mol, mo_coeff[1], mol))

    occidxa = numpy.where(mo_occ[0]>0)[0]
    occidxb = numpy.where(mo_occ[1]>0)[0]
    viridxa = numpy.where(mo_occ[0]==0)[0]
    viridxb = numpy.where(mo_occ[1]==0)[0]
    nocca = len(occidxa)
    noccb = len(occidxb)
    nvira = len(viridxa)
    nvirb = len(viridxb)
    orboa = mo_coeff[0][:,occidxa]
    orbob = mo_coeff[1][:,occidxb]
    orbva = mo_coeff[0][:,viridxa]
    orbvb = mo_coeff[1][:,viridxb]
    if with_symmetry and mol.symmetry:
        orbsyma, orbsymb = uhf_symm.get_orbsym(mol, mo_coeff)
        sym_forbida = orbsyma[viridxa,None] != orbsyma[occidxa]
        sym_forbidb = orbsymb[viridxb,None] != orbsymb[occidxb]
        sym_forbid = numpy.hstack((sym_forbida.ravel(), sym_forbidb.ravel()))

    if fock_ao is None:
        if getattr(mf, '_scf', None) and mf._scf.mol != mol:
            h1e = mf.get_hcore(mol)
        dm0 = mf.make_rdm1(mo_coeff, mo_occ)
        fock_ao = mf.get_fock(h1e, dm=dm0)
        focka = reduce(numpy.dot, (mo_coeff[0].conj().T, fock_ao[0], mo_coeff[0]))
        fockb = reduce(numpy.dot, (mo_coeff[1].conj().T, fock_ao[1], mo_coeff[1]))
    else:
        focka = reduce(numpy.dot, (mo_coeff0[0].conj().T, fock_ao[0], mo_coeff0[0]))
        fockb = reduce(numpy.dot, (mo_coeff0[1].conj().T, fock_ao[1], mo_coeff0[1]))
    fooa = focka[occidxa[:,None],occidxa]
    fvva = focka[viridxa[:,None],viridxa]
    foob = fockb[occidxb[:,None],occidxb]
    fvvb = fockb[viridxb[:,None],viridxb]

    g = numpy.hstack((focka[viridxa[:,None],occidxa].ravel(),
                      fockb[viridxb[:,None],occidxb].ravel()))

    h_diaga = fvva.diagonal().real[:,None] - fooa.diagonal().real
    h_diagb = fvvb.diagonal().real[:,None] - foob.diagonal().real
    h_diag = numpy.hstack((h_diaga.reshape(-1), h_diagb.reshape(-1)))

    if with_symmetry and mol.symmetry:
        g[sym_forbid] = 0
        h_diag[sym_forbid] = 0

    vind = mf.gen_response(mo_coeff, mo_occ, hermi=1)

    def h_op(x):
        if with_symmetry and mol.symmetry:
            x = x.copy()
            x[sym_forbid] = 0
        x1a = x[:nvira*nocca].reshape(nvira,nocca)
        x1b = x[nvira*nocca:].reshape(nvirb,noccb)
        x2a = numpy.einsum('pr,rq->pq', fvva, x1a)
        x2a-= numpy.einsum('sq,ps->pq', fooa, x1a)
        x2b = numpy.einsum('pr,rq->pq', fvvb, x1b)
        x2b-= numpy.einsum('sq,ps->pq', foob, x1b)

        d1a = reduce(numpy.dot, (orbva, x1a, orboa.conj().T))
        d1b = reduce(numpy.dot, (orbvb, x1b, orbob.conj().T))
        dm1 = numpy.array((d1a+d1a.conj().T,d1b+d1b.conj().T))
        v1 = vind(dm1)
        x2a += reduce(numpy.dot, (orbva.conj().T, v1[0], orboa))
        x2b += reduce(numpy.dot, (orbvb.conj().T, v1[1], orbob))

        x2 = numpy.hstack((x2a.ravel(), x2b.ravel()))
        if with_symmetry and mol.symmetry:
            x2[sym_forbid] = 0
        return x2

    return g, h_op, h_diag

def gen_g_hop_ghf(mf, mo_coeff, mo_occ, fock_ao=None, h1e=None,
                  with_symmetry=True):
    mol = mf.mol
    occidx = numpy.where(mo_occ==1)[0]
    viridx = numpy.where(mo_occ==0)[0]
    nocc = len(occidx)
    nvir = len(viridx)
    orbo = mo_coeff[:,occidx]
    orbv = mo_coeff[:,viridx]
    if with_symmetry and mol.symmetry:
        orbsym = ghf_symm.get_orbsym(mol, mo_coeff)
        sym_forbid = orbsym[viridx,None] != orbsym[occidx]

    if fock_ao is None:
        dm0 = mf.make_rdm1(mo_coeff, mo_occ)
        fock_ao = mf.get_fock(h1e, dm=dm0)
    fock = reduce(numpy.dot, (mo_coeff.conj().T, fock_ao, mo_coeff))

    g = fock[viridx[:,None],occidx]

    foo = fock[occidx[:,None],occidx]
    fvv = fock[viridx[:,None],viridx]

    h_diag = fvv.diagonal().real[:,None] - foo.diagonal().real

    if with_symmetry and mol.symmetry:
        g[sym_forbid] = 0
        h_diag[sym_forbid] = 0

    vind = mf.gen_response(mo_coeff, mo_occ, hermi=1)

    def h_op(x):
        x = x.reshape(nvir,nocc)
        if with_symmetry and mol.symmetry:
            x = x.copy()
            x[sym_forbid] = 0
        x2 = numpy.einsum('ps,sq->pq', fvv, x)
        x2-= numpy.einsum('ps,rp->rs', foo, x)

        d1 = reduce(numpy.dot, (orbv, x, orbo.conj().T))
        dm1 = d1 + d1.conj().T
        v1 = vind(dm1)
        x2 += reduce(numpy.dot, (orbv.conj().T, v1, orbo))
        if with_symmetry and mol.symmetry:
            x2[sym_forbid] = 0
        return x2.ravel()

    return g.reshape(-1), h_op, h_diag.reshape(-1)

def gen_g_hop_dhf(mf, mo_coeff, mo_occ, fock_ao=None, h1e=None,
                  with_symmetry=False):
    return gen_g_hop_ghf(mf, mo_coeff, mo_occ, fock_ao, h1e, with_symmetry)


# Dual basis for gradients and hessian
def project_mol(mol, dual_basis={}):
    from pyscf import df
    uniq_atoms = set([a[0] for a in mol._atom])
    newbasis = {}
    for symb in uniq_atoms:
        if gto.charge(symb) <= 10:
            newbasis[symb] = '321g'
        elif gto.charge(symb) <= 12:
            newbasis[symb] = 'dzp'
        elif gto.charge(symb) <= 18:
            newbasis[symb] = 'dz'
        elif gto.charge(symb) <= 86:
            newbasis[symb] = 'dzp'
        else:
            newbasis[symb] = 'sto3g'
    if isinstance(dual_basis, (dict, tuple, list)):
        newbasis.update(dual_basis)
    elif isinstance(dual_basis, str):
        for k in newbasis:
            newbasis[k] = dual_basis
    return df.addons.make_auxmol(mol, newbasis)


# TODO: check whether high order terms in (g_orb, h_op) affects optimization
# To include high order terms, we can generate mo_coeff every time u matrix
# changed and insert the mo_coeff to g_op, h_op.
# Seems the high order terms do not help optimization?
def _rotate_orb_cc(mf, h1e, s1e, conv_tol_grad=None, verbose=None):
    log = logger.new_logger(mf, verbose)

    if conv_tol_grad is None:
        conv_tol_grad = numpy.sqrt(mf.conv_tol*.1)
        #TODO: dynamically adjust max_stepsize, as done in mc1step.py

    def precond(x, e):
        hdiagd = h_diag-(e-mf.ah_level_shift)
        hdiagd[abs(hdiagd)<1e-8] = 1e-8
        x = x/hdiagd
        # Because of DFT, donot norm to 1 which leads 1st DM too large.
        #norm_x = numpy.linalg.norm(x)
        #if norm_x < 1e-2:
        #    x *= 1e-2/norm_x
        return x

    t3m = (logger.process_clock(), logger.perf_counter())
    u = g_kf = g_orb = norm_gorb = dxi = kfcount = jkcount = None
    dm0 = vhf0 = None
    g_op = lambda: g_orb
    while True:
        mo_coeff, mo_occ, dm0, vhf0, e_tot = (yield u, g_kf, kfcount, jkcount, dm0, vhf0)
        fock_ao = mf.get_fock(h1e, s1e, vhf0, dm0)

        g_kf, h_op, h_diag = mf.gen_g_hop(mo_coeff, mo_occ, fock_ao)
        norm_gkf = numpy.linalg.norm(g_kf)
        if g_orb is None:
            log.debug('    |g|= %4.3g (keyframe)', norm_gkf)
            kf_trust_region = mf.kf_trust_region
            x0_guess = g_kf
        else:
            norm_dg = numpy.linalg.norm(g_kf-g_orb)
            log.debug('    |g|= %4.3g (keyframe), |g-correction|= %4.3g',
                      norm_gkf, norm_dg)
            kf_trust_region = min(max(norm_gorb/(norm_dg+1e-9), mf.kf_trust_region), 10)
            log.debug1('Set  kf_trust_region = %g', kf_trust_region)
            x0_guess = dxi
        g_orb = g_kf
        norm_gorb = norm_gkf

        ah_conv_tol = min(norm_gorb**2, mf.ah_conv_tol)
        # increase the AH accuracy when approach convergence
        #ah_start_cycle = max(mf.ah_start_cycle, int(-numpy.log10(norm_gorb)))
        ah_start_cycle = mf.ah_start_cycle
        imic = 0
        dr = 0
        ukf = None
        jkcount = 0
        kfcount = 0
        ikf = 0

        for ah_end, ihop, w, dxi, hdxi, residual, seig \
                in ciah.davidson_cc(h_op, g_op, precond, x0_guess,
                                    tol=ah_conv_tol, max_cycle=mf.ah_max_cycle,
                                    lindep=mf.ah_lindep, verbose=log):
            norm_residual = numpy.linalg.norm(residual)
            ah_start_tol = min(norm_gorb*5, mf.ah_start_tol)
            if (ah_end or ihop == mf.ah_max_cycle or # make sure to use the last step
                ((norm_residual < ah_start_tol) and (ihop >= ah_start_cycle)) or
                (seig < mf.ah_lindep)):
                imic += 1
                dxmax = numpy.max(abs(dxi))
                if dxmax > mf.max_stepsize:
                    scale = mf.max_stepsize / dxmax
                    log.debug1('... scale rotation size %g', scale)
                    dxi *= scale
                    hdxi *= scale
                else:
                    scale = None

                dr = dr + dxi
                g_orb = g_orb + hdxi
                norm_dr = numpy.linalg.norm(dr)
                norm_gorb = numpy.linalg.norm(g_orb)
                norm_dxi = numpy.linalg.norm(dxi)
                log.debug('    imic %d(%d)  |g|= %4.3g  |dxi|= %4.3g  '
                          'max(|x|)= %4.3g  |dr|= %4.3g  eig= %4.3g  seig= %4.3g',
                          imic, ihop, norm_gorb, norm_dxi,
                          dxmax, norm_dr, w, seig)

                max_cycle = max(mf.max_cycle_inner,
                                mf.max_cycle_inner-int(numpy.log(norm_gkf+1e-9)*2))
                log.debug1('Set ah_start_tol %g, ah_start_cycle %d, max_cycle %d',
                           ah_start_tol, ah_start_cycle, max_cycle)
                ikf += 1
                if imic > 3 and norm_gorb > norm_gkf*mf.ah_grad_trust_region:
                    g_orb = g_orb - hdxi
                    dr -= dxi
                    norm_gorb = numpy.linalg.norm(g_orb)
                    log.debug('|g| >> keyframe, Restore previouse step')
                    break

                elif (imic >= max_cycle or norm_gorb < conv_tol_grad/mf.ah_grad_trust_region):
                    break

                elif (ikf > 2 and # avoid frequent keyframe
                      #TODO: replace it with keyframe_scheduler
                      (ikf >= max(mf.kf_interval, mf.kf_interval-numpy.log(norm_dr+1e-9)) or
                       # Insert keyframe if the keyframe and the esitimated g_orb are too different
                       norm_gorb < norm_gkf/kf_trust_region)):
                    ikf = 0
                    u = mf.update_rotate_matrix(dr, mo_occ, mo_coeff=mo_coeff)
                    if ukf is not None:
                        u = mf.rotate_mo(ukf, u)
                    ukf = u
                    dr[:] = 0
                    mo1 = mf.rotate_mo(mo_coeff, u)
                    dm = mf.make_rdm1(mo1, mo_occ)
                    # use mf._scf.get_veff to avoid density-fit mf polluting get_veff
                    vhf0 = mf._scf.get_veff(mf._scf.mol, dm, dm_last=dm0, vhf_last=vhf0)
                    dm0 = dm
                    # Use API to compute fock instead of "fock=h1e+vhf0". This is because get_fock
                    # is the hook being overloaded in many places.
                    fock_ao = mf.get_fock(h1e, s1e, vhf0, dm0)
                    g_kf1 = mf.get_grad(mo1, mo_occ, fock_ao)
                    norm_gkf1 = numpy.linalg.norm(g_kf1)
                    norm_dg = numpy.linalg.norm(g_kf1-g_orb)
                    jkcount += 1
                    kfcount += 1
                    if log.verbose >= logger.DEBUG:
                        e_tot, e_last = mf._scf.energy_tot(dm, h1e, vhf0), e_tot
                        log.debug('Adjust keyframe g_orb to |g|= %4.3g  '
                                  '|g-correction|=%4.3g  E=%.12g dE=%.5g',
                                  norm_gkf1, norm_dg, e_tot, e_tot-e_last)

                    if (norm_dg < norm_gorb*mf.ah_grad_trust_region  # kf not too diff
                        #or norm_gkf1 < norm_gkf  # grad is decaying
                        # close to solution
                        or norm_gkf1 < conv_tol_grad*mf.ah_grad_trust_region):
                        kf_trust_region = min(max(norm_gorb/(norm_dg+1e-9), mf.kf_trust_region), 10)
                        log.debug1('Set kf_trust_region = %g', kf_trust_region)
                        g_orb = g_kf = g_kf1
                        norm_gorb = norm_gkf = norm_gkf1
                    else:
                        g_orb = g_orb - hdxi
                        dr -= dxi
                        norm_gorb = numpy.linalg.norm(g_orb)
                        log.debug('Out of trust region. Restore previouse step')
                        break

        u = mf.update_rotate_matrix(dr, mo_occ, mo_coeff=mo_coeff)
        if ukf is not None:
            u = mf.rotate_mo(ukf, u)
        jkcount += ihop + 1
        log.debug('    tot inner=%d  %d JK  |g|= %4.3g  |u-1|= %4.3g',
                  imic, jkcount, norm_gorb, numpy.linalg.norm(dr))
        h_op = h_diag = None
        t3m = log.timer('aug_hess in %d inner iters' % imic, *t3m)


def kernel(mf, mo_coeff=None, mo_occ=None, dm=None,
           conv_tol=1e-10, conv_tol_grad=None, max_cycle=50, dump_chk=True,
           callback=None, verbose=logger.NOTE):
    cput0 = (logger.process_clock(), logger.perf_counter())
    log = logger.new_logger(mf, verbose)
    mol = mf._scf.mol
    if mol != mf.mol:
        logger.warn(mf, 'dual-basis SOSCF is an experimental feature. It is '
                    'still in testing.')

    if conv_tol_grad is None:
        conv_tol_grad = numpy.sqrt(conv_tol)
        log.info('Set conv_tol_grad to %g', conv_tol_grad)

    # call mf._scf.get_hcore, mf._scf.get_ovlp because they might be overloaded
    h1e = mf._scf.get_hcore(mol)
    s1e = mf._scf.get_ovlp(mol)

    if mo_coeff is not None and mo_occ is not None:
        dm = mf.make_rdm1(mo_coeff, mo_occ)
        # call mf._scf.get_veff, to avoid "newton().density_fit()" polluting get_veff
        vhf = mf._scf.get_veff(mol, dm)
        fock = mf.get_fock(h1e, s1e, vhf, dm, level_shift_factor=0)
        mo_energy, mo_tmp = mf.eig(fock, s1e)
        mf.get_occ(mo_energy, mo_tmp)
        mo_tmp = None

    else:
        if dm is None:
            logger.debug(mf, 'Initial guess density matrix is not given. '
                         'Generating initial guess from %s', mf.init_guess)
            dm = mf.get_init_guess(mf._scf.mol, mf.init_guess)
        vhf = mf._scf.get_veff(mol, dm)
        fock = mf.get_fock(h1e, s1e, vhf, dm, level_shift_factor=0)
        mo_energy, mo_coeff = mf.eig(fock, s1e)
        mo_occ = mf.get_occ(mo_energy, mo_coeff)
        dm, dm_last = mf.make_rdm1(mo_coeff, mo_occ), dm
        vhf = mf._scf.get_veff(mol, dm, dm_last=dm_last, vhf_last=vhf)

    # Save mo_coeff and mo_occ because they are needed by function rotate_mo
    mf.mo_coeff, mf.mo_occ = mo_coeff, mo_occ

    e_tot = mf._scf.energy_tot(dm, h1e, vhf)
    fock = mf.get_fock(h1e, s1e, vhf, dm, level_shift_factor=0)
    log.info('Initial guess E= %.15g  |g|= %g', e_tot,
             numpy.linalg.norm(mf._scf.get_grad(mo_coeff, mo_occ, fock)))

    if dump_chk and mf.chkfile:
        chkfile.save_mol(mol, mf.chkfile)

    # Copy the integral file to soscf object to avoid the integrals being
    # cached twice.
    if mol is mf.mol and not getattr(mf, 'with_df', None):
        mf._eri = mf._scf._eri
        # If different direct_scf_cutoff is assigned to newton_ah mf.opt
        # object, mf.opt should be different to mf._scf.opt
        #mf.opt = mf._scf.opt

    rotaiter = _rotate_orb_cc(mf, h1e, s1e, conv_tol_grad, verbose=log)
    next(rotaiter)  # start the iterator
    kftot = jktot = 0
    scf_conv = False
    cput1 = log.timer('initializing second order scf', *cput0)

    for imacro in range(max_cycle):
        u, g_orb, kfcount, jkcount, dm_last, vhf = \
                rotaiter.send((mo_coeff, mo_occ, dm, vhf, e_tot))
        kftot += kfcount + 1
        jktot += jkcount + 1

        last_hf_e = e_tot
        norm_gorb = numpy.linalg.norm(g_orb)
        mo_coeff = mf.rotate_mo(mo_coeff, u, log)
        dm = mf.make_rdm1(mo_coeff, mo_occ)
        vhf = mf._scf.get_veff(mol, dm, dm_last=dm_last, vhf_last=vhf)
        fock = mf.get_fock(h1e, s1e, vhf, dm, level_shift_factor=0)
        # NOTE: DO NOT change the initial guess mo_occ, mo_coeff
        if mf.verbose >= logger.DEBUG:
            mo_energy, mo_tmp = mf.eig(fock, s1e)
            mf.get_occ(mo_energy, mo_tmp)
            # call mf._scf.energy_tot for dft, because the (dft).get_veff step saved _exc in mf._scf
        e_tot = mf._scf.energy_tot(dm, h1e, vhf)

        log.info('macro= %d  E= %.15g  delta_E= %g  |g|= %g  %d KF %d JK',
                 imacro, e_tot, e_tot-last_hf_e, norm_gorb,
                 kfcount+1, jkcount)
        cput1 = log.timer('cycle= %d'%(imacro+1), *cput1)

        if callable(mf.check_convergence):
            scf_conv = mf.check_convergence(locals())
        elif abs(e_tot-last_hf_e) < conv_tol and norm_gorb < conv_tol_grad:
            scf_conv = True

        if dump_chk:
            mf.dump_chk(locals())

        if callable(callback):
            callback(locals())

        if scf_conv:
            break

    if callable(callback):
        callback(locals())

    rotaiter.close()
    mo_energy, mo_coeff1 = mf._scf.canonicalize(mo_coeff, mo_occ, fock)
    if mf.canonicalization:
        log.info('Canonicalize SCF orbitals')
        mo_coeff = mo_coeff1
        if dump_chk:
            mf.dump_chk(locals())
    log.info('macro X = %d  E=%.15g  |g|= %g  total %d KF %d JK',
             imacro+1, e_tot, norm_gorb, kftot+1, jktot+1)
    if (numpy.any(mo_occ==0) and
        mo_energy[mo_occ>0].max() > mo_energy[mo_occ==0].min()):
        log.warn('HOMO %s > LUMO %s was found in the canonicalized orbitals.',
                 mo_energy[mo_occ>0].max(), mo_energy[mo_occ==0].min())
    return scf_conv, e_tot, mo_energy, mo_coeff, mo_occ

# Note which function that "density_fit" decorated.  density-fit have been
# considered separatedly for newton_mf and newton_mf._scf, there are 3 cases
# 1. both grad and hessian by df, because the input mf obj is decorated by df
#       newton(scf.density_fit(scf.RHF(mol)))
# 2. both grad and hessian in df, because both newton_mf and the input mf objs
#    are decorated by df
#       scf.density_fit(newton(scf.density_fit(scf.RHF(mol))))
# 3. grad by explicit scheme, hessian by df, because only newton_mf obj is
#    decorated by df
#       scf.density_fit(newton(scf.RHF(mol)))
# The following function is not necessary
#def density_fit_(mf, auxbasis='weigend+etb'):
#    mfaux = mf.density_fit(auxbasis)
#    mf.gen_g_hop = mfaux.gen_g_hop
#    return mf
#def density_fit(mf, auxbasis='weigend+etb'):
#    return density_fit_(copy.copy(mf), auxbasis)


# A tag to label the derived SCF class
class _CIAH_SOSCF(object):
    '''
    Attributes for Newton solver:
        max_cycle_inner : int
            AH iterations within eacy macro iterations. Default is 10
        max_stepsize : int
            The step size for orbital rotation.  Small step is prefered.  Default is 0.05.
        canonicalization : bool
            To control whether to canonicalize the orbitals optimized by
            Newton solver.  Default is True.
    '''

    max_cycle_inner = getattr(__config__, 'soscf_newton_ah_SOSCF_max_cycle_inner', 12)
    max_stepsize = getattr(__config__, 'soscf_newton_ah_SOSCF_max_stepsize', .05)
    canonicalization = getattr(__config__, 'soscf_newton_ah_SOSCF_canonicalization', True)

    ah_start_tol = getattr(__config__, 'soscf_newton_ah_SOSCF_ah_start_tol', 1e9)
    ah_start_cycle = getattr(__config__, 'soscf_newton_ah_SOSCF_ah_start_cycle', 1)
    ah_level_shift = getattr(__config__, 'soscf_newton_ah_SOSCF_ah_level_shift', 0)
    ah_conv_tol = getattr(__config__, 'soscf_newton_ah_SOSCF_ah_conv_tol', 1e-12)
    ah_lindep = getattr(__config__, 'soscf_newton_ah_SOSCF_ah_lindep', 1e-14)
    ah_max_cycle = getattr(__config__, 'soscf_newton_ah_SOSCF_ah_max_cycle', 40)
    ah_grad_trust_region = getattr(__config__, 'soscf_newton_ah_SOSCF_ah_grad_trust_region', 2.5)
    kf_interval = getattr(__config__, 'soscf_newton_ah_SOSCF_kf_interval', 4)
    kf_trust_region = getattr(__config__, 'soscf_newton_ah_SOSCF_kf_trust_region', 5)

    def __init__(self, mf):
        self.__dict__.update(mf.__dict__)
        self._scf = mf
        self._keys.update(('max_cycle_inner', 'max_stepsize',
                           'canonicalization', 'ah_start_tol', 'ah_start_cycle',
                           'ah_level_shift', 'ah_conv_tol', 'ah_lindep',
                           'ah_max_cycle', 'ah_grad_trust_region', 'kf_interval',
                           'kf_trust_region'))

    def dump_flags(self, verbose=None):
        log = logger.new_logger(self, verbose)
        log.info('\n')
        self._scf.dump_flags(verbose)
        log.info('******** %s Newton solver flags ********', self._scf.__class__)
        log.info('SCF tol = %g', self.conv_tol)
        log.info('conv_tol_grad = %s',    self.conv_tol_grad)
        log.info('max. SCF cycles = %d', self.max_cycle)
        log.info('direct_scf = %s', self._scf.direct_scf)
        if self._scf.direct_scf:
            log.info('direct_scf_tol = %g', self._scf.direct_scf_tol)
        if self.chkfile:
            log.info('chkfile to save SCF result = %s', self.chkfile)
        log.info('max_cycle_inner = %d',  self.max_cycle_inner)
        log.info('max_stepsize = %g', self.max_stepsize)
        log.info('ah_start_tol = %g',     self.ah_start_tol)
        log.info('ah_level_shift = %g',   self.ah_level_shift)
        log.info('ah_conv_tol = %g',      self.ah_conv_tol)
        log.info('ah_lindep = %g',        self.ah_lindep)
        log.info('ah_start_cycle = %d',   self.ah_start_cycle)
        log.info('ah_max_cycle = %d',     self.ah_max_cycle)
        log.info('ah_grad_trust_region = %g', self.ah_grad_trust_region)
        log.info('kf_interval = %d', self.kf_interval)
        log.info('kf_trust_region = %d', self.kf_trust_region)
        log.info('canonicalization = %s', self.canonicalization)
        log.info('max_memory %d MB (current use %d MB)',
                 self.max_memory, lib.current_memory()[0])
        return self

    def build(self, mol=None):
        if mol is None: mol = self.mol
        if self.verbose >= logger.WARN:
            self.check_sanity()
        self._scf.build(mol)
        self.opt = None
        self._eri = None
        return self

    def reset(self, mol=None):
        if mol is not None:
            self.mol = mol
        return self._scf.reset(mol)

    def kernel(self, mo_coeff=None, mo_occ=None, dm0=None):
        cput0 = (logger.process_clock(), logger.perf_counter())
        if dm0 is not None:
            if isinstance(dm0, str):
                sys.stderr.write('Newton solver reads density matrix from chkfile %s\n' % dm0)
                dm0 = self.from_chk(dm0)

        elif mo_coeff is not None and mo_occ is None:
            logger.warn(self, 'Newton solver expects mo_coeff with '
                        'mo_occ as initial guess but mo_occ is not found in '
                        'the arguments.\n      The given '
                        'argument is treated as density matrix.')
            dm0 = mo_coeff
            mo_coeff = mo_occ = None

        else:
            if mo_coeff is None: mo_coeff = self.mo_coeff
            if mo_occ is None: mo_occ = self.mo_occ

            # TODO: assert mo_coeff orth-normality. If not orth-normal,
            # build dm from mo_coeff and mo_occ then unset mo_coeff and mo_occ.

        self.build(self.mol)
        self.dump_flags()

        self.converged, self.e_tot, \
                self.mo_energy, self.mo_coeff, self.mo_occ = \
                kernel(self, mo_coeff, mo_occ, dm0, conv_tol=self.conv_tol,
                       conv_tol_grad=self.conv_tol_grad,
                       max_cycle=self.max_cycle,
                       callback=self.callback, verbose=self.verbose)

        logger.timer(self, 'Second order SCF', *cput0)
        self._finalize()
        return self.e_tot

    def from_dm(self, dm):
        '''Transform the initial guess density matrix to initial orbital
        coefficients.

        Note kernel function can handle initial guess properly in pyscf-1.7 or
        newer versions. This function is kept for backward compatibility.
        '''
        mf = self._scf
        mol = mf.mol
        h1e = mf.get_hcore(mol)
        s1e = mf.get_ovlp(mol)
        vhf = mf.get_veff(mol, dm)
        fock = mf.get_fock(h1e, s1e, vhf, dm)
        mo_energy, mo_coeff = mf.eig(fock, s1e)
        mo_occ = mf.get_occ(mo_energy, mo_coeff)
        return mo_coeff, mo_occ

    gen_g_hop = gen_g_hop_rhf

    def update_rotate_matrix(self, dx, mo_occ, u0=1, mo_coeff=None):
        dr = hf.unpack_uniq_var(dx, mo_occ)

        if WITH_EX_EY_DEGENERACY:
            mol = self._scf.mol
            if mol.symmetry and mol.groupname in ('SO3', 'Dooh', 'Coov'):
                orbsym = hf_symm.get_orbsym(mol, mo_coeff)
                if mol.groupname == 'SO3':
                    _force_SO3_degeneracy_(dr, orbsym)
                else:
                    _force_Ex_Ey_degeneracy_(dr, orbsym)
        return numpy.dot(u0, expmat(dr))

    def rotate_mo(self, mo_coeff, u, log=None):
        mo = numpy.dot(mo_coeff, u)
        if self._scf.mol.symmetry:
            orbsym = hf_symm.get_orbsym(self._scf.mol, mo_coeff)
            mo = lib.tag_array(mo, orbsym=orbsym)

        if isinstance(log, logger.Logger) and log.verbose >= logger.DEBUG:
            idx = self.mo_occ > 0
            s = reduce(numpy.dot, (mo[:,idx].conj().T, self._scf.get_ovlp(),
                                   self.mo_coeff[:,idx]))
            log.debug('Overlap to initial guess, SVD = %s',
                      _effective_svd(s, 1e-5))
            log.debug('Overlap to last step, SVD = %s',
                      _effective_svd(u[idx][:,idx], 1e-5))
        return mo


def newton(mf):
    '''Co-iterative augmented hessian (CIAH) second order SCF solver

    Examples:

    >>> mol = gto.M(atom='H 0 0 0; H 0 0 1.1', basis='cc-pvdz')
    >>> mf = scf.RHF(mol).run(conv_tol=.5)
    >>> mf = scf.newton(mf).set(conv_tol=1e-9)
    >>> mf.kernel()
    -1.0811707843774987
    '''
    from pyscf import scf

    if isinstance(mf, _CIAH_SOSCF):
        return mf

    assert(isinstance(mf, hf.SCF))
    if mf.__doc__ is None:
        mf_doc = ''
    else:
        mf_doc = mf.__doc__

    if isinstance(mf, rohf.ROHF):
        class SecondOrderROHF(_CIAH_SOSCF, mf.__class__):
            __doc__ = mf_doc + _CIAH_SOSCF.__doc__
            gen_g_hop = gen_g_hop_rohf
        return SecondOrderROHF(mf)

    elif isinstance(mf, uhf.UHF):
        class SecondOrderUHF(_CIAH_SOSCF, mf.__class__):
            __doc__ = mf_doc + _CIAH_SOSCF.__doc__

            gen_g_hop = gen_g_hop_uhf

            def update_rotate_matrix(self, dx, mo_occ, u0=1, mo_coeff=None):
                occidxa = mo_occ[0] > 0
                occidxb = mo_occ[1] > 0
                viridxa = ~occidxa
                viridxb = ~occidxb

                nmo = len(occidxa)
                dr = numpy.zeros((2,nmo,nmo), dtype=dx.dtype)
                uniq = numpy.array((viridxa[:,None] & occidxa,
                                    viridxb[:,None] & occidxb))
                dr[uniq] = dx
                dr = dr - dr.conj().transpose(0,2,1)

                if WITH_EX_EY_DEGENERACY:
                    mol = self._scf.mol
                    if mol.symmetry and mol.groupname in ('SO3', 'Dooh', 'Coov'):
                        orbsyma, orbsymb = uhf_symm.get_orbsym(mol, mo_coeff)
                        if mol.groupname == 'SO3':
                            _force_SO3_degeneracy_(dr[0], orbsyma)
                            _force_SO3_degeneracy_(dr[1], orbsymb)
                        else:
                            _force_Ex_Ey_degeneracy_(dr[0], orbsyma)
                            _force_Ex_Ey_degeneracy_(dr[1], orbsymb)

                if isinstance(u0, int) and u0 == 1:
                    return numpy.asarray((expmat(dr[0]), expmat(dr[1])))
                else:
                    return numpy.asarray((numpy.dot(u0[0], expmat(dr[0])),
                                          numpy.dot(u0[1], expmat(dr[1]))))

            def rotate_mo(self, mo_coeff, u, log=None):
                mo = numpy.asarray((numpy.dot(mo_coeff[0], u[0]),
                                    numpy.dot(mo_coeff[1], u[1])))
                if self._scf.mol.symmetry:
                    orbsym = uhf_symm.get_orbsym(self._scf.mol, mo_coeff)
                    mo = lib.tag_array(mo, orbsym=orbsym)
                return mo

            def spin_square(self, mo_coeff=None, s=None):
                if mo_coeff is None:
                    mo_coeff = (self.mo_coeff[0][:,self.mo_occ[0]>0],
                                self.mo_coeff[1][:,self.mo_occ[1]>0])
                if getattr(self, '_scf', None) and self._scf.mol != self.mol:
                    s = self._scf.get_ovlp()
                return self._scf.spin_square(mo_coeff, s)

            def kernel(self, mo_coeff=None, mo_occ=None, dm0=None):
                if isinstance(mo_coeff, numpy.ndarray) and mo_coeff.ndim == 2:
                    mo_coeff = (mo_coeff, mo_coeff)
                if isinstance(mo_occ, numpy.ndarray) and mo_occ.ndim == 1:
                    mo_occ = (numpy.asarray(mo_occ >0, dtype=numpy.double),
                              numpy.asarray(mo_occ==2, dtype=numpy.double))
                return _CIAH_SOSCF.kernel(self, mo_coeff, mo_occ, dm0)

        return SecondOrderUHF(mf)

    elif isinstance(mf, scf.ghf.GHF):
        class SecondOrderGHF(_CIAH_SOSCF, mf.__class__):
            __doc__ = mf_doc + _CIAH_SOSCF.__doc__

            gen_g_hop = gen_g_hop_ghf

            def update_rotate_matrix(self, dx, mo_occ, u0=1, mo_coeff=None):
                dr = hf.unpack_uniq_var(dx, mo_occ)

                if WITH_EX_EY_DEGENERACY:
                    mol = self._scf.mol
                    if mol.symmetry and mol.groupname in ('SO3', 'Dooh', 'Coov'):
                        orbsym = scf.ghf_symm.get_orbsym(mol, mo_coeff)
                        if mol.groupname == 'SO3':
                            _force_SO3_degeneracy_(dr, orbsym)
                        else:
                            _force_Ex_Ey_degeneracy_(dr, orbsym)
                return numpy.dot(u0, expmat(dr))

            def rotate_mo(self, mo_coeff, u, log=None):
                mo = numpy.dot(mo_coeff, u)
                if self._scf.mol.symmetry:
                    orbsym = scf.ghf_symm.get_orbsym(self._scf.mol, mo_coeff)
                    mo = lib.tag_array(mo, orbsym=orbsym)
                return mo
        return SecondOrderGHF(mf)

    elif isinstance(mf, scf.dhf.RDHF):
        class SecondOrderRDHF(_CIAH_SOSCF, mf.__class__):
            __doc__ = mf_doc + _CIAH_SOSCF.__doc__

            gen_g_hop = gen_g_hop_dhf

            def update_rotate_matrix(self, dx, mo_occ, u0=1, mo_coeff=None):
                nmo = mo_occ.size
                nocc = numpy.count_nonzero(mo_occ)
                nvir = nmo - nocc
                dx = dx.reshape(nvir, nocc)
                dx_aa = dx[::2,::2]
                dr_aa = hf.unpack_uniq_var(dx_aa.ravel, mo_occ[::2])
                u = numpy.zeros((nmo, nmo), dtype=dr_aa.dtype)
                # Allows only the rotation within the up-up space and down-down space
                u[::2,::2] = u[1::2,1::2] = expmat(dr_aa)
                return numpy.dot(u0, u)

            def rotate_mo(self, mo_coeff, u, log=None):
                mo = numpy.dot(mo_coeff, u)
                return mo
        return SecondOrderRDHF(mf)

    elif isinstance(mf, scf.dhf.DHF):
        class SecondOrderDHF(_CIAH_SOSCF, mf.__class__):
            __doc__ = mf_doc + _CIAH_SOSCF.__doc__

            gen_g_hop = gen_g_hop_dhf

            def update_rotate_matrix(self, dx, mo_occ, u0=1, mo_coeff=None):
                dr = hf.unpack_uniq_var(dx, mo_occ)
                return numpy.dot(u0, expmat(dr))

            def rotate_mo(self, mo_coeff, u, log=None):
                mo = numpy.dot(mo_coeff, u)
                return mo
        return SecondOrderDHF(mf)

    else:
        class SecondOrderRHF(_CIAH_SOSCF, mf.__class__):
            __doc__ = mf_doc + _CIAH_SOSCF.__doc__
            gen_g_hop = gen_g_hop_rhf
        return SecondOrderRHF(mf)

SVD_TOL = getattr(__config__, 'soscf_newton_ah_effective_svd_tol', 1e-5)
def _effective_svd(a, tol=SVD_TOL):
    w = numpy.linalg.svd(a)[1]
    return w[(tol<w) & (w<1-tol)]
del(SVD_TOL)

def _force_SO3_degeneracy_(dr, orbsym):
    '''Force orbitals of same angular momentum to use the same rotation matrix'''
    orbsym = numpy.asarray(orbsym)
    orbsym_l = orbsym // 100
    lmax = max(orbsym_l)

    for l in range(lmax + 1):
        idx_l = numpy.where(orbsym_l == l)[0]
        nso_l = idx_l.size
        if nso_l > 0:
            degen = l * 2 + 1
            nso_m = nso_l // degen
            dr_l = dr[idx_l[:,None],idx_l].reshape(degen, nso_m, degen, nso_m)
            dr_avg = numpy.einsum('ipiq->pq', dr_l) / degen
            for m in range(degen):
                dr_l[m,:,m,:] = dr_avg
            dr[idx_l[:,None],idx_l] = dr_l.reshape(nso_l, nso_l)
    return dr

def _force_Ex_Ey_degeneracy_(dr, orbsym):
    '''Force the Ex and Ey orbitals to use the same rotation matrix'''
    # 0,1,4,5 are 1D irreps
    E_irrep_ids = set(orbsym).difference(set((0,1,4,5)))
    orbsym = numpy.asarray(orbsym)

    for ir in E_irrep_ids:
        if ir % 2 == 0:
            Ex = orbsym == ir
            Ey = orbsym ==(ir + 1)
            dr_x = dr[Ex[:,None] & Ex]
            dr_y = dr[Ey[:,None] & Ey]
            # In certain open-shell systems, the rotation amplitudes dr_x may
            # be equal to 0 while dr_y are not. In this case, we choose the
            # larger one to represent the rotation amplitudes for both.
            if numpy.linalg.norm(dr_x) > numpy.linalg.norm(dr_y):
                dr[Ey[:,None] & Ey] = dr_x
            else:
                dr[Ex[:,None] & Ex] = dr_y
    return dr


if __name__ == '__main__':
    from pyscf import scf

    mol = gto.Mole()
    mol.verbose = 0
    mol.atom = [
        ['H', ( 1.,-1.    , 0.   )],
        ['H', ( 0.,-1.    ,-1.   )],
        ['H', ( 1.,-0.5   ,-1.   )],
        ['H', ( 0., 1.    , 1.   )],
        ['H', ( 0., 0.5   , 1.   )],
        ['H', ( 1., 0.    ,-1.   )],
    ]

    mol.basis = '6-31g'
    mol.build()

    nmo = mol.nao_nr()
    m = newton(scf.RHF(mol))
    e0 = m.kernel()

#####################################
    mol.basis = '6-31g'
    mol.spin = 2
    mol.build(0, 0)
    m = scf.RHF(mol)
    m.max_cycle = 1
    #m.verbose = 5
    m.scf()
    e1 = kernel(newton(m), m.mo_coeff, m.mo_occ, max_cycle=50, verbose=5)[1]

    m = scf.UHF(mol)
    m.max_cycle = 1
    #m.verbose = 5
    m.scf()
    e2 = kernel(newton(m), m.mo_coeff, m.mo_occ, max_cycle=50, verbose=5)[1]

    m = scf.UHF(mol)
    m.max_cycle = 1
    #m.verbose = 5
    m.scf()
    nrmf = scf.density_fit(newton(m), 'weigend')
    nrmf.max_cycle = 50
    nrmf.conv_tol = 1e-8
    nrmf.conv_tol_grad = 1e-5
    #nrmf.verbose = 5
    e4 = nrmf.kernel()

    m = scf.density_fit(scf.UHF(mol), 'weigend')
    m.max_cycle = 1
    #m.verbose = 5
    m.scf()
    nrmf = scf.density_fit(newton(m), 'weigend')
    nrmf.max_cycle = 50
    nrmf.conv_tol_grad = 1e-5
    e5 = nrmf.kernel()

    m = newton(scf.GHF(mol))
    e6 = m.kernel()

    print(e0 - -2.93707955256)
    print(e1 - -2.99456398848)
    print(e2 - -2.99663808314)
    print(e4 - -2.99663808186)
    print(e5 - -2.99634506072)
    print(e6 - -3.002844505604826)


# pyscf/symm/Dmatrix.py

#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#


'''Wigner rotation D-matrix for real spherical harmonics'''

from math import sqrt
from functools import reduce
import numpy
from scipy.special import factorial
from pyscf.symm import sph


def Dmatrix(l, alpha, beta, gamma, reorder_p=False):
    '''Wigner rotation D-matrix

    D_{mm'} = <lm|R(alpha,beta,gamma)|lm'>
    alpha, beta, gamma are Euler angles (in z-y-z convention)

    Kwargs:
        reorder_p (bool): Whether to put the p functions in the (x,y,z) order.
    '''
    if l == 0:
        return numpy.eye(1)
    else:
        d = dmatrix(l, beta, reorder_p=False)
        ms = numpy.arange(-l, l+1)
        D = numpy.einsum('i,ij,j->ij', numpy.exp(-1j*alpha*ms), d,
                         numpy.exp(-1j*gamma*ms))

        D = _dmat_to_real(l, D, reorder_p=False)
        if reorder_p and l == 1:
            D = D[[2,0,1]][:,[2,0,1]]
        return D

def _dmat_to_real(l, d, reorder_p=False):
    ''' Transform the input D-matrix to make it compatible with the real
    spherical harmonic functions.

    Kwargs:
        reorder_p (bool): Whether to put the p functions in the (x,y,z) order.
    '''
    # The input D matrix works for pure spherical harmonics. The real
    # representation should be U^\dagger * D * U, where U is the unitary
    # matrix that transform the complex harmonics to the real harmonics.
    u = sph.sph_pure2real(l, reorder_p)
    return reduce(numpy.dot, (u.conj().T, d, u)).real

def dmatrix(l, beta, reorder_p=False):
    '''Wigner small-d matrix (in z-y-z convention)'''
    c = numpy.cos(beta/2)
    s = numpy.sin(beta/2)
    if l == 0:
        return numpy.eye(1)
    elif l == 1:
        mat = numpy.array(((c**2        , sqrt(2)*c*s , s**2       ),
                           (-sqrt(2)*c*s, c**2-s**2   , sqrt(2)*c*s),
                           (s**2        , -sqrt(2)*c*s, c**2       )))
        if reorder_p:
            mat = mat[[2,0,1]][:,[2,0,1]]
        return mat
    elif l == 2:
        c3s = c**3*s
        s3c = s**3*c
        c2s2 = (c*s)**2
        c4 = c**4
        s4 = s**4
        s631 = sqrt(6)*(c3s-s3c)
        s622 = sqrt(6)*c2s2
        c4s2 = c4-3*c2s2
        c2s4 = 3*c2s2-s4
        c4s4 = c4-4*c2s2+s4
        return numpy.array((( c4    , 2*c3s, s622, 2*s3c, s4   ),
                            (-2*c3s , c4s2 , s631, c2s4 , 2*s3c),
                            ( s622  ,-s631 , c4s4, s631 , s622 ),
                            (-2*s3c , c2s4 ,-s631, c4s2 , 2*c3s),
                            ( s4    ,-2*s3c, s622,-2*c3s, c4   )))
    else:
        facs = factorial(numpy.arange(2*l+1))
        cs = c**numpy.arange(2*l+1)
        ss = s**numpy.arange(2*l+1)

        mat = numpy.zeros((2*l+1,2*l+1))
        for i,m1 in enumerate(range(-l, l+1)):
            for j,m2 in enumerate(range(-l, l+1)):
                #:fac = sqrt( factorial(l+m1)*factorial(l-m1) \
                #:           *factorial(l+m2)*factorial(l-m2))
                #:for k in range(max(m2-m1,0), min(l+m2, l-m1)+1):
                #:    mat[i,j] += (-1)**(m1+m2+k) \
                #:            * c**(2*l+m2-m1-2*k) * s**(m1-m2+2*k) \
                #:            / (factorial(l+m2-k) * factorial(k) \
                #:               * factorial(m1-m2+k) * factorial(l-m1-k))
                #:mat[i,j] *= fac
                k = numpy.arange(max(m2-m1,0), min(l+m2, l-m1)+1)
                tmp = (cs[2*l+m2-m1-2*k] * ss[m1-m2+2*k] /
                       (facs[l+m2-k] * facs[k] * facs[m1-m2+k] * facs[l-m1-k]))

                mask = ((m1+m2+k) & 0b1).astype(bool)
                mat[i,j] -= tmp[ mask].sum()
                mat[i,j] += tmp[~mask].sum()

        ms = numpy.arange(-l, l+1)
        msfac = numpy.sqrt(facs[l+ms] * facs[l-ms])
        mat *= numpy.einsum('i,j->ij', msfac, msfac)
    return mat


def get_euler_angles(c1, c2):
    '''Find the three Euler angles (alpha, beta, gamma in z-y-z convention)
    that rotates coordinates c1 to coordinates c2.

    yp  = numpy.einsum('j,kj->k', c1[1], geom.rotation_mat(c1[2], beta))
    tmp = numpy.einsum('ij,kj->ik', c1 , geom.rotation_mat(c1[2], alpha))
    tmp = numpy.einsum('ij,kj->ik', tmp, geom.rotation_mat(yp   , beta ))
    c2  = numpy.einsum('ij,kj->ik', tmp, geom.rotation_mat(c2[2], gamma))

    (For backward compatibility) if c1 and c2 are two points in the real
    space, the Euler angles define the rotation transforms the old coordinates
    to the new coordinates (new_x, new_y, new_z) in which c1 is identical to c2.

    tmp = numpy.einsum('j,kj->k', c1 , geom.rotation_mat((0,0,1), gamma))
    tmp = numpy.einsum('j,kj->k', tmp, geom.rotation_mat((0,1,0), beta) )
    c2  = numpy.einsum('j,kj->k', tmp, geom.rotation_mat((0,0,1), alpha))
    '''
    c1 = numpy.asarray(c1)
    c2 = numpy.asarray(c2)
    if c1.ndim == 2 and c2.ndim == 2:
        zz = c1[2].dot(c2[2])
        beta = numpy.arccos(zz)
        if abs(zz) < 1 - 1e-12:
            yp = numpy.cross(c1[2], c2[2])
            yp /= numpy.linalg.norm(yp)
        else:
            yp = c1[1]

        yy = yp.dot(c1[1])
        alpha = numpy.arccos(yy)
        if numpy.cross(c1[1], yp).dot(c1[2]) < 0:
            alpha = -alpha

        gamma = numpy.arccos(yp.dot(c2[1]))
        if numpy.cross(yp, c2[1]).dot(c2[2]) < 0:
            gamma = -gamma

    else: # For backward compatibility, c1 and c2 are two points
        norm1 = numpy.linalg.norm(c1)
        norm2 = numpy.linalg.norm(c2)
        assert(abs(norm1 - norm2) < 1e-12)

        xy_norm = numpy.linalg.norm(c1[:2])
        if xy_norm > 1e-12:
            gamma = -numpy.arccos(c1[0] / xy_norm)
        else:
            gamma = 0

        xy_norm = numpy.linalg.norm(c2[:2])
        if xy_norm > 1e-12:
            alpha = numpy.arccos(c2[0] / xy_norm)
        else:
            alpha = 0

        beta = numpy.arccos(c2[2]/norm1) - numpy.arccos(c1[2]/norm2)

    return alpha, beta, gamma



# pyscf/symm/__init__.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
This module offers the functions to detect point group symmetry, basis
symmetriziation, Clebsch-Gordon coefficients.  This module works as a plugin of
PySCF package.  Symmetry is not hard coded in each method.
'''

from pyscf.symm import param
from pyscf.symm import geom
from pyscf.symm import basis
from pyscf.symm import cg

from pyscf.symm.geom import *
from pyscf.symm.basis import *
from pyscf.symm.addons import *


# pyscf/symm/addons.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

from functools import reduce
import numpy
import scipy.linalg
from pyscf import lib
from pyscf.lib import logger
from pyscf.lib.exceptions import PointGroupSymmetryError
from pyscf.symm import basis
from pyscf.symm import param
from pyscf import __config__


MULTI_IRREPS = -1

def label_orb_symm(mol, irrep_name, symm_orb, mo, s=None,
                   check=getattr(__config__, 'symm_addons_label_orb_symm_check', True),
                   tol=getattr(__config__, 'symm_addons_label_orb_symm_tol', 1e-9)):
    '''Label the symmetry of given orbitals

    irrep_name can be either the symbol or the ID of the irreducible
    representation.  If the ID is provided, it returns the numeric code
    associated with XOR operator, see :py:meth:`symm.param.IRREP_ID_TABLE`

    Args:
        mol : an instance of :class:`Mole`

        irrep_name : list of str or int
            A list of irrep ID or name,  it can be either mol.irrep_id or
            mol.irrep_name.  It can affect the return "label".
        symm_orb : list of 2d array
            the symmetry adapted basis
        mo : 2d array
            the orbitals to label

    Returns:
        list of symbols or integers to represent the irreps for the given
        orbitals

    Examples:

    >>> from pyscf import gto, scf, symm
    >>> mol = gto.M(atom='H 0 0 0; H 0 0 1', basis='ccpvdz',verbose=0, symmetry=1)
    >>> mf = scf.RHF(mol)
    >>> mf.kernel()
    >>> symm.label_orb_symm(mol, mol.irrep_name, mol.symm_orb, mf.mo_coeff)
    ['Ag', 'B1u', 'Ag', 'B1u', 'B2u', 'B3u', 'Ag', 'B2g', 'B3g', 'B1u']
    >>> symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, mf.mo_coeff)
    [0, 5, 0, 5, 6, 7, 0, 2, 3, 5]
    '''
    nmo = mo.shape[1]
    if s is None:
        s = mol.intor_symmetric('int1e_ovlp')
    s_mo = numpy.dot(s, mo)
    norm = numpy.zeros((len(irrep_name), nmo))
    for i, csym in enumerate(symm_orb):
        moso = numpy.dot(csym.T.conj(), s_mo)
        ovlpso = reduce(numpy.dot, (csym.T.conj(), s, csym))
        try:
            s_moso = lib.cho_solve(ovlpso, moso)
        except numpy.linalg.LinAlgError:
            ovlpso[numpy.diag_indices(csym.shape[1])] += 1e-12
            s_moso = lib.cho_solve(ovlpso, moso)
        norm[i] = numpy.einsum('ki,ki->i', moso.conj(), s_moso).real
    norm /= numpy.sum(norm, axis=0)  # for orbitals which are not normalized
    iridx = numpy.argmax(norm, axis=0)
    orbsym = numpy.asarray([irrep_name[i] for i in iridx])
    logger.debug(mol, 'irreps of each MO %s', orbsym)

    if check:
        largest_norm = norm[iridx,numpy.arange(nmo)]
        orbidx = numpy.where(largest_norm < 1-tol)[0]
        if orbidx.size > 0:
            idx = numpy.where(largest_norm < 1-tol*1e2)[0]
            if idx.size > 0:
                raise ValueError('orbitals %s not symmetrized, norm = %s' %
                                 (idx, largest_norm[idx]))
            else:
                logger.warn(mol, 'orbitals %s not strictly symmetrized.',
                            numpy.unique(orbidx))
                logger.warn(mol, 'They can be symmetrized with '
                            'pyscf.symm.symmetrize_space function.')
                logger.debug(mol, 'norm = %s', largest_norm[orbidx])
    return orbsym

def symmetrize_orb(mol, mo, orbsym=None, s=None,
                   check=getattr(__config__, 'symm_addons_symmetrize_orb_check', False)):
    '''Symmetrize the given orbitals.

    This function is different to the :func:`symmetrize_space`:  In this
    function, each orbital is symmetrized by removing non-symmetric components.
    :func:`symmetrize_space` symmetrizes the entire space by mixing different
    orbitals.

    Note this function might return non-orthorgonal orbitals.
    Call :func:`symmetrize_space` to find the symmetrized orbitals that are
    close to the given orbitals.

    Args:
        mo : 2D float array
            The orbital space to symmetrize

    Kwargs:
        orbsym : integer list
            Irrep id for each orbital.  If not given, the irreps are guessed
            by calling :func:`label_orb_symm`.
        s : 2D float array
            Overlap matrix.  If given, use this overlap than the the overlap
            of the input mol.

    Returns:
        2D orbital coefficients

    Examples:

    >>> from pyscf import gto, symm, scf
    >>> mol = gto.M(atom = 'C  0  0  0; H  1  1  1; H -1 -1  1; H  1 -1 -1; H -1  1 -1',
    ...             basis = 'sto3g')
    >>> mf = scf.RHF(mol).run()
    >>> mol.build(0, 0, symmetry='D2')
    >>> mo = symm.symmetrize_orb(mol, mf.mo_coeff)
    >>> print(symm.label_orb_symm(mol, mol.irrep_name, mol.symm_orb, mo))
    ['A', 'A', 'B1', 'B2', 'B3', 'A', 'B1', 'B2', 'B3']
    '''
    if s is None:
        s = mol.intor_symmetric('int1e_ovlp')
    if orbsym is None:
        orbsym = label_orb_symm(mol, mol.irrep_id, mol.symm_orb,
                                mo, s=s, check=check)
    orbsym = numpy.asarray(orbsym)
    s_mo = numpy.dot(s, mo)
    mo1 = numpy.empty_like(mo)

    if orbsym[0] in mol.irrep_name:
        irrep_id = mol.irrep_name
    else:
        irrep_id = mol.irrep_id

    for i, ir in enumerate(irrep_id):
        idx = orbsym == ir
        csym = mol.symm_orb[i]
        ovlpso = reduce(numpy.dot, (csym.T.conj(), s, csym))
        sc = lib.cho_solve(ovlpso, numpy.dot(csym.T.conj(), s_mo[:,idx]))
        mo1[:,idx] = numpy.dot(csym, sc)
    return mo1

def symmetrize_space(mol, mo, s=None,
                     check=getattr(__config__, 'symm_addons_symmetrize_space_check', True),
                     tol=getattr(__config__, 'symm_addons_symmetrize_space_tol', 1e-7)):
    '''Symmetrize the given orbital space.

    This function is different to the :func:`symmetrize_orb`:  In this function,
    the given orbitals are mixed to reveal the symmtery; :func:`symmetrize_orb`
    projects out non-symmetric components for each orbital.

    Args:
        mo : 2D float array
            The orbital space to symmetrize

    Kwargs:
        s : 2D float array
            Overlap matrix.  If not given, overlap is computed with the input mol.

    Returns:
        2D orbital coefficients

    Examples:

    >>> from pyscf import gto, symm, scf
    >>> mol = gto.M(atom = 'C  0  0  0; H  1  1  1; H -1 -1  1; H  1 -1 -1; H -1  1 -1',
    ...             basis = 'sto3g')
    >>> mf = scf.RHF(mol).run()
    >>> mol.build(0, 0, symmetry='D2')
    >>> mo = symm.symmetrize_space(mol, mf.mo_coeff)
    >>> print(symm.label_orb_symm(mol, mol.irrep_name, mol.symm_orb, mo))
    ['A', 'A', 'A', 'B1', 'B1', 'B2', 'B2', 'B3', 'B3']
    '''
    from pyscf.tools import mo_mapping
    if s is None:
        s = mol.intor_symmetric('int1e_ovlp')
    nmo = mo.shape[1]
    s_mo = numpy.dot(s, mo)
    if check and abs(numpy.dot(mo.conj().T, s_mo) - numpy.eye(nmo)).max() > tol:
        raise ValueError('Orbitals are not orthogonalized')

    mo1 = []
    for i, csym in enumerate(mol.symm_orb):
        moso = numpy.dot(csym.T.conj(), s_mo)
        ovlpso = reduce(numpy.dot, (csym.T.conj(), s, csym))

        # excluding orbitals which are already symmetrized
        try:
            diag = numpy.einsum('ki,ki->i', moso.conj(), lib.cho_solve(ovlpso, moso))
        except numpy.linalg.LinAlgError:
            ovlpso[numpy.diag_indices(csym.shape[1])] += 1e-12
            diag = numpy.einsum('ki,ki->i', moso.conj(), lib.cho_solve(ovlpso, moso))
        idx = abs(1-diag) < 1e-8
        orb_exclude = mo[:,idx]
        mo1.append(orb_exclude)
        moso1 = moso[:,~idx]
        dm = numpy.dot(moso1, moso1.T.conj())

        if dm.trace() > 1e-8:
            e, u = scipy.linalg.eigh(dm, ovlpso)
            mo1.append(numpy.dot(csym, u[:,abs(1-e) < 1e-6]))
    mo1 = numpy.hstack(mo1)
    if mo1.shape[1] != nmo:
        raise ValueError('The input orbital space is not symmetrized.\n One '
                         'possible reason is that the input mol and orbitals '
                         'are of different orientation.')
    if (check and
        abs(reduce(numpy.dot, (mo1.conj().T, s, mo1)) - numpy.eye(nmo)).max() > tol):
        raise ValueError('Orbitals are not orthogonalized')
    idx = mo_mapping.mo_1to1map(reduce(numpy.dot, (mo.T.conj(), s, mo1)))
    return mo1[:,idx]

def std_symb(gpname):
    '''std_symb('d2h') returns D2h; std_symb('D2H') returns D2h'''
    if gpname == 'SO3':
        return gpname
    else:
        return str(gpname[0].upper() + gpname[1:].lower())

def irrep_name2id(gpname, symb):
    '''Convert the irrep symbol to internal irrep ID

    Args:
        gpname : str
            The point group symbol
        symb : str
            Irrep symbol

    Returns:
        Irrep ID, int
    '''
    gpname = std_symb(gpname)
    symb = std_symb(symb)
    if gpname == 'SO3':
        return basis.so3_irrep_symb2id(symb)
    elif gpname in ('Dooh', 'Coov'):
        return basis.linearmole_irrep_symb2id(gpname, symb)
    else:
        return param.IRREP_ID_TABLE[gpname][symb]

def irrep_id2name(gpname, irrep_id):
    '''Convert the internal irrep ID to irrep symbol

    Args:
        gpname : str
            The point group symbol
        irrep_id : int
            See IRREP_ID_TABLE in pyscf/symm/param.py

    Returns:
        Irrep sybmol, str
    '''
    gpname = std_symb(gpname)
    if gpname == 'SO3':
        return basis.so3_irrep_id2symb(irrep_id)
    elif gpname in ('Dooh', 'Coov'):
        return basis.linearmole_irrep_id2symb(gpname, irrep_id)
    else:
        # irrep_id may be obtained from high symmetry (Dooh, Coov)
        irrep_id_in_d2h = irrep_id % 10
        return param.CHARACTER_TABLE[gpname][irrep_id_in_d2h][0]

def irrep_name(pgname, irrep_id):
    raise PointGroupSymmetryError('This function was obsoleted. Use irrep_id2name')

def route(target, nelec, orbsym):
    '''Pick orbitals to form a determinant which has the right symmetry.
    If solution is not found, return []
    '''
    def riter(target, nelec, orbsym):
        if nelec == 1:
            if target in orbsym:
                return [orbsym.index(target)]
            else:
                return []
        else:
            for i, ir in enumerate(orbsym):
                off = i + 1
                orb_left = orbsym[off:]
                res = riter(target ^ ir, nelec-1, orb_left)
                if res:
                    return [i] + [off+x for x in res]
            return []
    if isinstance(orbsym, numpy.ndarray):
        orbsym = orbsym.tolist()
    return riter(target, nelec, orbsym)

def eigh(h, orbsym):
    '''Solve eigenvalue problem based on the symmetry information for basis.
    See also pyscf/lib/linalg_helper.py :func:`eigh_by_blocks`

    Examples:

    >>> from pyscf import gto, symm
    >>> mol = gto.M(atom='H 0 0 0; H 0 0 1', basis='ccpvdz', symmetry=True)
    >>> c = numpy.hstack(mol.symm_orb)
    >>> vnuc_so = reduce(numpy.dot, (c.T, mol.intor('int1e_nuc_sph'), c))
    >>> orbsym = symm.label_orb_symm(mol, mol.irrep_name, mol.symm_orb, c)
    >>> symm.eigh(vnuc_so, orbsym)
    (array([-4.50766885, -1.80666351, -1.7808565 , -1.7808565 , -1.74189134,
            -0.98998583, -0.98998583, -0.40322226, -0.30242374, -0.07608981]),
     ...)
    '''
    return lib.eigh_by_blocks(h, labels=orbsym)

def direct_prod(orbsym1, orbsym2, groupname='D2h'):
    if groupname == 'SO3':
        prod = orbsym1[:,None] ^ orbsym2
        orbsym1_not_s = orbsym1 != 0
        orbsym2_not_s = orbsym2 != 0
        prod[orbsym1_not_s[:,None] & orbsym2_not_s != 0] = MULTI_IRREPS
        prod[orbsym1[:,None] == orbsym2] = 0
    elif groupname == 'Dooh':
        orbsym1_octa = (orbsym1 // 10) * 8 + orbsym1 % 10
        orbsym2_octa = (orbsym2 // 10) * 8 + orbsym2 % 10
        prod = orbsym1_octa[:,None] ^ orbsym2_octa
        prod = (prod % 8) + (prod // 8) * 10
        orbsym1_irrepE = (orbsym1 >= 2) & (orbsym1 != 4) & (orbsym1 != 5)
        orbsym2_irrepE = (orbsym2 >= 2) & (orbsym2 != 4) & (orbsym2 != 5)
        prod[orbsym1_irrepE[:,None] & orbsym2_irrepE] = MULTI_IRREPS
        prod[orbsym1[:,None] == orbsym2] = 0
    elif groupname == 'Coov':
        prod = orbsym1[:,None] ^ orbsym2
        orbsym1_irrepE = orbsym1 >= 2
        orbsym2_irrepE = orbsym2 >= 2
        prod[orbsym1_irrepE[:,None] & orbsym2_irrepE] = MULTI_IRREPS
        prod[orbsym1[:,None] == orbsym2] = 0
    else:  # D2h and subgroup
        prod = orbsym1[:,None] ^ orbsym2
    return prod


if __name__ == "__main__":
    from pyscf import gto
    from pyscf import scf
    mol = gto.Mole()
    mol.build(
        atom = [['H', (0,0,0)], ['H', (0,0,1)]],
        basis = {'H': 'cc-pvdz'},
        symmetry = 1
    )
    mf = scf.RHF(mol)
    mf.scf()

    nao, nmo = mf.mo_coeff.shape
    print(label_orb_symm(mol, mol.irrep_name, mol.symm_orb, mf.mo_coeff))
    numpy.random.seed(1)
    u = numpy.random.random((nmo,nmo))*1e-2
    u = scipy.linalg.expm(u - u.T)
    mo = symmetrize_orb(mol, numpy.dot(mf.mo_coeff, u))
    print(label_orb_symm(mol, mol.irrep_name, mol.symm_orb, mo))

    orbsym = [0, 3, 0, 2, 5, 6]
    res = route(7, 3, orbsym)
    print(res, reduce(lambda x, y: x ^ y, [orbsym[i] for i in res]))


# pyscf/symm/basis.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Generate symmetry adapted basis
'''

from functools import reduce
import numpy
from pyscf.data.elements import _symbol, _rm_digit
from pyscf import gto
from pyscf.lib.exceptions import PointGroupSymmetryError
from pyscf.symm import geom
from pyscf.symm import param

__all__ = ['tot_parity_odd',
           'symm_adapted_basis',
           'dump_symm_adapted_basis',
           'symmetrize_matrix',
           'linearmole_symm_descent',
           'linearmole_irrep_symb2id',
           'linearmole_irrep_id2symb',
           'linearmole_symm_adapted_basis',
           'so3_irrep_symb2id',
           'so3_irrep_id2symb',]

OP_PARITY_ODD = {
    'E'  : (0, 0, 0),
    'C2x': (0, 1, 1),
    'C2y': (1, 0, 1),
    'C2z': (1, 1, 0),
    'i'  : (1, 1, 1),
    'sx' : (1, 0, 0),
    'sy' : (0, 1, 0),
    'sz' : (0, 0, 1),
}

def tot_parity_odd(op, l, m):
    if op == 'E':
        return 0
    else:
        ox,oy,oz = OP_PARITY_ODD[op]
        gx,gy,gz = param.SPHERIC_GTO_PARITY_ODD[l][l+m]
        return (ox and gx) ^ (oy and gy) ^ (oz and gz)

def symm_adapted_basis(mol, gpname, orig=0, coordinates=None):
    if gpname == 'SO3':
        return so3_symm_adapted_basis(mol, gpname, orig, coordinates)
    elif gpname in ('Dooh', 'Coov'):
        return linearmole_symm_adapted_basis(mol, gpname, orig, coordinates)

    # prop_atoms are the atoms relocated wrt the charge center with proper
    # orientation
    if coordinates is None:
        coordinates = numpy.eye(3)
    prop_atoms = mol.format_atom(mol._atom, orig, coordinates, 'Bohr')
    eql_atom_ids = geom.symm_identical_atoms(gpname, prop_atoms)

    ops = numpy.asarray([param.D2H_OPS[op] for op in param.OPERATOR_TABLE[gpname]])
    chartab = numpy.array([x[1:] for x in param.CHARACTER_TABLE[gpname]])
    nirrep = chartab.__len__()
    aoslice = mol.aoslice_by_atom()
    nao = mol.nao_nr()
    atom_coords = numpy.array([a[1] for a in prop_atoms])

    sodic = [[] for i in range(8)]
    for atom_ids in eql_atom_ids:
        r0 = atom_coords[atom_ids[0]]
        op_coords = numpy.einsum('x,nxy->ny', r0, ops)
# Using ops to generate other atoms from atom_ids[0]
        coords0 = atom_coords[atom_ids]
        dc = abs(op_coords.reshape(-1,1,3) - coords0).max(axis=2)
        op_relate_idx = numpy.argwhere(dc < geom.TOLERANCE)[:,1]
        ao_loc = numpy.array([aoslice[atom_ids[i],2] for i in op_relate_idx])

        b0, b1 = aoslice[atom_ids[0],:2]
        ip = 0
        for ib in range(b0, b1):
            l = mol.bas_angular(ib)
            if mol.cart:
                degen = (l + 1) * (l + 2) // 2
                cbase = numpy.zeros((degen,nirrep,nao))
                for op_id, op in enumerate(ops):
                    n = 0
                    for x in range(l, -1, -1):
                        for y in range(l-x, -1, -1):
                            z = l-x-y
                            idx = ao_loc[op_id] + n
                            sign = op[0,0]**x * op[1,1]**y * op[2,2]**z
                            cbase[n,:,idx] += sign * chartab[:,op_id]
                            n += 1
            else:
                degen = l * 2 + 1
                cbase = numpy.zeros((degen,nirrep,nao))
                for op_id, op in enumerate(param.OPERATOR_TABLE[gpname]):
                    for n, m in enumerate(range(-l, l+1)):
                        idx = ao_loc[op_id] + n
                        if tot_parity_odd(op, l, m):
                            cbase[n,:,idx] -= chartab[:,op_id]
                        else:
                            cbase[n,:,idx] += chartab[:,op_id]
            norms = numpy.sqrt(numpy.einsum('mij,mij->mi', cbase, cbase))

            for i in range(mol.bas_nctr(ib)):
                for n, ir in numpy.argwhere(norms > 1e-12):
                    c = numpy.zeros(nao)
                    c[ip:] = cbase[n,ir,:nao-ip] / norms[n,ir]
                    sodic[ir].append(c)
                ip += degen

    ao_loc = mol.ao_loc_nr()
    l_idx = {}
    ANG_OF = 1
    for l in range(mol._bas[:,ANG_OF].max()+1):
        idx = [numpy.arange(ao_loc[ib], ao_loc[ib+1])
               for ib in numpy.where(mol._bas[:,ANG_OF] == l)[0]]
        if idx:
            l_idx[l] = numpy.hstack(idx)

    Ds = _ao_rotation_matrices(mol, coordinates)
    so = []
    irrep_ids = []
    for ir, c in enumerate(sodic):
        if len(c) > 0:
            irrep_ids.append(ir)
            c_ir = numpy.vstack(c).T
            nso = c_ir.shape[1]
            for l, idx in l_idx.items():
                c = c_ir[idx].reshape(-1,Ds[l].shape[1],nso)
                c_ir[idx] = numpy.einsum('nm,smp->snp', Ds[l], c).reshape(-1,nso)

            so.append(c_ir)

    return so, irrep_ids

def _ao_rotation_matrices(mol, axes):
    '''Cache the rotation matrices'''
    from pyscf import lib
    from pyscf.symm.Dmatrix import Dmatrix, get_euler_angles
    alpha, beta, gamma = get_euler_angles(numpy.eye(3), axes)
    ANG_OF = 1
    l_max = mol._bas[:,ANG_OF].max()

    if not mol.cart:
        return [Dmatrix(l, alpha, beta, gamma, reorder_p=True)
                for l in range(l_max+1)]

    pp = Dmatrix(1, alpha, beta, gamma, reorder_p=True)

    Ds = [numpy.ones((1,1))]
    for l in range(1, l_max+1):
        # All possible x,y,z combinations
        cidx = numpy.sort(lib.cartesian_prod([(0, 1, 2)] * l), axis=1)

        addr = 0
        affine = numpy.ones((1,1))
        for i in range(l):
            nd = affine.shape[0] * 3
            affine = numpy.einsum('ik,jl->ijkl', affine, pp).reshape(nd, nd)
            addr = addr * 3 + cidx[:,i]

        uniq_addr, rev_addr = numpy.unique(addr, return_inverse=True)
        ncart = (l + 1) * (l + 2) // 2
        assert ncart == uniq_addr.size
        trans = numpy.zeros((ncart,ncart))
        for i, k in enumerate(rev_addr):
            trans[k] += affine[i,uniq_addr]
        Ds.append(trans)
    return Ds

def dump_symm_adapted_basis(mol, so):
    raise PointGroupSymmetryError('TODO')

def symmetrize_matrix(mat, so):
    return [reduce(numpy.dot, (c.conj().T,mat,c)) for c in so]

def _basis_offset_for_atoms(atoms, basis_tab):
    basoff = [0]
    n = 0
    for at in atoms:
        symb = _symbol(at[0])
        if symb in basis_tab:
            bas0 = basis_tab[symb]
        else:
            bas0 = basis_tab[_rm_digit(symb)]
        for b in bas0:
            angl = b[0]
            n += _num_contract(b) * (angl*2+1)
        basoff.append(n)
    return n, basoff

def _num_contract(basis):
    if isinstance(basis[1], int):
        # This branch should never be reached if basis_tab is formated by
        # function mole.format_basis
        nctr = len(basis[2]) - 1
    else:
        nctr = len(basis[1]) - 1
    return nctr

###############################
# SO3 (real spherical harmonics)
# Irreps ID maps
# SO3       ->  Dooh
# s       0 ->  A1g   0

# pz    105 ->  A1u   5
# py    106 ->  E1uy  6
# px    107 ->  E1ux  7

# dz2   200 ->  A1g   0
# dyz   203 ->  E1gy  3
# dxz   202 ->  E1gx  2
# dxy   211 ->  E2gy  11
# dx2y2 210 ->  E2gx  10

# f0    305 ->  A1u   5
# f-1   306 ->  E1uy  6
# f+1   307 ->  E1ux  7
# f-2   314 ->  E2uy  14
# f+2   315 ->  E2ux  15
# f-3   316 ->  E3uy  16
# f+3   317 ->  E3ux  17

# g0    400 ->  A1g   0
# g-1   403 ->  E1gy  3
# g+1   402 ->  E1gx  2
# g-2   411 ->  E2gy  11
# g+2   410 ->  E2gx  10
# g-3   413 ->  E3gy  13
# g+3   412 ->  E3gx  12
# g-4   421 ->  E4gy  21
# g+4   420 ->  E4gx  20
_SO3_SYMB2ID = {
    's+0' :     0,
    'p-1':    106,
    'p+0':    105,
    'p+1':    107,
    'd-2':    211,
    'd-1':    203,
    'd+0':    200,
    'd+1':    202,
    'd+2':    210,
    'f-3':    316,
    'f-2':    314,
    'f-1':    306,
    'f+0':    305,
    'f+1':    307,
    'f+2':    315,
    'f+3':    317,
    'g-4':    421,
    'g-3':    413,
    'g-2':    411,
    'g-1':    403,
    'g+0':    400,
    'g+1':    402,
    'g+2':    410,
    'g+3':    412,
    'g+4':    420,
    'h-5':    526,
    'h-4':    524,
    'h-3':    516,
    'h-2':    514,
    'h-1':    506,
    'h+0':    505,
    'h+1':    507,
    'h+2':    515,
    'h+3':    517,
    'h+4':    525,
    'h+5':    527,
    'i-6':    631,
    'i-5':    623,
    'i-4':    621,
    'i-3':    613,
    'i-2':    611,
    'i-1':    603,
    'i+0':    600,
    'i+1':    602,
    'i+2':    610,
    'i+3':    612,
    'i+4':    620,
    'i+5':    622,
    'i+6':    630,
}
_SO3_ID2SYMB = dict([(v, k) for k, v in _SO3_SYMB2ID.items()])
_ANGULAR = 'spdfghiklmnortu'

def so3_irrep_symb2id(symb):
    symb = symb.lower()
    if symb in _SO3_SYMB2ID:
        return _SO3_SYMB2ID[symb]
    else:
        s = ''.join([i for i in symb if i.isalpha()])
        n = int(''.join([i for i in symb if not i.isalpha()]))
        e2 = abs(n) // 2
        l = _ANGULAR.index(s)
        if abs(n) > l:
            raise KeyError(symb)
        gu = 'u' if l % 2 else 'g'
        xy = 'y' if n < 0 else 'x'
        parity = '_odd' if n % 2 else '_even'
        return l * 100 + e2 * 10 + DOOH_IRREP_ID_TABLE[parity + gu + xy]

def so3_irrep_id2symb(irrep_id):
    if irrep_id in _SO3_ID2SYMB:
        return _SO3_ID2SYMB[irrep_id]
    else:
        l, e2 = divmod(irrep_id, 100)
        e2, n = divmod(e2, 10)
        if n in (0, 1, 5, 4):
            rn = e2 * 2
        elif n in (2, 3, 6, 7):
            rn = e2 * 2 + 1
        else:
            raise KeyError(f'Unknown Irrep ID {irrep_id} for SO3 group')
        if l <= 1 or rn > l:
            raise KeyError(f'Unknown Irrep ID {irrep_id} for SO3 group')
        if n in (6, 4, 3, 1):
            rn = -rn
        return f'{_ANGULAR[l]}{rn:+d}'

def so3_symm_adapted_basis(mol, gpname, orig=0, coordinates=None):
    assert gpname == 'SO3'
    assert mol.natm == 1

    ao_loc = mol.ao_loc_nr(cart=False)
    nao_sph = ao_loc[-1]

    if mol.cart:
        coeff = mol.cart2sph(normalized='sp')
    else:
        coeff = numpy.eye(nao_sph)
    nao = coeff.shape[0]

    lmax = max(mol._bas[:,gto.ANG_OF])
    so = []
    irrep_names = []
    for l in range(lmax+1):
        bas_idx = mol._bas[:,gto.ANG_OF] == l
        if sum(bas_idx) == 0: # Skip any unused angular momentum channels
            continue
        cs = [coeff[:,p0:p1]
              for p0, p1 in zip(ao_loc[:-1][bas_idx], ao_loc[1:][bas_idx])]
        c_groups = numpy.hstack(cs).reshape(nao, -1, l*2+1)
        if l == 1:
            so.extend([c_groups[:,:,1], c_groups[:,:,2], c_groups[:,:,0]])
            irrep_names.extend(['p-1', 'p+0', 'p+1'])
        else:
            for m in range(-l, l+1):
                so.append(c_groups[:,:,l+m])
                irrep_names.append('%s%+d' % (_ANGULAR[l], m))

    irrep_ids = [so3_irrep_symb2id(ir) for ir in irrep_names]
    return so, irrep_ids


###############################
# Linear molecule
# Irreps ID maps
# Dooh     ->  D2h        |   Coov      -> C2v
# A1g   0      Ag    0    |   A1    0      A1    0
# A2g   1      B1g   1    |   A2    1      A2    1
# A1u   5      B1u   5    |   E1x   2      B1    2
# A2u   4      Au    4    |   E1y   3      B2    3
# E1gx  2      B2g   2    |   E2x   10     A1    0
# E1gy  3      B3g   3    |   E2y   11     A2    1
# E1ux  7      B3u   7    |   E3x   12     B1    2
# E1uy  6      B2u   6    |   E3y   13     B2    3
# E2gx  10     Ag    0    |   E4x   20     A1    0
# E2gy  11     B1g   1    |   E4y   21     A2    1
# E2ux  15     B1u   5    |   E5x   22     B1    2
# E2uy  14     Au    4    |   E5y   23     B2    3
# E3gx  12     B2g   2    |
# E3gy  13     B3g   3    |
# E3ux  17     B3u   7    |
# E3uy  16     B2u   6    |
# E4gx  20     Ag    0    |
# E4gy  21     B1g   1    |
# E4ux  25     B1u   5    |
# E4uy  24     Au    4    |
# E5gx  22     B2g   2    |
# E5gy  23     B3g   3    |
# E5ux  27     B3u   7    |
# E5uy  26     B2u   6    |

DOOH_IRREP_ID_TABLE = {
    'A1g' : 0,
    'A2g' : 1,
    'A1u' : 5,
    'A2u' : 4,
    'E1gx': 2,
    'E1gy': 3,
    'E1ux': 7,
    'E1uy': 6,
    '_evengx': 0,
    '_evengy': 1,
    '_evenux': 5,
    '_evenuy': 4,
    '_oddgx': 2,
    '_oddgy': 3,
    '_oddux': 7,
    '_odduy': 6,
}
COOV_IRREP_ID_TABLE = {
    'A1' : 0,
    'A2' : 1,
    'E1x': 2,
    'E1y': 3,
    '_evenx': 0,
    '_eveny': 1,
    '_oddx': 2,
    '_oddy': 3,
}

def linearmole_symm_descent(gpname, irrep_id):
    '''Map irreps to D2h or C2v'''
    if gpname in ('Dooh', 'Coov'):
        return irrep_id % 10
    else:
        raise PointGroupSymmetryError('%s is not proper for linear molecule.' % gpname)

def linearmole_irrep_symb2id(gpname, symb):
    if gpname == 'Dooh':
        if symb in DOOH_IRREP_ID_TABLE:
            return DOOH_IRREP_ID_TABLE[symb]
        else:
            n = int(''.join([i for i in symb if i.isdigit()]))
            if n % 2:
                return (n//2)*10 + DOOH_IRREP_ID_TABLE['_odd'+symb[-2:]]
            else:
                return (n//2)*10 + DOOH_IRREP_ID_TABLE['_even'+symb[-2:]]
    elif gpname == 'Coov':
        if symb in COOV_IRREP_ID_TABLE:
            return COOV_IRREP_ID_TABLE[symb]
        else:
            n = int(''.join([i for i in symb if i.isdigit()]))
            if n % 2:
                return (n//2)*10 + COOV_IRREP_ID_TABLE['_odd'+symb[-1]]
            else:
                return (n//2)*10 + COOV_IRREP_ID_TABLE['_even'+symb[-1]]
    else:
        raise PointGroupSymmetryError('%s is not proper for linear molecule.' % gpname)

DOOH_IRREP_SYMBS = ('A1g' , 'A2g' , 'E1gx', 'E1gy' , 'A2u', 'A1u' , 'E1uy', 'E1ux')
DOOH_IRREP_SYMBS_EXT = ('gx' , 'gy' , 'gx', 'gy' , 'uy', 'ux' , 'uy', 'ux')
COOV_IRREP_SYMBS = ('A1' , 'A2' , 'E1x', 'E1y')
def linearmole_irrep_id2symb(gpname, irrep_id):
    if gpname == 'Dooh':
        if irrep_id < 10:
            return DOOH_IRREP_SYMBS[irrep_id]
        else:
            n = irrep_id % 10
            m = irrep_id // 10
            if n in (0, 1, 5, 4):
                rn = m*2
            else:
                rn = m*2+1
            return 'E%d%s' % (rn, DOOH_IRREP_SYMBS_EXT[n])
    elif gpname == 'Coov':
        if irrep_id < 10:
            return COOV_IRREP_SYMBS[irrep_id]
        else:
            n = irrep_id % 10
            m = irrep_id // 10
            if n < 2:
                rn = m*2
            else:
                rn = m*2+1
            if n % 2:
                xy = 'y'
            else:
                xy = 'x'
            return 'E%d%s' % (rn, xy)
    else:
        raise PointGroupSymmetryError('%s is not proper for linear molecule.' % gpname)

def linearmole_symm_adapted_basis(mol, gpname, orig=0, coordinates=None):
    assert(gpname in ('Dooh', 'Coov'))
    assert(not mol.cart)

    if coordinates is None:
        coordinates = numpy.eye(3)
    prop_atoms = mol.format_atom(mol._atom, orig, coordinates, 'Bohr')
    eql_atom_ids = geom.symm_identical_atoms(gpname, prop_atoms)

    aoslice = mol.aoslice_by_atom()
    basoff = aoslice[:,2]
    nao = mol.nao_nr()
    sodic = {}
    shalf = numpy.sqrt(.5)
    def plus(i0, i1):
        c = numpy.zeros(nao)
        c[i0] = c[i1] = shalf
        return c
    def minus(i0, i1):
        c = numpy.zeros(nao)
        c[i0] = shalf
        c[i1] =-shalf
        return c
    def identity(i0):
        c = numpy.zeros(nao)
        c[i0] = 1
        return c
    def add_so(irrep_name, c):
        if irrep_name in sodic:
            sodic[irrep_name].append(c)
        else:
            sodic[irrep_name] = [c]

    if gpname == 'Dooh':
        for atom_ids in eql_atom_ids:
            if len(atom_ids) == 2:
                at0 = atom_ids[0]
                at1 = atom_ids[1]
                ip = 0
                b0, b1, p0, p1 = aoslice[at0]
                for ib in range(b0, b1):
                    angl = mol.bas_angular(ib)
                    nc = mol.bas_nctr(ib)
                    degen = angl * 2 + 1
                    if angl == 1:
                        for i in range(nc):
                            aoff = ip + i*degen + angl
# m = 0
                            idx0 = basoff[at0] + aoff + 1
                            idx1 = basoff[at1] + aoff + 1
                            add_so('A1g', minus(idx0, idx1))
                            add_so('A1u', plus (idx0, idx1))
# m = +/- 1
                            idx0 = basoff[at0] + aoff - 1
                            idy0 = basoff[at0] + aoff
                            idx1 = basoff[at1] + aoff - 1
                            idy1 = basoff[at1] + aoff
                            add_so('E1ux', plus (idx0, idx1))
                            add_so('E1uy', plus (idy0, idy1))
                            add_so('E1gx', minus(idx0, idx1))
                            add_so('E1gy', minus(idy0, idy1))
                    else:
                        for i in range(nc):
                            aoff = ip + i*degen + angl
# m = 0
                            idx0 = basoff[at0] + aoff
                            idx1 = basoff[at1] + aoff
                            if angl % 2: # p-sigma, f-sigma
                                add_so('A1g', minus(idx0, idx1))
                                add_so('A1u', plus (idx0, idx1))
                            else: # s-sigma, d-sigma
                                add_so('A1g', plus (idx0, idx1))
                                add_so('A1u', minus(idx0, idx1))
# +/-m
                            for m in range(1,angl+1):
                                idx0 = basoff[at0] + aoff + m
                                idy0 = basoff[at0] + aoff - m
                                idx1 = basoff[at1] + aoff + m
                                idy1 = basoff[at1] + aoff - m
                                if angl % 2: # odd parity
                                    add_so('E%dux'%m, plus (idx0, idx1))
                                    add_so('E%duy'%m, plus (idy0, idy1))
                                    add_so('E%dgx'%m, minus(idx0, idx1))
                                    add_so('E%dgy'%m, minus(idy0, idy1))
                                else:
                                    add_so('E%dgy'%m, plus (idy0, idy1))
                                    add_so('E%dgx'%m, plus (idx0, idx1))
                                    add_so('E%duy'%m, minus(idy0, idy1))
                                    add_so('E%dux'%m, minus(idx0, idx1))
                    ip += nc * degen
            elif len(atom_ids) == 1:
                at0 = atom_ids[0]
                ip = 0
                b0, b1, p0, p1 = aoslice[at0]
                for ib in range(b0, b1):
                    angl = mol.bas_angular(ib)
                    nc = mol.bas_nctr(ib)
                    degen = angl * 2 + 1
                    if angl == 1:
                        for i in range(nc):
                            aoff = ip + i*degen + angl
# m = 0
                            idx0 = basoff[at0] + aoff + 1
                            add_so('A1u', identity(idx0))
# m = +/- 1
                            idx0 = basoff[at0] + aoff - 1
                            idy0 = basoff[at0] + aoff
                            add_so('E1uy', identity(idy0))
                            add_so('E1ux', identity(idx0))
                    else:
                        for i in range(nc):
                            aoff = ip + i*degen + angl
                            idx0 = basoff[at0] + aoff
# m = 0
                            if angl % 2:
                                add_so('A1u', identity(idx0))
                            else:
                                add_so('A1g', identity(idx0))
# +/-m
                            for m in range(1,angl+1):
                                idx0 = basoff[at0] + aoff + m
                                idy0 = basoff[at0] + aoff - m
                                if angl % 2: # p, f functions
                                    add_so('E%dux'%m, identity(idx0))
                                    add_so('E%duy'%m, identity(idy0))
                                else: # d, g functions
                                    add_so('E%dgy'%m, identity(idy0))
                                    add_so('E%dgx'%m, identity(idx0))
                    ip += nc * degen
    elif gpname == 'Coov':
        for atom_ids in eql_atom_ids:
            at0 = atom_ids[0]
            ip = 0
            b0, b1, p0, p1 = aoslice[at0]
            for ib in range(b0, b1):
                angl = mol.bas_angular(ib)
                nc = mol.bas_nctr(ib)
                degen = angl * 2 + 1
                if angl == 1:
                    for i in range(nc):
                        aoff = ip + i*degen + angl
# m = 0
                        idx0 = basoff[at0] + aoff + 1
                        add_so('A1', identity(idx0))
# m = +/- 1
                        idx0 = basoff[at0] + aoff - 1
                        idy0 = basoff[at0] + aoff
                        add_so('E1x', identity(idx0))
                        add_so('E1y', identity(idy0))
                else:
                    for i in range(nc):
                        aoff = ip + i*degen + angl
                        idx0 = basoff[at0] + aoff
# m = 0
                        add_so('A1', identity(idx0))
# +/-m
                        for m in range(1,angl+1):
                            idx0 = basoff[at0] + aoff + m
                            idy0 = basoff[at0] + aoff - m
                            add_so('E%dx'%m, identity(idx0))
                            add_so('E%dy'%m, identity(idy0))
                ip += nc * degen

    irrep_ids = []
    irrep_names = list(sodic.keys())
    for irname in irrep_names:
        irrep_ids.append(linearmole_irrep_symb2id(gpname, irname))
    irrep_idx = numpy.argsort(irrep_ids)
    irrep_ids = [irrep_ids[i] for i in irrep_idx]

    ao_loc = mol.ao_loc_nr()
    l_idx = {}
    ANG_OF = 1
    for l in range(mol._bas[:,ANG_OF].max()+1):
        idx = [numpy.arange(ao_loc[ib], ao_loc[ib+1])
               for ib in numpy.where(mol._bas[:,ANG_OF] == l)[0]]
        if idx:
            l_idx[l] = numpy.hstack(idx)

    Ds = _ao_rotation_matrices(mol, coordinates)
    so = []
    for i in irrep_idx:
        c_ir = numpy.vstack(sodic[irrep_names[i]]).T
        nso = c_ir.shape[1]
        for l, idx in l_idx.items():
            c = c_ir[idx].reshape(-1,Ds[l].shape[1],nso)
            c_ir[idx] = numpy.einsum('nm,smp->snp', Ds[l], c).reshape(-1,nso)

        so.append(c_ir)

    return so, irrep_ids


if __name__ == "__main__":
    h2o = gto.Mole()
    h2o.verbose = 0
    h2o.output = None
    h2o.atom = [['O' , (1. , 0.    , 0.   ,)],
                [1   , (0. , -.757 , 0.587,)],
                [1   , (0. , 0.757 , 0.587,)] ]
    h2o.basis = {'H': 'cc-pvdz',
                 'O': 'cc-pvdz',}
    h2o.build()
    gpname, origin, axes = geom.detect_symm(h2o._atom)
    atoms = gto.format_atom(h2o._atom, origin, axes)
    h2o.build(False, False, atom=atoms)
    print(gpname)
    eql_atoms = geom.symm_identical_atoms(gpname, atoms)
    print(symm_adapted_basis(h2o, gpname, eql_atoms)[1])

    mol = gto.M(
        atom = [['H', (0,0,0)], ['H', (0,0,-1)], ['H', (0,0,1)]],
        basis = 'ccpvtz', charge=1)
    gpname, orig, axes = geom.detect_symm(mol._atom)
    atoms = gto.format_atom(mol._atom, orig, axes)
    mol.build(False, False, atom=atoms)
    print(gpname)
    eql_atoms = geom.symm_identical_atoms(gpname, atoms)
    print(symm_adapted_basis(mol, gpname, eql_atoms)[1])

    mol = gto.M(
        atom = [['H', (0,0,0)], ['H', (0,0,-1)], ['He', (0,0,1)]],
        basis = 'ccpvtz')
    gpname, orig, axes = geom.detect_symm(mol._atom)
    atoms = gto.format_atom(mol._atom, orig, axes)
    mol.build(False, False, atom=atoms)
    print(gpname)
    eql_atoms = geom.symm_identical_atoms(gpname, atoms)
    print(symm_adapted_basis(mol, gpname, eql_atoms)[1])


# pyscf/symm/cg.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy

def cg_spin(l, jdouble, mjdouble, spin):
    '''Clebsch Gordon coefficient of <l,m,1/2,spin|j,mj>'''
    ll1 = 2 * l + 1
    if jdouble == 2*l+1:
        if spin > 0:
            c = numpy.sqrt(.5*(ll1+mjdouble)/ll1)
        else:
            c = numpy.sqrt(.5*(ll1-mjdouble)/ll1)
    elif jdouble == 2*l-1:
        if spin > 0:
            c =-numpy.sqrt(.5*(ll1-mjdouble)/ll1)
        else:
            c = numpy.sqrt(.5*(ll1+mjdouble)/ll1)
    else:
        c = 0
    return c


if __name__ == '__main__':
    for kappa in list(range(-4,0)) + list(range(1,4)):
        if kappa < 0:
            l = -kappa - 1
            j = l * 2 + 1
        else:
            l = kappa
            j = l * 2 - 1
        print(kappa,l,j)
        for mj in range(-j, j+1, 2):
            print(cg_spin(l, j, mj, 1), cg_spin(l, j, mj, -1))



# pyscf/symm/geom.py

#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#
#
# The symmetry detection method implemented here is not strictly follow the
# point group detection flowchart.  The detection is based on the degeneracy
# of cartesian basis of multipole momentum, eg
# http://symmetry.jacobs-university.de/cgi-bin/group.cgi?group=604&option=4
# see the column of "linear functions, quadratic functions and cubic functions".
#
# Different point groups have different combinations of degeneracy for each
# type of cartesian functions.  Based on the degeneracy of cartesian function
# basis, one can quickly filter out a few candidates of point groups for the
# given molecule.  Regular operations (rotation, mirror etc) can be applied
# then to identify the symmetry.  Current implementation only checks the
# rotation functions and it's roughly enough for D2h and subgroups.
#
# There are special cases this detection method may break down, eg two H8 cube
# molecules sitting on the same center but with random orientation.  The
# system is in C1 while this detection method gives O group because the
# 3 rotation bases are degenerated.  In this case, the code use the regular
# method (point group detection flowchart) to detect the point group.
#

import sys
import re
import numpy
import scipy.linalg
from pyscf.gto import mole
from pyscf.lib import norm
from pyscf.lib import logger
from pyscf.lib.exceptions import PointGroupSymmetryError
from pyscf.symm.param import OPERATOR_TABLE
from pyscf import __config__

TOLERANCE = getattr(__config__, 'symm_geom_tol', 1e-5)

# For code compatiblity in python-2 and python-3
if sys.version_info >= (3,):
    unicode = str

def parallel_vectors(v1, v2, tol=TOLERANCE):
    if numpy.allclose(v1, 0, atol=tol) or numpy.allclose(v2, 0, atol=tol):
        return True
    else:
        cos = numpy.dot(_normalize(v1), _normalize(v2))
        return (abs(cos-1) < TOLERANCE) | (abs(cos+1) < TOLERANCE)

def argsort_coords(coords, decimals=None):
    if decimals is None:
        decimals = int(-numpy.log10(TOLERANCE)) - 1
    coords = numpy.around(coords, decimals=decimals)
    idx = numpy.lexsort((coords[:,2], coords[:,1], coords[:,0]))
    return idx

def sort_coords(coords, decimals=None):
    if decimals is None:
        decimals = int(-numpy.log10(TOLERANCE)) - 1
    coords = numpy.asarray(coords)
    idx = argsort_coords(coords, decimals=decimals)
    return coords[idx]

# ref. http://en.wikipedia.org/wiki/Rotation_matrix
def rotation_mat(vec, theta):
    '''rotate angle theta along vec
    new(x,y,z) = R * old(x,y,z)'''
    vec = _normalize(vec)
    uu = vec.reshape(-1,1) * vec.reshape(1,-1)
    ux = numpy.array((
        ( 0     ,-vec[2], vec[1]),
        ( vec[2], 0     ,-vec[0]),
        (-vec[1], vec[0], 0     )))
    c = numpy.cos(theta)
    s = numpy.sin(theta)
    r = c * numpy.eye(3) + s * ux + (1-c) * uu
    return r

# reflection operation with householder
def householder(vec):
    vec = _normalize(vec)
    return numpy.eye(3) - vec[:,None]*vec*2

def closest_axes(axes, ref):
    xcomp, ycomp, zcomp = numpy.einsum('ix,jx->ji', axes, ref)
    zmax = numpy.amax(abs(zcomp))
    zmax_idx = numpy.where(abs(abs(zcomp)-zmax)<TOLERANCE)[0]
    z_id = numpy.amax(zmax_idx)
    #z_id = numpy.argmax(abs(zcomp))
    xcomp[z_id] = ycomp[z_id] = 0       # remove z
    xmax = numpy.amax(abs(xcomp))
    xmax_idx = numpy.where(abs(abs(xcomp)-xmax)<TOLERANCE)[0]
    x_id = numpy.amax(xmax_idx)
    #x_id = numpy.argmax(abs(xcomp))
    ycomp[x_id] = 0                     # remove x
    y_id = numpy.argmax(abs(ycomp))
    return x_id, y_id, z_id

def alias_axes(axes, ref):
    '''Rename axes, make it as close as possible to the ref axes
    '''
    x_id, y_id, z_id = closest_axes(axes, ref)
    new_axes = axes[[x_id,y_id,z_id]]
    if numpy.linalg.det(new_axes) < 0:
        new_axes = axes[[y_id,x_id,z_id]]
    return new_axes


def detect_symm(atoms, basis=None, verbose=logger.WARN):
    '''Detect the point group symmetry for given molecule.

    Return group name, charge center, and nex_axis (three rows for x,y,z)
    '''
    if isinstance(verbose, logger.Logger):
        log = verbose
    else:
        log = logger.Logger(sys.stdout, verbose)

    tol = TOLERANCE / numpy.sqrt(1+len(atoms))
    decimals = int(-numpy.log10(tol))
    log.debug('geometry tol = %g', tol)

    rawsys = SymmSys(atoms, basis)
    w1, u1 = rawsys.cartesian_tensor(1)
    axes = u1.T
    log.debug('principal inertia moments %s', w1)
    charge_center = rawsys.charge_center

    if numpy.allclose(w1, 0, atol=tol):
        gpname = 'SO3'
        return gpname, charge_center, numpy.eye(3)

    elif numpy.allclose(w1[:2], 0, atol=tol): # linear molecule
        if rawsys.has_icenter():
            gpname = 'Dooh'
        else:
            gpname = 'Coov'
        return gpname, charge_center, axes

    else:
        w1_degeneracy, w1_degen_values = _degeneracy(w1, decimals)
        w2, u2 = rawsys.cartesian_tensor(2)
        w2_degeneracy, w2_degen_values = _degeneracy(w2, decimals)
        log.debug('2d tensor %s', w2)

        n = None
        c2x = None
        mirrorx = None
        if 3 in w1_degeneracy: # T, O, I
            # Because rotation vectors Rx Ry Rz are 3-degenerated representation
            # See http://www.webqc.org/symmetrypointgroup-td.html

            w3, u3 = rawsys.cartesian_tensor(3)
            w3_degeneracy, w3_degen_values = _degeneracy(w3, decimals)
            log.debug('3d tensor %s', w3)
            if (5 in w2_degeneracy and
                4 in w3_degeneracy and len(w3_degeneracy) == 3):  # I group
                gpname, new_axes = _search_i_group(rawsys)
                if gpname is not None:
                    return gpname, charge_center, _refine(new_axes)

            elif 3 in w2_degeneracy and len(w2_degeneracy) <= 3:  # T/O group
                gpname, new_axes = _search_ot_group(rawsys)
                if gpname is not None:
                    return gpname, charge_center, _refine(new_axes)

        elif (2 in w1_degeneracy and
              numpy.any(w2_degeneracy[w2_degen_values>0] >= 2)):
            if numpy.allclose(w1[1], w1[2], atol=tol):
                axes = axes[[1,2,0]]
            n = rawsys.search_c_highest(axes[2])[1]
            if n == 1:
                n = None
            else:
                c2x = rawsys.search_c2x(axes[2], n)
                mirrorx = rawsys.search_mirrorx(axes[2], n)

        else:
            n = -1  # tag as D2h and subgroup

# They must not be I/O/T group, at most one C3 or higher rotation axis
        if n is None:
            zaxis, n = rawsys.search_c_highest()
            if n > 1:
                c2x = rawsys.search_c2x(zaxis, n)
                mirrorx = rawsys.search_mirrorx(zaxis, n)
                if c2x is not None:
                    axes = _make_axes(zaxis, c2x)
                elif mirrorx is not None:
                    axes = _make_axes(zaxis, mirrorx)
                else:
                    for axis in numpy.eye(3):
                        if not parallel_vectors(axis, zaxis):
                            axes = _make_axes(zaxis, axis)
                            break
            else:  # Ci or Cs or C1 with degenerated w1
                mirror = rawsys.search_mirrorx(None, 1)
                if mirror is not None:
                    xaxis = numpy.array((1.,0.,0.))
                    axes = _make_axes(mirror, xaxis)
                else:
                    axes = numpy.eye(3)

        log.debug('Highest C_n = C%d', n)
        if n >= 2:
            if c2x is not None:
                if rawsys.has_mirror(axes[2]):
                    gpname = 'D%dh' % n
                elif rawsys.has_improper_rotation(axes[2], n):
                    gpname = 'D%dd' % n
                else:
                    gpname = 'D%d' % n
                # yaxis = numpy.cross(axes[2], c2x)
                axes = _make_axes(axes[2], c2x)
            elif mirrorx is not None:
                gpname = 'C%dv' % n
                axes = _make_axes(axes[2], mirrorx)
            elif rawsys.has_mirror(axes[2]):
                gpname = 'C%dh' % n
            elif rawsys.has_improper_rotation(axes[2], n):
                gpname = 'S%d' % (n*2)
            else:
                gpname = 'C%d' % n
            return gpname, charge_center, _refine(axes)

        else:
            is_c2x = rawsys.has_rotation(axes[0], 2)
            is_c2y = rawsys.has_rotation(axes[1], 2)
            is_c2z = rawsys.has_rotation(axes[2], 2)
# rotate to old axes, as close as possible?
            if is_c2z and is_c2x and is_c2y:
                if rawsys.has_icenter():
                    gpname = 'D2h'
                else:
                    gpname = 'D2'
                axes = alias_axes(axes, numpy.eye(3))
            elif is_c2z or is_c2x or is_c2y:
                if is_c2x:
                    axes = axes[[1,2,0]]
                if is_c2y:
                    axes = axes[[2,0,1]]
                if rawsys.has_mirror(axes[2]):
                    gpname = 'C2h'
                elif rawsys.has_mirror(axes[0]):
                    gpname = 'C2v'
                else:
                    gpname = 'C2'
            else:
                if rawsys.has_icenter():
                    gpname = 'Ci'
                elif rawsys.has_mirror(axes[0]):
                    gpname = 'Cs'
                    axes = axes[[1,2,0]]
                elif rawsys.has_mirror(axes[1]):
                    gpname = 'Cs'
                    axes = axes[[2,0,1]]
                elif rawsys.has_mirror(axes[2]):
                    gpname = 'Cs'
                else:
                    gpname = 'C1'
                    axes = numpy.eye(3)
                    charge_center = numpy.zeros(3)
    return gpname, charge_center, axes


# reduce to D2h and its subgroups
# FIXME, CPL, 209, 506
def get_subgroup(gpname, axes):
    if gpname in ('D2h', 'D2' , 'C2h', 'C2v', 'C2' , 'Ci' , 'Cs' , 'C1'):
        return gpname, axes
    elif gpname in ('SO3',):
        #return 'D2h', alias_axes(axes, numpy.eye(3))
        return 'SO3', axes
    elif gpname in ('Dooh',):
        #return 'D2h', alias_axes(axes, numpy.eye(3))
        return 'Dooh', axes
    elif gpname in ('Coov',):
        #return 'C2v', axes
        return 'Coov', axes
    elif gpname in ('Oh',):
        return 'D2h', alias_axes(axes, numpy.eye(3))
    elif gpname in ('O',):
        return 'D2', alias_axes(axes, numpy.eye(3))
    elif gpname in ('Ih',):
        return 'Ci', alias_axes(axes, numpy.eye(3))
    elif gpname in ('I',):
        return 'C1', axes
    elif gpname in ('Td', 'T', 'Th'):
        return 'D2', alias_axes(axes, numpy.eye(3))
    elif re.search(r'S\d+', gpname):
        n = int(re.search(r'\d+', gpname).group(0))
        if n % 2 == 0:
            return 'C%d'%(n//2), axes
        else:
            return 'Ci', axes
    else:
        n = int(re.search(r'\d+', gpname).group(0))
        if n % 2 == 0:
            if re.search(r'D\d+d', gpname):
                subname = 'D2'
            elif re.search(r'D\d+h', gpname):
                subname = 'D2h'
            elif re.search(r'D\d+', gpname):
                subname = 'D2'
            elif re.search(r'C\d+h', gpname):
                subname = 'C2h'
            elif re.search(r'C\d+v', gpname):
                subname = 'C2v'
            else:
                subname = 'C2'
        else:
            # rotate axes and
            # Dnh -> C2v
            # Dn  -> C2
            # Dnd -> Ci
            # Cnh -> Cs
            # Cnv -> Cs
            if re.search(r'D\d+h', gpname):
                subname = 'C2v'
                axes = axes[[1,2,0]]
            elif re.search(r'D\d+d', gpname):
                subname = 'C2h'
                axes = axes[[1,2,0]]
            elif re.search(r'D\d+', gpname):
                subname = 'C2'
                axes = axes[[1,2,0]]
            elif re.search(r'C\d+h', gpname):
                subname = 'Cs'
            elif re.search(r'C\d+v', gpname):
                subname = 'Cs'
                axes = axes[[1,2,0]]
            else:
                subname = 'C1'
        return subname, axes
subgroup = get_subgroup

def as_subgroup(topgroup, axes, subgroup=None):
    from pyscf.symm import std_symb
    from pyscf.symm.param import SUBGROUP

    groupname, axes = get_subgroup(topgroup, axes)

    if isinstance(subgroup, (str, unicode)):
        subgroup = std_symb(subgroup)
        if groupname == 'C2v' and subgroup == 'Cs':
            axes = numpy.einsum('ij,kj->ki', rotation_mat(axes[1], numpy.pi/2), axes)

        elif (groupname == 'D2' and re.search(r'D\d+d', topgroup) and
              subgroup in ('C2v', 'Cs')):
            # Special treatment for D2d, D4d, .... get_subgroup gives D2 by
            # default while C2v is also D2d's subgroup.
            groupname = 'C2v'
            axes = numpy.einsum('ij,kj->ki', rotation_mat(axes[2], numpy.pi/4), axes)

        elif topgroup in ('Td', 'T', 'Th') and subgroup == 'C2v':
            x, y, z = axes
            x = _normalize(x+y)
            y = numpy.cross(z, x)
            axes = numpy.array((x,y,z))

        elif subgroup not in SUBGROUP[groupname]:
            raise PointGroupSymmetryError('%s not in Ablien subgroup of %s' %
                                          (subgroup, topgroup))

        groupname = subgroup
    return groupname, axes

def symm_ops(gpname, axes=None):
    if axes is not None:
        raise PointGroupSymmetryError('TODO: non-standard orientation')
    op1 = numpy.eye(3)
    opi = -1

    opc2z = -numpy.eye(3)
    opc2z[2,2] = 1
    opc2x = -numpy.eye(3)
    opc2x[0,0] = 1
    opc2y = -numpy.eye(3)
    opc2y[1,1] = 1

    opcsz = numpy.dot(opc2z, opi)
    opcsx = numpy.dot(opc2x, opi)
    opcsy = numpy.dot(opc2y, opi)
    opdic = {'E'  : op1,
             'C2z': opc2z,
             'C2x': opc2x,
             'C2y': opc2y,
             'i'  : opi,
             'sz' : opcsz,
             'sx' : opcsx,
             'sy' : opcsy,}
    return opdic

def symm_identical_atoms(gpname, atoms):
    '''Symmetry identical atoms'''
    # from pyscf import gto
    # Dooh Coov for linear molecule
    if gpname == 'Dooh':
        coords = numpy.array([a[1] for a in atoms], dtype=float)
        idx0 = argsort_coords(coords)
        coords0 = coords[idx0]
        opdic = symm_ops(gpname)
        newc = numpy.dot(coords, opdic['sz'])
        idx1 = argsort_coords(newc)
        dup_atom_ids = numpy.sort((idx0,idx1), axis=0).T
        uniq_idx = numpy.unique(dup_atom_ids[:,0], return_index=True)[1]
        eql_atom_ids = dup_atom_ids[uniq_idx]
        eql_atom_ids = [list(sorted(set(i))) for i in eql_atom_ids]
        return eql_atom_ids
    elif gpname == 'Coov':
        eql_atom_ids = [[i] for i,a in enumerate(atoms)]
        return eql_atom_ids

    coords = numpy.array([a[1] for a in atoms])

#    charges = numpy.array([gto.charge(a[0]) for a in atoms])
#    center = numpy.einsum('z,zr->r', charges, coords)/charges.sum()
#    if not numpy.allclose(center, 0, atol=TOLERANCE):
#        sys.stderr.write('WARN: Molecular charge center %s is not on (0,0,0)\n'
#                        % center)
    opdic = symm_ops(gpname)
    ops = [opdic[op] for op in OPERATOR_TABLE[gpname]]
    idx = argsort_coords(coords)
    coords0 = coords[idx]

    dup_atom_ids = []
    for op in ops:
        newc = numpy.dot(coords, op)
        idx = argsort_coords(newc)
        if not numpy.allclose(coords0, newc[idx], atol=TOLERANCE):
            raise PointGroupSymmetryError('Symmetry identical atoms not found')
        dup_atom_ids.append(idx)

    dup_atom_ids = numpy.sort(dup_atom_ids, axis=0).T
    uniq_idx = numpy.unique(dup_atom_ids[:,0], return_index=True)[1]
    eql_atom_ids = dup_atom_ids[uniq_idx]
    eql_atom_ids = [list(sorted(set(i))) for i in eql_atom_ids]
    return eql_atom_ids

def check_symm(gpname, atoms, basis=None):
    '''
    Check whether the declared symmetry (gpname) exists in the system

    If basis is specified, this function checks also the basis functions have
    the required symmetry.

    Args:
        gpname: str
            point group name
        atoms: list
            [[symbol, [x, y, z]], [symbol, [x, y, z]], ...]
    '''

    #FIXME: compare the basis set when basis is given
    if gpname == 'Dooh':
        coords = numpy.array([a[1] for a in atoms], dtype=float)
        if numpy.allclose(coords[:,:2], 0, atol=TOLERANCE):
            opdic = symm_ops(gpname)
            rawsys = SymmSys(atoms, basis)
            return rawsys.has_icenter() and numpy.allclose(rawsys.charge_center, 0)
        else:
            return False
    elif gpname == 'Coov':
        coords = numpy.array([a[1] for a in atoms], dtype=float)
        return numpy.allclose(coords[:,:2], 0, atol=TOLERANCE)

    opdic = symm_ops(gpname)
    ops = [opdic[op] for op in OPERATOR_TABLE[gpname]]
    rawsys = SymmSys(atoms, basis)
    for lst in rawsys.atomtypes.values():
        coords = rawsys.atoms[lst,1:]
        idx = argsort_coords(coords)
        coords0 = coords[idx]

        for op in ops:
            newc = numpy.dot(coords, op)
            idx = argsort_coords(newc)
            if not numpy.allclose(coords0, newc[idx], atol=TOLERANCE):
                return False
    return True

check_given_symm = check_symm

def shift_atom(atoms, orig, axis):
    c = numpy.array([a[1] for a in atoms])
    c = numpy.dot(c - orig, numpy.array(axis).T)
    return [[atoms[i][0], c[i]] for i in range(len(atoms))]

class RotationAxisNotFound(PointGroupSymmetryError):
    pass

class SymmSys(object):
    def __init__(self, atoms, basis=None):
        self.atomtypes = mole.atom_types(atoms, basis)
        # fake systems, which treates the atoms of different basis as different atoms.
        # the fake systems do not have the same symmetry as the potential
        # it's only used to determine the main (Z-)axis
        chg1 = numpy.pi - 2
        coords = []
        fake_chgs = []
        idx = []
        for k, lst in self.atomtypes.items():
            idx.append(lst)
            coords.append([atoms[i][1] for i in lst])
            ksymb = mole._rm_digit(k)
            if ksymb != k:
                # Put random charges on the decorated atoms
                fake_chgs.append([chg1] * len(lst))
                chg1 *= numpy.pi-2
            elif mole.is_ghost_atom(k):
                if ksymb == 'X' or ksymb.upper() == 'GHOST':
                    fake_chgs.append([.3] * len(lst))
                elif ksymb[0] == 'X':
                    fake_chgs.append([mole.charge(ksymb[1:])+.3] * len(lst))
                elif ksymb[:5] == 'GHOST':
                    fake_chgs.append([mole.charge(ksymb[5:])+.3] * len(lst))
            else:
                fake_chgs.append([mole.charge(ksymb)] * len(lst))
        coords = numpy.array(numpy.vstack(coords), dtype=float)
        fake_chgs = numpy.hstack(fake_chgs)
        self.charge_center = numpy.einsum('i,ij->j', fake_chgs, coords)/fake_chgs.sum()
        coords = coords - self.charge_center

        idx = numpy.argsort(numpy.hstack(idx))
        self.atoms = numpy.hstack((fake_chgs.reshape(-1,1), coords))[idx]

        self.group_atoms_by_distance = []
        decimals = int(-numpy.log10(TOLERANCE)) - 1
        for index in self.atomtypes.values():
            index = numpy.asarray(index)
            c = self.atoms[index,1:]
            dists = numpy.around(norm(c, axis=1), decimals)
            u, idx = numpy.unique(dists, return_inverse=True)
            for i, s in enumerate(u):
                self.group_atoms_by_distance.append(index[idx == i])

    def cartesian_tensor(self, n):
        z = self.atoms[:,0]
        r = self.atoms[:,1:]
        ncart = (n+1)*(n+2)//2
        natm = len(z)
        tensor = numpy.sqrt(numpy.copy(z).reshape(natm,-1) / z.sum())
        for i in range(n):
            tensor = numpy.einsum('zi,zj->zij', tensor, r).reshape(natm,-1)
        e, c = scipy.linalg.eigh(numpy.dot(tensor.T,tensor))
        return e[-ncart:], c[:,-ncart:]

    def symmetric_for(self, op):
        for lst in self.group_atoms_by_distance:
            r0 = self.atoms[lst,1:]
            r1 = numpy.dot(r0, op)
# FIXME: compare whehter two sets of coordinates are identical
            yield all((_vec_in_vecs(x, r0) for x in r1))

    def has_icenter(self):
        return all(self.symmetric_for(-1))

    def has_rotation(self, axis, n):
        op = rotation_mat(axis, numpy.pi*2/n).T
        return all(self.symmetric_for(op))

    def has_mirror(self, perp_vec):
        return all(self.symmetric_for(householder(perp_vec).T))

    def has_improper_rotation(self, axis, n):
        s_op = numpy.dot(householder(axis), rotation_mat(axis, numpy.pi/n)).T
        return all(self.symmetric_for(s_op))

    def search_possible_rotations(self, zaxis=None):
        '''If zaxis is given, the rotation axis is parallel to zaxis'''
        maybe_cn = []
        for lst in self.group_atoms_by_distance:
            natm = len(lst)
            if natm > 1:
                coords = self.atoms[lst,1:]
# possible C2 axis
                for i in range(1, natm):
                    if abs(coords[0]+coords[i]).sum() > TOLERANCE:
                        maybe_cn.append((coords[0]+coords[i], 2))
                    else: # abs(coords[0]-coords[i]).sum() > TOLERANCE:
                        maybe_cn.append((coords[0]-coords[i], 2))

# atoms of equal distances may be associated with rotation axis > C2.
                r0 = coords - coords[0]
                distance = norm(r0, axis=1)
                eq_distance = abs(distance[:,None] - distance) < TOLERANCE
                for i in range(2, natm):
                    for j in numpy.where(eq_distance[i,:i])[0]:
                        cos = numpy.dot(r0[i],r0[j]) / (distance[i]*distance[j])
                        ang = numpy.arccos(cos)
                        nfrac = numpy.pi*2 / (numpy.pi-ang)
                        n = int(numpy.around(nfrac))
                        if abs(nfrac-n) < TOLERANCE:
                            maybe_cn.append((numpy.cross(r0[i],r0[j]),n))

        # remove zero-vectors and duplicated vectors
        vecs = numpy.vstack([x[0] for x in maybe_cn])
        idx = norm(vecs, axis=1) > TOLERANCE
        ns = numpy.hstack([x[1] for x in maybe_cn])
        vecs = _normalize(vecs[idx])
        ns = ns[idx]

        if zaxis is not None:  # Keep parallel rotation axes
            cos = numpy.dot(vecs, _normalize(zaxis))
            vecs = vecs[(abs(cos-1) < TOLERANCE) | (abs(cos+1) < TOLERANCE)]
            ns = ns[(abs(cos-1) < TOLERANCE) | (abs(cos+1) < TOLERANCE)]

        possible_cn = []
        seen = numpy.zeros(len(vecs), dtype=bool)
        for k, v in enumerate(vecs):
            if not seen[k]:
                where1 = numpy.einsum('ix->i', abs(vecs[k:] - v)) < TOLERANCE
                where1 = numpy.where(where1)[0] + k
                where2 = numpy.einsum('ix->i', abs(vecs[k:] + v)) < TOLERANCE
                where2 = numpy.where(where2)[0] + k
                seen[where1] = True
                seen[where2] = True

                vk = _normalize((numpy.einsum('ix->x', vecs[where1]) -
                                 numpy.einsum('ix->x', vecs[where2])))
                for n in (set(ns[where1]) | set(ns[where2])):
                    possible_cn.append((vk,n))
        return possible_cn

    def search_c2x(self, zaxis, n):
        '''C2 axis which is perpendicular to z-axis'''
        decimals = int(-numpy.log10(TOLERANCE)) - 1
        for lst in self.group_atoms_by_distance:
            if len(lst) > 1:
                r0 = self.atoms[lst,1:]
                zcos = numpy.around(numpy.einsum('ij,j->i', r0, zaxis),
                                    decimals=decimals)
                uniq_zcos = numpy.unique(zcos)
                maybe_c2x = []
                for d in uniq_zcos:
                    if d > TOLERANCE:
                        mirrord = abs(zcos+d)<TOLERANCE
                        if mirrord.sum() == (zcos==d).sum():
                            above = r0[zcos==d]
                            below = r0[mirrord]
                            nelem = len(below)
                            maybe_c2x.extend([above[0] + below[i]
                                              for i in range(nelem)])
                    elif abs(d) < TOLERANCE: # plane which crosses the orig
                        r1 = r0[zcos==d][0]
                        maybe_c2x.append(r1)
                        r2 = numpy.dot(rotation_mat(zaxis, numpy.pi*2/n), r1)
                        if abs(r1+r2).sum() > TOLERANCE:
                            maybe_c2x.append(r1+r2)
                        else:
                            maybe_c2x.append(r2-r1)

                if len(maybe_c2x) > 0:
                    idx = norm(maybe_c2x, axis=1) > TOLERANCE
                    maybe_c2x = _normalize(maybe_c2x)[idx]
                    maybe_c2x = _remove_dupvec(maybe_c2x)
                    for c2x in maybe_c2x:
                        if (not parallel_vectors(c2x, zaxis) and
                            self.has_rotation(c2x, 2)):
                            return c2x

    def search_mirrorx(self, zaxis, n):
        '''mirror which is parallel to z-axis'''
        if n > 1:
            for lst in self.group_atoms_by_distance:
                natm = len(lst)
                r0 = self.atoms[lst[0],1:]
                if natm > 1 and not parallel_vectors(r0, zaxis):
                    r1 = numpy.dot(rotation_mat(zaxis, numpy.pi*2/n), r0)
                    mirrorx = _normalize(r1-r0)
                    if self.has_mirror(mirrorx):
                        return mirrorx
        else:
            for lst in self.group_atoms_by_distance:
                natm = len(lst)
                r0 = self.atoms[lst,1:]
                if natm > 1:
                    maybe_mirror = [r0[i]-r0[0] for i in range(1, natm)]
                    for mirror in _normalize(maybe_mirror):
                        if self.has_mirror(mirror):
                            return mirror

    def search_c_highest(self, zaxis=None):
        possible_cn = self.search_possible_rotations(zaxis)
        nmax = 1
        cmax = numpy.array([0.,0.,1.])
        for cn, n in possible_cn:
            if n > nmax and self.has_rotation(cn, n):
                nmax = n
                cmax = cn
        return cmax, nmax


def _normalize(vecs):
    vecs = numpy.asarray(vecs)
    if vecs.ndim == 1:
        return vecs / (numpy.linalg.norm(vecs) + 1e-200)
    else:
        return vecs / (norm(vecs, axis=1).reshape(-1,1) + 1e-200)

def _vec_in_vecs(vec, vecs):
    norm = numpy.sqrt(len(vecs))
    return min(numpy.einsum('ix->i', abs(vecs-vec))/norm) < TOLERANCE

def _search_i_group(rawsys):
    possible_cn = rawsys.search_possible_rotations()
    c5_axes = [c5 for c5, n in possible_cn
               if n == 5 and rawsys.has_rotation(c5, 5)]
    if len(c5_axes) <= 1:
        return None,None

    zaxis = c5_axes[0]
    cos = numpy.dot(c5_axes, zaxis)
    assert(numpy.all((abs(cos[1:]+1/numpy.sqrt(5)) < TOLERANCE) |
                     (abs(cos[1:]-1/numpy.sqrt(5)) < TOLERANCE)))

    if rawsys.has_icenter():
        gpname = 'Ih'
    else:
        gpname = 'I'

    c5 = c5_axes[1]
    if numpy.dot(c5, zaxis) < 0:
        c5 = -c5
    c5a = numpy.dot(rotation_mat(zaxis, numpy.pi*6/5), c5)
    xaxis = c5a + c5
    return gpname, _make_axes(zaxis, xaxis)

def _search_ot_group(rawsys):
    possible_cn = rawsys.search_possible_rotations()
    c4_axes = [c4 for c4, n in possible_cn
               if n == 4 and rawsys.has_rotation(c4, 4)]

    if len(c4_axes) > 0:  # T group
        assert(len(c4_axes) > 1)
        if rawsys.has_icenter():
            gpname = 'Oh'
        else:
            gpname = 'O'
        return gpname, _make_axes(c4_axes[0], c4_axes[1])

    else:  # T group
        c3_axes = [c3 for c3, n in possible_cn
                   if n == 3 and rawsys.has_rotation(c3, 3)]
        if len(c3_axes) <= 1:
            return None, None

        cos = numpy.dot(c3_axes, c3_axes[0])
        assert(numpy.all((abs(cos[1:]+1./3) < TOLERANCE) |
                         (abs(cos[1:]-1./3) < TOLERANCE)))

        if rawsys.has_icenter():
            gpname = 'Th'
# Because C3 axes are on the mirror of Td, two C3 can determine a mirror.
        elif rawsys.has_mirror(numpy.cross(c3_axes[0], c3_axes[1])):
            gpname = 'Td'
        else:
            gpname = 'T'

        c3a = c3_axes[0]
        if numpy.dot(c3a, c3_axes[1]) > 0:
            c3a = -c3a
        c3b = numpy.dot(rotation_mat(c3a,-numpy.pi*2/3), c3_axes[1])
        c3c = numpy.dot(rotation_mat(c3a, numpy.pi*2/3), c3_axes[1])
        zaxis, xaxis = c3a+c3b, c3a+c3c
        return gpname, _make_axes(zaxis, xaxis)

def _degeneracy(e, decimals):
    e = numpy.around(e, decimals)
    u, idx = numpy.unique(e, return_inverse=True)
    degen = numpy.array([numpy.count_nonzero(idx==i) for i in range(len(u))])
    return degen, u

def _pseudo_vectors(vs):
    idy0 = abs(vs[:,1])<TOLERANCE
    idz0 = abs(vs[:,2])<TOLERANCE
    vs = vs.copy()
    # ensure z component > 0
    vs[vs[:,2]<0] *= -1
    # if z component == 0, ensure y component > 0
    vs[(vs[:,1]<0) & idz0] *= -1
    # if y and z component == 0, ensure x component > 0
    vs[(vs[:,0]<0) & idy0 & idz0] *= -1
    return vs

def _remove_dupvec(vs):
    def rm_iter(vs):
        if len(vs) <= 1:
            return vs
        else:
            x = numpy.sum(abs(vs[1:]-vs[0]), axis=1)
            rest = rm_iter(vs[1:][x>TOLERANCE])
            return numpy.vstack((vs[0], rest))
    return rm_iter(_pseudo_vectors(vs))

def _make_axes(z, x):
    y = numpy.cross(z, x)
    x = numpy.cross(y, z)  # because x might not perp to z
    return _normalize(numpy.array((x,y,z)))

def _refine(axes):
    # Make sure the axes can be rotated from continuous unitary transformation
    if axes[2,2] < 0:
        axes[2] *= -1
    if abs(axes[0,0]) > abs(axes[1,0]):
        x_id, y_id = 0, 1
    else:
        x_id, y_id = 1, 0
    if axes[x_id,0] < 0:
        axes[x_id] *= -1
    if numpy.linalg.det(axes) < 0:
        axes[y_id] *= -1
    return axes


if __name__ == "__main__":
    atom = [["O" , (1. , 0.    , 0.   ,)],
            ['H' , (0. , -.757 , 0.587,)],
            ['H' , (0. , 0.757 , 0.587,)] ]
    gpname, orig, axes = detect_symm(atom)
    atom = shift_atom(atom, orig, axes)
    print(gpname, symm_identical_atoms(gpname, atom))

    atom = [['H', (0,0,0)], ['H', (0,0,-1)], ['H', (0,0,1)]]
    gpname, orig, axes = detect_symm(atom)
    print(gpname, orig, axes)
    atom = shift_atom(atom, orig, axes)
    print(gpname, symm_identical_atoms(gpname, atom))

    atom = [['H', (0., 0., 0.)],
            ['H', (0., 0., 1.)],
            ['H', (0., 1., 0.)],
            ['H', (1., 0., 0.)],
            ['H', (-1, 0., 0.)],
            ['H', (0.,-1., 0.)],
            ['H', (0., 0.,-1.)]]
    gpname, orig, axes = detect_symm(atom)
    print(gpname, orig, axes)
    atom = shift_atom(atom, orig, axes)
    print(gpname, symm_identical_atoms(subgroup(gpname, axes)[0], atom))


# pyscf/symm/sph.py

#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''
Spherical harmonics
'''

import numpy
import scipy.linalg
from pyscf.symm.cg import cg_spin

def real_sph_vec(r, lmax, reorder_p=False):
    '''Computes (all) real spherical harmonics up to the angular momentum lmax'''
    #:import scipy.special
    #:ngrid = r.shape[0]
    #:cosphi = r[:,2]
    #:sinphi = (1-cosphi**2)**.5
    #:costheta = numpy.ones(ngrid)
    #:sintheta = numpy.zeros(ngrid)
    #:costheta[sinphi!=0] = r[sinphi!=0,0] / sinphi[sinphi!=0]
    #:sintheta[sinphi!=0] = r[sinphi!=0,1] / sinphi[sinphi!=0]
    #:costheta[costheta> 1] = 1
    #:costheta[costheta<-1] =-1
    #:sintheta[sintheta> 1] = 1
    #:sintheta[sintheta<-1] =-1
    #:varphi = numpy.arccos(cosphi)
    #:theta = numpy.arccos(costheta)
    #:theta[sintheta<0] = 2*numpy.pi - theta[sintheta<0]
    #:ylms = []
    #:for l in range(lmax+1):
    #:    ylm = numpy.empty((l*2+1,ngrid))
    #:    ylm[l] = scipy.special.sph_harm(0, l, theta, varphi).real
    #:    for m in range(1, l+1):
    #:        f1 = scipy.special.sph_harm(-m, l, theta, varphi)
    #:        f2 = scipy.special.sph_harm( m, l, theta, varphi)
    #:        # complex to real spherical functions
    #:        if m % 2 == 1:
    #:            ylm[l-m] = (-f1.imag - f2.imag) / numpy.sqrt(2)
    #:            ylm[l+m] = ( f1.real - f2.real) / numpy.sqrt(2)
    #:        else:
    #:            ylm[l-m] = (-f1.imag + f2.imag) / numpy.sqrt(2)
    #:            ylm[l+m] = ( f1.real + f2.real) / numpy.sqrt(2)
    #:    ylms.append(ylm)
    #:return ylms

    # When r is a normalized vector:
    norm = 1./numpy.linalg.norm(r, axis=1)
    return multipoles(r*norm.reshape(-1,1), lmax, reorder_p)


def multipoles(r, lmax, reorder_dipole=True):
    '''
    Compute all multipoles upto lmax

    rad = numpy.linalg.norm(r, axis=1)
    ylms = real_ylm(r/rad.reshape(-1,1), lmax)
    pol = [rad**l*y for l, y in enumerate(ylms)]

    Kwargs:
        reorder_p : bool
            sort dipole to the order (x,y,z)
    '''
    from pyscf import gto

# libcint cart2sph transformation provide the capability to compute
# multipole directly.  cart2sph function is fast for low angular moment.
    ngrid = r.shape[0]
    xs = numpy.ones((lmax+1,ngrid))
    ys = numpy.ones((lmax+1,ngrid))
    zs = numpy.ones((lmax+1,ngrid))
    for i in range(1,lmax+1):
        xs[i] = xs[i-1] * r[:,0]
        ys[i] = ys[i-1] * r[:,1]
        zs[i] = zs[i-1] * r[:,2]
    ylms = []
    for l in range(lmax+1):
        nd = (l+1)*(l+2)//2
        c = numpy.empty((nd,ngrid))
        k = 0
        for lx in reversed(range(0, l+1)):
            for ly in reversed(range(0, l-lx+1)):
                lz = l - lx - ly
                c[k] = xs[lx] * ys[ly] * zs[lz]
                k += 1
        ylm = gto.cart2sph(l, c.T).T
        ylms.append(ylm)

# when call libcint, p functions are ordered as px,py,pz
# reorder px,py,pz to p(-1),p(0),p(1)
    if (not reorder_dipole) and lmax >= 1:
        ylms[1] = ylms[1][[1,2,0]]
    return ylms

def sph_pure2real(l, reorder_p=True):
    r'''
    Transformation matrix: from the pure spherical harmonic functions Y_m to
    the real spherical harmonic functions O_m.
          O_m = \sum Y_m' * U(m',m)
    Y(-1) = 1/\sqrt(2){-iO(-1) + O(1)}; Y(1) = 1/\sqrt(2){-iO(-1) - O(1)}
    Y(-2) = 1/\sqrt(2){-iO(-2) + O(2)}; Y(2) = 1/\sqrt(2){iO(-2) + O(2)}
    O(-1) = i/\sqrt(2){Y(-1) + Y(1)};   O(1) = 1/\sqrt(2){Y(-1) - Y(1)}
    O(-2) = i/\sqrt(2){Y(-2) - Y(2)};   O(2) = 1/\sqrt(2){Y(-2) + Y(2)}

    Kwargs:
        reorder_p (bool): Whether the p functions are in the (x,y,z) order.

    Returns:
        2D array U_{complex,real}
    '''
    n = 2 * l + 1
    u = numpy.zeros((n,n), dtype=complex)
    sqrthfr = numpy.sqrt(.5)
    sqrthfi = numpy.sqrt(.5)*1j

    if reorder_p and l == 1:
        u[1,2] = 1
        u[0,1] =  sqrthfi
        u[2,1] =  sqrthfi
        u[0,0] =  sqrthfr
        u[2,0] = -sqrthfr
    else:
        u[l,l] = 1
        for m in range(1, l+1, 2):
            u[l-m,l-m] =  sqrthfi
            u[l+m,l-m] =  sqrthfi
            u[l-m,l+m] =  sqrthfr
            u[l+m,l+m] = -sqrthfr
        for m in range(2, l+1, 2):
            u[l-m,l-m] =  sqrthfi
            u[l+m,l-m] = -sqrthfi
            u[l-m,l+m] =  sqrthfr
            u[l+m,l+m] =  sqrthfr

    return u

def sph_real2pure(l, reorder_p=True):
    '''
    Transformation matrix: from real spherical harmonic functions to the pure
    spherical harmonic functions.

    Kwargs:
        reorder_p (bool): Whether the real p functions are in the (x,y,z) order.
    '''
    # numpy.linalg.inv(sph_pure2real(l))
    return sph_pure2real(l, reorder_p).conj().T

# |spinor> = (|real_sph>, |real_sph>) * / u_alpha \
#                                       \ u_beta  /
# Return 2D array U_{sph,spinor}
def sph2spinor(l, reorder_p=True):
    if l == 0:
        return numpy.array((0., 1.)).reshape(1,-1), \
               numpy.array((1., 0.)).reshape(1,-1)
    else:
        u1 = sph_real2pure(l, reorder_p)
        ua = numpy.zeros((2*l+1,4*l+2),dtype=complex)
        ub = numpy.zeros((2*l+1,4*l+2),dtype=complex)
        j = l * 2 - 1
        mla = l + (-j-1)//2
        mlb = l + (-j+1)//2
        for k,mj in enumerate(range(-j, j+1, 2)):
            ua[:,k] = u1[:,mla] * cg_spin(l, j, mj, 1)
            ub[:,k] = u1[:,mlb] * cg_spin(l, j, mj,-1)
            mla += 1
            mlb += 1
        j = l * 2 + 1
        mla = l + (-j-1)//2
        mlb = l + (-j+1)//2
        for k,mj in enumerate(range(-j, j+1, 2)):
            if mla < 0:
                ua[:,l*2+k] = 0
            else:
                ua[:,l*2+k] = u1[:,mla] * cg_spin(l, j, mj, 1)
            if mlb >= 2*l+1:
                ub[:,l*2+k] = 0
            else:
                ub[:,l*2+k] = u1[:,mlb] * cg_spin(l, j, mj,-1)
            mla += 1
            mlb += 1
    return ua, ub
real2spinor = sph2spinor

# Returns 2D array U_{sph,spinor}
def sph2spinor_coeff(mol):
    '''Transformation matrix that transforms real-spherical GTOs to spinor
    GTOs for all basis functions

    Examples::

    >>> from pyscf import gto
    >>> from pyscf.symm import sph
    >>> mol = gto.M(atom='H 0 0 0; F 0 0 1', basis='ccpvtz')
    >>> ca, cb = sph.sph2spinor_coeff(mol)
    >>> s0 = mol.intor('int1e_ovlp_spinor')
    >>> s1 = ca.conj().T.dot(mol.intor('int1e_ovlp_sph')).dot(ca)
    >>> s1+= cb.conj().T.dot(mol.intor('int1e_ovlp_sph')).dot(cb)
    >>> print(abs(s1-s0).max())
    >>> 6.66133814775e-16
    '''
    lmax = max([mol.bas_angular(i) for i in range(mol.nbas)])
    ualst = []
    ublst = []
    for l in range(lmax+1):
        u1, u2 = sph2spinor(l, reorder_p=True)
        ualst.append(u1)
        ublst.append(u2)

    ca = []
    cb = []
    for ib in range(mol.nbas):
        l = mol.bas_angular(ib)
        kappa = mol.bas_kappa(ib)
        if kappa == 0:
            ua = ualst[l]
            ub = ublst[l]
        elif kappa < 0:
            ua = ualst[l][:,l*2:]
            ub = ublst[l][:,l*2:]
        else:
            ua = ualst[l][:,:l*2]
            ub = ublst[l][:,:l*2]
        nctr = mol.bas_nctr(ib)
        ca.extend([ua]*nctr)
        cb.extend([ub]*nctr)
    return numpy.stack([scipy.linalg.block_diag(*ca),
                        scipy.linalg.block_diag(*cb)])
real2spinor_whole = sph2spinor_coeff

def cart2spinor(l):
    '''Cartesian to spinor for angular moment l'''
    from pyscf import gto
    return gto.cart2spinor_l(l)


if __name__ == '__main__':
    for l in range(3):
        print(sph_pure2real(l))
        print(sph_real2pure(l))

    for l in range(3):
        print(sph2spinor(l)[0])
        print(sph2spinor(l)[1])


# pyscf/tddft/__init__.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pyscf.tdscf import *


# pyscf/tdscf/__init__.py

# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

from pyscf import scf
from pyscf.tdscf import rhf
from pyscf.tdscf import uhf
from pyscf.tdscf import ghf
from pyscf.tdscf import dhf
from pyscf.tdscf.rhf import TDRHF
from pyscf.tdscf.uhf import TDUHF
from pyscf.tdscf.ghf import TDGHF

try:
    from pyscf import dft
    from pyscf.tdscf import rks
    from pyscf.tdscf import uks
    from pyscf.tdscf import gks
    from pyscf.tdscf import dks
    from pyscf.tdscf.rks import TDRKS
    from pyscf.tdscf.uks import TDUKS
    from pyscf.tdscf.gks import TDGKS
except (ImportError, IOError):
    pass


def TDHF(mf):
    if isinstance(mf, scf.hf.KohnShamDFT):
        raise RuntimeError('TDHF does not support DFT object %s' % mf)
    if isinstance(mf, scf.uhf.UHF):
        # Call to_*hf() to remove newton decoration
        return uhf.TDHF(mf.to_uhf())
    else:
        return rhf.TDHF(mf.to_rhf())

def TDA(mf):
    if isinstance(mf, scf.uhf.UHF):
        if isinstance(mf, dft.rks.KohnShamDFT):
            return uks.TDA(mf.to_uks())
        else:
            return uhf.TDA(mf.to_uhf())
    else:
        if isinstance(mf, dft.rks.KohnShamDFT):
            return rks.TDA(mf.to_rks())
        else:
            return rhf.TDA(mf.to_rhf())

def TDDFT(mf):
    if isinstance(mf, scf.uhf.UHF):
        if isinstance(mf, dft.rks.KohnShamDFT):
            return uks.tddft(mf.to_uks())
        else:
            return uhf.TDHF(mf.to_uhf())
    else:
        if isinstance(mf, dft.rks.KohnShamDFT):
            return rks.tddft(mf.to_rks())
        else:
            return rhf.TDHF(mf.to_rhf())

TD = TDDFT


def RPA(mf):
    return TDDFT(mf)

def dRPA(mf):
    if isinstance(mf, scf.uhf.UHF):
        return uks.dRPA(mf)
    else:
        return rks.dRPA(mf)

def dTDA(mf):
    if isinstance(mf, scf.uhf.UHF):
        return uks.dTDA(mf)
    else:
        return rks.dTDA(mf)


# pyscf/tdscf/common_slow.py

#  Author: Artem Pulkin
"""
This and other `_slow` modules implement the time-dependent procedure. The primary performance drawback is
that, unlike other 'fast' routines with an implicit construction of the eigenvalue problem, these modules construct
TDHF matrices explicitly. As a result, regular `numpy.linalg.eig` can be used to retrieve TDHF roots in a reliable
fashion without any issues related to the Davidson procedure.

This is a helper module defining basic interfaces.
"""

import sys
from pyscf.lib import logger

from pyscf.pbc.tools import get_kconserv

import numpy
from scipy.linalg import solve

from itertools import count, groupby

if sys.version_info >= (3,):
    unicode = str

def msize(m):
    """
    Checks whether the matrix is square and returns its size.
    Args:
        m (numpy.ndarray): the matrix to measure;

    Returns:
        An integer with the size.
    """
    s = m.shape[0]
    if m.shape != (s, s):
        raise ValueError("Do not recognize the shape (must be a square matrix): {}".format(m.shape))
    return s


def full2ab(full, tolerance=1e-12):
    """
    Transforms a full TD matrix into A and B parts.
    Args:
        full (numpy.ndarray): the full TD matrix;
        tolerance (float): a tolerance for checking whether the full matrix is in the ABBA-form;

    Returns:
        A and B submatrices.
    """
    s = msize(full)
    if s % 2 != 0:
        raise ValueError("Not an even matrix size: {:d}".format(s))

    s2 = s // 2
    a, b = full[:s2, :s2], full[:s2, s2:]
    b_, a_ = full[s2:, :s2].conj(), full[s2:, s2:].conj()
    delta = max(abs(a + a_).max(), abs(b + b_).max())
    if delta > tolerance:
        raise ValueError("The full matrix is not in the ABBA-form, delta: {:.3e}".format(delta))

    return full[:s2, :s2], full[:s2, s2:]


def ab2full(a, b):
    """
    Transforms A and B TD matrices into a full matrix.
    Args:
        a (numpy.ndarray): TD A-matrix;
        b (numpy.ndarray): TD B-matrix;

    Returns:
        The full TD matrix.
    """
    sa = msize(a)
    sb = msize(b)
    if sa != sb:
        raise ValueError("Input matrix dimensions do not match: {:d} vs {:d}".format(sa, sb))
    return numpy.block([[a, b], [-b.conj(), -a.conj()]])


def ab2mkk(a, b, tolerance=1e-12):
    """
    Transforms A and B TD matrices into MK and K matrices.
    Args:
        a (numpy.ndarray): TD A-matrix;
        b (numpy.ndarray): TD B-matrix;
        tolerance (float): a tolerance for checking whether the input matrices are real;

    Returns:
        MK and K submatrices.
    """
    if max(abs(a.imag).max(), abs(b.imag).max()) > tolerance:
        raise ValueError("A- and/or B-matrixes are complex-valued: no transform is possible")
    a, b = a.real, b.real
    tdhf_k, tdhf_m = a - b, a + b
    tdhf_mk = tdhf_m.dot(tdhf_k)
    return tdhf_mk, tdhf_k


def mkk2ab(mk, k):
    """
    Transforms MK and M TD matrices into A and B matrices.
    Args:
        mk (numpy.ndarray): TD MK-matrix;
        k (numpy.ndarray): TD K-matrix;

    Returns:
        A and B submatrices.
    """
    if numpy.iscomplexobj(mk) or numpy.iscomplexobj(k):
        raise ValueError("MK- and/or K-matrixes are complex-valued: no transform is possible")
    m = solve(k.T, mk.T).T
    a = 0.5 * (m + k)
    b = 0.5 * (m - k)
    return a, b


def full2mkk(full):
    """
    Transforms a full TD matrix into MK and K parts.
    Args:
        full (numpy.ndarray): the full TD matrix;

    Returns:
        MK and K submatrices.
    """
    return ab2mkk(*full2ab(full))


def mkk2full(mk, k):
    """
    Transforms MK and M TD matrices into a full TD matrix.
    Args:
        mk (numpy.ndarray): TD MK-matrix;
        k (numpy.ndarray): TD K-matrix;

    Returns:
        The full TD matrix.
    """
    return ab2full(*mkk2ab(mk, k))


class TDMatrixBlocks(object):
    def tdhf_primary_form(self, *args, **kwargs):
        """
        A primary form of TDHF matrixes.

        Returns:
            Output type: "full", "ab", or "mk" and the corresponding matrix(es).
        """
        raise NotImplementedError

    @staticmethod
    def __check_primary_form__(m):
        if not isinstance(m, tuple):
            raise ValueError("The value returned by `tdhf_primary_form` is not a tuple")
        if len(m) < 1:
            raise ValueError("Empty tuple returned by `tdhf_primary_form`")
        if not isinstance(m[0], (str, unicode)):
            raise ValueError("The first item returned by `tdhf_primary_form` must be a string")
        forms = dict(ab=3, mk=3, full=2)
        if m[0] in forms:
            if len(m) != forms[m[0]]:
                raise ValueError("The {} form returned by `tdhf_primary_form` must contain {:d} values".format(
                    m[0].upper(), forms[m[0]],
                ))
        else:
            raise ValueError("Unknown form specification returned by `tdhf_primary_form`: {}".format(m[0]))

    def tdhf_ab_form(self, *args, **kwargs):
        """
        The A-B form of the TD problem.

        Returns:
            A and B TD matrices.
        """
        m = self.tdhf_primary_form(*args, **kwargs)
        self.__check_primary_form__(m)
        if m[0] == "ab":
            return m[1:]
        elif m[0] == "full":
            return full2ab(m[1])
        elif m[0] == "mk":
            return mkk2ab(*m[1:])

    def tdhf_full_form(self, *args, **kwargs):
        """
        The full form of the TD problem.

        Returns:
            The full TD matrix.
        """
        m = self.tdhf_primary_form(*args, **kwargs)
        self.__check_primary_form__(m)
        if m[0] == "ab":
            return ab2full(*m[1:])
        elif m[0] == "full":
            return m[1]
        elif m[0] == "mk":
            return mkk2full(*m[1:])

    def tdhf_mk_form(self, *args, **kwargs):
        """
        The MK form of the TD problem.

        Returns:
            MK and K TD matrixes.
        """
        m = self.tdhf_primary_form(*args, **kwargs)
        self.__check_primary_form__(m)
        if m[0] == "ab":
            return ab2mkk(*m[1:])
        elif m[0] == "full":
            return full2mkk(m[1])
        elif m[0] == "mk":
            return m[1:]


def mknj2i(item):
    """
    Transforms "mknj" notation into tensor index order for the ERI.
    Args:
        item (str): an arbitrary transpose of "mknj" letters;

    Returns:
        4 indexes.
    """
    notation = "mknj"
    notation = dict(zip(notation, range(len(notation))))
    return tuple(notation[i] for i in item)


class TDERIMatrixBlocks(TDMatrixBlocks):
    symmetries = [
        ((0, 1, 2, 3), False),
    ]

    def __init__(self):
        """
        This a prototype class for TD calculations based on ERI (TD-HF). It handles integral blocks and
        the diagonal part, see Eq. 7.5 of RevModPhys.36.844.
        """
        # Caching
        self.__eri__ = {}

    def __get_mo_energies__(self, *args, **kwargs):
        """This routine collects occupied and virtual MO energies."""
        raise NotImplementedError

    def __calc_block__(self, item, *args):
        raise NotImplementedError

    def tdhf_diag(self, *args):
        """
        Retrieves the diagonal block.
        Args:
            *args: args passed to `__get_mo_energies__`;

        Returns:
            The diagonal block.
        """
        e_occ, e_virt = self.__get_mo_energies__(*args)
        diag = (- e_occ[:, numpy.newaxis] + e_virt[numpy.newaxis, :]).reshape(-1)
        return numpy.diag(diag).reshape((len(e_occ) * len(e_virt), len(e_occ) * len(e_virt)))

    def eri_ov(self, item, *args):
        """
        Retrieves ERI block using 'ov' notation.
        Args:
            item (str): a 4-character string of 'o' and 'v' letters;
            *args: other args passed to `__calc_block__`;

        Returns:
            The corresponding block of ERI (4-tensor, phys notation).
        """
        if len(item) != 4 or not isinstance(item, str) or not set(item).issubset('ov'):
            raise ValueError("Unknown item: {}".format(repr(item)))

        args = (tuple(item), ) + args
        if args in self.__eri__:
            return self.__eri__[args]

        result = self.__calc_block__(*args)

        for permutation, conjugation in self.symmetries:
            permuted_args = tuple(
                tuple(arg[_i] for _i in permutation)
                for arg in args
            )
            if conjugation:
                self.__eri__[permuted_args] = result.transpose(*permutation).conj()
            else:
                self.__eri__[permuted_args] = result.transpose(*permutation)
        return result

    def eri_mknj(self, item, *args):
        """
        Retrieves ERI block using 'mknj' notation.
        Args:
            item (str): a 4-character string of 'mknj' letters;
            *args: other arguments passed to `get_block_ov_notation`;

        Returns:
            The corresponding block of ERI (matrix with paired dimensions).
        """
        if len(item) != 4 or not isinstance(item, str) or set(item) != set('mknj'):
            raise ValueError("Unknown item: {}".format(repr(item)))

        item = mknj2i(item)
        n_ov = ''.join('o' if i % 2 == 0 else 'v' for i in item)
        args = tuple(
            tuple(arg[i] for i in item)
            for arg in args
        )
        result = self.eri_ov(n_ov, *args).transpose(*numpy.argsort(item))
        i, j, k, l = result.shape
        result = result.reshape((i * j, k * l))
        return result

    def __getitem__(self, item):
        if isinstance(item, str):
            spec, args = item, tuple()
        else:
            spec, args = item[0], item[1:]
        if set(spec) == set("mknj"):
            return self.eri_mknj(spec, *args)
        elif set(spec).issubset("ov"):
            return self.eri_ov(spec, *args)
        else:
            raise ValueError("Unknown item: {}".format(repr(item)))

    def tdhf_primary_form(self, *args, **kwargs):
        """
        A primary form of TDHF matrixes (AB).

        Returns:
            Output type: "ab", and the corresponding matrixes.
        """
        d = self.tdhf_diag(*args, **kwargs)
        a = d + 2 * self["knmj"] - self["knjm"]
        b = 2 * self["kjmn"] - self["kjnm"]
        return "ab", a, b


class TDProxyMatrixBlocks(TDMatrixBlocks):
    def __init__(self, model):
        """
        This a prototype class for TD calculations based on proxying pyscf classes such as TDDFT. It is a work-around
        class. It accepts a `pyscf.tdscf.*` class and uses its matvec to construct a full-sized TD matrix.
        Args:
            model: a pyscf base model to extract TD matrix from;
        """
        super(TDProxyMatrixBlocks, self).__init__()
        self.proxy_model = model
        self.proxy_vind, self.proxy_diag = self.proxy_model.gen_vind(self.proxy_model._scf)
        self.proxy_vind = VindTracker(self.proxy_vind)

    def tdhf_primary_form(self, *args, **kwargs):
        raise NotImplementedError


def format_frozen_mol(frozen, nmo):
    """
    Formats the argument into a mask array of bools where False values correspond to frozen molecular orbitals.
    Args:
        frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals;
        nmo (int): the total number of molecular orbitals;

    Returns:
        The mask array.
    """
    space = numpy.ones(nmo, dtype=bool)
    if frozen is None:
        pass
    elif isinstance(frozen, int):
        space[:frozen] = False
    elif isinstance(frozen, (tuple, list, numpy.ndarray)):
        space[frozen] = False
    else:
        raise ValueError("Cannot recognize the 'frozen' argument: expected None, int or Iterable")
    return space


class MolecularMFMixin(object):
    def __init__(self, model, frozen=None):
        """
        A mixin to support custom slices of mean-field attributes: `mo_coeff`, `mo_energy`, ...

        Molecular version. Also supports single k-point inputs.

        Args:
            model: the base model;
            frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals;
        """
        self.__is_k__ = False
        if "kpts" in dir(model):
            self.__is_k__ = True
            if len(model.kpts) != 1:
                raise ValueError("Only a single k-point supported, found: model.kpts = {}".format(model.kpts))
        self.model = model
        self.space = format_frozen_mol(frozen, len(self.squeeze(model.mo_energy)))

    def squeeze(self, x):
        """Squeezes quantities in the case of a PBC model."""
        return x[0] if self.__is_k__ else x

    @property
    def mo_coeff(self):
        """MO coefficients."""
        return self.squeeze(self.model.mo_coeff)[:, self.space]

    @property
    def mo_energy(self):
        """MO energies."""
        return self.squeeze(self.model.mo_energy)[self.space]

    @property
    def mo_occ(self):
        """MO occupation numbers."""
        return self.squeeze(self.model.mo_occ)[self.space]

    @property
    def nocc(self):
        """The number of occupied orbitals."""
        return int(self.squeeze(self.model.mo_occ)[self.space].sum() // 2)

    @property
    def nmo(self):
        """The total number of molecular orbitals."""
        return self.space.sum()

    @property
    def mo_coeff_full(self):
        """MO coefficients."""
        return self.squeeze(self.model.mo_coeff)

    @property
    def nocc_full(self):
        """The true (including frozen degrees of freedom) number of occupied orbitals."""
        return int(self.squeeze(self.model.mo_occ).sum() // 2)

    @property
    def nmo_full(self):
        """The true (including frozen degrees of freedom) total number of molecular orbitals."""
        return len(self.space)


def format_frozen_k(frozen, nmo, nk):
    """
    Formats the argument into a mask array of bools where False values correspond to frozen orbitals for each k-point.
    Args:
        frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals for all k-points or
        multiple lists of frozen orbitals for each k-point;
        nmo (int): the total number of molecular orbitals;
        nk (int): the total number of k-points;

    Returns:
        The mask array.
    """
    space = numpy.ones((nk, nmo), dtype=bool)
    if frozen is None:
        pass
    elif isinstance(frozen, int):
        space[:, :frozen] = False
    elif isinstance(frozen, (tuple, list, numpy.ndarray)):
        if len(frozen) > 0:
            if isinstance(frozen[0], int):
                space[:, frozen] = False
            else:
                for i in range(nk):
                    space[i, frozen[i]] = False
    else:
        raise ValueError("Cannot recognize the 'frozen' argument: expected None, int or Iterable")
    return space


def k_nocc(model):
    """
    Retrieves occupation numbers.
    Args:
        model (RHF): the model;

    Returns:
        Numbers of occupied orbitals in the model.
    """
    return tuple(int(i.sum() // 2) for i in model.mo_occ)


def k_nmo(model):
    """
    Retrieves number of AOs per k-point.
    Args:
        model (RHF): the model;

    Returns:
        Numbers of AOs in the model.
    """
    return tuple(i.shape[1] for i in model.mo_coeff)


class PeriodicMFMixin(object):
    def __init__(self, model, frozen=None):
        """
        A mixin to support custom slices of mean-field attributes: `mo_coeff`, `mo_energy`, ...

        PBC version.

        Args:
            model: the base model;
            frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals;
        """
        self.model = model
        self.space = format_frozen_k(frozen, len(model.mo_energy[0]), len(model.kpts))
        self.kconserv = get_kconserv(self.model.cell, self.model.kpts).swapaxes(1, 2)

    @property
    def mo_coeff(self):
        """MO coefficients."""
        return tuple(i[:, j] for i, j in zip(self.model.mo_coeff, self.space))

    @property
    def mo_energy(self):
        """MO energies."""
        return tuple(i[j] for i, j in zip(self.model.mo_energy, self.space))

    @property
    def mo_occ(self):
        """MO occupation numbers."""
        return tuple(i[j] for i, j in zip(self.model.mo_occ, self.space))

    @property
    def nocc(self):
        """The number of occupied orbitals."""
        return k_nocc(self)

    @property
    def nmo(self):
        """The total number of molecular orbitals."""
        return k_nmo(self)

    @property
    def mo_coeff_full(self):
        """MO coefficients."""
        return self.model.mo_coeff

    @property
    def nocc_full(self):
        """The true (including frozen degrees of freedom) number of occupied orbitals."""
        return k_nocc(self.model)

    @property
    def nmo_full(self):
        """The true (including frozen degrees of freedom) total number of molecular orbitals."""
        return k_nmo(self.model)


class VindTracker(object):
    def __init__(self, vind):
        """
        Tracks calls to `vind` (a matrix-vector multiplication density response routine).
        Args:
            vind (Callable): a matvec product routine;
        """
        self.vind = vind
        self.args = self.results = self.errors = None
        self.reset()

    def reset(self):
        """
        Resets statistics.
        """
        self.args = []
        self.results = []
        self.errors = []

    def __call__(self, v):
        if not isinstance(v, numpy.ndarray):
            raise ValueError("The input is not an array")
        self.args.append(v.shape)
        try:
            r = self.vind(v)
        except Exception as e:
            self.results.append(None)
            self.errors.append(e)
            raise
        r = numpy.array(r)
        self.results.append(r.shape)
        self.errors.append(None)
        return r

    def __iter__(self):
        for i, o, e in zip(self.args, self.results, self.errors):
            yield i, o, e

    @property
    def ncalls(self):
        return len(self.args)

    @property
    def msize(self):
        for i in self.results:
            if i is not None:
                return i[1]
        return None

    @property
    def elements_total(self):
        return self.msize ** 2

    @property
    def elements_calc(self):
        return sum(map(lambda i: i[0] * i[1] if i is not None else 0, self.results))

    @property
    def ratio(self):
        return 1.0 * self.elements_calc / self.elements_total

    def text_stats(self):
        return "--------------------\nVind call statistics\n--------------------\n" \
            "  calls: {total_calls:d}\n" \
            "  elements total: {total_elems:d} ({size})\n" \
            "  elements calculated: {total_calc:d}\n" \
            "  ratio: {ratio:.3f}".format(
                total_calls=self.ncalls,
                total_elems=self.elements_total,
                size="x".join((str(self.msize),) * 2),
                total_calc=self.elements_calc,
                ratio=self.ratio,
            )


def eig(m, driver=None, nroots=None, half=True):
    """
    Eigenvalue problem solver.
    Args:
        m (numpy.ndarray): the matrix to diagonalize;
        driver (str): one of the drivers;
        nroots (int): the number of roots ot calculate (ignored for `driver` == 'eig');
        half (bool): if True, implies spectrum symmetry and takes only a half of eigenvalues;

    Returns:

    """
    if driver is None:
        driver = 'eig'
    if driver == 'eig':
        vals, vecs = numpy.linalg.eig(m)
        order = numpy.argsort(vals)
        vals, vecs = vals[order], vecs[:, order]
        if half:
            vals, vecs = vals[len(vals) // 2:], vecs[:, vecs.shape[1] // 2:]
            vecs = vecs[:, ]
        vals, vecs = vals[:nroots], vecs[:, :nroots]
    else:
        raise ValueError("Unknown driver: {}".format(driver))
    return vals, vecs


def kernel(eri, driver=None, fast=True, nroots=None, **kwargs):
    """
    Calculates eigenstates and eigenvalues of the TDHF problem.
    Args:
        eri (TDDFTMatrixBlocks): ERI;
        driver (str): one of the eigenvalue problem drivers;
        fast (bool): whether to run diagonalization on smaller matrixes;
        nroots (int): the number of roots to calculate;
        **kwargs: arguments to `eri.tdhf_matrix`;

    Returns:
        Positive eigenvalues and eigenvectors.
    """
    if not isinstance(eri, TDMatrixBlocks):
        raise ValueError("The argument must be ERI object")

    if fast:

        logger.debug1(eri.model, "Preparing TDHF matrix (fast) ...")
        tdhf_mk, tdhf_k = eri.tdhf_mk_form(**kwargs)
        logger.debug1(eri.model, "Diagonalizing a {} matrix with {} ...".format(
            'x'.join(map(str, tdhf_mk.shape)),
            "'{}'".format(driver) if driver is not None else "a default method",
        ))
        vals, vecs_x = eig(tdhf_mk, driver=driver, nroots=nroots, half=False)

        vals = vals ** .5
        vecs_y = (1. / vals)[numpy.newaxis, :] * tdhf_k.dot(vecs_x)
        vecs_u, vecs_v = vecs_y + vecs_x, vecs_y - vecs_x
        return vals, numpy.concatenate((vecs_u, vecs_v), axis=0)

    else:

        logger.debug1(eri.model, "Preparing TDHF matrix ...")
        m = eri.tdhf_full_form(**kwargs)
        logger.debug1(eri.model, "Diagonalizing a {} matrix with {} ...".format(
            'x'.join(map(str, m.shape)),
            "'{}'".format(driver) if driver is not None else "a default method",
        ))
        return eig(m, driver=driver, nroots=nroots)


class TDBase(object):
    v2a = None

    def __init__(self, mf, frozen=None):
        """
        Performs TD calculation. Roots and eigenvectors are stored in `self.e`, `self.xy`.
        Args:
            mf: the mean-field model;
            frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals;
        """
        self._scf = mf
        self.driver = None
        self.nroots = None
        self.eri = None
        self.xy = None
        self.e = None
        self.frozen = frozen
        self.fast = not numpy.iscomplexobj(numpy.asanyarray(mf.mo_coeff))

    def __kernel__(self, **kwargs):
        """Silent implementation of kernel which does not change attributes."""
        if self.eri is None:
            self.eri = self.ao2mo()

        e, xy = kernel(
            self.eri,
            driver=self.driver,
            nroots=self.nroots,
            fast=self.fast,
            **kwargs
        )
        xy = self.vector_to_amplitudes(xy)
        return e, xy

    def kernel(self):
        """
        Calculates eigenstates and eigenvalues of the TDHF problem.

        Returns:
            Positive eigenvalues and eigenvectors.
        """
        self.e, self.xy = self.__kernel__()
        return self.e, self.xy

    def ao2mo(self):
        """
        Picks ERI: either 4-fold or 8-fold symmetric.

        Returns:
            A suitable ERI.
        """
        raise NotImplementedError

    def vector_to_amplitudes(self, vectors):
        """
        Transforms (reshapes) and normalizes vectors into amplitudes.
        Args:
            vectors (numpy.ndarray): raw eigenvectors to transform;

        Returns:
            Amplitudes with the following shape: (# of roots, 2 (x or y), # of
            occupied orbitals, # of virtual orbitals).
        """
        return self.v2a(vectors, self.eri.nocc, self.eri.nmo)


def format_mask(x):
    """
    Formats a mask into a readable string.
    Args:
        x (ndarray): an array with the mask;

    Returns:
        A readable string with the mask.
    """
    x = numpy.asanyarray(x)
    if len(x) == 0:
        return "(empty)"
    if x.dtype == bool:
        x = numpy.argwhere(x)[:, 0]
    grps = tuple(list(g) for _, g in groupby(x, lambda n, c=count(): n-next(c)))
    return ",".join("{:d}-{:d}".format(i[0], i[-1]) if len(i) > 1 else "{:d}".format(i[0]) for i in grps)


# pyscf/tdscf/dhf.py

#!/usr/bin/env python
# Copyright 2021-2022 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
TDA and TDHF for no-pair DKS Hamiltonian
'''

from functools import reduce
import numpy
from pyscf import lib
from pyscf import dft
from pyscf import ao2mo
from pyscf.lib import logger
from pyscf.tdscf import rhf
from pyscf.data import nist
from pyscf import __config__

OUTPUT_THRESHOLD = getattr(__config__, 'tdscf_uhf_get_nto_threshold', 0.3)
REAL_EIG_THRESHOLD = getattr(__config__, 'tdscf_uhf_TDDFT_pick_eig_threshold', 1e-4)

# Low excitation filter to avoid numerical instability
POSTIVE_EIG_THRESHOLD = getattr(__config__, 'tdscf_uhf_TDDFT_positive_eig_threshold', 1e-3)

def gen_tda_operation(mf, fock_ao=None):
    '''A x
    '''
    mo_coeff = mf.mo_coeff
    mo_energy = mf.mo_energy
    mo_occ = mf.mo_occ
    nao, nmo = mo_coeff.shape
    # Remove all negative states
    n2c = nmo // 2
    occidx = n2c + numpy.where(mo_occ[n2c:] == 1)[0]
    viridx = n2c + numpy.where(mo_occ[n2c:] == 0)[0]
    nocc = len(occidx)
    nvir = len(viridx)
    orbv = mo_coeff[:,viridx]
    orbo = mo_coeff[:,occidx]

    if fock_ao is None:
        #dm0 = mf.make_rdm1(mo_coeff, mo_occ)
        #fock_ao = mf.get_hcore() + mf.get_veff(mol, dm0)
        foo = numpy.diag(mo_energy[occidx])
        fvv = numpy.diag(mo_energy[viridx])
    else:
        fock = reduce(numpy.dot, (mo_coeff.conj().T, fock_ao, mo_coeff))
        foo = fock[occidx[:,None],occidx]
        fvv = fock[viridx[:,None],viridx]
    hdiag = (fvv.diagonal() - foo.diagonal()[:,None]).ravel()

    mo_coeff = numpy.asarray(numpy.hstack((orbo,orbv)), order='F')
    vresp = mf.gen_response(hermi=0)

    def vind(zs):
        zs = numpy.asarray(zs).reshape(-1,nocc,nvir)
        dmov = lib.einsum('xov,qv,po->xpq', zs, orbv.conj(), orbo)
        v1ao = vresp(dmov)
        v1ov = lib.einsum('xpq,qv,po->xov', v1ao, orbv, orbo.conj())
        v1ov += lib.einsum('xqs,sp->xqp', zs, fvv)
        v1ov -= lib.einsum('xpr,sp->xsr', zs, foo)
        return v1ov.reshape(v1ov.shape[0], -1)

    return vind, hdiag
gen_tda_hop = gen_tda_operation

def get_ab(mf, mo_energy=None, mo_coeff=None, mo_occ=None):
    r'''A and B matrices for TDDFT response function.

    A[i,a,j,b] = \delta_{ab}\delta_{ij}(E_a - E_i) + (ia||bj)
    B[i,a,j,b] = (ia||jb)
    '''
    if mo_energy is None: mo_energy = mf.mo_energy
    if mo_coeff is None: mo_coeff = mf.mo_coeff
    if mo_occ is None: mo_occ = mf.mo_occ
    assert(mo_coeff.dtype == numpy.double)

    mol = mf.mol
    nao, nmo = mo_coeff.shape
    n2c = nmo // 2
    occidx = n2c + numpy.where(mo_occ[n2c:] == 1)[0]
    viridx = n2c + numpy.where(mo_occ[n2c:] == 0)[0]
    orbv = mo_coeff[:,viridx]
    orbo = mo_coeff[:,occidx]
    nvir = orbv.shape[1]
    nocc = orbo.shape[1]
    nmo = nocc + nvir
    mo = numpy.hstack((orbo, orbv))
    moL = numpy.asarray(mo[:n2c], order='F')
    moS = numpy.asarray(mo[n2c:], order='F')
    orboL = moL[:,:nocc]
    orboS = moS[:,:nocc]
    orbvL = moL[:,nocc:]
    orbvS = moS[:,nocc:]

    e_ia = lib.direct_sum('a-i->ia', mo_energy[viridx], mo_energy[occidx])
    a = numpy.diag(e_ia.ravel()).reshape(nocc,nvir,nocc,nvir)
    b = numpy.zeros_like(a)

    def add_hf_(a, b, hyb=1):
        eri_mo = ao2mo.kernel(mol, [orboL, moL, moL, moL], intor='int2e_spinor')
        eri_mo+= ao2mo.kernel(mol, [orboS, moS, moS, moS], intor='int2e_spsp1_spsp2_spinor')
        eri_mo+= ao2mo.kernel(mol, [orboS, moS, moL, moL], intor='int2e_spsp1_spinor')
        eri_mo+= ao2mo.kernel(mol, [moS, moS, orboL, moL], intor='int2e_spsp1_spinor').T
        eri_mo = eri_mo.reshape(nocc,nmo,nmo,nmo)

        a += numpy.einsum('iabj->iajb', eri_mo[:nocc,nocc:,nocc:,:nocc])
        a -= numpy.einsum('ijba->iajb', eri_mo[:nocc,:nocc,nocc:,nocc:]) * hyb
        b += numpy.einsum('iajb->iajb', eri_mo[:nocc,nocc:,:nocc,nocc:])
        b -= numpy.einsum('jaib->iajb', eri_mo[:nocc,nocc:,:nocc,nocc:]) * hyb

    if isinstance(mf, dft.KohnShamDFT):
        ni = mf._numint
        ni.libxc.test_deriv_order(mf.xc, 2, raise_error=True)
        if getattr(mf, 'nlc', '') != '':
            raise NotImplementedError

        omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin)

        add_hf_(a, b, hyb)

        xctype = ni._xc_type(mf.xc)
        dm0 = mf.make_rdm1(mo_coeff, mo_occ)
        make_rho = ni._gen_rho_evaluator(mol, dm0, hermi=1)[0]
        mem_now = lib.current_memory()[0]
        max_memory = max(2000, mf.max_memory*.8-mem_now)

        if xctype == 'LDA':
            ao_deriv = 0
            for ao, mask, weight, coords \
                    in ni.block_loop(mol, mf.grids, nao, ao_deriv, max_memory):
                rho = make_rho(0, ao, mask, xctype)
                fxc = ni.eval_xc(mf.xc, rho, 1, deriv=2)[2]
                u_u, u_d, d_d = fxc[0].T

                aoLa, aoLb, aoSa, aoSb = ao
                moL_oa = lib.einsum('rp,pi->ri', aoLa, orboL)
                moL_va = lib.einsum('rp,pi->ri', aoLa, orbvL)
                moL_ob = lib.einsum('rp,pi->ri', aoLb, orboL)
                moL_vb = lib.einsum('rp,pi->ri', aoLb, orbvL)
                moS_oa = lib.einsum('rp,pi->ri', aoSa, orboS)
                moS_va = lib.einsum('rp,pi->ri', aoSa, orbvS)
                moS_ob = lib.einsum('rp,pi->ri', aoSb, orboS)
                moS_vb = lib.einsum('rp,pi->ri', aoSb, orbvS)
                if mf.collinear:
                    rhoL_ov_a = numpy.einsum('ri,ra->ria', moL_oa.conj(), moL_va)
                    rhoL_ov_b = numpy.einsum('ri,ra->ria', moL_ob.conj(), moL_vb)
                    rhoS_ov_a = numpy.einsum('ri,ra->ria', moS_oa.conj(), moS_va)
                    rhoS_ov_b = numpy.einsum('ri,ra->ria', moS_ob.conj(), moS_vb)

                    # M = \beta\Sigam leads to the flip of up-down in small
                    # component densities
                    rho_ov_a = rhoL_ov_a + rhoS_ov_b
                    rho_ov_b = rhoL_ov_b + rhoS_ov_a
                    rho_vo_a = rho_ov_a.conj()
                    rho_vo_b = rho_ov_b.conj()

                    w_ov_a = numpy.einsum('ria,r->ria', u_u * rho_ov_a + u_d * rho_ov_b, weight)
                    a += lib.einsum('ria,rjb->iajb', w_ov_a, rho_vo_a)
                    b += lib.einsum('ria,rjb->iajb', w_ov_a, rho_ov_a)

                    w_ov_b = numpy.einsum('ria,r->ria', u_d * rho_ov_a + d_d * rho_ov_b, weight)
                    a += lib.einsum('ria,rjb->iajb', w_ov_b, rho_vo_b)
                    b += lib.einsum('ria,rjb->iajb', w_ov_b, rho_ov_b)
                else:
                    raise NotImplementedError

        elif xctype == 'GGA':
            ao_deriv = 1
            for ao, mask, weight, coords \
                    in ni.block_loop(mol, mf.grids, nao, ao_deriv, max_memory):
                rho0 = make_rho(0, ao, mask, xctype)
                r0, (mx0, my0, mz0) = rho0
                rho0a = (r0 + mz0) * .5
                rho0b = (r0 - mz0) * .5
                vxc, fxc = ni.eval_xc(mf.xc, rho0, 1, deriv=2)[1:3]
                uu, ud, dd = vxc[1].T
                u_u, u_d, d_d = fxc[0].T
                u_uu, u_ud, u_dd, d_uu, d_ud, d_dd = fxc[1].T
                uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd = fxc[2].T

                aoLa, aoLb, aoSa, aoSb = ao
                moL_oa = lib.einsum('xrp,pi->xri', aoLa, orboL)
                moL_va = lib.einsum('xrp,pi->xri', aoLa, orbvL)
                moL_ob = lib.einsum('xrp,pi->xri', aoLb, orboL)
                moL_vb = lib.einsum('xrp,pi->xri', aoLb, orbvL)
                moS_oa = lib.einsum('xrp,pi->xri', aoSa, orboS)
                moS_va = lib.einsum('xrp,pi->xri', aoSa, orbvS)
                moS_ob = lib.einsum('xrp,pi->xri', aoSb, orboS)
                moS_vb = lib.einsum('xrp,pi->xri', aoSb, orbvS)
                if mf.collinear:
                    rhoL_ov_a = numpy.einsum('xri,ra->xria', moL_oa.conj(), moL_va[0])
                    rhoL_ov_b = numpy.einsum('xri,ra->xria', moL_ob.conj(), moL_vb[0])
                    rhoL_ov_a[1:4] += numpy.einsum('ri,xra->xria', moL_oa[0].conj(), moL_va[1:4])
                    rhoL_ov_b[1:4] += numpy.einsum('ri,xra->xria', moL_ob[0].conj(), moL_vb[1:4])
                    rhoS_ov_a = numpy.einsum('xri,ra->xria', moS_oa.conj(), moS_va[0])
                    rhoS_ov_b = numpy.einsum('xri,ra->xria', moS_ob.conj(), moS_vb[0])
                    rhoS_ov_a[1:4] += numpy.einsum('ri,xra->xria', moS_oa[0].conj(), moS_va[1:4])
                    rhoS_ov_b[1:4] += numpy.einsum('ri,xra->xria', moS_ob[0].conj(), moS_vb[1:4])

                    # M = \beta\Sigam leads to the flip of up-down in small
                    # component densities
                    rho_ov_a = rhoL_ov_a + rhoS_ov_b
                    rho_ov_b = rhoL_ov_b + rhoS_ov_a
                    rho_vo_a = rho_ov_a.conj()
                    rho_vo_b = rho_vo_a.conj()

                    a0a1 = numpy.einsum('xr,xria->ria', rho0a[1:4], rho_ov_a[1:4])
                    a0b1 = numpy.einsum('xr,xria->ria', rho0a[1:4], rho_ov_b[1:4])
                    b0a1 = numpy.einsum('xr,xria->ria', rho0b[1:4], rho_ov_a[1:4])
                    b0b1 = numpy.einsum('xr,xria->ria', rho0b[1:4], rho_ov_b[1:4])

                    # aaaa
                    w_ov = numpy.empty_like(rho_ov_a)
                    w_ov[0]  = numpy.einsum('r,ria->ria', u_u, rho_ov_a[0])
                    w_ov[0] += numpy.einsum('r,ria->ria', 2*u_uu, a0a1)
                    w_ov[0] += numpy.einsum('r,ria->ria',   u_ud, b0a1)
                    f_ov_a  = numpy.einsum('r,ria->ria', 4*uu_uu, a0a1)
                    f_ov_b  = numpy.einsum('r,ria->ria', 2*uu_ud, a0a1)
                    f_ov_a += numpy.einsum('r,ria->ria', 2*uu_ud, b0a1)
                    f_ov_b += numpy.einsum('r,ria->ria',   ud_ud, b0a1)
                    f_ov_a += numpy.einsum('r,ria->ria', 2*u_uu, rho_ov_a[0])
                    f_ov_b += numpy.einsum('r,ria->ria',   u_ud, rho_ov_a[0])
                    w_ov[1:] = numpy.einsum('ria,xr->xria', f_ov_a, rho0a[1:4])
                    w_ov[1:]+= numpy.einsum('ria,xr->xria', f_ov_b, rho0b[1:4])
                    w_ov[1:]+= numpy.einsum('r,xria->xria', 2*uu, rho_ov_a[1:4])
                    w_ov *= weight[:,None,None]
                    a += lib.einsum('xria,xrjb->iajb', w_ov, rho_vo_a)
                    b += lib.einsum('xria,xrjb->iajb', w_ov, rho_ov_a)

                    # bbbb
                    w_ov = numpy.empty_like(rho_ov_b)
                    w_ov[0]  = numpy.einsum('r,ria->ria', d_d, rho_ov_b[0])
                    w_ov[0] += numpy.einsum('r,ria->ria', 2*d_dd, b0b1)
                    w_ov[0] += numpy.einsum('r,ria->ria',   d_ud, a0b1)
                    f_ov_b  = numpy.einsum('r,ria->ria', 4*dd_dd, b0b1)
                    f_ov_a  = numpy.einsum('r,ria->ria', 2*ud_dd, b0b1)
                    f_ov_b += numpy.einsum('r,ria->ria', 2*ud_dd, a0b1)
                    f_ov_a += numpy.einsum('r,ria->ria',   ud_ud, a0b1)
                    f_ov_b += numpy.einsum('r,ria->ria', 2*d_dd, rho_ov_b[0])
                    f_ov_a += numpy.einsum('r,ria->ria',   d_ud, rho_ov_b[0])
                    w_ov[1:] = numpy.einsum('ria,xr->xria', f_ov_a, rho0a[1:4])
                    w_ov[1:]+= numpy.einsum('ria,xr->xria', f_ov_b, rho0b[1:4])
                    w_ov[1:]+= numpy.einsum('r,xria->xria', 2*dd, rho_ov_b[1:4])
                    w_ov *= weight[:,None,None]
                    a += lib.einsum('xria,xrjb->iajb', w_ov, rho_vo_b)
                    b += lib.einsum('xria,xrjb->iajb', w_ov, rho_ov_b)

                    # aabb and bbaa
                    w_ov = numpy.empty_like(rho_ov_b)
                    w_ov[0]  = numpy.einsum('r,ria->ria', u_d, rho_ov_b[0])
                    w_ov[0] += numpy.einsum('r,ria->ria', 2*u_dd, b0b1)
                    w_ov[0] += numpy.einsum('r,ria->ria',   u_ud, a0b1)
                    f_ov_a  = numpy.einsum('r,ria->ria', 4*uu_dd, b0b1)
                    f_ov_b  = numpy.einsum('r,ria->ria', 2*ud_dd, b0b1)
                    f_ov_a += numpy.einsum('r,ria->ria', 2*uu_ud, a0b1)
                    f_ov_b += numpy.einsum('r,ria->ria',   ud_ud, a0b1)
                    f_ov_a += numpy.einsum('r,ria->ria', 2*d_uu, rho_ov_b[0])
                    f_ov_b += numpy.einsum('r,ria->ria',   d_ud, rho_ov_b[0])
                    w_ov[1:] = numpy.einsum('ria,xr->xria', f_ov_a, rho0a[1:4])
                    w_ov[1:]+= numpy.einsum('ria,xr->xria', f_ov_b, rho0b[1:4])
                    w_ov[1:]+= numpy.einsum('r,xria->xria', ud, rho_ov_b[1:4])
                    w_ov *= weight[:,None,None]
                    a_iajb = lib.einsum('xria,xrjb->iajb', w_ov, rho_vo_a)
                    b_iajb = lib.einsum('xria,xrjb->iajb', w_ov, rho_ov_a)
                    a += a_iajb
                    a += a_iajb.conj().transpose(2,3,0,1)
                    b += b_iajb * 2
                else:
                    raise NotImplementedError

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

    else:
        add_hf_(a, b)

    return a, b

def get_nto(tdobj, state=1, threshold=OUTPUT_THRESHOLD, verbose=None):
    raise NotImplementedError('get_nto')

def analyze(tdobj, verbose=None):
    raise NotImplementedError('analyze')

def _contract_multipole(tdobj, ints, hermi=True, xy=None):
    raise NotImplementedError


class TDMixin(rhf.TDMixin):

    @lib.with_doc(get_ab.__doc__)
    def get_ab(self, mf=None):
        if mf is None: mf = self._scf
        return get_ab(mf)

    analyze = analyze
    get_nto = get_nto
    _contract_multipole = _contract_multipole  # needed by transition dipoles

    def nuc_grad_method(self):
        raise NotImplementedError


@lib.with_doc(rhf.TDA.__doc__)
class TDA(TDMixin):
    def gen_vind(self, mf):
        '''Compute Ax'''
        return gen_tda_hop(mf)

    def init_guess(self, mf, nstates=None, wfnsym=None):
        if nstates is None: nstates = self.nstates
        mo_energy = mf.mo_energy
        mo_occ = mf.mo_occ
        # Remove all negative states
        n2c = mf.mo_occ.size // 2
        occidx = n2c + numpy.where(mo_occ[n2c:] == 1)[0]
        viridx = n2c + numpy.where(mo_occ[n2c:] == 0)[0]
        e_ia = mo_energy[viridx] - mo_energy[occidx,None]
        e_ia_max = e_ia.max()

        nov = e_ia.size
        nstates = min(nstates, nov)
        e_ia = e_ia.ravel()
        e_threshold = min(e_ia_max, e_ia[numpy.argsort(e_ia)[nstates-1]])
        # Handle degeneracy, include all degenerated states in initial guess
        e_threshold += 1e-6

        idx = numpy.where(e_ia <= e_threshold)[0]
        x0 = numpy.zeros((idx.size, nov))
        for i, j in enumerate(idx):
            x0[i, j] = 1  # Koopmans' excitations
        return x0

    def kernel(self, x0=None, nstates=None):
        '''TDA diagonalization solver
        '''
        cpu0 = (logger.process_clock(), logger.perf_counter())
        self.check_sanity()
        self.dump_flags()
        if nstates is None:
            nstates = self.nstates
        else:
            self.nstates = nstates

        log = logger.Logger(self.stdout, self.verbose)

        vind, hdiag = self.gen_vind(self._scf)
        precond = self.get_precond(hdiag)

        if x0 is None:
            x0 = self.init_guess(self._scf, self.nstates)

        def pickeig(w, v, nroots, envs):
            idx = numpy.where(w > POSTIVE_EIG_THRESHOLD**2)[0]
            return w[idx], v[:,idx], idx

        # FIXME: Is it correct to call davidson1 for complex integrals
        self.converged, self.e, x1 = \
                lib.davidson1(vind, x0, precond,
                              tol=self.conv_tol,
                              nroots=nstates, lindep=self.lindep,
                              max_cycle=self.max_cycle,
                              max_space=self.max_space, pick=pickeig,
                              verbose=log)

        nocc = (self._scf.mo_occ>0).sum()
        nmo = self._scf.mo_occ.size
        n2c = nmo // 2
        nvir = n2c - nocc
        self.xy = [(xi.reshape(nocc,nvir), 0) for xi in x1]

        if self.chkfile:
            lib.chkfile.save(self.chkfile, 'tddft/e', self.e)
            lib.chkfile.save(self.chkfile, 'tddft/xy', self.xy)

        log.timer('TDA', *cpu0)
        self._finalize()
        return self.e, self.xy


def gen_tdhf_operation(mf, fock_ao=None):
    '''Generate function to compute

    [ A   B ][X]
    [-B* -A*][Y]
    '''
    mo_coeff = mf.mo_coeff
    mo_energy = mf.mo_energy
    mo_occ = mf.mo_occ
    nao, nmo = mo_coeff.shape
    n2c = nmo // 2
    occidx = n2c + numpy.where(mo_occ[n2c:] == 1)[0]
    viridx = n2c + numpy.where(mo_occ[n2c:] == 0)[0]
    nocc = len(occidx)
    nvir = len(viridx)
    orbv = mo_coeff[:,viridx]
    orbo = mo_coeff[:,occidx]

    foo = numpy.diag(mo_energy[occidx])
    fvv = numpy.diag(mo_energy[viridx])
    hdiag = fvv.diagonal() - foo.diagonal()[:,None]
    hdiag = numpy.hstack((hdiag.ravel(), -hdiag.ravel())).real

    mo_coeff = numpy.asarray(numpy.hstack((orbo,orbv)), order='F')
    vresp = mf.gen_response(hermi=0)

    def vind(xys):
        xys = numpy.asarray(xys).reshape(-1,2,nocc,nvir)
        xs, ys = xys.transpose(1,0,2,3)
        # dms = AX + BY
        dms  = lib.einsum('xov,qv,po->xpq', xs, orbv.conj(), orbo)
        dms += lib.einsum('xov,pv,qo->xpq', ys, orbv, orbo.conj())

        v1ao = vresp(dms)
        v1ov = lib.einsum('xpq,po,qv->xov', v1ao, orbo.conj(), orbv)
        v1vo = lib.einsum('xpq,qo,pv->xov', v1ao, orbo, orbv.conj())
        v1ov += lib.einsum('xqs,sp->xqp', xs, fvv)  # AX
        v1ov -= lib.einsum('xpr,sp->xsr', xs, foo)  # AX
        v1vo += lib.einsum('xqs,sp->xqp', ys, fvv)  # AY
        v1vo -= lib.einsum('xpr,sp->xsr', ys, foo)  # AY

        # (AX, (-A*)Y)
        nz = xys.shape[0]
        hx = numpy.hstack((v1ov.reshape(nz,-1), -v1vo.reshape(nz,-1)))
        return hx

    return vind, hdiag


class TDHF(TDMixin):
    @lib.with_doc(gen_tdhf_operation.__doc__)
    def gen_vind(self, mf):
        return gen_tdhf_operation(mf)

    def init_guess(self, mf, nstates=None, wfnsym=None):
        x0 = TDA.init_guess(self, mf, nstates, wfnsym)
        y0 = numpy.zeros_like(x0)
        return numpy.hstack((x0, y0))

    def kernel(self, x0=None, nstates=None):
        '''TDHF diagonalization with non-Hermitian eigenvalue solver
        '''
        cpu0 = (logger.process_clock(), logger.perf_counter())
        self.check_sanity()
        self.dump_flags()
        if nstates is None:
            nstates = self.nstates
        else:
            self.nstates = nstates

        log = logger.Logger(self.stdout, self.verbose)

        vind, hdiag = self.gen_vind(self._scf)
        precond = self.get_precond(hdiag)
        if x0 is None:
            x0 = self.init_guess(self._scf, self.nstates)

        def pickeig(w, v, nroots, envs):
            realidx = numpy.where((abs(w.imag) < REAL_EIG_THRESHOLD) &
                                  (w.real > POSTIVE_EIG_THRESHOLD))[0]
            # FIXME: Should the amplitudes be real?
            return lib.linalg_helper._eigs_cmplx2real(w, v, realidx,
                                                      real_eigenvectors=False)

        self.converged, w, x1 = \
                lib.davidson_nosym1(vind, x0, precond,
                                    tol=self.conv_tol,
                                    nroots=nstates, lindep=self.lindep,
                                    max_cycle=self.max_cycle,
                                    max_space=self.max_space, pick=pickeig,
                                    verbose=log)

        nocc = (self._scf.mo_occ>0).sum()
        nmo = self._scf.mo_occ.size
        n2c = nmo // 2
        nvir = n2c - nocc
        self.e = w
        def norm_xy(z):
            x, y = z.reshape(2,nocc,nvir)
            norm = lib.norm(x)**2 - lib.norm(y)**2
            norm = numpy.sqrt(1./norm)
            return x*norm, y*norm
        self.xy = [norm_xy(z) for z in x1]

        if self.chkfile:
            lib.chkfile.save(self.chkfile, 'tddft/e', self.e)
            lib.chkfile.save(self.chkfile, 'tddft/xy', self.xy)

        log.timer('TDDFT', *cpu0)
        self._finalize()
        return self.e, self.xy

from pyscf import scf
scf.dhf.DHF.TDA = lib.class_as_method(TDA)
scf.dhf.DHF.TDHF = lib.class_as_method(TDHF)

del(OUTPUT_THRESHOLD)


# pyscf/tdscf/dks.py

#!/usr/bin/env python
# Copyright 2021-2022 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#


import numpy
from pyscf import lib
from pyscf.tdscf import dhf
from pyscf.data import nist
from pyscf.dft.rks import KohnShamDFT
from pyscf import __config__

# Low excitation filter to avoid numerical instability
POSTIVE_EIG_THRESHOLD = getattr(__config__, 'tdscf_rhf_TDDFT_positive_eig_threshold', 1e-3)


class TDA(dhf.TDA):
    pass

class TDDFT(dhf.TDHF):
    pass

RPA = TDDKS = TDDFT

from pyscf import dft
dft.dks.DKS.TDA   = lib.class_as_method(TDA)
dft.dks.DKS.TDHF  = None
dft.dks.DKS.TDDFT = lib.class_as_method(TDDFT)


# pyscf/tdscf/ghf.py

#!/usr/bin/env python
# Copyright 2021-2022 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#
# Ref:
# Chem Phys Lett, 256, 454
# J. Mol. Struct. THEOCHEM, 914, 3
# Recent Advances in Density Functional Methods, Chapter 5, M. E. Casida
#


from functools import reduce
import numpy
from pyscf import lib
from pyscf import dft
from pyscf.dft import numint
from pyscf import ao2mo
from pyscf import symm
from pyscf.lib import logger
from pyscf.tdscf import rhf
from pyscf.scf import ghf_symm
from pyscf.data import nist
from pyscf import __config__

OUTPUT_THRESHOLD = getattr(__config__, 'tdscf_rhf_get_nto_threshold', 0.3)
REAL_EIG_THRESHOLD = getattr(__config__, 'tdscf_rhf_TDDFT_pick_eig_threshold', 1e-4)

# Low excitation filter to avoid numerical instability
POSTIVE_EIG_THRESHOLD = getattr(__config__, 'tdscf_rhf_TDDFT_positive_eig_threshold', 1e-3)


def gen_tda_operation(mf, fock_ao=None, wfnsym=None):
    '''A x

    Kwargs:
        wfnsym : int or str
            Point group symmetry irrep symbol or ID for excited CIS wavefunction.
    '''
    mol = mf.mol
    mo_coeff = mf.mo_coeff
    mo_energy = mf.mo_energy
    mo_occ = mf.mo_occ
    nao, nmo = mo_coeff.shape
    occidx = numpy.where(mo_occ == 1)[0]
    viridx = numpy.where(mo_occ == 0)[0]
    nocc = len(occidx)
    nvir = len(viridx)
    orbv = mo_coeff[:,viridx]
    orbo = mo_coeff[:,occidx]

    if wfnsym is not None and mol.symmetry:
        if isinstance(wfnsym, str):
            wfnsym = symm.irrep_name2id(mol.groupname, wfnsym)
        wfnsym = wfnsym % 10  # convert to D2h subgroup
        orbsym = ghf_symm.get_orbsym(mol, mo_coeff)
        orbsym_in_d2h = numpy.asarray(orbsym) % 10  # convert to D2h irreps
        sym_forbid = (orbsym_in_d2h[occidx,None] ^ orbsym_in_d2h[viridx]) != wfnsym

    if fock_ao is None:
        #dm0 = mf.make_rdm1(mo_coeff, mo_occ)
        #fock_ao = mf.get_hcore() + mf.get_veff(mol, dm0)
        foo = numpy.diag(mo_energy[occidx])
        fvv = numpy.diag(mo_energy[viridx])
    else:
        fock = reduce(numpy.dot, (mo_coeff.conj().T, fock_ao, mo_coeff))
        foo = fock[occidx[:,None],occidx]
        fvv = fock[viridx[:,None],viridx]

    hdiag = fvv.diagonal() - foo.diagonal()[:,None]
    if wfnsym is not None and mol.symmetry:
        hdiag[sym_forbid] = 0
    hdiag = hdiag.ravel().real

    mo_coeff = numpy.asarray(numpy.hstack((orbo,orbv)), order='F')
    vresp = mf.gen_response(hermi=0)

    def vind(zs):
        zs = numpy.asarray(zs).reshape(-1,nocc,nvir)
        if wfnsym is not None and mol.symmetry:
            zs = numpy.copy(zs)
            zs[:,sym_forbid] = 0

        dmov = lib.einsum('xov,qv,po->xpq', zs, orbv.conj(), orbo)
        v1ao = vresp(dmov)
        v1ov = lib.einsum('xpq,qv,po->xov', v1ao, orbv, orbo.conj())
        v1ov += lib.einsum('xqs,sp->xqp', zs, fvv)
        v1ov -= lib.einsum('xpr,sp->xsr', zs, foo)
        if wfnsym is not None and mol.symmetry:
            v1ov[:,sym_forbid] = 0
        return v1ov.reshape(v1ov.shape[0],-1)

    return vind, hdiag
gen_tda_hop = gen_tda_operation

def get_ab(mf, mo_energy=None, mo_coeff=None, mo_occ=None):
    r'''A and B matrices for TDDFT response function.

    A[i,a,j,b] = \delta_{ab}\delta_{ij}(E_a - E_i) + (ia||bj)
    B[i,a,j,b] = (ia||jb)
    '''
    if mo_energy is None: mo_energy = mf.mo_energy
    if mo_coeff is None: mo_coeff = mf.mo_coeff
    if mo_occ is None: mo_occ = mf.mo_occ
    assert mo_coeff.dtype == numpy.double

    mol = mf.mol
    nao, nmo = mo_coeff.shape
    occidx = numpy.where(mo_occ==1)[0]
    viridx = numpy.where(mo_occ==0)[0]
    orbv = mo_coeff[:,viridx]
    orbo = mo_coeff[:,occidx]
    nvir = orbv.shape[1]
    nocc = orbo.shape[1]
    mo = numpy.hstack((orbo,orbv))
    moa = mo[:nao].copy()
    mob = mo[nao:].copy()
    orboa = orbo[:nao]
    orbob = orbo[nao:]
    nmo = nocc + nvir

    e_ia = lib.direct_sum('a-i->ia', mo_energy[viridx], mo_energy[occidx])
    a = numpy.diag(e_ia.ravel()).reshape(nocc,nvir,nocc,nvir)
    b = numpy.zeros_like(a)

    def add_hf_(a, b, hyb=1):
        eri_mo_aa  = ao2mo.general(mol, [orboa,moa,moa,moa], compact=False)
        eri_mo_aa += ao2mo.general(mol, [orbob,mob,mob,mob], compact=False)
        eri_mo_aa = eri_mo_aa.reshape(nocc,nmo,nmo,nmo)
        eri_mo  = ao2mo.general(mol, [orboa,moa,mob,mob], compact=False)
        eri_mo += ao2mo.general(mol, [orbob,mob,moa,moa], compact=False)
        eri_mo = eri_mo.reshape(nocc,nmo,nmo,nmo)
        eri_mo += eri_mo_aa
        a += numpy.einsum('iabj->iajb', eri_mo   [:nocc,nocc:,nocc:,:nocc])
        a -= numpy.einsum('ijba->iajb', eri_mo_aa[:nocc,:nocc,nocc:,nocc:]) * hyb
        b += numpy.einsum('iajb->iajb', eri_mo   [:nocc,nocc:,:nocc,nocc:])
        b -= numpy.einsum('jaib->iajb', eri_mo_aa[:nocc,nocc:,:nocc,nocc:]) * hyb

    if isinstance(mf, dft.KohnShamDFT):
        ni = mf._numint
        ni.libxc.test_deriv_order(mf.xc, 2, raise_error=True)
        if getattr(mf, 'nlc', '') != '':
            raise NotImplementedError

        if not mf.collinear:
            raise NotImplementedError

        omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin)

        add_hf_(a, b, hyb)

        xctype = ni._xc_type(mf.xc)
        dm0 = mf.make_rdm1(mo_coeff, mo_occ)
        dm0a = dm0[:nao,:nao].real.copy()
        dm0b = dm0[nao:,nao:].real.copy()
        mem_now = lib.current_memory()[0]
        max_memory = max(2000, mf.max_memory*.8-mem_now)

        def get_mo_value(ao):
            if ao.ndim == 2:
                mo_a = lib.einsum('rp,pi->ri', ao, moa)
                mo_b = lib.einsum('rp,pi->ri', ao, mob)
                return mo_a[:,:nocc], mo_a[:,nocc:], mo_b[:,:nocc], mo_b[:,nocc:]
            else:
                mo_a = lib.einsum('xrp,pi->xri', ao, moa)
                mo_b = lib.einsum('xrp,pi->xri', ao, mob)
                return mo_a[:,:,:nocc], mo_a[:,:,nocc:], mo_b[:,:,:nocc], mo_b[:,:,nocc:]

        if xctype == 'LDA':
            ao_deriv = 0
            for ao, mask, weight, coords \
                    in ni.block_loop(mol, mf.grids, nao, ao_deriv, max_memory):
                rho0a = numint.eval_rho(mol, ao, dm0a, mask, xctype)
                rho0b = numint.eval_rho(mol, ao, dm0b, mask, xctype)
                fxc = ni.eval_xc(mf.xc, (rho0a,rho0b), 1, deriv=2)[2]
                u_u, u_d, d_d = fxc[0].T

                mo_oa, mo_va, mo_ob, mo_vb = get_mo_value(ao)
                rho_ov_a = numpy.einsum('ri,ra->ria', mo_oa.conj(), mo_va)
                rho_ov_b = numpy.einsum('ri,ra->ria', mo_ob.conj(), mo_vb)
                rho_vo_a = rho_ov_a.conj()
                rho_vo_b = rho_ov_b.conj()

                w_ov_a = numpy.einsum('ria,r->ria', u_u * rho_ov_a + u_d * rho_ov_b, weight)
                a += lib.einsum('ria,rjb->iajb', w_ov_a, rho_vo_a)
                b += lib.einsum('ria,rjb->iajb', w_ov_a, rho_ov_a)

                w_ov_b = numpy.einsum('ria,r->ria', u_d * rho_ov_a + d_d * rho_ov_b, weight)
                a += lib.einsum('ria,rjb->iajb', w_ov_b, rho_vo_b)
                b += lib.einsum('ria,rjb->iajb', w_ov_b, rho_ov_b)

        elif xctype == 'GGA':
            ao_deriv = 1
            for ao, mask, weight, coords \
                    in ni.block_loop(mol, mf.grids, nao, ao_deriv, max_memory):
                rho0a = numint.eval_rho(mol, ao, dm0a, mask, xctype)
                rho0b = numint.eval_rho(mol, ao, dm0b, mask, xctype)
                vxc, fxc = ni.eval_xc(mf.xc, (rho0a,rho0b), 0, deriv=2)[1:3]
                uu, ud, dd = vxc[1].T
                u_u, u_d, d_d = fxc[0].T
                u_uu, u_ud, u_dd, d_uu, d_ud, d_dd = fxc[1].T
                uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd = fxc[2].T

                mo_oa, mo_va, mo_ob, mo_vb = get_mo_value(ao)
                rho_ov_a = numpy.einsum('xri,ra->xria', mo_oa.conj(), mo_va[0])
                rho_ov_b = numpy.einsum('xri,ra->xria', mo_ob.conj(), mo_vb[0])
                rho_ov_a[1:4] += numpy.einsum('ri,xra->xria', mo_oa[0].conj(), mo_va[1:4])
                rho_ov_b[1:4] += numpy.einsum('ri,xra->xria', mo_ob[0].conj(), mo_vb[1:4])
                rho_vo_a = rho_ov_a.conj()
                rho_vo_b = rho_vo_a.conj()
                # sigma1 ~ \nabla(\rho_\alpha+\rho_\beta) dot \nabla(|b><j|) z_{bj}
                a0a1 = numpy.einsum('xr,xria->ria', rho0a[1:4], rho_ov_a[1:4])
                a0b1 = numpy.einsum('xr,xria->ria', rho0a[1:4], rho_ov_b[1:4])
                b0a1 = numpy.einsum('xr,xria->ria', rho0b[1:4], rho_ov_a[1:4])
                b0b1 = numpy.einsum('xr,xria->ria', rho0b[1:4], rho_ov_b[1:4])

                # aaaa
                w_ov = numpy.empty_like(rho_ov_a)
                w_ov[0]  = numpy.einsum('r,ria->ria', u_u, rho_ov_a[0])
                w_ov[0] += numpy.einsum('r,ria->ria', 2*u_uu, a0a1)
                w_ov[0] += numpy.einsum('r,ria->ria',   u_ud, b0a1)
                f_ov_a  = numpy.einsum('r,ria->ria', 4*uu_uu, a0a1)
                f_ov_b  = numpy.einsum('r,ria->ria', 2*uu_ud, a0a1)
                f_ov_a += numpy.einsum('r,ria->ria', 2*uu_ud, b0a1)
                f_ov_b += numpy.einsum('r,ria->ria',   ud_ud, b0a1)
                f_ov_a += numpy.einsum('r,ria->ria', 2*u_uu, rho_ov_a[0])
                f_ov_b += numpy.einsum('r,ria->ria',   u_ud, rho_ov_a[0])
                w_ov[1:] = numpy.einsum('ria,xr->xria', f_ov_a, rho0a[1:4])
                w_ov[1:]+= numpy.einsum('ria,xr->xria', f_ov_b, rho0b[1:4])
                w_ov[1:]+= numpy.einsum('r,xria->xria', 2*uu, rho_ov_a[1:4])
                w_ov *= weight[:,None,None]
                a += lib.einsum('xria,xrjb->iajb', w_ov, rho_vo_a)
                b += lib.einsum('xria,xrjb->iajb', w_ov, rho_ov_a)

                # bbbb
                w_ov = numpy.empty_like(rho_ov_b)
                w_ov[0]  = numpy.einsum('r,ria->ria', d_d, rho_ov_b[0])
                w_ov[0] += numpy.einsum('r,ria->ria', 2*d_dd, b0b1)
                w_ov[0] += numpy.einsum('r,ria->ria',   d_ud, a0b1)
                f_ov_b  = numpy.einsum('r,ria->ria', 4*dd_dd, b0b1)
                f_ov_a  = numpy.einsum('r,ria->ria', 2*ud_dd, b0b1)
                f_ov_b += numpy.einsum('r,ria->ria', 2*ud_dd, a0b1)
                f_ov_a += numpy.einsum('r,ria->ria',   ud_ud, a0b1)
                f_ov_b += numpy.einsum('r,ria->ria', 2*d_dd, rho_ov_b[0])
                f_ov_a += numpy.einsum('r,ria->ria',   d_ud, rho_ov_b[0])
                w_ov[1:] = numpy.einsum('ria,xr->xria', f_ov_a, rho0a[1:4])
                w_ov[1:]+= numpy.einsum('ria,xr->xria', f_ov_b, rho0b[1:4])
                w_ov[1:]+= numpy.einsum('r,xria->xria', 2*dd, rho_ov_b[1:4])
                w_ov *= weight[:,None,None]
                a += lib.einsum('xria,xrjb->iajb', w_ov, rho_vo_b)
                b += lib.einsum('xria,xrjb->iajb', w_ov, rho_ov_b)

                # aabb and bbaa
                w_ov = numpy.empty_like(rho_ov_b)
                w_ov[0]  = numpy.einsum('r,ria->ria', u_d, rho_ov_b[0])
                w_ov[0] += numpy.einsum('r,ria->ria', 2*u_dd, b0b1)
                w_ov[0] += numpy.einsum('r,ria->ria',   u_ud, a0b1)
                f_ov_a  = numpy.einsum('r,ria->ria', 4*uu_dd, b0b1)
                f_ov_b  = numpy.einsum('r,ria->ria', 2*ud_dd, b0b1)
                f_ov_a += numpy.einsum('r,ria->ria', 2*uu_ud, a0b1)
                f_ov_b += numpy.einsum('r,ria->ria',   ud_ud, a0b1)
                f_ov_a += numpy.einsum('r,ria->ria', 2*d_uu, rho_ov_b[0])
                f_ov_b += numpy.einsum('r,ria->ria',   d_ud, rho_ov_b[0])
                w_ov[1:] = numpy.einsum('ria,xr->xria', f_ov_a, rho0a[1:4])
                w_ov[1:]+= numpy.einsum('ria,xr->xria', f_ov_b, rho0b[1:4])
                w_ov[1:]+= numpy.einsum('r,xria->xria', ud, rho_ov_b[1:4])
                w_ov *= weight[:,None,None]
                a_iajb = lib.einsum('xria,xrjb->iajb', w_ov, rho_vo_a)
                b_iajb = lib.einsum('xria,xrjb->iajb', w_ov, rho_ov_a)
                a += a_iajb
                a += a_iajb.conj().transpose(2,3,0,1)
                b += b_iajb * 2

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

    else:
        add_hf_(a, b)

    return a, b

def get_nto(tdobj, state=1, threshold=OUTPUT_THRESHOLD, verbose=None):
    raise NotImplementedError('get_nto')

def analyze(tdobj, verbose=None):
    raise NotImplementedError('analyze')

def _contract_multipole(tdobj, ints, hermi=True, xy=None):
    raise NotImplementedError


class TDMixin(rhf.TDMixin):

    @lib.with_doc(get_ab.__doc__)
    def get_ab(self, mf=None):
        if mf is None: mf = self._scf
        return get_ab(mf)

    analyze = analyze
    get_nto = get_nto
    _contract_multipole = _contract_multipole  # needed by transition dipoles

    def nuc_grad_method(self):
        raise NotImplementedError


@lib.with_doc(rhf.TDA.__doc__)
class TDA(TDMixin):

    singlet = None

    def gen_vind(self, mf):
        '''Compute Ax'''
        return gen_tda_hop(mf, wfnsym=self.wfnsym)

    def init_guess(self, mf, nstates=None, wfnsym=None):
        if nstates is None: nstates = self.nstates
        if wfnsym is None: wfnsym = self.wfnsym

        mo_energy = mf.mo_energy
        mo_occ = mf.mo_occ
        occidx = numpy.where(mo_occ==2)[0]
        viridx = numpy.where(mo_occ==0)[0]
        e_ia = mo_energy[viridx] - mo_energy[occidx,None]
        e_ia_max = e_ia.max()

        if wfnsym is not None and mf.mol.symmetry:
            if isinstance(wfnsym, str):
                wfnsym = symm.irrep_name2id(mf.mol.groupname, wfnsym)
            wfnsym = wfnsym % 10  # convert to D2h subgroup
            orbsym = ghf_symm.get_orbsym(mf.mol, mf.mo_coeff)
            orbsym_in_d2h = numpy.asarray(orbsym) % 10  # convert to D2h irreps
            e_ia[(orbsym_in_d2h[occidx,None] ^ orbsym_in_d2h[viridx]) != wfnsym] = 1e99

        nov = e_ia.size
        nstates = min(nstates, nov)
        e_ia = e_ia.ravel()
        e_threshold = min(e_ia_max, e_ia[numpy.argsort(e_ia)[nstates-1]])
        # Handle degeneracy, include all degenerated states in initial guess
        e_threshold += 1e-6

        idx = numpy.where(e_ia <= e_threshold)[0]
        x0 = numpy.zeros((idx.size, nov))
        for i, j in enumerate(idx):
            x0[i, j] = 1  # Koopmans' excitations
        return x0

    def kernel(self, x0=None, nstates=None):
        '''TDA diagonalization solver
        '''
        cpu0 = (logger.process_clock(), logger.perf_counter())
        self.check_sanity()
        self.dump_flags()
        if nstates is None:
            nstates = self.nstates
        else:
            self.nstates = nstates

        log = logger.Logger(self.stdout, self.verbose)

        vind, hdiag = self.gen_vind(self._scf)
        precond = self.get_precond(hdiag)

        if x0 is None:
            x0 = self.init_guess(self._scf, self.nstates)

        def pickeig(w, v, nroots, envs):
            idx = numpy.where(w > POSTIVE_EIG_THRESHOLD**2)[0]
            return w[idx], v[:,idx], idx

        # FIXME: Is it correct to call davidson1 for complex integrals
        self.converged, self.e, x1 = \
                lib.davidson1(vind, x0, precond,
                              tol=self.conv_tol,
                              nroots=nstates, lindep=self.lindep,
                              max_cycle=self.max_cycle,
                              max_space=self.max_space, pick=pickeig,
                              verbose=log)

        nocc = (self._scf.mo_occ>0).sum()
        nmo = self._scf.mo_occ.size
        nvir = nmo - nocc
        self.xy = [(xi.reshape(nocc,nvir), 0) for xi in x1]

        if self.chkfile:
            lib.chkfile.save(self.chkfile, 'tddft/e', self.e)
            lib.chkfile.save(self.chkfile, 'tddft/xy', self.xy)

        log.timer('TDA', *cpu0)
        self._finalize()
        return self.e, self.xy

CIS = TDA


def gen_tdhf_operation(mf, fock_ao=None, wfnsym=None):
    '''Generate function to compute

    [ A   B ][X]
    [-B* -A*][Y]
    '''
    mol = mf.mol
    mo_coeff = mf.mo_coeff
    mo_energy = mf.mo_energy
    mo_occ = mf.mo_occ
    nao, nmo = mo_coeff.shape
    occidx = numpy.where(mo_occ == 1)[0]
    viridx = numpy.where(mo_occ == 0)[0]
    nocc = len(occidx)
    nvir = len(viridx)
    orbv = mo_coeff[:,viridx]
    orbo = mo_coeff[:,occidx]

    if wfnsym is not None and mol.symmetry:
        if isinstance(wfnsym, str):
            wfnsym = symm.irrep_name2id(mol.groupname, wfnsym)
        wfnsym = wfnsym % 10  # convert to D2h subgroup
        orbsym = ghf_symm.get_orbsym(mol, mo_coeff)
        orbsym_in_d2h = numpy.asarray(orbsym) % 10  # convert to D2h irreps
        sym_forbid = (orbsym_in_d2h[occidx,None] ^ orbsym_in_d2h[viridx]) != wfnsym

    #dm0 = mf.make_rdm1(mo_coeff, mo_occ)
    #fock_ao = mf.get_hcore() + mf.get_veff(mol, dm0)
    #fock = reduce(numpy.dot, (mo_coeff.T, fock_ao, mo_coeff))
    #foo = fock[occidx[:,None],occidx]
    #fvv = fock[viridx[:,None],viridx]
    foo = numpy.diag(mo_energy[occidx])
    fvv = numpy.diag(mo_energy[viridx])

    hdiag = fvv.diagonal() - foo.diagonal()[:,None]
    if wfnsym is not None and mol.symmetry:
        hdiag[sym_forbid] = 0
    hdiag = numpy.hstack((hdiag.ravel(), -hdiag.ravel())).real

    mo_coeff = numpy.asarray(numpy.hstack((orbo,orbv)), order='F')
    vresp = mf.gen_response(hermi=0)

    def vind(xys):
        xys = numpy.asarray(xys).reshape(-1,2,nocc,nvir)
        if wfnsym is not None and mol.symmetry:
            # shape(nz,2,nocc,nvir): 2 ~ X,Y
            xys = numpy.copy(xys)
            xys[:,:,sym_forbid] = 0

        xs, ys = xys.transpose(1,0,2,3)
        # dms = AX + BY
        dms  = lib.einsum('xov,qv,po->xpq', xs, orbv.conj(), orbo)
        dms += lib.einsum('xov,pv,qo->xpq', ys, orbv, orbo.conj())

        v1ao = vresp(dms)
        v1ov = lib.einsum('xpq,po,qv->xov', v1ao, orbo.conj(), orbv)
        v1vo = lib.einsum('xpq,qo,pv->xov', v1ao, orbo, orbv.conj())
        v1ov += lib.einsum('xqs,sp->xqp', xs, fvv)  # AX
        v1ov -= lib.einsum('xpr,sp->xsr', xs, foo)  # AX
        v1vo += lib.einsum('xqs,sp->xqp', ys, fvv)  # AY
        v1vo -= lib.einsum('xpr,sp->xsr', ys, foo)  # AY

        if wfnsym is not None and mol.symmetry:
            v1ov[:,sym_forbid] = 0
            v1vo[:,sym_forbid] = 0

        # (AX, (-A*)Y)
        nz = xys.shape[0]
        hx = numpy.hstack((v1ov.reshape(nz,-1), -v1vo.reshape(nz,-1)))
        return hx

    return vind, hdiag


class TDHF(TDMixin):

    singlet = None

    @lib.with_doc(gen_tdhf_operation.__doc__)
    def gen_vind(self, mf):
        return gen_tdhf_operation(mf, wfnsym=self.wfnsym)

    def init_guess(self, mf, nstates=None, wfnsym=None):
        x0 = TDA.init_guess(self, mf, nstates, wfnsym)
        y0 = numpy.zeros_like(x0)
        return numpy.hstack((x0,y0))

    def kernel(self, x0=None, nstates=None):
        '''TDHF diagonalization with non-Hermitian eigenvalue solver
        '''
        cpu0 = (logger.process_clock(), logger.perf_counter())
        self.check_sanity()
        self.dump_flags()
        if nstates is None:
            nstates = self.nstates
        else:
            self.nstates = nstates

        log = logger.Logger(self.stdout, self.verbose)

        vind, hdiag = self.gen_vind(self._scf)
        precond = self.get_precond(hdiag)
        if x0 is None:
            x0 = self.init_guess(self._scf, self.nstates)

        def pickeig(w, v, nroots, envs):
            realidx = numpy.where((abs(w.imag) < REAL_EIG_THRESHOLD) &
                                  (w.real > POSTIVE_EIG_THRESHOLD))[0]
            # FIXME: Should the amplitudes be real? It also affects x2c-tdscf
            return lib.linalg_helper._eigs_cmplx2real(w, v, realidx,
                                                      real_eigenvectors=False)

        self.converged, w, x1 = \
                lib.davidson_nosym1(vind, x0, precond,
                                    tol=self.conv_tol,
                                    nroots=nstates, lindep=self.lindep,
                                    max_cycle=self.max_cycle,
                                    max_space=self.max_space, pick=pickeig,
                                    verbose=log)

        nocc = (self._scf.mo_occ>0).sum()
        nmo = self._scf.mo_occ.size
        nvir = nmo - nocc
        self.e = w
        def norm_xy(z):
            x, y = z.reshape(2,nocc,nvir)
            norm = lib.norm(x)**2 - lib.norm(y)**2
            norm = numpy.sqrt(1./norm)
            return x*norm, y*norm
        self.xy = [norm_xy(z) for z in x1]

        if self.chkfile:
            lib.chkfile.save(self.chkfile, 'tddft/e', self.e)
            lib.chkfile.save(self.chkfile, 'tddft/xy', self.xy)

        log.timer('TDDFT', *cpu0)
        self._finalize()
        return self.e, self.xy

RPA = TDGHF = TDHF

from pyscf import scf
scf.ghf.GHF.TDA = lib.class_as_method(TDA)
scf.ghf.GHF.TDHF = lib.class_as_method(TDHF)


# pyscf/tdscf/gks.py

#!/usr/bin/env python
# Copyright 2021-2022 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#


import numpy
from pyscf import lib
from pyscf import symm
from pyscf.tdscf import ghf
from pyscf.scf import ghf_symm
from pyscf.scf import _response_functions  # noqa
from pyscf.data import nist
from pyscf.dft.rks import KohnShamDFT
from pyscf import __config__

# Low excitation filter to avoid numerical instability
POSTIVE_EIG_THRESHOLD = getattr(__config__, 'tdscf_rhf_TDDFT_positive_eig_threshold', 1e-3)


class TDA(ghf.TDA):
    pass

class TDDFT(ghf.TDHF):
    pass

RPA = TDGKS = TDDFT

class CasidaTDDFT(TDDFT, TDA):
    '''Solve the Casida TDDFT formula (A-B)(A+B)(X+Y) = (X+Y)w^2
    '''
    init_guess = TDA.init_guess

    def gen_vind(self, mf):
        wfnsym = self.wfnsym
        mol = mf.mol
        mo_coeff = mf.mo_coeff
        assert mo_coeff.dtype == numpy.double
        mo_energy = mf.mo_energy
        mo_occ = mf.mo_occ
        nao, nmo = mo_coeff.shape
        occidx = numpy.where(mo_occ==1)[0]
        viridx = numpy.where(mo_occ==0)[0]
        nocc = len(occidx)
        nvir = len(viridx)
        orbv = mo_coeff[:,viridx]
        orbo = mo_coeff[:,occidx]

        if wfnsym is not None and mol.symmetry:
            if isinstance(wfnsym, str):
                wfnsym = symm.irrep_name2id(mol.groupname, wfnsym)
            wfnsym = wfnsym % 10  # convert to D2h subgroup
            orbsym = ghf_symm.get_orbsym(mol, mo_coeff)
            orbsym_in_d2h = numpy.asarray(orbsym) % 10  # convert to D2h irreps
            sym_forbid = (orbsym_in_d2h[occidx,None] ^ orbsym_in_d2h[viridx]) != wfnsym

        e_ia = (mo_energy[viridx].reshape(-1,1) - mo_energy[occidx]).T
        if wfnsym is not None and mol.symmetry:
            e_ia[sym_forbid] = 0
        d_ia = numpy.sqrt(e_ia)
        ed_ia = e_ia * d_ia
        hdiag = e_ia.ravel() ** 2

        vresp = mf.gen_response(mo_coeff, mo_occ, hermi=1)

        def vind(zs):
            zs = numpy.asarray(zs).reshape(-1,nocc,nvir)
            if wfnsym is not None and mol.symmetry:
                zs = numpy.copy(zs)
                zs[:,sym_forbid] = 0

            dmov = lib.einsum('xov,ov,po,qv->xpq', zs, d_ia, orbo, orbv)
            # +cc for A+B and K_{ai,jb} in A == K_{ai,bj} in B
            dmov = dmov + dmov.transpose(0,2,1)

            v1ao = vresp(dmov)
            v1ov = lib.einsum('xpq,po,qv->xov', v1ao, orbo, orbv)

            # numpy.sqrt(e_ia) * (e_ia*d_ia*z + v1ov)
            v1ov += numpy.einsum('xov,ov->xov', zs, ed_ia)
            v1ov *= d_ia
            if wfnsym is not None and mol.symmetry:
                v1ov[:,sym_forbid] = 0
            return v1ov.reshape(v1ov.shape[0],-1)

        return vind, hdiag

    def kernel(self, x0=None, nstates=None):
        '''TDDFT diagonalization solver
        '''
        cpu0 = (lib.logger.process_clock(), lib.logger.perf_counter())
        mf = self._scf
        if mf._numint.libxc.is_hybrid_xc(mf.xc):
            raise RuntimeError('%s cannot be used with hybrid functional'
                               % self.__class__)
        self.check_sanity()
        self.dump_flags()
        if nstates is None:
            nstates = self.nstates
        else:
            self.nstates = nstates

        log = lib.logger.Logger(self.stdout, self.verbose)

        vind, hdiag = self.gen_vind(self._scf)
        precond = self.get_precond(hdiag)
        if x0 is None:
            x0 = self.init_guess(self._scf, self.nstates)

        def pickeig(w, v, nroots, envs):
            idx = numpy.where(w > POSTIVE_EIG_THRESHOLD**2)[0]
            return w[idx], v[:,idx], idx

        self.converged, w2, x1 = \
                lib.davidson1(vind, x0, precond,
                              tol=self.conv_tol,
                              nroots=nstates, lindep=self.lindep,
                              max_cycle=self.max_cycle,
                              max_space=self.max_space, pick=pickeig,
                              verbose=log)

        mo_energy = self._scf.mo_energy
        mo_occ = self._scf.mo_occ
        occidx = numpy.where(mo_occ==1)[0]
        viridx = numpy.where(mo_occ==0)[0]
        e_ia = (mo_energy[viridx,None] - mo_energy[occidx]).T
        e_ia = numpy.sqrt(e_ia)
        def norm_xy(w, z):
            zp = e_ia * z.reshape(e_ia.shape)
            zm = w/e_ia * z.reshape(e_ia.shape)
            x = (zp + zm) * .5
            y = (zp - zm) * .5
            norm = lib.norm(x)**2 - lib.norm(y)**2
            norm = numpy.sqrt(1./norm)
            return (x*norm, y*norm)

        idx = numpy.where(w2 > POSTIVE_EIG_THRESHOLD**2)[0]
        self.e = numpy.sqrt(w2[idx])
        self.xy = [norm_xy(self.e[i], x1[i]) for i in idx]

        if self.chkfile:
            lib.chkfile.save(self.chkfile, 'tddft/e', self.e)
            lib.chkfile.save(self.chkfile, 'tddft/xy', self.xy)

        log.timer('TDDFT', *cpu0)
        self._finalize()
        return self.e, self.xy

    def nuc_grad_method(self):
        raise NotImplementedError

TDDFTNoHybrid = CasidaTDDFT


def tddft(mf):
    '''Driver to create TDDFT or CasidaTDDFT object'''
    if (not mf._numint.libxc.is_hybrid_xc(mf.xc) and
        # Casida formula can be applied for real orbitals only
        mf.mo_coeff.dtype == numpy.double):
        return CasidaTDDFT(mf)
    else:
        return TDDFT(mf)

from pyscf import dft
dft.gks.GKS.TDA           = dft.gks_symm.GKS.TDA           = lib.class_as_method(TDA)
dft.gks.GKS.TDHF          = dft.gks_symm.GKS.TDHF          = None
dft.gks.GKS.TDDFTNoHybrid = dft.gks_symm.GKS.TDDFTNoHybrid = lib.class_as_method(TDDFTNoHybrid)
dft.gks.GKS.CasidaTDDFT   = dft.gks_symm.GKS.CasidaTDDFT   = lib.class_as_method(CasidaTDDFT)
dft.gks.GKS.TDDFT         = dft.gks_symm.GKS.TDDFT         = tddft


# pyscf/tdscf/proxy.py

#  Author: Artem Pulkin
"""
This and other `proxy` modules implement the time-dependent mean-field procedure using the existing pyscf
implementations as a black box. The main purpose of these modules is to overcome the existing limitations in pyscf
(i.e. real-only orbitals, davidson diagonalizer, incomplete Bloch space, etc). The primary performance drawback is that,
unlike the original pyscf routines with an implicit construction of the eigenvalue problem, these modules construct TD
matrices explicitly by proxying to pyscf density response routines with a O(N^4) complexity scaling. As a result,
regular `numpy.linalg.eig` can be used to retrieve TD roots. Several variants of proxy-TD are available:

 * (this module) `pyscf.tdscf.proxy`: the molecular implementation;
 * `pyscf.pbc.tdscf.proxy`: PBC (periodic boundary condition) Gamma-point-only implementation;
 * `pyscf.pbc.tdscf.kproxy_supercell`: PBC implementation constructing supercells. Works with an arbitrary number of
   k-points but has an overhead due to ignoring the momentum conservation law. In addition, works only with
   time reversal invariant (TRI) models: i.e. the k-point grid has to be aligned and contain at least one TRI momentum.
 * `pyscf.pbc.tdscf.kproxy`: same as the above but respect the momentum conservation and, thus, diagonlizes smaller
   matrices (the performance gain is the total number of k-points in the model).
"""

# Convention for these modules:
# * PhysERI is the proxying class constructing time-dependent matrices
# * vector_to_amplitudes reshapes and normalizes the solution
# * TDProxy provides a container

from pyscf.tdscf.common_slow import TDProxyMatrixBlocks, MolecularMFMixin, TDBase, format_mask
from pyscf.tdscf import rhf_slow, TDDFT, TDHF
from pyscf.lib import logger

import numpy


def molecular_response_ov(vind, space_ov, nocc, nmo, double, log_dest):
    """
    Retrieves a raw response matrix.
    Args:
        vind (Callable): a pyscf matvec routine;
        space_ov (ndarray): the active `ov` space mask: either the same mask for both rows and columns (1D array) or
        separate `ov` masks for rows and columns (2D array);
        nocc (int): the number of occupied orbitals (frozen and active);
        nmo (int): the total number of orbitals;
        double (bool): set to True if `vind` returns the double-sized (i.e. full) matrix;
        log_dest (object): pyscf logging;

    Returns:
        The TD matrix.

    Note:
        The runtime scales with the size of the column mask `space_ov[1]` but not the row mask `space_ov[0]`.
    """
    space_ov = numpy.array(space_ov)

    nocc_full, nvirt_full = nocc, nmo - nocc
    size_full = nocc_full * nvirt_full

    if space_ov.shape not in ((size_full,), (2, size_full)):
        raise ValueError("The 'space_ov' argument should be a 1D array with dimension {size_full:d} or a 2D array with"
                         " dimensions 2x{size_full:d}, found: {actual}".format(
                             size_full=size_full,
                             actual=space_ov.shape,))
    ov1, ov2 = space_ov

    size = sum(ov2)

    probe = numpy.zeros((size, 2 * size_full if double else size_full))
    probe[numpy.arange(probe.shape[0]), numpy.argwhere(ov2)[:, 0]] = 1
    logger.debug1(log_dest, "Requesting response against {} matrix (column space: {})".format(
        "x".join(str(i) for i in probe.shape), format_mask(ov2),
    ))
    result = vind(probe).T
    logger.debug1(log_dest, "  output: {}".format(result.shape))

    if double:
        result = result[numpy.tile(ov1, 2)]
        half = sum(ov1)
        result_a = result[:half]
        result_b = result[half:]
        return result_a, -result_b.conj()
    else:
        return result[ov1]


def orb2ov(space, nocc):
    """
    Converts orbital active space specification into ov-pairs space spec.
    Args:
        space (ndarray): the obital space;
        nocc (int): the number of occupied orbitals;

    Returns:
        The ov space specification.
    """
    space = numpy.array(space)
    o = space[..., :nocc]
    v = space[..., nocc:]
    return (o[..., numpy.newaxis] * v[..., numpy.newaxis, :]).reshape(space.shape[:-1] + (-1,))


def molecular_response(vind, space, nocc, nmo, double, log_dest):
    """
    Retrieves a raw response matrix.
    Args:
        vind (Callable): a pyscf matvec routine;
        space (ndarray): the active orbital space mask: either the same mask for both rows and columns (1D array) or
        separate orbital masks for rows and columns (2D array);
        nocc (int): the number of occupied orbitals (frozen and active);
        nmo (int): the total number of orbitals;
        double (bool): set to True if `vind` returns the double-sized (i.e. full) matrix;
        log_dest (object): pyscf logging;

    Returns:
        The TD matrix.
    """
    space = numpy.array(space)

    if space.shape == (nmo,):
        space = numpy.repeat(space[numpy.newaxis, :], 2, axis=0)
    elif space.shape != (2, nmo):
        raise ValueError(
            "The 'space' argument should be a 1D array with dimension {size_full:d} or a 2D array with"
            " dimensions 2x{size_full:d}, found: {actual}".format(
                size_full=nocc,
                actual=space.shape,
            ))
    return molecular_response_ov(vind, orb2ov(space, nocc), nocc, nmo, double, log_dest)


def mk_make_canonic(m, o, v, return_ov=False, space_ov=None):
    """
    Makes the output of pyscf TDDFT matrix (MK form) to be canonic.
    Args:
        m (ndarray): the TDDFT matrix;
        o (ndarray): occupied orbital energies;
        v (ndarray): virtual orbital energies;
        return_ov (bool): if True, returns the K-matrix as well;
        space_ov (ndarray): an optional ov space;

    Returns:
        The rotated matrix as well as an optional K-matrix.
    """
    e_ov = (v[numpy.newaxis, :] - o[:, numpy.newaxis]).reshape(-1)
    if space_ov:
        e_ov = e_ov[space_ov]
    e_ov_sqr = e_ov ** .5
    result = m * (e_ov_sqr[numpy.newaxis, :] / e_ov_sqr[:, numpy.newaxis])
    if return_ov:
        return result, numpy.diag(e_ov)
    else:
        return result


class PhysERI(MolecularMFMixin, TDProxyMatrixBlocks):
    proxy_choices = {
        "hf": TDHF,
        "dft": TDDFT,
    }

    def __init__(self, model, proxy, frozen=None):
        """
        A proxy class for calculating TD matrix blocks (molecular version).

        Args:
            model: the base model;
            proxy: a pyscf proxy with TD response function, one of 'hf', 'dft';
            frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals;
        """
        # Note: the "proxy" argument does not matter at all since pyscf returns same objects for both HF and DFT
        # mean-field solutions. The argument is left for consistency and explicitness
        if proxy not in self.proxy_choices:
            raise ValueError("Unknown proxy: {}".format(repr(proxy)))
        TDProxyMatrixBlocks.__init__(self, self.proxy_choices[proxy](model))
        MolecularMFMixin.__init__(self, model, frozen=frozen)

    def proxy_is_double(self):
        """
        Determines if double-sized matrices are proxied.
        Returns:
            True if double-sized matrices are proxied.
        """
        nocc_full = self.nocc_full
        nmo_full = self.nmo_full
        size_full = nocc_full * (nmo_full - nocc_full)
        size_hdiag = len(self.proxy_diag)

        if size_full == size_hdiag:
            return False

        elif 2 * size_full == size_hdiag:
            return True

        else:
            raise RuntimeError("Do not recognize the size of TD diagonal: {:d}. The size of ov-space is {:d}".format(
                size_hdiag, size_full
            ))

    def proxy_response(self):
        """
        A raw response matrix.
        Returns:
            A raw response matrix.
        """
        return molecular_response(
            self.proxy_vind,
            self.space,
            self.nocc_full,
            self.nmo_full,
            self.proxy_is_double(),
            self.model,
        )

    def tdhf_primary_form(self, *args, **kwargs):
        """
        A primary form of TD matrixes.

        Returns:
            Output type: "full", "ab", or "mk" and the corresponding matrix(es).
        """
        if not self.proxy_is_double():
            # The MK case
            e_occ, e_virt = self.mo_energy[:self.nocc], self.mo_energy[self.nocc:]
            return ("mk",) + mk_make_canonic(self.proxy_response(), e_occ, e_virt, return_ov=True)

        else:
            # Full case
            return ("ab",) + self.proxy_response()


vector_to_amplitudes = rhf_slow.vector_to_amplitudes


class TDProxy(TDBase):
    proxy_eri = PhysERI
    v2a = staticmethod(vector_to_amplitudes)

    def __init__(self, mf, proxy, frozen=None):
        """
        Performs TD calculation. Roots and eigenvectors are stored in `self.e`, `self.xy`.
        Args:
            mf: the base restricted mean-field model;
            proxy: a pyscf proxy with TD response function, one of 'hf', 'dft';
            frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals;
        """
        super(TDProxy, self).__init__(mf, frozen=frozen)
        self.__proxy__ = proxy

    def ao2mo(self):
        """
        Prepares ERI.

        Returns:
            A suitable ERI.
        """
        return self.proxy_eri(self._scf, self.__proxy__, frozen=self.frozen)


# pyscf/tdscf/rhf.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#
# Ref:
# Chem Phys Lett, 256, 454
# J. Mol. Struct. THEOCHEM, 914, 3
# Recent Advances in Density Functional Methods, Chapter 5, M. E. Casida
#


from functools import reduce
import numpy
from pyscf import lib
from pyscf import gto
from pyscf import scf
from pyscf import ao2mo
from pyscf import symm
from pyscf.lib import logger
from pyscf.scf import hf_symm
from pyscf.scf import _response_functions
from pyscf.data import nist
from pyscf import __config__

OUTPUT_THRESHOLD = getattr(__config__, 'tdscf_rhf_get_nto_threshold', 0.3)
REAL_EIG_THRESHOLD = getattr(__config__, 'tdscf_rhf_TDDFT_pick_eig_threshold', 1e-4)
MO_BASE = getattr(__config__, 'MO_BASE', 1)

# Low excitation filter to avoid numerical instability
POSTIVE_EIG_THRESHOLD = getattr(__config__, 'tdscf_rhf_TDDFT_positive_eig_threshold', 1e-3)


def gen_tda_operation(mf, fock_ao=None, singlet=True, wfnsym=None):
    '''Generate function to compute A x

    Kwargs:
        wfnsym : int or str
            Point group symmetry irrep symbol or ID for excited CIS wavefunction.
    '''
    mol = mf.mol
    mo_coeff = mf.mo_coeff
    assert(mo_coeff.dtype == numpy.double)
    mo_energy = mf.mo_energy
    mo_occ = mf.mo_occ
    nao, nmo = mo_coeff.shape
    occidx = numpy.where(mo_occ==2)[0]
    viridx = numpy.where(mo_occ==0)[0]
    nocc = len(occidx)
    nvir = len(viridx)
    orbv = mo_coeff[:,viridx]
    orbo = mo_coeff[:,occidx]

    if wfnsym is not None and mol.symmetry:
        if isinstance(wfnsym, str):
            wfnsym = symm.irrep_name2id(mol.groupname, wfnsym)
        wfnsym = wfnsym % 10  # convert to D2h subgroup
        orbsym = hf_symm.get_orbsym(mol, mo_coeff)
        orbsym_in_d2h = numpy.asarray(orbsym) % 10  # convert to D2h irreps
        sym_forbid = (orbsym_in_d2h[occidx,None] ^ orbsym_in_d2h[viridx]) != wfnsym

    if fock_ao is None:
        #dm0 = mf.make_rdm1(mo_coeff, mo_occ)
        #fock_ao = mf.get_hcore() + mf.get_veff(mol, dm0)
        foo = numpy.diag(mo_energy[occidx])
        fvv = numpy.diag(mo_energy[viridx])
    else:
        fock = reduce(numpy.dot, (mo_coeff.conj().T, fock_ao, mo_coeff))
        foo = fock[occidx[:,None],occidx]
        fvv = fock[viridx[:,None],viridx]

    hdiag = fvv.diagonal() - foo.diagonal()[:,None]
    if wfnsym is not None and mol.symmetry:
        hdiag[sym_forbid] = 0
    hdiag = hdiag.ravel()

    mo_coeff = numpy.asarray(numpy.hstack((orbo,orbv)), order='F')
    vresp = mf.gen_response(singlet=singlet, hermi=0)

    def vind(zs):
        zs = numpy.asarray(zs).reshape(-1,nocc,nvir)
        if wfnsym is not None and mol.symmetry:
            zs = numpy.copy(zs)
            zs[:,sym_forbid] = 0

        # *2 for double occupancy
        dmov = lib.einsum('xov,po,qv->xpq', zs*2, orbo, orbv.conj())
        v1ao = vresp(dmov)
        v1ov = lib.einsum('xpq,po,qv->xov', v1ao, orbo.conj(), orbv)
        v1ov += lib.einsum('xqs,sp->xqp', zs, fvv)
        v1ov -= lib.einsum('xpr,sp->xsr', zs, foo)
        if wfnsym is not None and mol.symmetry:
            v1ov[:,sym_forbid] = 0
        return v1ov.reshape(v1ov.shape[0],-1)

    return vind, hdiag
gen_tda_hop = gen_tda_operation

def get_ab(mf, mo_energy=None, mo_coeff=None, mo_occ=None):
    r'''A and B matrices for TDDFT response function.

    A[i,a,j,b] = \delta_{ab}\delta_{ij}(E_a - E_i) + (ia||bj)
    B[i,a,j,b] = (ia||jb)
    '''
    if mo_energy is None: mo_energy = mf.mo_energy
    if mo_coeff is None: mo_coeff = mf.mo_coeff
    if mo_occ is None: mo_occ = mf.mo_occ
    assert(mo_coeff.dtype == numpy.double)

    mol = mf.mol
    nao, nmo = mo_coeff.shape
    occidx = numpy.where(mo_occ==2)[0]
    viridx = numpy.where(mo_occ==0)[0]
    orbv = mo_coeff[:,viridx]
    orbo = mo_coeff[:,occidx]
    nvir = orbv.shape[1]
    nocc = orbo.shape[1]
    mo = numpy.hstack((orbo,orbv))
    nmo = nocc + nvir

    e_ia = lib.direct_sum('a-i->ia', mo_energy[viridx], mo_energy[occidx])
    a = numpy.diag(e_ia.ravel()).reshape(nocc,nvir,nocc,nvir)
    b = numpy.zeros_like(a)

    def add_hf_(a, b, hyb=1):
        eri_mo = ao2mo.general(mol, [orbo,mo,mo,mo], compact=False)
        eri_mo = eri_mo.reshape(nocc,nmo,nmo,nmo)
        a += numpy.einsum('iabj->iajb', eri_mo[:nocc,nocc:,nocc:,:nocc]) * 2
        a -= numpy.einsum('ijba->iajb', eri_mo[:nocc,:nocc,nocc:,nocc:]) * hyb

        b += numpy.einsum('iajb->iajb', eri_mo[:nocc,nocc:,:nocc,nocc:]) * 2
        b -= numpy.einsum('jaib->iajb', eri_mo[:nocc,nocc:,:nocc,nocc:]) * hyb

    if isinstance(mf, scf.hf.KohnShamDFT):
        ni = mf._numint
        ni.libxc.test_deriv_order(mf.xc, 2, raise_error=True)
        if getattr(mf, 'nlc', '') != '':
            logger.warn(mf, 'NLC functional found in DFT object.  Its second '
                        'deriviative is not available. Its contribution is '
                        'not included in the response function.')
        omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin)

        add_hf_(a, b, hyb)

        xctype = ni._xc_type(mf.xc)
        dm0 = mf.make_rdm1(mo_coeff, mo_occ)
        make_rho = ni._gen_rho_evaluator(mol, dm0, hermi=1)[0]
        mem_now = lib.current_memory()[0]
        max_memory = max(2000, mf.max_memory*.8-mem_now)

        if xctype == 'LDA':
            ao_deriv = 0
            for ao, mask, weight, coords \
                    in ni.block_loop(mol, mf.grids, nao, ao_deriv, max_memory):
                rho = make_rho(0, ao, mask, xctype)
                fxc = ni.eval_xc(mf.xc, rho, 0, deriv=2)[2]
                frr = fxc[0]

                rho_o = lib.einsum('rp,pi->ri', ao, orbo)
                rho_v = lib.einsum('rp,pi->ri', ao, orbv)
                rho_ov = numpy.einsum('ri,ra->ria', rho_o, rho_v)
                w_ov = numpy.einsum('ria,r->ria', rho_ov, weight*frr)
                iajb = lib.einsum('ria,rjb->iajb', rho_ov, w_ov) * 2
                a += iajb
                b += iajb

        elif xctype == 'GGA':
            ao_deriv = 1
            for ao, mask, weight, coords \
                    in ni.block_loop(mol, mf.grids, nao, ao_deriv, max_memory):
                rho = make_rho(0, ao, mask, xctype)
                vxc, fxc = ni.eval_xc(mf.xc, rho, 0, deriv=2)[1:3]
                vgamma = vxc[1]
                frho, frhogamma, fgg = fxc[:3]

                rho_o = lib.einsum('xrp,pi->xri', ao, orbo)
                rho_v = lib.einsum('xrp,pi->xri', ao, orbv)
                rho_ov = numpy.einsum('xri,ra->xria', rho_o, rho_v[0])
                rho_ov[1:4] += numpy.einsum('ri,xra->xria', rho_o[0], rho_v[1:4])
                # sigma1 ~ \nabla(\rho_\alpha+\rho_\beta) dot \nabla(|b><j|) z_{bj}
                sigma1 = numpy.einsum('xr,xria->ria', rho[1:4], rho_ov[1:4])

                w_ov = numpy.empty_like(rho_ov)
                w_ov[0]  = numpy.einsum('r,ria->ria', frho, rho_ov[0])
                w_ov[0] += numpy.einsum('r,ria->ria', 2*frhogamma, sigma1)
                f_ov = numpy.einsum('r,ria->ria', 4*fgg, sigma1)
                f_ov+= numpy.einsum('r,ria->ria', 2*frhogamma, rho_ov[0])
                w_ov[1:] = numpy.einsum('ria,xr->xria', f_ov, rho[1:4])
                w_ov[1:]+= numpy.einsum('r,xria->xria', 2*vgamma, rho_ov[1:4])
                w_ov *= weight[:,None,None]
                iajb = lib.einsum('xria,xrjb->iajb', rho_ov, w_ov) * 2
                a += iajb
                b += iajb

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

    else:
        add_hf_(a, b)

    return a, b

def get_nto(tdobj, state=1, threshold=OUTPUT_THRESHOLD, verbose=None):
    r'''
    Natural transition orbital analysis.

    The natural transition density matrix between ground state and excited
    state :math:`Tia = \langle \Psi_{ex} | i a^\dagger | \Psi_0 \rangle` can
    be transformed to diagonal form through SVD
    :math:`T = O \sqrt{\lambda} V^\dagger`. O and V are occupied and virtual
    natural transition orbitals. The diagonal elements :math:`\lambda` are the
    weights of the occupied-virtual orbital pair in the excitation.

    Ref: Martin, R. L., JCP, 118, 4775-4777

    Note in the TDHF/TDDFT calculations, the excitation part (X) is
    interpreted as the CIS coefficients and normalized to 1. The de-excitation
    part (Y) is ignored.

    Args:
        tdobj : TDA, or TDHF, or TDDFT object

        state : int
            Excited state ID.  state = 1 means the first excited state.
            If state < 0, state ID is counted from the last excited state.

    Kwargs:
        threshold : float
            Above which the NTO coefficients will be printed in the output.

    Returns:
        A list (weights, NTOs).  NTOs are natural orbitals represented in AO
        basis. The first N_occ NTOs are occupied NTOs and the rest are virtual
        NTOs.
    '''
    if state == 0:
        logger.warn(tdobj, 'Excited state starts from 1. '
                    'Set state=1 for first excited state.')
        state_id = state
    elif state < 0:
        state_id = state
    else:
        state_id = state - 1

    mol = tdobj.mol
    mo_coeff = tdobj._scf.mo_coeff
    mo_occ = tdobj._scf.mo_occ
    orbo = mo_coeff[:,mo_occ==2]
    orbv = mo_coeff[:,mo_occ==0]
    nocc = orbo.shape[1]
    nvir = orbv.shape[1]

    cis_t1 = tdobj.xy[state_id][0]
    # TDDFT (X,Y) has X^2-Y^2=1.
    # Renormalizing X (X^2=1) to map it to CIS coefficients
    cis_t1 *= 1. / numpy.linalg.norm(cis_t1)

# TODO: Comparing to the NTOs defined in JCP, 142, 244103.  JCP, 142, 244103
# provides a method to incorporate the Y matrix in the transition density
# matrix.  However, it may break the point group symmetry of the NTO orbitals
# when the system has degenerated irreducible representations.

    if mol.symmetry:
        orbsym = hf_symm.get_orbsym(mol, mo_coeff)
        orbsym_in_d2h = numpy.asarray(orbsym) % 10  # convert to D2h irreps
        o_sym = orbsym_in_d2h[mo_occ==2]
        v_sym = orbsym_in_d2h[mo_occ==0]
        nto_o = numpy.eye(nocc)
        nto_v = numpy.eye(nvir)
        weights_o = numpy.zeros(nocc)
        weights_v = numpy.zeros(nvir)
        for ir in set(orbsym_in_d2h):
            idx = numpy.where(o_sym == ir)[0]
            if idx.size > 0:
                dm_oo = numpy.dot(cis_t1[idx], cis_t1[idx].T)
                weights_o[idx], nto_o[idx[:,None],idx] = numpy.linalg.eigh(dm_oo)

            idx = numpy.where(v_sym == ir)[0]
            if idx.size > 0:
                dm_vv = numpy.dot(cis_t1[:,idx].T, cis_t1[:,idx])
                weights_v[idx], nto_v[idx[:,None],idx] = numpy.linalg.eigh(dm_vv)

        # weights in descending order
        idx = numpy.argsort(-weights_o)
        weights_o = weights_o[idx]
        nto_o = nto_o[:,idx]
        o_sym = o_sym[idx]

        idx = numpy.argsort(-weights_v)
        weights_v = weights_v[idx]
        nto_v = nto_v[:,idx]
        v_sym = v_sym[idx]

        nto_orbsym = numpy.hstack((o_sym, v_sym))

        if nocc < nvir:
            weights = weights_o
        else:
            weights = weights_v

    else:
        nto_o, w, nto_vT = numpy.linalg.svd(cis_t1)
        nto_v = nto_vT.conj().T
        weights = w**2
        nto_orbsym = None

    idx = numpy.argmax(abs(nto_o.real), axis=0)
    nto_o[:,nto_o[idx,numpy.arange(nocc)].real<0] *= -1
    idx = numpy.argmax(abs(nto_v.real), axis=0)
    nto_v[:,nto_v[idx,numpy.arange(nvir)].real<0] *= -1

    occupied_nto = numpy.dot(orbo, nto_o)
    virtual_nto = numpy.dot(orbv, nto_v)
    nto_coeff = numpy.hstack((occupied_nto, virtual_nto))

    if mol.symmetry:
        nto_coeff = lib.tag_array(nto_coeff, orbsym=nto_orbsym)

    log = logger.new_logger(tdobj, verbose)
    if log.verbose >= logger.INFO:
        log.info('State %d: %g eV  NTO largest component %s',
                 state_id+1, tdobj.e[state_id]*nist.HARTREE2EV, weights[0])
        o_idx = numpy.where(abs(nto_o[:,0]) > threshold)[0]
        v_idx = numpy.where(abs(nto_v[:,0]) > threshold)[0]
        fmt = '%' + str(lib.param.OUTPUT_DIGITS) + 'f (MO #%d)'
        log.info('    occ-NTO: ' +
                 ' + '.join([(fmt % (nto_o[i,0], i+MO_BASE))
                             for i in o_idx]))
        log.info('    vir-NTO: ' +
                 ' + '.join([(fmt % (nto_v[i,0], i+MO_BASE+nocc))
                             for i in v_idx]))
    return weights, nto_coeff


def analyze(tdobj, verbose=None):
    log = logger.new_logger(tdobj, verbose)
    mol = tdobj.mol
    mo_coeff = tdobj._scf.mo_coeff
    mo_occ = tdobj._scf.mo_occ
    nocc = numpy.count_nonzero(mo_occ == 2)

    e_ev = numpy.asarray(tdobj.e) * nist.HARTREE2EV
    e_wn = numpy.asarray(tdobj.e) * nist.HARTREE2WAVENUMBER
    wave_length = 1e7/e_wn

    if tdobj.singlet:
        log.note('\n** Singlet excitation energies and oscillator strengths **')
    else:
        log.note('\n** Triplet excitation energies and oscillator strengths **')

    if mol.symmetry:
        orbsym = hf_symm.get_orbsym(mol, mo_coeff)
        x_sym = symm.direct_prod(orbsym[mo_occ==2], orbsym[mo_occ==0], mol.groupname)
    else:
        x_sym = None

    f_oscillator = tdobj.oscillator_strength()
    for i, ei in enumerate(tdobj.e):
        x, y = tdobj.xy[i]
        if x_sym is None:
            log.note('Excited State %3d: %12.5f eV %9.2f nm  f=%.4f',
                     i+1, e_ev[i], wave_length[i], f_oscillator[i])
        else:
            wfnsym = analyze_wfnsym(tdobj, x_sym, x)
            log.note('Excited State %3d: %4s %12.5f eV %9.2f nm  f=%.4f',
                     i+1, wfnsym, e_ev[i], wave_length[i], f_oscillator[i])

        if log.verbose >= logger.INFO:
            o_idx, v_idx = numpy.where(abs(x) > 0.1)
            for o, v in zip(o_idx, v_idx):
                log.info('    %4d -> %-4d %12.5f',
                         o+MO_BASE, v+MO_BASE+nocc, x[o,v])

    if log.verbose >= logger.INFO:
        log.info('\n** Transition electric dipole moments (AU) **')
        log.info('state          X           Y           Z        Dip. S.      Osc.')
        trans_dip = tdobj.transition_dipole()
        for i, ei in enumerate(tdobj.e):
            dip = trans_dip[i]
            log.info('%3d    %11.4f %11.4f %11.4f %11.4f %11.4f',
                     i+1, dip[0], dip[1], dip[2], numpy.dot(dip, dip),
                     f_oscillator[i])

        log.info('\n** Transition velocity dipole moments (imaginary part, AU) **')
        log.info('state          X           Y           Z        Dip. S.      Osc.')
        trans_v = tdobj.transition_velocity_dipole()
        f_v = tdobj.oscillator_strength(gauge='velocity', order=0)
        for i, ei in enumerate(tdobj.e):
            v = trans_v[i]
            log.info('%3d    %11.4f %11.4f %11.4f %11.4f %11.4f',
                     i+1, v[0], v[1], v[2], numpy.dot(v, v), f_v[i])

        log.info('\n** Transition magnetic dipole moments (imaginary part, AU) **')
        log.info('state          X           Y           Z')
        trans_m = tdobj.transition_magnetic_dipole()
        for i, ei in enumerate(tdobj.e):
            m = trans_m[i]
            log.info('%3d    %11.4f %11.4f %11.4f',
                     i+1, m[0], m[1], m[2])
    return tdobj

def analyze_wfnsym(tdobj, x_sym, x):
    '''Guess the wfn symmetry of TDDFT X amplitude'''
    possible_sym = x_sym[(x > 1e-7) | (x < -1e-7)]
    if numpy.all(possible_sym == symm.MULTI_IRREPS):
        if tdobj.wfnsym is None:
            wfnsym = '???'
        else:
            wfnsym = tdobj.wfnsym
    else:
        ids = possible_sym[possible_sym != symm.MULTI_IRREPS]
        ids = numpy.unique(ids)
        if ids.size == 1:
            wfnsym = symm.irrep_id2name(tdobj.mol.groupname, ids[0])
        else:
            wfnsym = '???'
    return wfnsym


def transition_dipole(tdobj, xy=None):
    '''Transition dipole moments in the length gauge'''
    mol = tdobj.mol
    with mol.with_common_orig(_charge_center(mol)):
        ints = mol.intor_symmetric('int1e_r', comp=3)
    return tdobj._contract_multipole(ints, hermi=True, xy=xy)

def transition_velocity_dipole(tdobj, xy=None):
    '''Transition dipole moments in the velocity gauge (imaginary part only)
    '''
    ints = tdobj.mol.intor('int1e_ipovlp', comp=3, hermi=2)
    v = tdobj._contract_multipole(ints, hermi=False, xy=xy)
    return -v

def transition_magnetic_dipole(tdobj, xy=None):
    '''Transition magnetic dipole moments (imaginary part only)'''
    mol = tdobj.mol
    with mol.with_common_orig(_charge_center(mol)):
        ints = mol.intor('int1e_cg_irxp', comp=3, hermi=2)
    m_pol = tdobj._contract_multipole(ints, hermi=False, xy=xy)
    return -m_pol

def transition_quadrupole(tdobj, xy=None):
    '''Transition quadrupole moments in the length gauge'''
    mol = tdobj.mol
    nao = mol.nao_nr()
    with mol.with_common_orig(_charge_center(mol)):
        ints = mol.intor('int1e_rr', comp=9, hermi=0).reshape(3,3,nao,nao)
    quad = tdobj._contract_multipole(ints, hermi=True, xy=xy)
    return quad

def transition_velocity_quadrupole(tdobj, xy=None):
    '''Transition quadrupole moments in the velocity gauge (imaginary part only)
    '''
    mol = tdobj.mol
    nao = mol.nao_nr()
    with mol.with_common_orig(_charge_center(mol)):
        ints = mol.intor('int1e_irp', comp=9, hermi=0).reshape(3,3,nao,nao)
    ints = ints + ints.transpose(1,0,3,2)
    quad = tdobj._contract_multipole(ints, hermi=True, xy=xy)
    return -quad

def transition_magnetic_quadrupole(tdobj, xy=None):
    '''Transition magnetic quadrupole moments (imaginary part only)'''
    XX, XY, XZ, YX, YY, YZ, ZX, ZY, ZZ = range(9)
    mol = tdobj.mol
    nao = mol.nao_nr()
    with mol.with_common_orig(_charge_center(mol)):
        ints = mol.intor('int1e_irrp', comp=27, hermi=0).reshape(3,9,nao,nao)
    m_ints = (ints[:,[YZ,ZX,XY]] - ints[:,[ZY,XZ,YX]]).transpose(1,0,2,3)
    with mol.with_common_orig(_charge_center(mol)):
        ints = mol.intor('int1e_irpr', comp=27, hermi=0).reshape(9,3,nao,nao)
    m_ints += ints[[YZ,ZX,XY]] - ints[[ZY,XZ,YX]]
    m_quad = tdobj._contract_multipole(m_ints, hermi=True, xy=xy)
    return -m_quad

def transition_octupole(tdobj, xy=None):
    '''Transition octupole moments in the length gauge'''
    mol = tdobj.mol
    nao = mol.nao_nr()
    with mol.with_common_orig(_charge_center(mol)):
        ints = mol.intor('int1e_rrr', comp=27, hermi=0).reshape(3,3,3,nao,nao)
    o_pol = tdobj._contract_multipole(ints, hermi=True, xy=xy)
    return o_pol

def transition_velocity_octupole(tdobj, xy=None):
    '''Transition octupole moments in the velocity gauge (imaginary part only)
    '''
    mol = tdobj.mol
    nao = mol.nao_nr()
    with mol.with_common_orig(_charge_center(mol)):
        ints = mol.intor('int1e_irrp', comp=27, hermi=0).reshape(3,3,3,nao,nao)
    ints = ints + ints.transpose(2,1,0,4,3)
    with mol.with_common_orig(_charge_center(mol)):
        ints += mol.intor('int1e_irpr', comp=27, hermi=0).reshape(3,3,3,nao,nao)
    o_pol = tdobj._contract_multipole(ints, hermi=True, xy=xy)
    return -o_pol

def _charge_center(mol):
    charges = mol.atom_charges()
    coords  = mol.atom_coords()
    return numpy.einsum('z,zr->r', charges, coords)/charges.sum()

def _contract_multipole(tdobj, ints, hermi=True, xy=None):
    if xy is None: xy = tdobj.xy
    mo_coeff = tdobj._scf.mo_coeff
    mo_occ = tdobj._scf.mo_occ
    orbo = mo_coeff[:,mo_occ==2]
    orbv = mo_coeff[:,mo_occ==0]

    nstates = len(xy)
    pol_shape = ints.shape[:-2]
    nao = ints.shape[-1]

    #Incompatible to old numpy version
    #ints = numpy.einsum('...pq,pi,qj->...ij', ints, orbo.conj(), orbv)
    ints = lib.einsum('xpq,pi,qj->xij', ints.reshape(-1,nao,nao), orbo.conj(), orbv)
    pol = numpy.array([numpy.einsum('xij,ij->x', ints, x) * 2 for x,y in xy])
    if isinstance(xy[0][1], numpy.ndarray):
        if hermi:
            pol += [numpy.einsum('xij,ij->x', ints, y) * 2 for x,y in xy]
        else:  # anti-Hermitian
            pol -= [numpy.einsum('xij,ij->x', ints, y) * 2 for x,y in xy]
    pol = pol.reshape((nstates,)+pol_shape)
    return pol

def oscillator_strength(tdobj, e=None, xy=None, gauge='length', order=0):
    if e is None: e = tdobj.e

    if gauge == 'length':
        trans_dip = transition_dipole(tdobj, xy)
        f = 2./3. * numpy.einsum('s,sx,sx->s', e, trans_dip, trans_dip)
        return f

    else:  # velocity gauge
        # Ref. JCP, 143, 234103
        trans_dip = transition_velocity_dipole(tdobj, xy)
        f = 2./3. * numpy.einsum('s,sx,sx->s', 1./e, trans_dip, trans_dip)

        if order > 0:
            m_dip = .5 * transition_magnetic_dipole(tdobj, xy)
            f_m = numpy.einsum('s,sx,sx->s', e, m_dip, m_dip)
            f_m = nist.ALPHA**2/6 * f_m.real
            f += f_m

            quad = .5 * transition_velocity_quadrupole(tdobj, xy)
            f_quad = numpy.einsum('s,sxy,sxy->s', e, quad, quad)
            f_quad-= 1./3 * numpy.einsum('s,sxx,sxx->s', e, quad, quad)
            f_quad = nist.ALPHA**2/20 * f_quad.real
            f += f_quad
            logger.debug(tdobj, '    First order correction to oscillator '
                         'strength (velocity gague)')
            logger.debug(tdobj, '    %s', f_m+f_quad)

        if order > 1:
            m_quad = -1./6 * 1j*transition_magnetic_quadrupole(tdobj, xy)
            f_m = numpy.einsum('s,sy,szx,xyz->s', e, trans_dip*1j, m_quad,
                               lib.LeviCivita)
            f_m = nist.ALPHA**3/9 * f_m.real
            f += f_m

            o_pol = -1./6 * 1j*transition_velocity_octupole(tdobj, xy)
            f_o = numpy.einsum('s,sy,sxxy->s', e, trans_dip*1j, o_pol)
            f_o = -2*nist.ALPHA**2/45 * f_o.real
            f += f_o
            logger.debug(tdobj, '    Second order correction to oscillator '
                         'strength (velocity gague)')
            logger.debug(tdobj, '    %s', f_m+f_o)

    return f


def as_scanner(td):
    '''Generating a scanner/solver for TDA/TDHF/TDDFT PES.

    The returned solver is a function. This function requires one argument
    "mol" as input and returns total TDA/TDHF/TDDFT energy.

    The solver will automatically use the results of last calculation as the
    initial guess of the new calculation.  All parameters assigned in the
    TDA/TDDFT and the underlying SCF objects (conv_tol, max_memory etc) are
    automatically applied in the solver.

    Note scanner has side effects.  It may change many underlying objects
    (_scf, with_df, with_x2c, ...) during calculation.

    Examples::

    >>> from pyscf import gto, scf, tdscf
    >>> mol = gto.M(atom='H 0 0 0; F 0 0 1')
    >>> td_scanner = tdscf.TDHF(scf.RHF(mol)).as_scanner()
    >>> de = td_scanner(gto.M(atom='H 0 0 0; F 0 0 1.1'))
    [ 0.34460866  0.34460866  0.7131453 ]
    >>> de = td_scanner(gto.M(atom='H 0 0 0; F 0 0 1.5'))
    [ 0.14844013  0.14844013  0.47641829]
    '''
    if isinstance(td, lib.SinglePointScanner):
        return td

    logger.info(td, 'Set %s as a scanner', td.__class__)

    class TD_Scanner(td.__class__, lib.SinglePointScanner):
        def __init__(self, td):
            self.__dict__.update(td.__dict__)
            self._scf = td._scf.as_scanner()
        def __call__(self, mol_or_geom, **kwargs):
            if isinstance(mol_or_geom, gto.Mole):
                mol = mol_or_geom
            else:
                mol = self.mol.set_geom_(mol_or_geom, inplace=False)

            self.reset(mol)

            mf_scanner = self._scf
            mf_e = mf_scanner(mol)
            self.kernel(**kwargs)
            return mf_e + self.e
    return TD_Scanner(td)


class TDMixin(lib.StreamObject):
    conv_tol = getattr(__config__, 'tdscf_rhf_TDA_conv_tol', 1e-9)
    nstates = getattr(__config__, 'tdscf_rhf_TDA_nstates', 3)
    singlet = getattr(__config__, 'tdscf_rhf_TDA_singlet', True)
    lindep = getattr(__config__, 'tdscf_rhf_TDA_lindep', 1e-12)
    level_shift = getattr(__config__, 'tdscf_rhf_TDA_level_shift', 0)
    max_space = getattr(__config__, 'tdscf_rhf_TDA_max_space', 50)
    max_cycle = getattr(__config__, 'tdscf_rhf_TDA_max_cycle', 100)

    def __init__(self, mf):
        self.verbose = mf.verbose
        self.stdout = mf.stdout
        self.mol = mf.mol
        self._scf = mf
        self.max_memory = mf.max_memory
        self.chkfile = mf.chkfile

        self.wfnsym = None

        # xy = (X,Y), normalized to 1/2: 2(XX-YY) = 1
        # In TDA, Y = 0
        self.converged = None
        self.e = None
        self.xy = None

        keys = set(('conv_tol', 'nstates', 'singlet', 'lindep', 'level_shift',
                    'max_space', 'max_cycle'))
        self._keys = set(self.__dict__.keys()).union(keys)

    @property
    def nroots(self):
        return self.nstates
    @nroots.setter
    def nroots(self, x):
        self.nstates = x

    @property
    def e_tot(self):
        '''Excited state energies'''
        return self._scf.e_tot + self.e

    def dump_flags(self, verbose=None):
        log = logger.new_logger(self, verbose)
        log.info('\n')
        log.info('******** %s for %s ********',
                 self.__class__, self._scf.__class__)
        if self.singlet is None:
            log.info('nstates = %d', self.nstates)
        elif self.singlet:
            log.info('nstates = %d singlet', self.nstates)
        else:
            log.info('nstates = %d triplet', self.nstates)
        log.info('wfnsym = %s', self.wfnsym)
        log.info('conv_tol = %g', self.conv_tol)
        log.info('eigh lindep = %g', self.lindep)
        log.info('eigh level_shift = %g', self.level_shift)
        log.info('eigh max_space = %d', self.max_space)
        log.info('eigh max_cycle = %d', self.max_cycle)
        log.info('chkfile = %s', self.chkfile)
        log.info('max_memory %d MB (current use %d MB)',
                 self.max_memory, lib.current_memory()[0])
        if not self._scf.converged:
            log.warn('Ground state SCF is not converged')
        log.info('\n')

    def check_sanity(self):
        if self._scf.mo_coeff is None:
            raise RuntimeError('SCF object is not initialized')
        lib.StreamObject.check_sanity(self)

    def reset(self, mol=None):
        if mol is not None:
            self.mol = mol
        self._scf.reset(mol)
        return self

    def gen_vind(self, mf):
        raise NotImplementedError

    @lib.with_doc(get_ab.__doc__)
    def get_ab(self, mf=None):
        if mf is None: mf = self._scf
        return get_ab(mf)

    def get_precond(self, hdiag):
        def precond(x, e, x0):
            diagd = hdiag - (e-self.level_shift)
            diagd[abs(diagd)<1e-8] = 1e-8
            return x/diagd
        return precond

    analyze = analyze
    get_nto = get_nto
    oscillator_strength = oscillator_strength

    _contract_multipole = _contract_multipole  # needed by following methods
    transition_dipole              = transition_dipole
    transition_quadrupole          = transition_quadrupole
    transition_octupole            = transition_octupole
    transition_velocity_dipole     = transition_velocity_dipole
    transition_velocity_quadrupole = transition_velocity_quadrupole
    transition_velocity_octupole   = transition_velocity_octupole
    transition_magnetic_dipole     = transition_magnetic_dipole
    transition_magnetic_quadrupole = transition_magnetic_quadrupole

    as_scanner = as_scanner

    def nuc_grad_method(self):
        from pyscf.grad import tdrhf
        return tdrhf.Gradients(self)

    def _finalize(self):
        '''Hook for dumping results and clearing up the object.'''
        if not all(self.converged):
            logger.note(self, 'TD-SCF states %s not converged.',
                        [i for i, x in enumerate(self.converged) if not x])
        logger.note(self, 'Excited State energies (eV)\n%s', self.e * nist.HARTREE2EV)
        return self

class TDA(TDMixin):
    '''Tamm-Dancoff approximation

    Attributes:
        conv_tol : float
            Diagonalization convergence tolerance.  Default is 1e-9.
        nstates : int
            Number of TD states to be computed. Default is 3.

    Saved results:

        converged : bool
            Diagonalization converged or not
        e : 1D array
            excitation energy for each excited state.
        xy : A list of two 2D arrays
            The two 2D arrays are Excitation coefficients X (shape [nocc,nvir])
            and de-excitation coefficients Y (shape [nocc,nvir]) for each
            excited state.  (X,Y) are normalized to 1/2 in RHF/RKS methods and
            normalized to 1 for UHF/UKS methods. In the TDA calculation, Y = 0.
    '''
    def gen_vind(self, mf):
        '''Compute Ax'''
        return gen_tda_hop(mf, singlet=self.singlet, wfnsym=self.wfnsym)

    def init_guess(self, mf, nstates=None, wfnsym=None):
        if nstates is None: nstates = self.nstates
        if wfnsym is None: wfnsym = self.wfnsym

        mo_energy = mf.mo_energy
        mo_occ = mf.mo_occ
        occidx = numpy.where(mo_occ==2)[0]
        viridx = numpy.where(mo_occ==0)[0]
        e_ia = mo_energy[viridx] - mo_energy[occidx,None]
        e_ia_max = e_ia.max()

        if wfnsym is not None and mf.mol.symmetry:
            if isinstance(wfnsym, str):
                wfnsym = symm.irrep_name2id(mf.mol.groupname, wfnsym)
            wfnsym = wfnsym % 10  # convert to D2h subgroup
            orbsym = hf_symm.get_orbsym(mf.mol, mf.mo_coeff)
            orbsym_in_d2h = numpy.asarray(orbsym) % 10  # convert to D2h irreps
            e_ia[(orbsym_in_d2h[occidx,None] ^ orbsym_in_d2h[viridx]) != wfnsym] = 1e99

        nov = e_ia.size
        nstates = min(nstates, nov)
        e_ia = e_ia.ravel()
        e_threshold = min(e_ia_max, e_ia[numpy.argsort(e_ia)[nstates-1]])
        # Handle degeneracy, include all degenerated states in initial guess
        e_threshold += 1e-6

        idx = numpy.where(e_ia <= e_threshold)[0]
        x0 = numpy.zeros((idx.size, nov))
        for i, j in enumerate(idx):
            x0[i, j] = 1  # Koopmans' excitations
        return x0

    def kernel(self, x0=None, nstates=None):
        '''TDA diagonalization solver
        '''
        cpu0 = (logger.process_clock(), logger.perf_counter())
        self.check_sanity()
        self.dump_flags()
        if nstates is None:
            nstates = self.nstates
        else:
            self.nstates = nstates

        log = logger.Logger(self.stdout, self.verbose)

        vind, hdiag = self.gen_vind(self._scf)
        precond = self.get_precond(hdiag)

        if x0 is None:
            x0 = self.init_guess(self._scf, self.nstates)

        def pickeig(w, v, nroots, envs):
            idx = numpy.where(w > POSTIVE_EIG_THRESHOLD**2)[0]
            return w[idx], v[:,idx], idx

        self.converged, self.e, x1 = \
                lib.davidson1(vind, x0, precond,
                              tol=self.conv_tol,
                              nroots=nstates, lindep=self.lindep,
                              max_cycle=self.max_cycle,
                              max_space=self.max_space, pick=pickeig,
                              verbose=log)

        nocc = (self._scf.mo_occ>0).sum()
        nmo = self._scf.mo_occ.size
        nvir = nmo - nocc
# 1/sqrt(2) because self.x is for alpha excitation amplitude and 2(X^+*X) = 1
        self.xy = [(xi.reshape(nocc,nvir)*numpy.sqrt(.5),0) for xi in x1]

        if self.chkfile:
            lib.chkfile.save(self.chkfile, 'tddft/e', self.e)
            lib.chkfile.save(self.chkfile, 'tddft/xy', self.xy)

        log.timer('TDA', *cpu0)
        self._finalize()
        return self.e, self.xy

CIS = TDA


def gen_tdhf_operation(mf, fock_ao=None, singlet=True, wfnsym=None):
    '''Generate function to compute

    [ A  B][X]
    [-B -A][Y]
    '''
    mol = mf.mol
    mo_coeff = mf.mo_coeff
    assert(mo_coeff.dtype == numpy.double)
    mo_energy = mf.mo_energy
    mo_occ = mf.mo_occ
    nao, nmo = mo_coeff.shape
    occidx = numpy.where(mo_occ==2)[0]
    viridx = numpy.where(mo_occ==0)[0]
    nocc = len(occidx)
    nvir = len(viridx)
    orbv = mo_coeff[:,viridx]
    orbo = mo_coeff[:,occidx]

    if wfnsym is not None and mol.symmetry:
        if isinstance(wfnsym, str):
            wfnsym = symm.irrep_name2id(mol.groupname, wfnsym)
        wfnsym = wfnsym % 10  # convert to D2h subgroup
        orbsym = hf_symm.get_orbsym(mol, mo_coeff)
        orbsym_in_d2h = numpy.asarray(orbsym) % 10  # convert to D2h irreps
        sym_forbid = (orbsym_in_d2h[occidx,None] ^ orbsym_in_d2h[viridx]) != wfnsym

    #dm0 = mf.make_rdm1(mo_coeff, mo_occ)
    #fock_ao = mf.get_hcore() + mf.get_veff(mol, dm0)
    #fock = reduce(numpy.dot, (mo_coeff.T, fock_ao, mo_coeff))
    #foo = fock[occidx[:,None],occidx]
    #fvv = fock[viridx[:,None],viridx]
    foo = numpy.diag(mo_energy[occidx])
    fvv = numpy.diag(mo_energy[viridx])

    hdiag = fvv.diagonal() - foo.diagonal()[:,None]
    if wfnsym is not None and mol.symmetry:
        hdiag[sym_forbid] = 0
    hdiag = numpy.hstack((hdiag.ravel(), -hdiag.ravel()))

    mo_coeff = numpy.asarray(numpy.hstack((orbo,orbv)), order='F')
    vresp = mf.gen_response(singlet=singlet, hermi=0)

    def vind(xys):
        xys = numpy.asarray(xys).reshape(-1,2,nocc,nvir)
        if wfnsym is not None and mol.symmetry:
            # shape(nz,2,nocc,nvir): 2 ~ X,Y
            xys = numpy.copy(xys)
            xys[:,:,sym_forbid] = 0

        xs, ys = xys.transpose(1,0,2,3)
        # dms = AX + BY
        # *2 for double occupancy
        dms  = lib.einsum('xov,po,qv->xpq', xs*2, orbo, orbv.conj())
        dms += lib.einsum('xov,pv,qo->xpq', ys*2, orbv, orbo.conj())

        v1ao = vresp(dms)
        v1ov = lib.einsum('xpq,po,qv->xov', v1ao, orbo.conj(), orbv)
        v1vo = lib.einsum('xpq,pv,qo->xov', v1ao, orbv.conj(), orbo)
        v1ov += lib.einsum('xqs,sp->xqp', xs, fvv)  # AX
        v1ov -= lib.einsum('xpr,sp->xsr', xs, foo)  # AX
        v1vo += lib.einsum('xqs,sp->xqp', ys, fvv)  # AY
        v1vo -= lib.einsum('xpr,sp->xsr', ys, foo)  # AY

        if wfnsym is not None and mol.symmetry:
            v1ov[:,sym_forbid] = 0
            v1vo[:,sym_forbid] = 0

        # (AX, -AY)
        nz = xys.shape[0]
        hx = numpy.hstack((v1ov.reshape(nz,-1), -v1vo.reshape(nz,-1)))
        return hx

    return vind, hdiag


class TDHF(TDA):
    '''Time-dependent Hartree-Fock

    Attributes:
        conv_tol : float
            Diagonalization convergence tolerance.  Default is 1e-9.
        nstates : int
            Number of TD states to be computed. Default is 3.

    Saved results:

        converged : bool
            Diagonalization converged or not
        e : 1D array
            excitation energy for each excited state.
        xy : A list of two 2D arrays
            The two 2D arrays are Excitation coefficients X (shape [nocc,nvir])
            and de-excitation coefficients Y (shape [nocc,nvir]) for each
            excited state.  (X,Y) are normalized to 1/2 in RHF/RKS methods and
            normalized to 1 for UHF/UKS methods. In the TDA calculation, Y = 0.
    '''
    @lib.with_doc(gen_tdhf_operation.__doc__)
    def gen_vind(self, mf):
        return gen_tdhf_operation(mf, singlet=self.singlet, wfnsym=self.wfnsym)

    def init_guess(self, mf, nstates=None, wfnsym=None):
        x0 = TDA.init_guess(self, mf, nstates, wfnsym)
        y0 = numpy.zeros_like(x0)
        return numpy.hstack((x0,y0))

    def kernel(self, x0=None, nstates=None):
        '''TDHF diagonalization with non-Hermitian eigenvalue solver
        '''
        cpu0 = (logger.process_clock(), logger.perf_counter())
        self.check_sanity()
        self.dump_flags()
        if nstates is None:
            nstates = self.nstates
        else:
            self.nstates = nstates

        log = logger.Logger(self.stdout, self.verbose)

        vind, hdiag = self.gen_vind(self._scf)
        precond = self.get_precond(hdiag)
        if x0 is None:
            x0 = self.init_guess(self._scf, self.nstates)

        # We only need positive eigenvalues
        def pickeig(w, v, nroots, envs):
            realidx = numpy.where((abs(w.imag) < REAL_EIG_THRESHOLD) &
                                  (w.real > POSTIVE_EIG_THRESHOLD))[0]
            # If the complex eigenvalue has small imaginary part, both the
            # real part and the imaginary part of the eigenvector can
            # approximately be used as the "real" eigen solutions.
            return lib.linalg_helper._eigs_cmplx2real(w, v, realidx,
                                                      real_eigenvectors=True)

        self.converged, w, x1 = \
                lib.davidson_nosym1(vind, x0, precond,
                                    tol=self.conv_tol,
                                    nroots=nstates, lindep=self.lindep,
                                    max_cycle=self.max_cycle,
                                    max_space=self.max_space, pick=pickeig,
                                    verbose=log)

        nocc = (self._scf.mo_occ>0).sum()
        nmo = self._scf.mo_occ.size
        nvir = nmo - nocc
        self.e = w
        def norm_xy(z):
            x, y = z.reshape(2,nocc,nvir)
            norm = lib.norm(x)**2 - lib.norm(y)**2
            norm = numpy.sqrt(.5/norm)  # normalize to 0.5 for alpha spin
            return x*norm, y*norm
        self.xy = [norm_xy(z) for z in x1]

        if self.chkfile:
            lib.chkfile.save(self.chkfile, 'tddft/e', self.e)
            lib.chkfile.save(self.chkfile, 'tddft/xy', self.xy)

        log.timer('TDDFT', *cpu0)
        self._finalize()
        return self.e, self.xy

    def nuc_grad_method(self):
        from pyscf.grad import tdrhf
        return tdrhf.Gradients(self)

RPA = TDRHF = TDHF

scf.hf.RHF.TDA = lib.class_as_method(TDA)
scf.hf.RHF.TDHF = lib.class_as_method(TDHF)
scf.rohf.ROHF.TDA = None
scf.rohf.ROHF.TDHF = None
scf.hf_symm.ROHF.TDA = None
scf.hf_symm.ROHF.TDHF = None

del(OUTPUT_THRESHOLD)


# pyscf/tdscf/rhf_slow.py

#  Author: Artem Pulkin
"""
This and other `_slow` modules implement the time-dependent Hartree-Fock procedure. The primary performance drawback is
that, unlike other 'fast' routines with an implicit construction of the eigenvalue problem, these modules construct
TDHF matrices explicitly via an AO-MO transformation, i.e. with a O(N^5) complexity scaling. As a result, regular
`numpy.linalg.eig` can be used to retrieve TDHF roots in a reliable fashion without any issues related to the Davidson
procedure. Several variants of TDHF are available:

 * (this module) `pyscf.tdscf.rhf_slow`: the molecular implementation;
 * `pyscf.pbc.tdscf.rhf_slow`: PBC (periodic boundary condition) implementation for RHF objects of `pyscf.pbc.scf`
   modules;
 * `pyscf.pbc.tdscf.krhf_slow_supercell`: PBC implementation for KRHF objects of `pyscf.pbc.scf` modules. Works with
   an arbitrary number of k-points but has a overhead due to an effective construction of a supercell.
 * `pyscf.pbc.tdscf.krhf_slow_gamma`: A Gamma-point calculation resembling the original `pyscf.pbc.tdscf.krhf`
   module. Despite its name, it accepts KRHF objects with an arbitrary number of k-points but finds only few TDHF roots
   corresponding to collective oscillations without momentum transfer;
 * `pyscf.pbc.tdscf.krhf_slow`: PBC implementation for KRHF objects of `pyscf.pbc.scf` modules. Works with
   an arbitrary number of k-points and employs k-point conservation (diagonalizes matrix blocks separately).
"""

from pyscf import ao2mo
from pyscf.lib import logger
from pyscf.tdscf.common_slow import TDERIMatrixBlocks, MolecularMFMixin, TDBase

import numpy

# Convention for these modules:
# * PhysERI, PhysERI4, PhysERI8 are 2-electron integral routines computed directly (for debug purposes), with a 4-fold
#   symmetry and with an 8-fold symmetry
# * vector_to_amplitudes reshapes and normalizes the solution
# * TDRHF provides a container


class PhysERI(MolecularMFMixin, TDERIMatrixBlocks):
    def __init__(self, model, frozen=None):
        """
        The TDHF ERI implementation performing a full AO-MO transformation of integrals. No symmetries are employed in
        this class.

        Args:
            model (RHF): the base model;
            frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals;
        """
        TDERIMatrixBlocks.__init__(self)
        MolecularMFMixin.__init__(self, model, frozen=frozen)
        self.__full_eri__ = self.ao2mo((self.mo_coeff,) * 4)

    def ao2mo(self, coeff):
        """
        Phys ERI in MO basis.
        Args:
            coeff (Iterable): MO orbitals;

        Returns:
            ERI in MO basis.
        """
        coeff = (coeff[0], coeff[2], coeff[1], coeff[3])

        if "with_df" in dir(self.model):
            if "kpt" in dir(self.model):
                result = self.model.with_df.ao2mo(coeff, (self.model.kpt,) * 4, compact=False)
            else:
                result = self.model.with_df.ao2mo(coeff, compact=False)
        else:
            result = ao2mo.general(self.model.mol, coeff, compact=False)

        return result.reshape(
            tuple(i.shape[1] for i in coeff)
        ).swapaxes(1, 2)

    def __get_mo_energies__(self):
        return self.mo_energy[:self.nocc], self.mo_energy[self.nocc:]

    def __calc_block__(self, item):
        slc = tuple(slice(self.nocc) if i == 'o' else slice(self.nocc, None) for i in item)
        return self.__full_eri__[slc]


class PhysERI4(PhysERI):
    symmetries = [
        ((0, 1, 2, 3), False),
        ((1, 0, 3, 2), False),
        ((2, 3, 0, 1), True),
        ((3, 2, 1, 0), True),
    ]

    def __init__(self, model, frozen=None):
        """
        The TDHF ERI implementation performing a partial AO-MO transformation of integrals of a molecular system. A
        4-fold symmetry of complex-valued orbitals is used.

        Args:
            model (RHF): the base model;
            frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals;
        """
        TDERIMatrixBlocks.__init__(self)
        MolecularMFMixin.__init__(self, model, frozen=frozen)

    def __calc_block__(self, item):
        o = self.mo_coeff[:, :self.nocc]
        v = self.mo_coeff[:, self.nocc:]
        logger.info(self.model, "Computing {} ...".format(''.join(item)))
        return self.ao2mo(tuple(o if i == "o" else v for i in item))


class PhysERI8(PhysERI4):
    symmetries = [
        ((0, 1, 2, 3), False),
        ((1, 0, 3, 2), False),
        ((2, 3, 0, 1), False),
        ((3, 2, 1, 0), False),

        ((2, 1, 0, 3), False),
        ((3, 0, 1, 2), False),
        ((0, 3, 2, 1), False),
        ((1, 2, 3, 0), False),
    ]

    def __init__(self, model, frozen=None):
        """
        The TDHF ERI implementation performing a partial AO-MO transformation of integrals of a molecular system. An
        8-fold symmetry of real-valued orbitals is used.

        Args:
            model (RHF): the base model;
        """
        super(PhysERI8, self).__init__(model, frozen=frozen)


def vector_to_amplitudes(vectors, nocc, nmo):
    """
    Transforms (reshapes) and normalizes vectors into amplitudes.
    Args:
        vectors (numpy.ndarray): raw eigenvectors to transform;
        nocc (int): number of occupied orbitals;
        nmo (int): the total number of orbitals;

    Returns:
        Amplitudes with the following shape: (# of roots, 2 (x or y), # of occupied orbitals, # of virtual orbitals).
    """
    vectors = numpy.asanyarray(vectors)
    vectors = vectors.reshape(2, nocc, nmo-nocc, vectors.shape[1])
    norm = (abs(vectors) ** 2).sum(axis=(1, 2))
    norm = 2 * (norm[0] - norm[1])
    vectors /= norm ** .5
    return vectors.transpose(3, 0, 1, 2)


class TDRHF(TDBase):
    eri1 = PhysERI
    eri4 = PhysERI4
    eri8 = PhysERI8
    v2a = staticmethod(vector_to_amplitudes)

    def ao2mo(self):
        """
        Picks ERI: either 4-fold or 8-fold symmetric.

        Returns:
            A suitable ERI.
        """
        if numpy.iscomplexobj(self._scf.mo_coeff):
            if self.eri4 is not None:
                logger.debug1(self._scf, "4-fold symmetry used (complex orbitals)")
                return self.eri4(self._scf, frozen=self.frozen)
            elif self.eri1 is not None:
                logger.debug1(self._scf, "fallback: no symmetry used (complex orbitals)")
                return self.eri1(self._scf, frozen=self.frozen)
            else:
                raise RuntimeError("Failed to pick ERI for complex MOs: both eri1 and eri4 are None")
        else:
            if self.eri8 is not None:
                logger.debug1(self._scf, "8-fold symmetry used (real orbitals)")
                return self.eri8(self._scf, frozen=self.frozen)
            elif self.eri1 is not None:
                logger.debug1(self._scf, "fallback: no symmetry used (real orbitals)")
                return self.eri1(self._scf, frozen=self.frozen)
            else:
                raise RuntimeError("Failed to pick ERI for real MOs: both eri1 and eri8 are None")


# pyscf/tdscf/rks.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#
# Ref:
# Chem Phys Lett, 256, 454
# J. Mol. Struct. THEOCHEM, 914, 3
#


import numpy
from pyscf import lib
from pyscf import symm
from pyscf.tdscf import rhf
from pyscf.scf import hf_symm
from pyscf.data import nist
from pyscf.dft.rks import KohnShamDFT
from pyscf import __config__

# Low excitation filter to avoid numerical instability
POSTIVE_EIG_THRESHOLD = getattr(__config__, 'tdscf_rhf_TDDFT_positive_eig_threshold', 1e-3)


class TDA(rhf.TDA):
    def nuc_grad_method(self):
        from pyscf.grad import tdrks
        return tdrks.Gradients(self)

class TDDFT(rhf.TDHF):
    def nuc_grad_method(self):
        from pyscf.grad import tdrks
        return tdrks.Gradients(self)

RPA = TDRKS = TDDFT

class CasidaTDDFT(TDDFT, TDA):
    '''Solve the Casida TDDFT formula (A-B)(A+B)(X+Y) = (X+Y)w^2
    '''
    init_guess = TDA.init_guess

    def gen_vind(self, mf):
        wfnsym = self.wfnsym
        singlet = self.singlet

        mol = mf.mol
        mo_coeff = mf.mo_coeff
        assert mo_coeff.dtype == numpy.double
        mo_energy = mf.mo_energy
        mo_occ = mf.mo_occ
        nao, nmo = mo_coeff.shape
        occidx = numpy.where(mo_occ==2)[0]
        viridx = numpy.where(mo_occ==0)[0]
        nocc = len(occidx)
        nvir = len(viridx)
        orbv = mo_coeff[:,viridx]
        orbo = mo_coeff[:,occidx]

        if wfnsym is not None and mol.symmetry:
            if isinstance(wfnsym, str):
                wfnsym = symm.irrep_name2id(mol.groupname, wfnsym)
            wfnsym = wfnsym % 10  # convert to D2h subgroup
            orbsym = hf_symm.get_orbsym(mol, mo_coeff)
            orbsym_in_d2h = numpy.asarray(orbsym) % 10  # convert to D2h irreps
            sym_forbid = (orbsym_in_d2h[occidx,None] ^ orbsym_in_d2h[viridx]) != wfnsym

        e_ia = (mo_energy[viridx].reshape(-1,1) - mo_energy[occidx]).T
        if wfnsym is not None and mol.symmetry:
            e_ia[sym_forbid] = 0
        d_ia = numpy.sqrt(e_ia)
        ed_ia = e_ia * d_ia
        hdiag = e_ia.ravel() ** 2

        vresp = mf.gen_response(singlet=singlet, hermi=1)

        def vind(zs):
            zs = numpy.asarray(zs).reshape(-1,nocc,nvir)
            # *2 for double occupancy
            dmov = lib.einsum('xov,ov,po,qv->xpq', zs, d_ia*2, orbo, orbv.conj())
            # +cc for A+B and K_{ai,jb} in A == K_{ai,bj} in B
            dmov = dmov + dmov.conj().transpose(0,2,1)

            v1ao = vresp(dmov)
            v1ov = lib.einsum('xpq,po,qv->xov', v1ao, orbo.conj(), orbv)

            # numpy.sqrt(e_ia) * (e_ia*d_ia*z + v1ov)
            v1ov += numpy.einsum('xov,ov->xov', zs, ed_ia)
            v1ov *= d_ia
            return v1ov.reshape(v1ov.shape[0],-1)

        return vind, hdiag

    def kernel(self, x0=None, nstates=None):
        '''TDDFT diagonalization solver
        '''
        cpu0 = (lib.logger.process_clock(), lib.logger.perf_counter())
        mf = self._scf
        if mf._numint.libxc.is_hybrid_xc(mf.xc):
            raise RuntimeError('%s cannot be used with hybrid functional'
                               % self.__class__)
        self.check_sanity()
        self.dump_flags()
        if nstates is None:
            nstates = self.nstates
        else:
            self.nstates = nstates

        log = lib.logger.Logger(self.stdout, self.verbose)

        vind, hdiag = self.gen_vind(self._scf)
        precond = self.get_precond(hdiag)
        if x0 is None:
            x0 = self.init_guess(self._scf, self.nstates)

        def pickeig(w, v, nroots, envs):
            idx = numpy.where(w > POSTIVE_EIG_THRESHOLD**2)[0]
            return w[idx], v[:,idx], idx

        self.converged, w2, x1 = \
                lib.davidson1(vind, x0, precond,
                              tol=self.conv_tol,
                              nroots=nstates, lindep=self.lindep,
                              max_cycle=self.max_cycle,
                              max_space=self.max_space, pick=pickeig,
                              verbose=log)

        mo_energy = self._scf.mo_energy
        mo_occ = self._scf.mo_occ
        occidx = numpy.where(mo_occ==2)[0]
        viridx = numpy.where(mo_occ==0)[0]
        e_ia = (mo_energy[viridx,None] - mo_energy[occidx]).T
        e_ia = numpy.sqrt(e_ia)
        def norm_xy(w, z):
            zp = e_ia * z.reshape(e_ia.shape)
            zm = w/e_ia * z.reshape(e_ia.shape)
            x = (zp + zm) * .5
            y = (zp - zm) * .5
            norm = lib.norm(x)**2 - lib.norm(y)**2
            norm = numpy.sqrt(.5/norm)  # normalize to 0.5 for alpha spin
            return (x*norm, y*norm)

        idx = numpy.where(w2 > POSTIVE_EIG_THRESHOLD**2)[0]
        self.e = numpy.sqrt(w2[idx])
        self.xy = [norm_xy(self.e[i], x1[i]) for i in idx]

        if self.chkfile:
            lib.chkfile.save(self.chkfile, 'tddft/e', self.e)
            lib.chkfile.save(self.chkfile, 'tddft/xy', self.xy)

        log.timer('TDDFT', *cpu0)
        self._finalize()
        return self.e, self.xy

    def nuc_grad_method(self):
        from pyscf.grad import tdrks
        return tdrks.Gradients(self)

TDDFTNoHybrid = CasidaTDDFT

class dRPA(TDDFTNoHybrid):
    def __init__(self, mf):
        if not isinstance(mf, KohnShamDFT):
            raise RuntimeError("direct RPA can only be applied with DFT; for HF+dRPA, use .xc='hf'")
        mf = mf.to_rhf()
        # commit fc8d1967995b7e033b60d4428ddcca87aac78e4f handles xc='' .
        # xc='0*LDA' is equivalent to xc=''
        #mf.xc = '0.0*LDA'
        mf.xc = ''
        TDDFTNoHybrid.__init__(self, mf)

TDH = dRPA

class dTDA(TDA):
    def __init__(self, mf):
        if not isinstance(mf, KohnShamDFT):
            raise RuntimeError("direct TDA can only be applied with DFT; for HF+dTDA, use .xc='hf'")
        mf = mf.to_rhf()
        # commit fc8d1967995b7e033b60d4428ddcca87aac78e4f handles xc='' .
        # xc='0*LDA' is equivalent to xc=''
        #mf.xc = '0.0*LDA'
        mf.xc = ''
        TDA.__init__(self, mf)


def tddft(mf):
    '''Driver to create TDDFT or CasidaTDDFT object'''
    if mf._numint.libxc.is_hybrid_xc(mf.xc):
        return TDDFT(mf)
    else:
        return CasidaTDDFT(mf)

from pyscf import dft
dft.rks.RKS.TDA           = dft.rks_symm.RKS.TDA           = lib.class_as_method(TDA)
dft.rks.RKS.TDHF          = dft.rks_symm.RKS.TDHF          = None
#dft.rks.RKS.TDDFT         = dft.rks_symm.RKS.TDDFT         = lib.class_as_method(TDDFT)
dft.rks.RKS.TDDFTNoHybrid = dft.rks_symm.RKS.TDDFTNoHybrid = lib.class_as_method(TDDFTNoHybrid)
dft.rks.RKS.CasidaTDDFT   = dft.rks_symm.RKS.CasidaTDDFT   = lib.class_as_method(CasidaTDDFT)
dft.rks.RKS.TDDFT         = dft.rks_symm.RKS.TDDFT         = tddft
dft.rks.RKS.dTDA          = dft.rks_symm.RKS.dTDA          = lib.class_as_method(dTDA)
dft.rks.RKS.dRPA          = dft.rks_symm.RKS.dRPA          = lib.class_as_method(dRPA)
dft.roks.ROKS.TDA           = dft.rks_symm.ROKS.TDA           = None
dft.roks.ROKS.TDHF          = dft.rks_symm.ROKS.TDHF          = None
dft.roks.ROKS.TDDFT         = dft.rks_symm.ROKS.TDDFT         = None
dft.roks.ROKS.TDDFTNoHybrid = dft.rks_symm.ROKS.TDDFTNoHybrid = None
dft.roks.ROKS.dTDA          = dft.rks_symm.ROKS.dTDA          = None
dft.roks.ROKS.dRPA          = dft.rks_symm.ROKS.dRPA          = None


# pyscf/tdscf/uhf.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#


import numpy
from pyscf import lib
from pyscf import scf
from pyscf import symm
from pyscf import ao2mo
from pyscf.lib import logger
from pyscf.tdscf import rhf
from pyscf.scf import uhf_symm
from pyscf.scf import _response_functions
from pyscf.data import nist
from pyscf import __config__

OUTPUT_THRESHOLD = getattr(__config__, 'tdscf_rhf_get_nto_threshold', 0.3)
REAL_EIG_THRESHOLD = getattr(__config__, 'tdscf_rhf_TDDFT_pick_eig_threshold', 1e-4)
MO_BASE = getattr(__config__, 'MO_BASE', 1)

# Low excitation filter to avoid numerical instability
POSTIVE_EIG_THRESHOLD = getattr(__config__, 'tdscf_rhf_TDDFT_positive_eig_threshold', 1e-3)


def gen_tda_operation(mf, fock_ao=None, wfnsym=None):
    '''A x

    Kwargs:
        wfnsym : int or str
            Point group symmetry irrep symbol or ID for excited CIS wavefunction.
    '''
    mol = mf.mol
    mo_coeff = mf.mo_coeff
    assert(mo_coeff[0].dtype == numpy.double)
    mo_energy = mf.mo_energy
    mo_occ = mf.mo_occ
    nao, nmo = mo_coeff[0].shape
    occidxa = numpy.where(mo_occ[0]>0)[0]
    occidxb = numpy.where(mo_occ[1]>0)[0]
    viridxa = numpy.where(mo_occ[0]==0)[0]
    viridxb = numpy.where(mo_occ[1]==0)[0]
    nocca = len(occidxa)
    noccb = len(occidxb)
    nvira = len(viridxa)
    nvirb = len(viridxb)
    orboa = mo_coeff[0][:,occidxa]
    orbob = mo_coeff[1][:,occidxb]
    orbva = mo_coeff[0][:,viridxa]
    orbvb = mo_coeff[1][:,viridxb]

    if wfnsym is not None and mol.symmetry:
        if isinstance(wfnsym, str):
            wfnsym = symm.irrep_name2id(mol.groupname, wfnsym)
        orbsyma, orbsymb = uhf_symm.get_orbsym(mol, mo_coeff)
        wfnsym = wfnsym % 10  # convert to D2h subgroup
        orbsyma_in_d2h = numpy.asarray(orbsyma) % 10
        orbsymb_in_d2h = numpy.asarray(orbsymb) % 10
        sym_forbida = (orbsyma_in_d2h[occidxa,None] ^ orbsyma_in_d2h[viridxa]) != wfnsym
        sym_forbidb = (orbsymb_in_d2h[occidxb,None] ^ orbsymb_in_d2h[viridxb]) != wfnsym
        sym_forbid = numpy.hstack((sym_forbida.ravel(), sym_forbidb.ravel()))

    e_ia_a = mo_energy[0][viridxa] - mo_energy[0][occidxa,None]
    e_ia_b = mo_energy[1][viridxb] - mo_energy[1][occidxb,None]
    e_ia = numpy.hstack((e_ia_a.reshape(-1), e_ia_b.reshape(-1)))
    hdiag = e_ia
    if wfnsym is not None and mol.symmetry:
        hdiag[sym_forbid] = 0

    mem_now = lib.current_memory()[0]
    max_memory = max(2000, mf.max_memory*.8-mem_now)
    vresp = mf.gen_response(hermi=0, max_memory=max_memory)

    def vind(zs):
        zs = numpy.asarray(zs)
        if wfnsym is not None and mol.symmetry:
            zs = numpy.copy(zs)
            zs[:,sym_forbid] = 0

        za = zs[:,:nocca*nvira].reshape(-1,nocca,nvira)
        zb = zs[:,nocca*nvira:].reshape(-1,noccb,nvirb)
        dmova = lib.einsum('xov,po,qv->xpq', za, orboa, orbva.conj())
        dmovb = lib.einsum('xov,po,qv->xpq', zb, orbob, orbvb.conj())

        v1ao = vresp(numpy.asarray((dmova,dmovb)))

        v1a = lib.einsum('xpq,po,qv->xov', v1ao[0], orboa.conj(), orbva)
        v1b = lib.einsum('xpq,po,qv->xov', v1ao[1], orbob.conj(), orbvb)
        v1a += numpy.einsum('xia,ia->xia', za, e_ia_a)
        v1b += numpy.einsum('xia,ia->xia', zb, e_ia_b)

        nz = zs.shape[0]
        hx = numpy.hstack((v1a.reshape(nz,-1), v1b.reshape(nz,-1)))
        if wfnsym is not None and mol.symmetry:
            hx[:,sym_forbid] = 0
        return hx

    return vind, hdiag
gen_tda_hop = gen_tda_operation

def get_ab(mf, mo_energy=None, mo_coeff=None, mo_occ=None):
    r'''A and B matrices for TDDFT response function.

    A[i,a,j,b] = \delta_{ab}\delta_{ij}(E_a - E_i) + (ia||bj)
    B[i,a,j,b] = (ia||jb)

    Spin symmetry is considered in the returned A, B lists.  List A has three
    items: (A_aaaa, A_aabb, A_bbbb). A_bbaa = A_aabb.transpose(2,3,0,1).
    B has three items: (B_aaaa, B_aabb, B_bbbb).
    B_bbaa = B_aabb.transpose(2,3,0,1).
    '''
    if mo_energy is None: mo_energy = mf.mo_energy
    if mo_coeff is None: mo_coeff = mf.mo_coeff
    if mo_occ is None: mo_occ = mf.mo_occ

    mol = mf.mol
    nao = mol.nao_nr()
    occidx_a = numpy.where(mo_occ[0]==1)[0]
    viridx_a = numpy.where(mo_occ[0]==0)[0]
    occidx_b = numpy.where(mo_occ[1]==1)[0]
    viridx_b = numpy.where(mo_occ[1]==0)[0]
    orbo_a = mo_coeff[0][:,occidx_a]
    orbv_a = mo_coeff[0][:,viridx_a]
    orbo_b = mo_coeff[1][:,occidx_b]
    orbv_b = mo_coeff[1][:,viridx_b]
    nocc_a = orbo_a.shape[1]
    nvir_a = orbv_a.shape[1]
    nocc_b = orbo_b.shape[1]
    nvir_b = orbv_b.shape[1]
    mo_a = numpy.hstack((orbo_a,orbv_a))
    mo_b = numpy.hstack((orbo_b,orbv_b))
    nmo_a = nocc_a + nvir_a
    nmo_b = nocc_b + nvir_b

    e_ia_a = (mo_energy[0][viridx_a,None] - mo_energy[0][occidx_a]).T
    e_ia_b = (mo_energy[1][viridx_b,None] - mo_energy[1][occidx_b]).T
    a_aa = numpy.diag(e_ia_a.ravel()).reshape(nocc_a,nvir_a,nocc_a,nvir_a)
    a_bb = numpy.diag(e_ia_b.ravel()).reshape(nocc_b,nvir_b,nocc_b,nvir_b)
    a_ab = numpy.zeros((nocc_a,nvir_a,nocc_b,nvir_b))
    b_aa = numpy.zeros_like(a_aa)
    b_ab = numpy.zeros_like(a_ab)
    b_bb = numpy.zeros_like(a_bb)
    a = (a_aa, a_ab, a_bb)
    b = (b_aa, b_ab, b_bb)

    def add_hf_(a, b, hyb=1):
        eri_aa = ao2mo.general(mol, [orbo_a,mo_a,mo_a,mo_a], compact=False)
        eri_ab = ao2mo.general(mol, [orbo_a,mo_a,mo_b,mo_b], compact=False)
        eri_bb = ao2mo.general(mol, [orbo_b,mo_b,mo_b,mo_b], compact=False)
        eri_aa = eri_aa.reshape(nocc_a,nmo_a,nmo_a,nmo_a)
        eri_ab = eri_ab.reshape(nocc_a,nmo_a,nmo_b,nmo_b)
        eri_bb = eri_bb.reshape(nocc_b,nmo_b,nmo_b,nmo_b)
        a_aa, a_ab, a_bb = a
        b_aa, b_ab, b_bb = b

        a_aa += numpy.einsum('iabj->iajb', eri_aa[:nocc_a,nocc_a:,nocc_a:,:nocc_a])
        a_aa -= numpy.einsum('ijba->iajb', eri_aa[:nocc_a,:nocc_a,nocc_a:,nocc_a:]) * hyb
        b_aa += numpy.einsum('iajb->iajb', eri_aa[:nocc_a,nocc_a:,:nocc_a,nocc_a:])
        b_aa -= numpy.einsum('jaib->iajb', eri_aa[:nocc_a,nocc_a:,:nocc_a,nocc_a:]) * hyb

        a_bb += numpy.einsum('iabj->iajb', eri_bb[:nocc_b,nocc_b:,nocc_b:,:nocc_b])
        a_bb -= numpy.einsum('ijba->iajb', eri_bb[:nocc_b,:nocc_b,nocc_b:,nocc_b:]) * hyb
        b_bb += numpy.einsum('iajb->iajb', eri_bb[:nocc_b,nocc_b:,:nocc_b,nocc_b:])
        b_bb -= numpy.einsum('jaib->iajb', eri_bb[:nocc_b,nocc_b:,:nocc_b,nocc_b:]) * hyb

        a_ab += numpy.einsum('iabj->iajb', eri_ab[:nocc_a,nocc_a:,nocc_b:,:nocc_b])
        b_ab += numpy.einsum('iajb->iajb', eri_ab[:nocc_a,nocc_a:,:nocc_b,nocc_b:])

    if isinstance(mf, scf.hf.KohnShamDFT):
        ni = mf._numint
        ni.libxc.test_deriv_order(mf.xc, 2, raise_error=True)
        if getattr(mf, 'nlc', '') != '':
            logger.warn(mf, 'NLC functional found in DFT object.  Its second '
                        'deriviative is not available. Its contribution is '
                        'not included in the response function.')
        omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin)

        add_hf_(a, b, hyb)

        xctype = ni._xc_type(mf.xc)
        dm0 = mf.make_rdm1(mo_coeff, mo_occ)
        make_rho = ni._gen_rho_evaluator(mol, dm0, hermi=1)[0]
        mem_now = lib.current_memory()[0]
        max_memory = max(2000, mf.max_memory*.8-mem_now)

        if xctype == 'LDA':
            ao_deriv = 0
            for ao, mask, weight, coords \
                    in ni.block_loop(mol, mf.grids, nao, ao_deriv, max_memory):
                rho0a = make_rho(0, ao, mask, 'LDA')
                rho0b = make_rho(1, ao, mask, 'LDA')
                fxc = ni.eval_xc(mf.xc, (rho0a,rho0b), 1, deriv=2)[2]
                u_u, u_d, d_d = fxc[0].T

                rho_o_a = lib.einsum('rp,pi->ri', ao, orbo_a)
                rho_v_a = lib.einsum('rp,pi->ri', ao, orbv_a)
                rho_o_b = lib.einsum('rp,pi->ri', ao, orbo_b)
                rho_v_b = lib.einsum('rp,pi->ri', ao, orbv_b)
                rho_ov_a = numpy.einsum('ri,ra->ria', rho_o_a, rho_v_a)
                rho_ov_b = numpy.einsum('ri,ra->ria', rho_o_b, rho_v_b)

                w_ov = numpy.einsum('ria,r->ria', rho_ov_a, weight*u_u)
                iajb = lib.einsum('ria,rjb->iajb', rho_ov_a, w_ov)
                a_aa += iajb
                b_aa += iajb

                w_ov = numpy.einsum('ria,r->ria', rho_ov_b, weight*u_d)
                iajb = lib.einsum('ria,rjb->iajb', rho_ov_a, w_ov)
                a_ab += iajb
                b_ab += iajb

                w_ov = numpy.einsum('ria,r->ria', rho_ov_b, weight*d_d)
                iajb = lib.einsum('ria,rjb->iajb', rho_ov_b, w_ov)
                a_bb += iajb
                b_bb += iajb

        elif xctype == 'GGA':
            ao_deriv = 1
            for ao, mask, weight, coords \
                    in ni.block_loop(mol, mf.grids, nao, ao_deriv, max_memory):
                rho0a = make_rho(0, ao, mask, 'GGA')
                rho0b = make_rho(1, ao, mask, 'GGA')
                vxc, fxc = ni.eval_xc(mf.xc, (rho0a,rho0b), 1, deriv=2)[1:3]
                uu, ud, dd = vxc[1].T
                u_u, u_d, d_d = fxc[0].T
                u_uu, u_ud, u_dd, d_uu, d_ud, d_dd = fxc[1].T
                uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd = fxc[2].T

                rho_o_a = lib.einsum('xrp,pi->xri', ao, orbo_a)
                rho_v_a = lib.einsum('xrp,pi->xri', ao, orbv_a)
                rho_o_b = lib.einsum('xrp,pi->xri', ao, orbo_b)
                rho_v_b = lib.einsum('xrp,pi->xri', ao, orbv_b)
                rho_ov_a = numpy.einsum('xri,ra->xria', rho_o_a, rho_v_a[0])
                rho_ov_b = numpy.einsum('xri,ra->xria', rho_o_b, rho_v_b[0])
                rho_ov_a[1:4] += numpy.einsum('ri,xra->xria', rho_o_a[0], rho_v_a[1:4])
                rho_ov_b[1:4] += numpy.einsum('ri,xra->xria', rho_o_b[0], rho_v_b[1:4])
                # sigma1 ~ \nabla(\rho_\alpha+\rho_\beta) dot \nabla(|b><j|) z_{bj}
                a0a1 = numpy.einsum('xr,xria->ria', rho0a[1:4], rho_ov_a[1:4])
                a0b1 = numpy.einsum('xr,xria->ria', rho0a[1:4], rho_ov_b[1:4])
                b0a1 = numpy.einsum('xr,xria->ria', rho0b[1:4], rho_ov_a[1:4])
                b0b1 = numpy.einsum('xr,xria->ria', rho0b[1:4], rho_ov_b[1:4])

                w_ov = numpy.empty_like(rho_ov_a)
                w_ov[0]  = numpy.einsum('r,ria->ria', u_u, rho_ov_a[0])
                w_ov[0] += numpy.einsum('r,ria->ria', 2*u_uu, a0a1)
                w_ov[0] += numpy.einsum('r,ria->ria',   u_ud, b0a1)
                f_ov_a = numpy.einsum('r,ria->ria', 4*uu_uu, a0a1)
                f_ov_b = numpy.einsum('r,ria->ria', 2*uu_ud, a0a1)
                f_ov_a+= numpy.einsum('r,ria->ria', 2*uu_ud, b0a1)
                f_ov_b+= numpy.einsum('r,ria->ria',   ud_ud, b0a1)
                f_ov_a+= numpy.einsum('r,ria->ria', 2*u_uu, rho_ov_a[0])
                f_ov_b+= numpy.einsum('r,ria->ria',   u_ud, rho_ov_a[0])
                w_ov[1:] = numpy.einsum('ria,xr->xria', f_ov_a, rho0a[1:4])
                w_ov[1:]+= numpy.einsum('ria,xr->xria', f_ov_b, rho0b[1:4])
                w_ov[1:]+= numpy.einsum('r,xria->xria', 2*uu, rho_ov_a[1:4])
                w_ov *= weight[:,None,None]
                iajb = lib.einsum('xria,xrjb->iajb', rho_ov_a, w_ov)
                a_aa += iajb
                b_aa += iajb

                w_ov = numpy.empty_like(rho_ov_b)
                w_ov[0]  = numpy.einsum('r,ria->ria', d_d, rho_ov_b[0])
                w_ov[0] += numpy.einsum('r,ria->ria', 2*d_dd, b0b1)
                w_ov[0] += numpy.einsum('r,ria->ria',   d_ud, a0b1)
                f_ov_b = numpy.einsum('r,ria->ria', 4*dd_dd, b0b1)
                f_ov_a = numpy.einsum('r,ria->ria', 2*ud_dd, b0b1)
                f_ov_b+= numpy.einsum('r,ria->ria', 2*ud_dd, a0b1)
                f_ov_a+= numpy.einsum('r,ria->ria',   ud_ud, a0b1)
                f_ov_b+= numpy.einsum('r,ria->ria', 2*d_dd, rho_ov_b[0])
                f_ov_a+= numpy.einsum('r,ria->ria',   d_ud, rho_ov_b[0])
                w_ov[1:] = numpy.einsum('ria,xr->xria', f_ov_a, rho0a[1:4])
                w_ov[1:]+= numpy.einsum('ria,xr->xria', f_ov_b, rho0b[1:4])
                w_ov[1:]+= numpy.einsum('r,xria->xria', 2*dd, rho_ov_b[1:4])
                w_ov *= weight[:,None,None]
                iajb = lib.einsum('xria,xrjb->iajb', rho_ov_b, w_ov)
                a_bb += iajb
                b_bb += iajb

                w_ov = numpy.empty_like(rho_ov_b)
                w_ov[0]  = numpy.einsum('r,ria->ria', u_d, rho_ov_b[0])
                w_ov[0] += numpy.einsum('r,ria->ria', 2*u_dd, b0b1)
                w_ov[0] += numpy.einsum('r,ria->ria',   u_ud, a0b1)
                f_ov_a = numpy.einsum('r,ria->ria', 4*uu_dd, b0b1)
                f_ov_b = numpy.einsum('r,ria->ria', 2*ud_dd, b0b1)
                f_ov_a+= numpy.einsum('r,ria->ria', 2*uu_ud, a0b1)
                f_ov_b+= numpy.einsum('r,ria->ria',   ud_ud, a0b1)
                f_ov_a+= numpy.einsum('r,ria->ria', 2*d_uu, rho_ov_b[0])
                f_ov_b+= numpy.einsum('r,ria->ria',   d_ud, rho_ov_b[0])
                w_ov[1:] = numpy.einsum('ria,xr->xria', f_ov_a, rho0a[1:4])
                w_ov[1:]+= numpy.einsum('ria,xr->xria', f_ov_b, rho0b[1:4])
                w_ov[1:]+= numpy.einsum('r,xria->xria', ud, rho_ov_b[1:4])
                w_ov *= weight[:,None,None]
                iajb = lib.einsum('xria,xrjb->iajb', rho_ov_a, w_ov)
                a_ab += iajb
                b_ab += iajb

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

    else:
        add_hf_(a, b)

    return a, b

def get_nto(tdobj, state=1, threshold=OUTPUT_THRESHOLD, verbose=None):
    r'''
    Natural transition orbital analysis.

    The natural transition density matrix between ground state and excited
    state :math:`Tia = \langle \Psi_{ex} | i a^\dagger | \Psi_0 \rangle` can
    be transformed to diagonal form through SVD
    :math:`T = O \sqrt{\lambda} V^\dagger`. O and V are occupied and virtual
    natural transition orbitals. The diagonal elements :math:`\lambda` are the
    weights of the occupied-virtual orbital pair in the excitation.

    Ref: Martin, R. L., JCP, 118, 4775-4777

    Note in the TDHF/TDDFT calculations, the excitation part (X) is
    interpreted as the CIS coefficients and normalized to 1. The de-excitation
    part (Y) is ignored.

    Args:
        state : int
            Excited state ID.  state = 1 means the first excited state.
            If state < 0, state ID is counted from the last excited state.

    Kwargs:
        threshold : float
            Above which the NTO coefficients will be printed in the output.

    Returns:
        A list (weights, NTOs).  NTOs are natural orbitals represented in AO
        basis. The first N_occ NTOs are occupied NTOs and the rest are virtual
        NTOs.
    '''
    if state == 0:
        logger.warn(tdobj, 'Excited state starts from 1. '
                    'Set state=1 for first excited state.')
        state_id = state
    elif state < 0:
        state_id = state
    else:
        state_id = state - 1

    mol = tdobj.mol
    mo_coeff = tdobj._scf.mo_coeff
    mo_occ = tdobj._scf.mo_occ
    orbo_a = mo_coeff[0][:,mo_occ[0]==1]
    orbv_a = mo_coeff[0][:,mo_occ[0]==0]
    orbo_b = mo_coeff[1][:,mo_occ[1]==1]
    orbv_b = mo_coeff[1][:,mo_occ[1]==0]
    nocc_a = orbo_a.shape[1]
    nvir_a = orbv_a.shape[1]
    nocc_b = orbo_b.shape[1]
    nvir_b = orbv_b.shape[1]

    cis_t1a, cis_t1b = tdobj.xy[state_id][0]
    norm = numpy.linalg.norm(cis_t1a)**2 + numpy.linalg.norm(cis_t1b)**2
    cis_t1a *= 1. / norm
    cis_t1b *= 1. / norm

    if mol.symmetry:
        orbsyma, orbsymb = uhf_symm.get_orbsym(mol, mo_coeff)
        o_sym_a = orbsyma[mo_occ[0]==1]
        v_sym_a = orbsyma[mo_occ[0]==0]
        o_sym_b = orbsymb[mo_occ[1]==1]
        v_sym_b = orbsymb[mo_occ[1]==0]
        nto_o_a = numpy.eye(nocc_a)
        nto_v_a = numpy.eye(nvir_a)
        nto_o_b = numpy.eye(nocc_b)
        nto_v_b = numpy.eye(nvir_b)
        weights_o_a = numpy.zeros(nocc_a)
        weights_v_a = numpy.zeros(nvir_a)
        weights_o_b = numpy.zeros(nocc_b)
        weights_v_b = numpy.zeros(nvir_b)

        for ir in set(orbsyma):
            idx = numpy.where(o_sym_a == ir)[0]
            if idx.size > 0:
                dm_oo = numpy.dot(cis_t1a[idx], cis_t1a[idx].T)
                weights_o_a[idx], nto_o_a[idx[:,None],idx] = numpy.linalg.eigh(dm_oo)

            idx = numpy.where(v_sym_a == ir)[0]
            if idx.size > 0:
                dm_vv = numpy.dot(cis_t1a[:,idx].T, cis_t1a[:,idx])
                weights_v_a[idx], nto_v_a[idx[:,None],idx] = numpy.linalg.eigh(dm_vv)

        for ir in set(orbsymb):
            idx = numpy.where(o_sym_b == ir)[0]
            if idx.size > 0:
                dm_oo = numpy.dot(cis_t1b[idx], cis_t1b[idx].T)
                weights_o_b[idx], nto_o_b[idx[:,None],idx] = numpy.linalg.eigh(dm_oo)

            idx = numpy.where(v_sym_b == ir)[0]
            if idx.size > 0:
                dm_vv = numpy.dot(cis_t1b[:,idx].T, cis_t1b[:,idx])
                weights_v_b[idx], nto_v_b[idx[:,None],idx] = numpy.linalg.eigh(dm_vv)

        def sort(weights, nto, sym):
            # weights in descending order
            idx = numpy.argsort(-weights)
            weights = weights[idx]
            nto = nto[:,idx]
            sym = sym[idx]
            return weights, nto, sym

        weights_o_a, nto_o_a, o_sym_a = sort(weights_o_a, nto_o_a, o_sym_a)
        weights_v_a, nto_v_a, v_sym_a = sort(weights_v_a, nto_v_a, v_sym_a)
        weights_o_b, nto_o_b, o_sym_b = sort(weights_o_b, nto_o_b, o_sym_b)
        weights_v_b, nto_v_b, v_sym_b = sort(weights_v_b, nto_v_b, v_sym_b)

        nto_orbsyma = numpy.hstack((o_sym_a, v_sym_a))
        nto_orbsymb = numpy.hstack((o_sym_b, v_sym_b))

        if nocc_a < nvir_a:
            weights_a = weights_o_a
        else:
            weights_a = weights_v_a
        if nocc_b < nvir_b:
            weights_b = weights_o_b
        else:
            weights_b = weights_v_b

    else:
        nto_o_a, w_a, nto_v_aT = numpy.linalg.svd(cis_t1a)
        nto_o_b, w_b, nto_v_bT = numpy.linalg.svd(cis_t1b)
        nto_v_a = nto_v_aT.conj().T
        nto_v_b = nto_v_bT.conj().T
        weights_a = w_a**2
        weights_b = w_b**2
        nto_orbsyma = nto_orbsymb = None

    def _set_phase_(c):
        idx = numpy.argmax(abs(c.real), axis=0)
        c[:,c[idx,numpy.arange(c.shape[1])].real<0] *= -1
    _set_phase_(nto_o_a)
    _set_phase_(nto_o_b)
    _set_phase_(nto_v_a)
    _set_phase_(nto_v_b)

    occupied_nto_a = numpy.dot(orbo_a, nto_o_a)
    occupied_nto_b = numpy.dot(orbo_b, nto_o_b)
    virtual_nto_a = numpy.dot(orbv_a, nto_v_a)
    virtual_nto_b = numpy.dot(orbv_b, nto_v_b)
    nto_coeff = (numpy.hstack((occupied_nto_a, virtual_nto_a)),
                 numpy.hstack((occupied_nto_b, virtual_nto_b)))

    if mol.symmetry:
        nto_coeff = (lib.tag_array(nto_coeff[0], orbsym=nto_orbsyma),
                     lib.tag_array(nto_coeff[1], orbsym=nto_orbsymb))

    log = logger.new_logger(tdobj, verbose)
    if log.verbose >= logger.INFO:
        log.info('State %d: %g eV  NTO largest component %s',
                 state_id+1, tdobj.e[state_id]*nist.HARTREE2EV,
                 weights_a[0]+weights_b[0])
        fmt = '%' + str(lib.param.OUTPUT_DIGITS) + 'f (MO #%d)'
        o_idx_a = numpy.where(abs(nto_o_a[:,0]) > threshold)[0]
        v_idx_a = numpy.where(abs(nto_v_a[:,0]) > threshold)[0]
        o_idx_b = numpy.where(abs(nto_o_b[:,0]) > threshold)[0]
        v_idx_b = numpy.where(abs(nto_v_b[:,0]) > threshold)[0]
        log.info('    alpha occ-NTO: ' +
                 ' + '.join([(fmt % (nto_o_a[i,0], i+MO_BASE))
                             for i in o_idx_a]))
        log.info('    alpha vir-NTO: ' +
                 ' + '.join([(fmt % (nto_v_a[i,0], i+MO_BASE+nocc_a))
                             for i in v_idx_a]))
        log.info('    beta occ-NTO: ' +
                 ' + '.join([(fmt % (nto_o_b[i,0], i+MO_BASE))
                             for i in o_idx_b]))
        log.info('    beta vir-NTO: ' +
                 ' + '.join([(fmt % (nto_v_b[i,0], i+MO_BASE+nocc_b))
                             for i in v_idx_b]))
    return (weights_a, weights_b), nto_coeff


def analyze(tdobj, verbose=None):
    log = logger.new_logger(tdobj, verbose)
    mol = tdobj.mol
    mo_coeff = tdobj._scf.mo_coeff
    mo_occ = tdobj._scf.mo_occ
    nocc_a = numpy.count_nonzero(mo_occ[0] == 1)
    nocc_b = numpy.count_nonzero(mo_occ[1] == 1)

    e_ev = numpy.asarray(tdobj.e) * nist.HARTREE2EV
    e_wn = numpy.asarray(tdobj.e) * nist.HARTREE2WAVENUMBER
    wave_length = 1e7/e_wn

    log.note('\n** Excitation energies and oscillator strengths **')

    if mol.symmetry:
        orbsyma, orbsymb = uhf_symm.get_orbsym(mol, mo_coeff)
        x_syma = symm.direct_prod(orbsyma[mo_occ[0]==1], orbsyma[mo_occ[0]==0], mol.groupname)
        x_symb = symm.direct_prod(orbsymb[mo_occ[1]==1], orbsymb[mo_occ[1]==0], mol.groupname)
    else:
        x_syma = None

    f_oscillator = tdobj.oscillator_strength()
    for i, ei in enumerate(tdobj.e):
        x, y = tdobj.xy[i]
        if x_syma is None:
            log.note('Excited State %3d: %12.5f eV %9.2f nm  f=%.4f',
                     i+1, e_ev[i], wave_length[i], f_oscillator[i])
        else:
            wfnsyma = rhf.analyze_wfnsym(tdobj, x_syma, x[0])
            wfnsymb = rhf.analyze_wfnsym(tdobj, x_symb, x[1])
            if wfnsyma == wfnsymb:
                wfnsym = wfnsyma
            else:
                wfnsym = '???'
            log.note('Excited State %3d: %4s %12.5f eV %9.2f nm  f=%.4f',
                     i+1, wfnsym, e_ev[i], wave_length[i], f_oscillator[i])

        if log.verbose >= logger.INFO:
            for o, v in zip(* numpy.where(abs(x[0]) > 0.1)):
                log.info('    %4da -> %4da %12.5f',
                         o+MO_BASE, v+MO_BASE+nocc_a, x[0][o,v])
            for o, v in zip(* numpy.where(abs(x[1]) > 0.1)):
                log.info('    %4db -> %4db %12.5f',
                         o+MO_BASE, v+MO_BASE+nocc_b, x[1][o,v])

    if log.verbose >= logger.INFO:
        log.info('\n** Transition electric dipole moments (AU) **')
        log.info('state          X           Y           Z        Dip. S.      Osc.')
        trans_dip = tdobj.transition_dipole()
        for i, ei in enumerate(tdobj.e):
            dip = trans_dip[i]
            log.info('%3d    %11.4f %11.4f %11.4f %11.4f %11.4f',
                     i+1, dip[0], dip[1], dip[2], numpy.dot(dip, dip),
                     f_oscillator[i])

        log.info('\n** Transition velocity dipole moments (imaginary part, AU) **')
        log.info('state          X           Y           Z        Dip. S.      Osc.')
        trans_v = tdobj.transition_velocity_dipole()
        f_v = tdobj.oscillator_strength(gauge='velocity', order=0)
        for i, ei in enumerate(tdobj.e):
            v = trans_v[i]
            log.info('%3d    %11.4f %11.4f %11.4f %11.4f %11.4f',
                     i+1, v[0], v[1], v[2], numpy.dot(v, v), f_v[i])

        log.info('\n** Transition magnetic dipole moments (imaginary part, AU) **')
        log.info('state          X           Y           Z')
        trans_m = tdobj.transition_magnetic_dipole()
        for i, ei in enumerate(tdobj.e):
            m = trans_m[i]
            log.info('%3d    %11.4f %11.4f %11.4f',
                     i+1, m[0], m[1], m[2])
    return tdobj


def _contract_multipole(tdobj, ints, hermi=True, xy=None):
    if xy is None: xy = tdobj.xy
    mo_coeff = tdobj._scf.mo_coeff
    mo_occ = tdobj._scf.mo_occ
    orbo_a = mo_coeff[0][:,mo_occ[0]==1]
    orbv_a = mo_coeff[0][:,mo_occ[0]==0]
    orbo_b = mo_coeff[1][:,mo_occ[1]==1]
    orbv_b = mo_coeff[1][:,mo_occ[1]==0]

    ints_a = numpy.einsum('...pq,pi,qj->...ij', ints, orbo_a.conj(), orbv_a)
    ints_b = numpy.einsum('...pq,pi,qj->...ij', ints, orbo_b.conj(), orbv_b)
    pol = [(numpy.einsum('...ij,ij->...', ints_a, x[0]) +
            numpy.einsum('...ij,ij->...', ints_b, x[1])) for x,y in xy]
    pol = numpy.array(pol)
    y = xy[0][1]
    if isinstance(y[0], numpy.ndarray):
        pol_y = [(numpy.einsum('...ij,ij->...', ints_a, y[0]) +
                  numpy.einsum('...ij,ij->...', ints_b, y[1])) for x,y in xy]
        if hermi:
            pol += pol_y
        else:  # anti-Hermitian
            pol -= pol_y
    return pol


class TDMixin(rhf.TDMixin):

    @lib.with_doc(get_ab.__doc__)
    def get_ab(self, mf=None):
        if mf is None: mf = self._scf
        return get_ab(mf)

    analyze = analyze
    get_nto = get_nto
    _contract_multipole = _contract_multipole  # needed by transition dipoles

    def nuc_grad_method(self):
        from pyscf.grad import tduhf
        return tduhf.Gradients(self)


@lib.with_doc(rhf.TDA.__doc__)
class TDA(TDMixin):

    singlet = None

    def gen_vind(self, mf):
        '''Compute Ax'''
        return gen_tda_hop(mf, wfnsym=self.wfnsym)

    def init_guess(self, mf, nstates=None, wfnsym=None):
        if nstates is None: nstates = self.nstates
        if wfnsym is None: wfnsym = self.wfnsym

        mol = mf.mol
        mo_energy = mf.mo_energy
        mo_occ = mf.mo_occ
        occidxa = numpy.where(mo_occ[0]>0)[0]
        occidxb = numpy.where(mo_occ[1]>0)[0]
        viridxa = numpy.where(mo_occ[0]==0)[0]
        viridxb = numpy.where(mo_occ[1]==0)[0]
        e_ia_a = (mo_energy[0][viridxa,None] - mo_energy[0][occidxa]).T
        e_ia_b = (mo_energy[1][viridxb,None] - mo_energy[1][occidxb]).T

        if wfnsym is not None and mol.symmetry:
            if isinstance(wfnsym, str):
                wfnsym = symm.irrep_name2id(mol.groupname, wfnsym)
            orbsyma, orbsymb = uhf_symm.get_orbsym(mol, mf.mo_coeff)
            wfnsym = wfnsym % 10  # convert to D2h subgroup
            orbsyma_in_d2h = numpy.asarray(orbsyma) % 10
            orbsymb_in_d2h = numpy.asarray(orbsymb) % 10
            e_ia_a[(orbsyma_in_d2h[occidxa,None] ^ orbsyma_in_d2h[viridxa]) != wfnsym] = 1e99
            e_ia_b[(orbsymb_in_d2h[occidxb,None] ^ orbsymb_in_d2h[viridxb]) != wfnsym] = 1e99

        e_ia = numpy.hstack((e_ia_a.ravel(), e_ia_b.ravel()))
        e_ia_max = e_ia.max()
        nov = e_ia.size
        nstates = min(nstates, nov)
        e_threshold = min(e_ia_max, e_ia[numpy.argsort(e_ia)[nstates-1]])
        # Handle degeneracy
        e_threshold += 1e-6

        idx = numpy.where(e_ia <= e_threshold)[0]
        x0 = numpy.zeros((idx.size, nov))
        for i, j in enumerate(idx):
            x0[i, j] = 1  # Koopmans' excitations
        return x0

    def kernel(self, x0=None, nstates=None):
        '''TDA diagonalization solver
        '''
        cpu0 = (logger.process_clock(), logger.perf_counter())
        self.check_sanity()
        self.dump_flags()
        if nstates is None:
            nstates = self.nstates
        else:
            self.nstates = nstates
        log = logger.Logger(self.stdout, self.verbose)

        vind, hdiag = self.gen_vind(self._scf)
        precond = self.get_precond(hdiag)
        if x0 is None:
            x0 = self.init_guess(self._scf, self.nstates)

        def pickeig(w, v, nroots, envs):
            idx = numpy.where(w > POSTIVE_EIG_THRESHOLD)[0]
            return w[idx], v[:,idx], idx

        self.converged, self.e, x1 = \
                lib.davidson1(vind, x0, precond,
                              tol=self.conv_tol,
                              nroots=nstates, lindep=self.lindep,
                              max_cycle=self.max_cycle,
                              max_space=self.max_space, pick=pickeig,
                              verbose=log)

        nmo = self._scf.mo_occ[0].size
        nocca = (self._scf.mo_occ[0]>0).sum()
        noccb = (self._scf.mo_occ[1]>0).sum()
        nvira = nmo - nocca
        nvirb = nmo - noccb
        self.xy = [((xi[:nocca*nvira].reshape(nocca,nvira),  # X_alpha
                     xi[nocca*nvira:].reshape(noccb,nvirb)), # X_beta
                    (0, 0))  # (Y_alpha, Y_beta)
                   for xi in x1]

        if self.chkfile:
            lib.chkfile.save(self.chkfile, 'tddft/e', self.e)
            lib.chkfile.save(self.chkfile, 'tddft/xy', self.xy)

        log.timer('TDA', *cpu0)
        self._finalize()
        return self.e, self.xy

CIS = TDA


def gen_tdhf_operation(mf, fock_ao=None, singlet=True, wfnsym=None):
    '''Generate function to compute

    [ A  B][X]
    [-B -A][Y]
    '''
    mol = mf.mol
    mo_coeff = mf.mo_coeff
    assert(mo_coeff[0].dtype == numpy.double)
    mo_energy = mf.mo_energy
    mo_occ = mf.mo_occ
    nao, nmo = mo_coeff[0].shape
    occidxa = numpy.where(mo_occ[0]>0)[0]
    occidxb = numpy.where(mo_occ[1]>0)[0]
    viridxa = numpy.where(mo_occ[0]==0)[0]
    viridxb = numpy.where(mo_occ[1]==0)[0]
    nocca = len(occidxa)
    noccb = len(occidxb)
    nvira = len(viridxa)
    nvirb = len(viridxb)
    orboa = mo_coeff[0][:,occidxa]
    orbob = mo_coeff[1][:,occidxb]
    orbva = mo_coeff[0][:,viridxa]
    orbvb = mo_coeff[1][:,viridxb]

    if wfnsym is not None and mol.symmetry:
        if isinstance(wfnsym, str):
            wfnsym = symm.irrep_name2id(mol.groupname, wfnsym)
        orbsyma, orbsymb = uhf_symm.get_orbsym(mol, mo_coeff)
        wfnsym = wfnsym % 10  # convert to D2h subgroup
        orbsyma_in_d2h = numpy.asarray(orbsyma) % 10
        orbsymb_in_d2h = numpy.asarray(orbsymb) % 10
        sym_forbida = (orbsyma_in_d2h[occidxa,None] ^ orbsyma_in_d2h[viridxa]) != wfnsym
        sym_forbidb = (orbsymb_in_d2h[occidxb,None] ^ orbsymb_in_d2h[viridxb]) != wfnsym
        sym_forbid = numpy.hstack((sym_forbida.ravel(), sym_forbidb.ravel()))

    e_ia_a = mo_energy[0][viridxa] - mo_energy[0][occidxa,None]
    e_ia_b = mo_energy[1][viridxb] - mo_energy[1][occidxb,None]
    e_ia = hdiag = numpy.hstack((e_ia_a.ravel(), e_ia_b.ravel()))
    if wfnsym is not None and mol.symmetry:
        hdiag[sym_forbid] = 0
    hdiag = numpy.hstack((hdiag, -hdiag))

    mem_now = lib.current_memory()[0]
    max_memory = max(2000, mf.max_memory*.8-mem_now)
    vresp = mf.gen_response(hermi=0, max_memory=max_memory)

    def vind(xys):
        nz = len(xys)
        xys = numpy.asarray(xys).reshape(nz,2,-1)
        if wfnsym is not None and mol.symmetry:
            # shape(nz,2,-1): 2 ~ X,Y
            xys = numpy.copy(xys)
            xys[:,:,sym_forbid] = 0

        xs, ys = xys.transpose(1,0,2)
        xa = xs[:,:nocca*nvira].reshape(nz,nocca,nvira)
        xb = xs[:,nocca*nvira:].reshape(nz,noccb,nvirb)
        ya = ys[:,:nocca*nvira].reshape(nz,nocca,nvira)
        yb = ys[:,nocca*nvira:].reshape(nz,noccb,nvirb)
        # dms = AX + BY
        dmsa  = lib.einsum('xov,po,qv->xpq', xa, orboa, orbva.conj())
        dmsa += lib.einsum('xov,pv,qo->xpq', ya, orbva, orboa.conj())
        dmsb  = lib.einsum('xov,po,qv->xpq', xb, orbob, orbvb.conj())
        dmsb += lib.einsum('xov,pv,qo->xpq', yb, orbvb, orbob.conj())

        v1ao = vresp(numpy.asarray((dmsa,dmsb)))

        v1aov = lib.einsum('xpq,po,qv->xov', v1ao[0], orboa.conj(), orbva)
        v1bov = lib.einsum('xpq,po,qv->xov', v1ao[1], orbob.conj(), orbvb)
        v1avo = lib.einsum('xpq,pv,qo->xov', v1ao[0], orbva.conj(), orboa)
        v1bvo = lib.einsum('xpq,pv,qo->xov', v1ao[1], orbvb.conj(), orbob)

        v1ov = xs * e_ia  # AX
        v1vo = ys * e_ia  # AY
        v1ov[:,:nocca*nvira] += v1aov.reshape(nz,-1)
        v1vo[:,:nocca*nvira] += v1avo.reshape(nz,-1)
        v1ov[:,nocca*nvira:] += v1bov.reshape(nz,-1)
        v1vo[:,nocca*nvira:] += v1bvo.reshape(nz,-1)
        if wfnsym is not None and mol.symmetry:
            v1ov[:,sym_forbid] = 0
            v1vo[:,sym_forbid] = 0
        hx = numpy.hstack((v1ov, -v1vo))
        return hx

    return vind, hdiag


class TDHF(TDMixin):

    singlet = None

    @lib.with_doc(gen_tdhf_operation.__doc__)
    def gen_vind(self, mf):
        return gen_tdhf_operation(mf, singlet=self.singlet, wfnsym=self.wfnsym)

    def init_guess(self, mf, nstates=None, wfnsym=None):
        x0 = TDA.init_guess(self, mf, nstates, wfnsym)
        y0 = numpy.zeros_like(x0)
        return numpy.hstack((x0,y0))

    def kernel(self, x0=None, nstates=None):
        '''TDHF diagonalization with non-Hermitian eigenvalue solver
        '''
        cpu0 = (logger.process_clock(), logger.perf_counter())
        self.check_sanity()
        self.dump_flags()
        if nstates is None:
            nstates = self.nstates
        else:
            self.nstates = nstates

        log = logger.Logger(self.stdout, self.verbose)

        vind, hdiag = self.gen_vind(self._scf)
        precond = self.get_precond(hdiag)
        if x0 is None:
            x0 = self.init_guess(self._scf, self.nstates)

        # We only need positive eigenvalues
        def pickeig(w, v, nroots, envs):
            realidx = numpy.where((abs(w.imag) < REAL_EIG_THRESHOLD) &
                                  (w.real > POSTIVE_EIG_THRESHOLD))[0]
            return lib.linalg_helper._eigs_cmplx2real(w, v, realidx,
                                                      real_eigenvectors=True)

        self.converged, w, x1 = \
                lib.davidson_nosym1(vind, x0, precond,
                                    tol=self.conv_tol,
                                    nroots=nstates, lindep=self.lindep,
                                    max_cycle=self.max_cycle,
                                    max_space=self.max_space, pick=pickeig,
                                    verbose=log)

        nmo = self._scf.mo_occ[0].size
        nocca = (self._scf.mo_occ[0]>0).sum()
        noccb = (self._scf.mo_occ[1]>0).sum()
        nvira = nmo - nocca
        nvirb = nmo - noccb
        e = []
        xy = []
        for i, z in enumerate(x1):
            x, y = z.reshape(2,-1)
            norm = lib.norm(x)**2 - lib.norm(y)**2
            if norm > 0:
                norm = 1/numpy.sqrt(norm)
                e.append(w[i])
                xy.append(((x[:nocca*nvira].reshape(nocca,nvira) * norm,  # X_alpha
                            x[nocca*nvira:].reshape(noccb,nvirb) * norm), # X_beta
                           (y[:nocca*nvira].reshape(nocca,nvira) * norm,  # Y_alpha
                            y[nocca*nvira:].reshape(noccb,nvirb) * norm)))# Y_beta
        self.e = numpy.array(e)
        self.xy = xy

        if self.chkfile:
            lib.chkfile.save(self.chkfile, 'tddft/e', self.e)
            lib.chkfile.save(self.chkfile, 'tddft/xy', self.xy)

        log.timer('TDDFT', *cpu0)
        self._finalize()
        return self.e, self.xy

RPA = TDUHF = TDHF

scf.uhf.UHF.TDA = lib.class_as_method(TDA)
scf.uhf.UHF.TDHF = lib.class_as_method(TDHF)

del(OUTPUT_THRESHOLD)


# pyscf/tdscf/uks.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#


import numpy
from pyscf import symm
from pyscf import lib
from pyscf.lib import logger
from pyscf.tdscf import uhf
from pyscf.scf import uhf_symm
from pyscf.data import nist
from pyscf.dft.rks import KohnShamDFT
from pyscf import __config__

# Low excitation filter to avoid numerical instability
POSTIVE_EIG_THRESHOLD = getattr(__config__, 'tdscf_rhf_TDDFT_positive_eig_threshold', 1e-3)


class TDA(uhf.TDA):
    def nuc_grad_method(self):
        from pyscf.grad import tduks
        return tduks.Gradients(self)

class TDDFT(uhf.TDHF):
    def nuc_grad_method(self):
        from pyscf.grad import tduks
        return tduks.Gradients(self)

RPA = TDUKS = TDDFT


class CasidaTDDFT(TDDFT, TDA):
    '''Solve the Casida TDDFT formula (A-B)(A+B)(X+Y) = (X+Y)w^2
    '''

    init_guess = TDA.init_guess

    def get_vind(self, mf):
        wfnsym = self.wfnsym

        mol = mf.mol
        mo_coeff = mf.mo_coeff
        assert mo_coeff[0].dtype == numpy.double
        mo_energy = mf.mo_energy
        mo_occ = mf.mo_occ
        nao, nmo = mo_coeff[0].shape
        occidxa = numpy.where(mo_occ[0]>0)[0]
        occidxb = numpy.where(mo_occ[1]>0)[0]
        viridxa = numpy.where(mo_occ[0]==0)[0]
        viridxb = numpy.where(mo_occ[1]==0)[0]
        nocca = len(occidxa)
        noccb = len(occidxb)
        nvira = len(viridxa)
        nvirb = len(viridxb)
        orboa = mo_coeff[0][:,occidxa]
        orbob = mo_coeff[1][:,occidxb]
        orbva = mo_coeff[0][:,viridxa]
        orbvb = mo_coeff[1][:,viridxb]

        if wfnsym is not None and mol.symmetry:
            if isinstance(wfnsym, str):
                wfnsym = symm.irrep_name2id(mol.groupname, wfnsym)
            orbsyma, orbsymb = uhf_symm.get_orbsym(mol, mo_coeff)
            wfnsym = wfnsym % 10  # convert to D2h subgroup
            orbsyma_in_d2h = numpy.asarray(orbsyma) % 10
            orbsymb_in_d2h = numpy.asarray(orbsymb) % 10
            sym_forbida = (orbsyma_in_d2h[occidxa,None] ^ orbsyma_in_d2h[viridxa]) != wfnsym
            sym_forbidb = (orbsymb_in_d2h[occidxb,None] ^ orbsymb_in_d2h[viridxb]) != wfnsym
            sym_forbid = numpy.hstack((sym_forbida.ravel(), sym_forbidb.ravel()))

        e_ia_a = (mo_energy[0][viridxa,None] - mo_energy[0][occidxa]).T
        e_ia_b = (mo_energy[1][viridxb,None] - mo_energy[1][occidxb]).T
        e_ia = numpy.hstack((e_ia_a.reshape(-1), e_ia_b.reshape(-1)))
        if wfnsym is not None and mol.symmetry:
            e_ia[sym_forbid] = 0
        d_ia = numpy.sqrt(e_ia).ravel()
        ed_ia = e_ia.ravel() * d_ia
        hdiag = e_ia.ravel() ** 2

        vresp = mf.gen_response(mo_coeff, mo_occ, hermi=1)

        def vind(zs):
            nz = len(zs)
            zs = numpy.asarray(zs).reshape(nz,-1)
            if wfnsym is not None and mol.symmetry:
                zs = numpy.copy(zs)
                zs[:,sym_forbid] = 0

            dmsa = (zs[:,:nocca*nvira] * d_ia[:nocca*nvira]).reshape(nz,nocca,nvira)
            dmsb = (zs[:,nocca*nvira:] * d_ia[nocca*nvira:]).reshape(nz,noccb,nvirb)
            dmsa = lib.einsum('xov,po,qv->xpq', dmsa, orboa, orbva.conj())
            dmsb = lib.einsum('xov,po,qv->xpq', dmsb, orbob, orbvb.conj())
            dmsa = dmsa + dmsa.conj().transpose(0,2,1)
            dmsb = dmsb + dmsb.conj().transpose(0,2,1)

            v1ao = vresp(numpy.asarray((dmsa,dmsb)))

            v1a = lib.einsum('xpq,po,qv->xov', v1ao[0], orboa.conj(), orbva)
            v1b = lib.einsum('xpq,po,qv->xov', v1ao[1], orbob.conj(), orbvb)

            hx = numpy.hstack((v1a.reshape(nz,-1), v1b.reshape(nz,-1)))
            hx += ed_ia * zs
            hx *= d_ia
            return hx

        return vind, hdiag

    def kernel(self, x0=None, nstates=None):
        '''TDDFT diagonalization solver
        '''
        cpu0 = (logger.process_clock(), logger.perf_counter())
        mf = self._scf
        if mf._numint.libxc.is_hybrid_xc(mf.xc):
            raise RuntimeError('%s cannot be used with hybrid functional'
                               % self.__class__)
        self.check_sanity()
        self.dump_flags()
        if nstates is None:
            nstates = self.nstates
        else:
            self.nstates = nstates
        log = logger.Logger(self.stdout, self.verbose)

        vind, hdiag = self.get_vind(self._scf)
        precond = self.get_precond(hdiag)

        if x0 is None:
            x0 = self.init_guess(self._scf, self.nstates)

        def pickeig(w, v, nroots, envs):
            idx = numpy.where(w > POSTIVE_EIG_THRESHOLD**2)[0]
            return w[idx], v[:,idx], idx

        self.converged, w2, x1 = \
                lib.davidson1(vind, x0, precond,
                              tol=self.conv_tol,
                              nroots=nstates, lindep=self.lindep,
                              max_cycle=self.max_cycle,
                              max_space=self.max_space, pick=pickeig,
                              verbose=log)

        mo_energy = self._scf.mo_energy
        mo_occ = self._scf.mo_occ
        occidxa = numpy.where(mo_occ[0]>0)[0]
        occidxb = numpy.where(mo_occ[1]>0)[0]
        viridxa = numpy.where(mo_occ[0]==0)[0]
        viridxb = numpy.where(mo_occ[1]==0)[0]
        nocca = len(occidxa)
        noccb = len(occidxb)
        nvira = len(viridxa)
        nvirb = len(viridxb)
        e_ia_a = (mo_energy[0][viridxa,None] - mo_energy[0][occidxa]).T
        e_ia_b = (mo_energy[1][viridxb,None] - mo_energy[1][occidxb]).T
        e_ia = numpy.hstack((e_ia_a.reshape(-1), e_ia_b.reshape(-1)))
        e_ia = numpy.sqrt(e_ia)

        e = []
        xy = []
        for i, z in enumerate(x1):
            if w2[i] < POSTIVE_EIG_THRESHOLD**2:
                continue
            w = numpy.sqrt(w2[i])
            zp = e_ia * z
            zm = w/e_ia * z
            x = (zp + zm) * .5
            y = (zp - zm) * .5
            norm = lib.norm(x)**2 - lib.norm(y)**2
            if norm > 0:
                norm = 1/numpy.sqrt(norm)
                e.append(w)
                xy.append(((x[:nocca*nvira].reshape(nocca,nvira) * norm,  # X_alpha
                            x[nocca*nvira:].reshape(noccb,nvirb) * norm), # X_beta
                           (y[:nocca*nvira].reshape(nocca,nvira) * norm,  # Y_alpha
                            y[nocca*nvira:].reshape(noccb,nvirb) * norm)))# Y_beta
        self.e = numpy.array(e)
        self.xy = xy

        if self.chkfile:
            lib.chkfile.save(self.chkfile, 'tddft/e', self.e)
            lib.chkfile.save(self.chkfile, 'tddft/xy', self.xy)

        log.timer('TDDFT', *cpu0)
        self._finalize()
        return self.e, self.xy

    def nuc_grad_method(self):
        from pyscf.grad import tduks
        return tduks.Gradients(self)

TDDFTNoHybrid = CasidaTDDFT


class dRPA(TDDFTNoHybrid):
    def __init__(self, mf):
        if not isinstance(mf, KohnShamDFT):
            raise RuntimeError("direct RPA can only be applied with DFT; for HF+dRPA, use .xc='hf'")
        mf = mf.to_uhf()
        mf.xc = ''
        TDDFTNoHybrid.__init__(self, mf)

TDH = dRPA

class dTDA(TDA):
    def __init__(self, mf):
        if not isinstance(mf, KohnShamDFT):
            raise RuntimeError("direct TDA can only be applied with DFT; for HF+dTDA, use .xc='hf'")
        mf = mf.to_uhf()
        mf.xc = ''
        TDA.__init__(self, mf)


def tddft(mf):
    '''Driver to create TDDFT or CasidaTDDFT object'''
    if mf._numint.libxc.is_hybrid_xc(mf.xc):
        return TDDFT(mf)
    else:
        return CasidaTDDFT(mf)

from pyscf import dft
dft.uks.UKS.TDA           = dft.uks_symm.UKS.TDA           = lib.class_as_method(TDA)
dft.uks.UKS.TDHF          = dft.uks_symm.UKS.TDHF          = None
#dft.uks.UKS.TDDFT         = dft.uks_symm.UKS.TDDFT         = lib.class_as_method(TDDFT)
dft.uks.UKS.TDDFTNoHybrid = dft.uks_symm.UKS.TDDFTNoHybrid = lib.class_as_method(TDDFTNoHybrid)
dft.uks.UKS.CasidaTDDFT   = dft.uks_symm.UKS.CasidaTDDFT   = lib.class_as_method(CasidaTDDFT)
dft.uks.UKS.TDDFT         = dft.uks_symm.UKS.TDDFT         = tddft
dft.uks.UKS.dTDA          = dft.uks_symm.UKS.dTDA          = lib.class_as_method(dTDA)
dft.uks.UKS.dRPA          = dft.uks_symm.UKS.dRPA          = lib.class_as_method(dRPA)


# pyscf/tools/__init__.py

# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pyscf.tools import fcidump
from pyscf.tools import molden
from pyscf.tools import dump_mat
from pyscf.tools import cubegen


# pyscf/tools/c60struct.py

# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from functools import reduce
import numpy

def rotmatz(ang):
    c = numpy.cos(ang)
    s = numpy.sin(ang)
    return numpy.array((( c, s, 0),
                        (-s, c, 0),
                        ( 0, 0, 1),))
def rotmaty(ang):
    c = numpy.cos(ang)
    s = numpy.sin(ang)
    return numpy.array((( c, 0, s),
                        ( 0, 1, 0),
                        (-s, 0, c),))

def r2edge(ang, r):
    return 2*r*numpy.sin(ang/2)



def make60(b5, b6):
    theta1 = numpy.arccos(1/numpy.sqrt(5))
    theta2 = (numpy.pi - theta1) * .5
    r = (b5*2+b6)/2/numpy.sin(theta1/2)
    rot72 = rotmatz(numpy.pi*2/5)
    s2 = numpy.sin(theta2)
    c2 = numpy.cos(theta2)
    p1 = numpy.array(( s2*b5,  0, r-c2*b5))
    p9 = numpy.array((-s2*b5,  0,-r+c2*b5))
    p2 = numpy.array(( s2*(b5+b6),  0, r-c2*(b5+b6)))
    rot1 = reduce(numpy.dot, (rotmaty(theta1), rot72, rotmaty(-theta1)))
    p2s = []
    for i in range(5):
        p2s.append(p2)
        p2 = numpy.dot(p2, rot1)

    coord = []
    for i in range(5):
        coord.append(p1)
        p1 = numpy.dot(p1, rot72)
    for pj in p2s:
        pi = pj
        for i in range(5):
            coord.append(pi)
            pi = numpy.dot(pi, rot72)
    for pj in p2s:
        pi = pj
        for i in range(5):
            coord.append(-pi)
            pi = numpy.dot(pi, rot72)
    for i in range(5):
        coord.append(p9)
        p9 = numpy.dot(p9, rot72)
    return numpy.array(coord)


def make12(b):
    theta1 = numpy.arccos(1/numpy.sqrt(5))
    r = b/2/numpy.sin(theta1/2)
    rot72 = rotmatz(numpy.pi*2/5)
    s1 = numpy.sin(theta1)
    c1 = numpy.cos(theta1)
    p1 = numpy.array(( s1*r,  0,  c1*r))
    p2 = numpy.array((-s1*r,  0, -c1*r))
    coord = [(  0,  0,    r)]
    for i in range(5):
        coord.append(p1)
        p1 = numpy.dot(p1, rot72)
    for i in range(5):
        coord.append(p2)
        p2 = numpy.dot(p2, rot72)
    coord.append((  0,  0,  -r))
    return numpy.array(coord)


def make20(b):
    theta1 = numpy.arccos(numpy.sqrt(5)/3)
    theta2 = numpy.arcsin(r2edge(theta1,1)/2/numpy.sin(numpy.pi/5))
    r = b/2/numpy.sin(theta1/2)
    rot72 = rotmatz(numpy.pi*2/5)
    s2 = numpy.sin(theta2)
    c2 = numpy.cos(theta2)
    s3 = numpy.sin(theta1+theta2)
    c3 = numpy.cos(theta1+theta2)
    p1 = numpy.array(( s2*r,  0,  c2*r))
    p2 = numpy.array(( s3*r,  0,  c3*r))
    p3 = numpy.array((-s3*r,  0, -c3*r))
    p4 = numpy.array((-s2*r,  0, -c2*r))
    coord = []
    for i in range(5):
        coord.append(p1)
        p1 = numpy.dot(p1, rot72)
    for i in range(5):
        coord.append(p2)
        p2 = numpy.dot(p2, rot72)
    for i in range(5):
        coord.append(p3)
        p3 = numpy.dot(p3, rot72)
    for i in range(5):
        coord.append(p4)
        p4 = numpy.dot(p4, rot72)
    return numpy.array(coord)


if __name__ == '__main__':
    b5 = 1.46
    b6 = 1.38
    for c in make60(b5, b6):
        print(c)

    b = 1.4
    for c in make12(b):
        print(c)

    for c in make20(b):
        print(c)


# pyscf/tools/chgcar.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Paul J. Robinson <pjrobinson@ucla.edu>
#         Qiming Sun <osirpt.sun@gmail.com>
#

'''
Vasp CHGCAR file format

See also
https://cms.mpi.univie.ac.at/vasp/vasp/CHGCAR_file.html
'''

import sys
import collections
import time
import numpy
import pyscf
from pyscf import lib
from pyscf import gto
from pyscf.pbc import gto as pbcgto
from pyscf.tools import cubegen

if sys.version_info >= (3,):
    unicode = str

RESOLUTION = cubegen.RESOLUTION
BOX_MARGIN = cubegen.BOX_MARGIN


def density(cell, outfile, dm, nx=60, ny=60, nz=60, resolution=RESOLUTION):
    '''Calculates electron density and write out in CHGCAR format.

    Args:
        cell : Mole or Cell object
            Mole or pbc Cell. If Mole object is given, the program will guess
            a cubic lattice for the molecule.
        outfile : str
            Name of Cube file to be written.
        dm : ndarray
            Density matrix of molecule.

    Kwargs:
        nx : int
            Number of grid point divisions in x direction.
            Note this is function of the molecule's size; a larger molecule
            will have a coarser representation than a smaller one for the
            same value.
        ny : int
            Number of grid point divisions in y direction.
        nz : int
            Number of grid point divisions in z direction.

    Returns:
        No return value. This function outputs a VASP chgcarlike file
        (with phase if desired)...it can be opened in VESTA or VMD or
        many other softwares

    Examples:

        >>> # generates the first MO from the list of mo_coefficents
        >>> from pyscf.pbc import gto, scf
        >>> from pyscf.tools import chgcar
        >>> cell = gto.M(atom='H 0 0 0; H 0 0 1', a=numpy.eye(3)*3)
        >>> mf = scf.RHF(cell).run()
        >>> chgcar.density(cell, 'h2.CHGCAR', mf.make_rdm1())

    '''
    cc = CHGCAR(cell, nx=nx, ny=ny, nz=nz, resolution=resolution)

    coords = cc.get_coords()
    ngrids = cc.get_ngrids()
    blksize = min(8000, ngrids)
    rho = numpy.empty(ngrids)
    for ip0, ip1 in lib.prange(0, ngrids, blksize):
        if isinstance(cell, pbcgto.cell.Cell):
            ao = cell.pbc_eval_gto('GTOval', coords[ip0:ip1])
        else:
            ao = cell.eval_gto('GTOval', coords[ip0:ip1])
        rho[ip0:ip1] = lib.einsum('pi,ij,pj->p', ao, dm, ao)
    rho = rho.reshape(nx,ny,nz)

    cc.write(rho, outfile)


def orbital(cell, outfile, coeff, nx=60, ny=60, nz=60, resolution=RESOLUTION):
    '''Calculate orbital value on real space grid and write out in
    CHGCAR format.

    Args:
        cell : Mole or Cell object
            Mole or pbc Cell. If Mole object is given, the program will guess
            a cubic lattice for the molecule.
        outfile : str
            Name of Cube file to be written.
        dm : ndarray
            Density matrix of molecule.

    Kwargs:
        nx : int
            Number of grid point divisions in x direction.
            Note this is function of the molecule's size; a larger molecule
            will have a coarser representation than a smaller one for the
            same value.
        ny : int
            Number of grid point divisions in y direction.
        nz : int
            Number of grid point divisions in z direction.

    Returns:
        No return value. This function outputs a VASP chgcarlike file
        (with phase if desired)...it can be opened in VESTA or VMD or
        many other softwares

    Examples:

        >>> # generates the first MO from the list of mo_coefficents
        >>> from pyscf.pbc import gto, scf
        >>> from pyscf.tools import chgcar
        >>> cell = gto.M(atom='H 0 0 0; H 0 0 1', a=numpy.eye(3)*3)
        >>> mf = scf.RHF(cell).run()
        >>> chgcar.orbital(cell, 'h2_mo1.CHGCAR', mf.mo_coeff[:,0])

    '''
    cc = CHGCAR(cell, nx=nx, ny=ny, nz=nz, resolution=resolution)

    coords = cc.get_coords()
    ngrids = cc.get_ngrids()
    blksize = min(8000, ngrids)
    orb_on_grid = numpy.empty(ngrids)
    for ip0, ip1 in lib.prange(0, ngrids, blksize):
        if isinstance(cell, pbcgto.cell.Cell):
            ao = cell.pbc_eval_gto('GTOval', coords[ip0:ip1])
        else:
            ao = cell.eval_gto('GTOval', coords[ip0:ip1])
        orb_on_grid[ip0:ip1] = numpy.dot(ao, coeff)
    orb_on_grid = orb_on_grid.reshape(nx,ny,nz)

    cc.write(orb_on_grid, outfile, comment='Orbital value in real space (1/Bohr^3)')


class CHGCAR(cubegen.Cube):
    '''  Read-write of the Vasp CHGCAR files  '''
    def __init__(self, cell, nx=60, ny=60, nz=60, resolution=RESOLUTION,
                 margin=BOX_MARGIN):
        if not isinstance(cell, pbcgto.cell.Cell):
            coord = cell.atom_coords()
            box = numpy.max(coord,axis=0) - numpy.min(coord,axis=0) + margin*2
            boxorig = numpy.min(coord,axis=0) - margin
            if resolution is not None:
                nx, ny, nz = numpy.ceil(box / resolution).astype(int)
            self.box = numpy.diag(box)
            lib.logger.warn(cell, 'Molecular system is found. FFT-grid is not '
                            'available for Molecule. Lattice (in Bohr)\n'
                            '%s\nand FFT grids %s are applied.',
                            self.box, (nx,ny,nz))
            self.mol = cell
            cell = cell.view(pbcgto.Cell)
            if (isinstance(cell.unit, (str, unicode)) and
                cell.unit.startswith(('B','b','au','AU'))):
                cell.a = self.box
            else:
                cell.a = self.box * lib.param.BOHR
            ptr = cell._atm[:,gto.PTR_COORD]
            cell._env[ptr+0] = coord[:,0] - boxorig[0]
            cell._env[ptr+1] = coord[:,1] - boxorig[1]
            cell._env[ptr+2] = coord[:,2] - boxorig[2]

        self.nx = nx
        self.ny = ny
        self.nz = nz
        self.cell = cell
        self.box = cell.lattice_vectors()
        self.boxorig = numpy.zeros(3)
        self.vol = cell.vol

    def get_coords(self) :
        """  Result: set of coordinates to compute a field which is to be stored
        in the file.
        """
        xs = numpy.arange(self.nx) * (1./self.nx)
        ys = numpy.arange(self.ny) * (1./self.ny)
        zs = numpy.arange(self.nz) * (1./self.nz)
        xyz = lib.cartesian_prod((xs, ys, zs))
        coords = numpy.dot(xyz, self.box)
        return numpy.asarray(coords, order='C')

    def write(self, field, fname, comment=None):
        """  Result: .vasp file with the field in the file fname.  """
        assert(field.ndim == 3)
        assert(field.shape == (self.nx, self.ny, self.nz))
        if comment is None:
            comment = 'VASP file: Electron density in real space (e/Bohr^3)  '

        cell = self.cell

        # See CHGCAR format https://cms.mpi.univie.ac.at/vasp/vasp/CHGCAR_file.html
        # the value of (total density * volume) was dumped
        field = field * self.vol

        boxA = self.box * lib.param.BOHR
        atomList= [cell.atom_pure_symbol(i) for i in range(cell.natm)]
        Axyz = zip(atomList, cell.atom_coords().tolist())
        Axyz = sorted(Axyz, key = lambda x: x[0])
        swappedCoords = [(vec[1]+self.boxorig) * lib.param.BOHR for vec in Axyz]
        vaspAtomicInfo = collections.Counter([xyz[0] for xyz in Axyz])
        vaspAtomicInfo = sorted(vaspAtomicInfo.items())
        with open(fname, 'w') as f:
            f.write(comment)
            f.write('PySCF Version: %s  Date: %s\n' % (pyscf.__version__, time.ctime()))
            f.write('1.0000000000\n')
            f.write('%14.8f %14.8f %14.8f \n' % (boxA[0,0],boxA[0,1],boxA[0,2]))
            f.write('%14.8f %14.8f %14.8f \n' % (boxA[1,0],boxA[1,1],boxA[1,2]))
            f.write('%14.8f %14.8f %14.8f \n' % (boxA[2,0],boxA[2,1],boxA[2,2]))
            f.write(''.join(['%5.3s'%atomN[0] for atomN in vaspAtomicInfo]) + '\n')
            f.write(''.join(['%5d'%atomN[1] for atomN in vaspAtomicInfo]) + '\n')
            f.write('Cartesian \n')
            for ia in range(cell.natm):
                f.write(' %14.8f %14.8f %14.8f\n' % tuple(swappedCoords[ia]))
            f.write('\n')
            f.write('%6.5s %6.5s %6.5s \n' % (self.nx,self.ny,self.nz))
            fmt = ' %14.8e '
            for iz in range(self.nz):
                for iy in range(self.ny):
                    f.write('\n')
                    for ix in range(self.nx):
                        f.write(fmt % field[ix,iy,iz])

    def read(self, chgcar_file):
        raise NotImplementedError


if __name__ == '__main__':
    from pyscf.pbc import scf
    from pyscf.tools import chgcar
    cell = gto.M(atom='H 0 0 0; H 0 0 1', a=numpy.eye(3)*3)
    mf = scf.RHF(cell).run()
    chgcar.density(cell, 'h2.CHGCAR', mf.make_rdm1()) #makes total density
    chgcar.orbital(cell, 'h2_mo1.CHGCAR', mf.mo_coeff[:,0]) # makes mo#1 (sigma)
    chgcar.orbital(cell, 'h2_mo2.CHGCAR', mf.mo_coeff[:,1]) # makes mo#2 (sigma*)



# pyscf/tools/chkfile_util.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import sys
import numpy
from pyscf.scf import chkfile
from pyscf.scf import hf

def mulliken(filename, key='scf'):
    '''Reading scf/mcscf information from chkfile, then do Mulliken population
    analysis for the density matrix
    '''
    if key.lower() == 'mcscf':
        mol = chkfile.load_mol(filename)
        mo_coeff = chkfile.load(filename, 'mcscf/mo_coeff')
        mo_occ = chkfile.load(filename, 'mcscf/mo_occ')
    else:
        mol, mf = chkfile.load_scf(filename)
        mo_coeff = mf['mo_coeff']
        mo_occ = mf['mo_occ']
    dm = numpy.dot(mo_coeff*mo_occ, mo_coeff.T)
    hf.mulliken_meta(mol, dm)

def dump_mo(filename, key='scf'):
    '''Read scf/mcscf information from chkfile, then dump the orbital
    coefficients.
    '''
    from pyscf.tools import dump_mat
    if key.lower() == 'mcscf':
        mol = chkfile.load_mol(filename)
        mo_coeff = chkfile.load(filename, 'mcscf/mo_coeff')
    else:
        mol, mf = chkfile.load_scf(filename)
        mo_coeff = mf['mo_coeff']
    dump_mat.dump_mo(mol, mo_coeff)

def molden(filename, key='scf'):
    '''Read scf/mcscf information from chkfile, then convert the scf/mcscf
    orbitals to molden format.
    '''
    from pyscf.tools import molden
    molden.from_chkfile(filename+'.molden', filename, key+'/mo_coeff')

if __name__ == '__main__':
    if len(sys.argv) == 1:
        print('%s chkfile_name [pop|mo|molden] [scf|mcscf]' % sys.argv[0])
        exit()
    filename = sys.argv[1]
    fndic = {'pop': mulliken,
             'mo': dump_mo,
             'molden': molden}
    fn = fndic[sys.argv[2].lower()]
    if len(sys.argv) > 3:
        key = sys.argv[3]
        fn(filename, key)
    else:
        fn(filename)


# pyscf/tools/cubegen.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#         Peter Koval <koval.peter@gmail.com>
#         Paul J. Robinson <pjrobinson@ucla.edu>
#

'''
Gaussian cube file format.  Reference:
http://paulbourke.net/dataformats/cube/
https://h5cube-spec.readthedocs.io/en/latest/cubeformat.html
http://gaussian.com/cubegen/

The output cube file has the following format

Comment line
Comment line
N_atom Ox Oy Oz         # number of atoms, followed by the coordinates of the origin
N1 vx1 vy1 vz1          # number of grids along each axis, followed by the step size in x/y/z direction.
N2 vx2 vy2 vz2          # ...
N3 vx3 vy3 vz3          # ...
Atom1 Z1 x y z          # Atomic number, charge, and coordinates of the atom
...                     # ...
AtomN ZN x y z          # ...
Data on grids           # (N1*N2) lines of records, each line has N3 elements
'''

import time
import numpy
import pyscf
from pyscf import lib
from pyscf import gto
from pyscf import df
from pyscf.dft import numint
from pyscf import __config__

RESOLUTION = getattr(__config__, 'cubegen_resolution', None)
BOX_MARGIN = getattr(__config__, 'cubegen_box_margin', 3.0)
ORIGIN = getattr(__config__, 'cubegen_box_origin', None)
# If given, EXTENT should be a 3-element ndarray/list/tuple to represent the
# extension in x, y, z
EXTENT = getattr(__config__, 'cubegen_box_extent', None)


def density(mol, outfile, dm, nx=80, ny=80, nz=80, resolution=RESOLUTION,
            margin=BOX_MARGIN):
    """Calculates electron density and write out in cube format.

    Args:
        mol : Mole
            Molecule to calculate the electron density for.
        outfile : str
            Name of Cube file to be written.
        dm : ndarray
            Density matrix of molecule.

    Kwargs:
        nx : int
            Number of grid point divisions in x direction.
            Note this is function of the molecule's size; a larger molecule
            will have a coarser representation than a smaller one for the
            same value. Conflicts to keyword resolution.
        ny : int
            Number of grid point divisions in y direction.
        nz : int
            Number of grid point divisions in z direction.
        resolution: float
            Resolution of the mesh grid in the cube box. If resolution is
            given in the input, the input nx/ny/nz have no effects.  The value
            of nx/ny/nz will be determined by the resolution and the cube box
            size.
    """
    from pyscf.pbc.gto import Cell
    cc = Cube(mol, nx, ny, nz, resolution, margin)

    GTOval = 'GTOval'
    if isinstance(mol, Cell):
        GTOval = 'PBC' + GTOval

    # Compute density on the .cube grid
    coords = cc.get_coords()
    ngrids = cc.get_ngrids()
    blksize = min(8000, ngrids)
    rho = numpy.empty(ngrids)
    for ip0, ip1 in lib.prange(0, ngrids, blksize):
        ao = mol.eval_gto(GTOval, coords[ip0:ip1])
        rho[ip0:ip1] = numint.eval_rho(mol, ao, dm)
    rho = rho.reshape(cc.nx,cc.ny,cc.nz)

    # Write out density to the .cube file
    cc.write(rho, outfile, comment='Electron density in real space (e/Bohr^3)')
    return rho


def orbital(mol, outfile, coeff, nx=80, ny=80, nz=80, resolution=RESOLUTION,
            margin=BOX_MARGIN):
    """Calculate orbital value on real space grid and write out in cube format.

    Args:
        mol : Mole
            Molecule to calculate the electron density for.
        outfile : str
            Name of Cube file to be written.
        coeff : 1D array
            coeff coefficient.

    Kwargs:
        nx : int
            Number of grid point divisions in x direction.
            Note this is function of the molecule's size; a larger molecule
            will have a coarser representation than a smaller one for the
            same value. Conflicts to keyword resolution.
        ny : int
            Number of grid point divisions in y direction.
        nz : int
            Number of grid point divisions in z direction.
        resolution: float
            Resolution of the mesh grid in the cube box. If resolution is
            given in the input, the input nx/ny/nz have no effects.  The value
            of nx/ny/nz will be determined by the resolution and the cube box
            size.
    """
    from pyscf.pbc.gto import Cell
    cc = Cube(mol, nx, ny, nz, resolution, margin)

    GTOval = 'GTOval'
    if isinstance(mol, Cell):
        GTOval = 'PBC' + GTOval

    # Compute density on the .cube grid
    coords = cc.get_coords()
    ngrids = cc.get_ngrids()
    blksize = min(8000, ngrids)
    orb_on_grid = numpy.empty(ngrids)
    for ip0, ip1 in lib.prange(0, ngrids, blksize):
        ao = mol.eval_gto(GTOval, coords[ip0:ip1])
        orb_on_grid[ip0:ip1] = numpy.dot(ao, coeff)
    orb_on_grid = orb_on_grid.reshape(cc.nx,cc.ny,cc.nz)

    # Write out orbital to the .cube file
    cc.write(orb_on_grid, outfile, comment='Orbital value in real space (1/Bohr^3)')
    return orb_on_grid


def mep(mol, outfile, dm, nx=80, ny=80, nz=80, resolution=RESOLUTION,
        margin=BOX_MARGIN):
    """Calculates the molecular electrostatic potential (MEP) and write out in
    cube format.

    Args:
        mol : Mole
            Molecule to calculate the electron density for.
        outfile : str
            Name of Cube file to be written.
        dm : ndarray
            Density matrix of molecule.

    Kwargs:
        nx : int
            Number of grid point divisions in x direction.
            Note this is function of the molecule's size; a larger molecule
            will have a coarser representation than a smaller one for the
            same value. Conflicts to keyword resolution.
        ny : int
            Number of grid point divisions in y direction.
        nz : int
            Number of grid point divisions in z direction.
        resolution: float
            Resolution of the mesh grid in the cube box. If resolution is
            given in the input, the input nx/ny/nz have no effects.  The value
            of nx/ny/nz will be determined by the resolution and the cube box
            size.
    """
    cc = Cube(mol, nx, ny, nz, resolution, margin)

    coords = cc.get_coords()

    # Nuclear potential at given points
    Vnuc = 0
    for i in range(mol.natm):
        r = mol.atom_coord(i)
        Z = mol.atom_charge(i)
        rp = r - coords
        Vnuc += Z / numpy.einsum('xi,xi->x', rp, rp)**.5

    # Potential of electron density
    Vele = numpy.empty_like(Vnuc)
    for p0, p1 in lib.prange(0, Vele.size, 600):
        fakemol = gto.fakemol_for_charges(coords[p0:p1])
        ints = df.incore.aux_e2(mol, fakemol)
        Vele[p0:p1] = numpy.einsum('ijp,ij->p', ints, dm)

    MEP = Vnuc - Vele     # MEP at each point
    MEP = MEP.reshape(cc.nx,cc.ny,cc.nz)

    # Write the potential
    cc.write(MEP, outfile, 'Molecular electrostatic potential in real space')
    return MEP


class Cube(object):
    '''  Read-write of the Gaussian CUBE files

    Attributes:
        nx : int
            Number of grid point divisions in x direction.
            Note this is function of the molecule's size; a larger molecule
            will have a coarser representation than a smaller one for the
            same value. Conflicts to keyword resolution.
        ny : int
            Number of grid point divisions in y direction.
        nz : int
            Number of grid point divisions in z direction.
        resolution: float
            Resolution of the mesh grid in the cube box. If resolution is
            given in the input, the input nx/ny/nz have no effects.  The value
            of nx/ny/nz will be determined by the resolution and the cube box
            size. The unit is Bohr.
    '''
    def __init__(self, mol, nx=80, ny=80, nz=80, resolution=RESOLUTION,
                 margin=BOX_MARGIN, origin=ORIGIN, extent=EXTENT):
        from pyscf.pbc.gto import Cell
        self.mol = mol
        coord = mol.atom_coords()

        # If the molecule is periodic, use lattice vectors as the box
        # and automatically determine margin, origin, and extent
        if isinstance(mol, Cell):
            self.box = mol.lattice_vectors()
            atom_center = (numpy.max(coord, axis=0) + numpy.min(coord, axis=0))/2
            box_center = (self.box[0] + self.box[1] + self.box[2])/2
            self.boxorig = atom_center - box_center
        else:
            # Create a box based on its origin and extents/lengths (rectangular cuboid only)
            # If extent is not supplied, use the coordinates plus a margin on both sides
            if extent is None:
                extent = numpy.max(coord, axis=0) - numpy.min(coord, axis=0) + 2*margin
            self.box = numpy.diag(extent)

            # if origin is not supplied, set it as the minimum coordinate minus a margin.
            if origin is None:
                origin = numpy.min(coord, axis=0) - margin
            self.boxorig = numpy.asarray(origin)

        if resolution is not None:
            nx, ny, nz = numpy.ceil(numpy.diag(self.box) / resolution).astype(int)

        self.nx = nx
        self.ny = ny
        self.nz = nz

        if isinstance(mol, Cell):
            # Use an asymmetric mesh for tiling unit cells
            self.xs = numpy.linspace(0, 1, nx, endpoint=False)
            self.ys = numpy.linspace(0, 1, ny, endpoint=False)
            self.zs = numpy.linspace(0, 1, nz, endpoint=False)
        else:
            # Use endpoint=True to get a symmetric mesh
            # see also the discussion https://github.com/sunqm/pyscf/issues/154
            self.xs = numpy.linspace(0, 1, nx, endpoint=True)
            self.ys = numpy.linspace(0, 1, ny, endpoint=True)
            self.zs = numpy.linspace(0, 1, nz, endpoint=True)

    def get_coords(self) :
        """  Result: set of coordinates to compute a field which is to be stored
        in the file.
        """
        frac_coords = lib.cartesian_prod([self.xs, self.ys, self.zs])
        return frac_coords @ self.box + self.boxorig # Convert fractional coordinates to real-space coordinates

    def get_ngrids(self):
        return self.nx * self.ny * self.nz

    def get_volume_element(self):
        return (self.xs[1]-self.xs[0])*(self.ys[1]-self.ys[0])*(self.zs[1]-self.zs[0])

    def write(self, field, fname, comment=None):
        """  Result: .cube file with the field in the file fname.  """
        assert(field.ndim == 3)
        assert(field.shape == (self.nx, self.ny, self.nz))
        if comment is None:
            comment = 'Generic field? Supply the optional argument "comment" to define this line'

        mol = self.mol
        coord = mol.atom_coords()
        with open(fname, 'w') as f:
            f.write(comment+'\n')
            f.write(f'PySCF Version: {pyscf.__version__}  Date: {time.ctime()}\n')
            f.write(f'{mol.natm:5d}')
            f.write('%12.6f%12.6f%12.6f\n' % tuple(self.boxorig.tolist()))
            dx = self.xs[-1] if len(self.xs) == 1 else self.xs[1]
            dy = self.ys[-1] if len(self.ys) == 1 else self.ys[1]
            dz = self.zs[-1] if len(self.zs) == 1 else self.zs[1]
            delta = (self.box.T * [dx,dy,dz]).T
            f.write(f'{self.nx:5d}{delta[0,0]:12.6f}{delta[0,1]:12.6f}{delta[0,2]:12.6f}\n')
            f.write(f'{self.ny:5d}{delta[1,0]:12.6f}{delta[1,1]:12.6f}{delta[1,2]:12.6f}\n')
            f.write(f'{self.nz:5d}{delta[2,0]:12.6f}{delta[2,1]:12.6f}{delta[2,2]:12.6f}\n')
            for ia in range(mol.natm):
                atmsymb = mol.atom_symbol(ia)
                f.write('%5d%12.6f'% (gto.charge(atmsymb), 0.))
                f.write('%12.6f%12.6f%12.6f\n' % tuple(coord[ia]))

            for ix in range(self.nx):
                for iy in range(self.ny):
                    for iz0, iz1 in lib.prange(0, self.nz, 6):
                        fmt = '%13.5E' * (iz1-iz0) + '\n'
                        f.write(fmt % tuple(field[ix,iy,iz0:iz1].tolist()))

    def read(self, cube_file):
        with open(cube_file, 'r') as f:
            f.readline()
            f.readline()
            data = f.readline().split()
            natm = int(data[0])
            self.boxorig = numpy.array([float(x) for x in data[1:]])
            def parse_nx(data):
                d = data.split()
                return int(d[0]), numpy.array([float(x) for x in d[1:]])
            self.nx, self.xs = parse_nx(f.readline())
            self.ny, self.ys = parse_nx(f.readline())
            self.nz, self.zs = parse_nx(f.readline())
            atoms = []
            for ia in range(natm):
                d = f.readline().split()
                atoms.append([int(d[0]), [float(x) for x in d[2:]]])
            self.mol = gto.M(atom=atoms, unit='Bohr')

            data = f.read()
        cube_data = numpy.array([float(x) for x in data.split()])
        return cube_data.reshape([self.nx, self.ny, self.nz])


if __name__ == '__main__':
    from pyscf import scf
    from pyscf.tools import cubegen
    mol = gto.M(atom='''O 0.00000000,  0.000000,  0.000000
                H 0.761561, 0.478993, 0.00000000
                H -0.761561, 0.478993, 0.00000000''', basis='6-31g*')
    mf = scf.RHF(mol).run()
    cubegen.density(mol, 'h2o_den.cube', mf.make_rdm1()) #makes total density
    cubegen.mep(mol, 'h2o_pot.cube', mf.make_rdm1())
    cubegen.orbital(mol, 'h2o_mo1.cube', mf.mo_coeff[:,0])


# pyscf/tools/dump_mat.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pyscf.lib.parameters import OUTPUT_DIGITS, OUTPUT_COLS
from pyscf import __config__

BASE = getattr(__config__, 'BASE', 0)

def dump_tri(stdout, c, label=None,
             ncol=OUTPUT_COLS, digits=OUTPUT_DIGITS, start=BASE):
    ''' Format print for the lower triangular part of an array

    Args:
        stdout : file object
            eg sys.stdout, or stdout = open('/path/to/file') or
            mol.stdout if mol is an object initialized from :class:`gto.Mole`
        c : numpy.ndarray
            coefficients

    Kwargs:
        label : list of strings
            Row labels (default is 1,2,3,4,...)
        ncol : int
            Number of columns in the format output (default 5)
        digits : int
            Number of digits of precision for floating point output (default 5)
        start : int
            The number to start to count the index (default 0)

    Examples:

        >>> import sys, numpy
        >>> dm = numpy.eye(3)
        >>> dump_tri(sys.stdout, dm)
                #0        #1        #2
        0       1.00000
        1       0.00000   1.00000
        2       0.00000   0.00000   1.00000
        >>> from pyscf import gto
        >>> mol = gto.M(atom='C 0 0 0')
        >>> dm = numpy.eye(mol.nao_nr())
        >>> dump_tri(sys.stdout, dm, label=mol.ao_labels(), ncol=9, digits=2)
                    #0     #1     #2     #3     #4     #5     #6     #7     #8
        0  C 1s     1.00
        0  C 2s     0.00   1.00
        0  C 3s     0.00   0.00   1.00
        0  C 2px    0.00   0.00   0.00   1.00
        0  C 2py    0.00   0.00   0.00   0.00   1.00
        0  C 2pz    0.00   0.00   0.00   0.00   0.00   1.00
        0  C 3px    0.00   0.00   0.00   0.00   0.00   0.00   1.00
        0  C 3py    0.00   0.00   0.00   0.00   0.00   0.00   0.00   1.00
        0  C 3pz    0.00   0.00   0.00   0.00   0.00   0.00   0.00   0.00   1.00
    '''
    nc = c.shape[1]
    for ic in range(0, nc, ncol):
        dc = c[:,ic:ic+ncol]
        m = dc.shape[1]
        fmt = (' %%%d.%df'%(digits+4,digits))*m + '\n'
        if label is None:
            stdout.write(((' '*(digits+3))+'%s\n') %
                         (' '*(digits)).join(['#%-4d'%i for i in range(start+ic,start+ic+m)]))
            for k, v in enumerate(dc[ic:ic+m]):
                fmt = (' %%%d.%df'%(digits+4,digits))*(k+1) + '\n'
                stdout.write(('%-5d' % (ic+k+start)) + (fmt % tuple(v[:k+1])))
            for k, v in enumerate(dc[ic+m:]):
                stdout.write(('%-5d' % (ic+m+k+start)) + (fmt % tuple(v)))
        else:
            stdout.write(((' '*(digits+10))+'%s\n') %
                         (' '*(digits)).join(['#%-4d'%i for i in range(start+ic,start+ic+m)]))
            #stdout.write('           ')
            #stdout.write(((' '*(digits)+'#%-5d')*m) % tuple(range(ic+start,ic+m+start)) + '\n')
            for k, v in enumerate(dc[ic:ic+m]):
                fmt = (' %%%d.%df'%(digits+4,digits))*(k+1) + '\n'
                stdout.write(('%12s' % label[ic+k]) + (fmt % tuple(v[:k+1])))
            for k, v in enumerate(dc[ic+m:]):
                stdout.write(('%12s' % label[ic+m+k]) + (fmt % tuple(v)))

def dump_rec(stdout, c, label=None, label2=None,
             ncol=OUTPUT_COLS, digits=OUTPUT_DIGITS, start=BASE):
    ''' Print an array in rectangular format

    Args:
        stdout : file object
            eg sys.stdout, or stdout = open('/path/to/file') or
            mol.stdout if mol is an object initialized from :class:`gto.Mole`
        c : numpy.ndarray
            coefficients

    Kwargs:
        label : list of strings
            Row labels (default is 1,2,3,4,...)
        label2 : list of strings
            Col labels (default is 1,2,3,4,...)
        ncol : int
            Number of columns in the format output (default 5)
        digits : int
            Number of digits of precision for floating point output (default 5)
        start : int
            The number to start to count the index (default 0)

    Examples:

        >>> import sys, numpy
        >>> dm = numpy.eye(3)
        >>> dump_rec(sys.stdout, dm)
                #0        #1        #2
        0       1.00000   0.00000   0.00000
        1       0.00000   1.00000   0.00000
        2       0.00000   0.00000   1.00000
        >>> from pyscf import gto
        >>> mol = gto.M(atom='C 0 0 0')
        >>> dm = numpy.eye(mol.nao_nr())
        >>> dump_rec(sys.stdout, dm, label=mol.ao_labels(), ncol=9, digits=2)
                    #0     #1     #2     #3     #4     #5     #6     #7     #8
        0  C 1s     1.00   0.00   0.00   0.00   0.00   0.00   0.00   0.00   0.00
        0  C 2s     0.00   1.00   0.00   0.00   0.00   0.00   0.00   0.00   0.00
        0  C 3s     0.00   0.00   1.00   0.00   0.00   0.00   0.00   0.00   0.00
        0  C 2px    0.00   0.00   0.00   1.00   0.00   0.00   0.00   0.00   0.00
        0  C 2py    0.00   0.00   0.00   0.00   1.00   0.00   0.00   0.00   0.00
        0  C 2pz    0.00   0.00   0.00   0.00   0.00   1.00   0.00   0.00   0.00
        0  C 3px    0.00   0.00   0.00   0.00   0.00   0.00   1.00   0.00   0.00
        0  C 3py    0.00   0.00   0.00   0.00   0.00   0.00   0.00   1.00   0.00
        0  C 3pz    0.00   0.00   0.00   0.00   0.00   0.00   0.00   0.00   1.00
    '''
    nc = c.shape[1]
    if label2 is None:
        fmt = '#%%-%dd' % (digits+3)
        label2 = [fmt%i for i in range(start,nc+start)]
    else:
        fmt = '%%-%ds' % (digits+4)
        label2 = [fmt%i for i in label2]
    for ic in range(0, nc, ncol):
        dc = c[:,ic:ic+ncol]
        m = dc.shape[1]
        fmt = (' %%%d.%df'%(digits+4,digits))*m + '\n'
        if label is None:
            stdout.write(((' '*(digits+3))+'%s\n') % ' '.join(label2[ic:ic+m]))
            for k, v in enumerate(dc):
                stdout.write(('%-5d' % (k+start)) + (fmt % tuple(v)))
        else:
            stdout.write(((' '*(digits+10))+'%s\n') % ' '.join(label2[ic:ic+m]))
            for k, v in enumerate(dc):
                stdout.write(('%12s' % label[k]) + (fmt % tuple(v)))

def dump_mo(mol, c, label=None,
            ncol=OUTPUT_COLS, digits=OUTPUT_DIGITS, start=BASE):
    ''' Format print for orbitals

    Args:
        stdout : file object
            eg sys.stdout, or stdout = open('/path/to/file') or
            mol.stdout if mol is an object initialized from :class:`gto.Mole`
        c : numpy.ndarray
            Orbitals, each column is an orbital

    Kwargs:
        label : list of strings
            Row labels (default is AO labels)

    Examples:

        >>> from pyscf import gto
        >>> mol = gto.M(atom='C 0 0 0')
        >>> mo = numpy.eye(mol.nao_nr())
        >>> dump_mo(mol, mo)
                    #0     #1     #2     #3     #4     #5     #6     #7     #8
        0  C 1s     1.00   0.00   0.00   0.00   0.00   0.00   0.00   0.00   0.00
        0  C 2s     0.00   1.00   0.00   0.00   0.00   0.00   0.00   0.00   0.00
        0  C 3s     0.00   0.00   1.00   0.00   0.00   0.00   0.00   0.00   0.00
        0  C 2px    0.00   0.00   0.00   1.00   0.00   0.00   0.00   0.00   0.00
        0  C 2py    0.00   0.00   0.00   0.00   1.00   0.00   0.00   0.00   0.00
        0  C 2pz    0.00   0.00   0.00   0.00   0.00   1.00   0.00   0.00   0.00
        0  C 3px    0.00   0.00   0.00   0.00   0.00   0.00   1.00   0.00   0.00
        0  C 3py    0.00   0.00   0.00   0.00   0.00   0.00   0.00   1.00   0.00
        0  C 3pz    0.00   0.00   0.00   0.00   0.00   0.00   0.00   0.00   1.00
    '''
    if label is None:
        label = mol.ao_labels()
    dump_rec(mol.stdout, c, label, None, ncol, digits, start)

del(BASE)


if __name__ == '__main__':
    import sys
    import numpy
    c = numpy.random.random((16,16))
    label = ['A%5d' % i for i in range(16)]
    dump_tri(sys.stdout, c, label, 10, 2, 1)
    dump_rec(sys.stdout, c, None, label, start=1)


# pyscf/tools/fcidump.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''
FCIDUMP functions (write, read) for real Hamiltonian
'''

import re
from functools import reduce
import numpy
from pyscf import gto
from pyscf import scf
from pyscf import ao2mo
from pyscf import symm
from pyscf import __config__

DEFAULT_FLOAT_FORMAT = getattr(__config__, 'fcidump_float_format', ' %.16g')
TOL = getattr(__config__, 'fcidump_write_tol', 1e-15)

MOLPRO_ORBSYM = getattr(__config__, 'fcidump_molpro_orbsym', False)

# Mapping Pyscf symmetry numbering to Molpro symmetry numbering for each irrep.
# See also pyscf.symm.param.IRREP_ID_TABLE
# https://www.molpro.net/info/current/doc/manual/node36.html
ORBSYM_MAP = {
    'D2h': (1,         # Ag
            4,         # B1g
            6,         # B2g
            7,         # B3g
            8,         # Au
            5,         # B1u
            3,         # B2u
            2),        # B3u
    'C2v': (1,         # A1
            4,         # A2
            2,         # B1
            3),        # B2
    'C2h': (1,         # Ag
            4,         # Bg
            2,         # Au
            3),        # Bu
    'D2' : (1,         # A
            4,         # B1
            3,         # B2
            2),        # B3
    'Cs' : (1,         # A'
            2),        # A"
    'C2' : (1,         # A
            2),        # B
    'Ci' : (1,         # Ag
            2),        # Au
    'C1' : (1,)
}

def write_head(fout, nmo, nelec, ms=0, orbsym=None):
    if not isinstance(nelec, (int, numpy.number)):
        ms = abs(nelec[0] - nelec[1])
        nelec = nelec[0] + nelec[1]
    fout.write(' &FCI NORB=%4d,NELEC=%2d,MS2=%d,\n' % (nmo, nelec, ms))
    if orbsym is not None and len(orbsym) > 0:
        fout.write('  ORBSYM=%s\n' % ','.join([str(x) for x in orbsym]))
    else:
        fout.write('  ORBSYM=%s\n' % ('1,' * nmo))
    fout.write('  ISYM=1,\n')
    fout.write(' &END\n')


def write_eri(fout, eri, nmo, tol=TOL, float_format=DEFAULT_FLOAT_FORMAT):
    npair = nmo*(nmo+1)//2
    output_format = float_format + ' %4d %4d %4d %4d\n'
    if eri.size == nmo**4:
        eri = ao2mo.restore(8, eri, nmo)

    if eri.ndim == 2: # 4-fold symmetry
        assert(eri.size == npair**2)
        ij = 0
        for i in range(nmo):
            for j in range(0, i+1):
                kl = 0
                for k in range(0, nmo):
                    for l in range(0, k+1):
                        if abs(eri[ij,kl]) > tol:
                            fout.write(output_format % (eri[ij,kl], i+1, j+1, k+1, l+1))
                        kl += 1
                ij += 1
    else:  # 8-fold symmetry
        assert(eri.size == npair*(npair+1)//2)
        ij = 0
        ijkl = 0
        for i in range(nmo):
            for j in range(0, i+1):
                kl = 0
                for k in range(0, i+1):
                    for l in range(0, k+1):
                        if ij >= kl:
                            if abs(eri[ijkl]) > tol:
                                fout.write(output_format % (eri[ijkl], i+1, j+1, k+1, l+1))
                            ijkl += 1
                        kl += 1
                ij += 1

def write_hcore(fout, h, nmo, tol=TOL, float_format=DEFAULT_FLOAT_FORMAT):
    h = h.reshape(nmo,nmo)
    output_format = float_format + ' %4d %4d  0  0\n'
    for i in range(nmo):
        for j in range(0, i+1):
            if abs(h[i,j]) > tol:
                fout.write(output_format % (h[i,j], i+1, j+1))


def from_chkfile(filename, chkfile, tol=TOL, float_format=DEFAULT_FLOAT_FORMAT,
                 molpro_orbsym=MOLPRO_ORBSYM, orbsym=None):
    '''Read SCF results from PySCF chkfile and transform 1-electron,
    2-electron integrals using the SCF orbitals.  The transformed integrals is
    written to FCIDUMP

    Kwargs:
        molpro_orbsym (bool): Whether to dump the orbsym in Molpro orbsym
            convention as documented in
            https://www.molpro.net/info/current/doc/manual/node36.html
    '''
    mol, scf_rec = scf.chkfile.load_scf(chkfile)
    mo_coeff = numpy.array(scf_rec['mo_coeff'])
    nmo = mo_coeff.shape[1]

    s = reduce(numpy.dot, (mo_coeff.conj().T, mol.intor_symmetric('int1e_ovlp'), mo_coeff))
    if abs(s - numpy.eye(nmo)).max() > 1e-6:
        # Not support the chkfile from pbc calculation
        raise RuntimeError('Non-orthogonal orbitals found in chkfile')

    if mol.symmetry:
        orbsym = symm.label_orb_symm(mol, mol.irrep_id,
                                     mol.symm_orb, mo_coeff, check=False)
    from_mo(mol, filename, mo_coeff, orbsym=orbsym, tol=tol,
            float_format=float_format, molpro_orbsym=molpro_orbsym,
            ms=mol.spin)

def from_integrals(filename, h1e, h2e, nmo, nelec, nuc=0, ms=0, orbsym=None,
                   tol=TOL, float_format=DEFAULT_FLOAT_FORMAT):
    '''Convert the given 1-electron and 2-electron integrals to FCIDUMP format'''
    with open(filename, 'w') as fout:
        write_head(fout, nmo, nelec, ms, orbsym)
        write_eri(fout, h2e, nmo, tol=tol, float_format=float_format)
        write_hcore(fout, h1e, nmo, tol=tol, float_format=float_format)
        output_format = float_format + '  0  0  0  0\n'
        fout.write(output_format % nuc)

def from_mo(mol, filename, mo_coeff, orbsym=None,
            tol=TOL, float_format=DEFAULT_FLOAT_FORMAT,
            molpro_orbsym=MOLPRO_ORBSYM, ms=0):
    '''Use the given MOs to transfrom the 1-electron and 2-electron integrals
    then dump them to FCIDUMP.

    Kwargs:
        molpro_orbsym (bool): Whether to dump the orbsym in Molpro orbsym
            convention as documented in
            https://www.molpro.net/info/current/doc/manual/node36.html
    '''
    if getattr(mol, '_mesh', None):
        raise NotImplementedError('PBC system')

    if orbsym is None:
        orbsym = getattr(mo_coeff, 'orbsym', None)
        if molpro_orbsym and orbsym is not None:
            orbsym = [ORBSYM_MAP[mol.groupname][i] for i in orbsym]
    h1ao = scf.hf.get_hcore(mol)
    h1e = reduce(numpy.dot, (mo_coeff.T, h1ao, mo_coeff))
    eri = ao2mo.full(mol, mo_coeff, verbose=0)
    nuc = mol.energy_nuc()
    from_integrals(filename, h1e, eri, h1e.shape[0], mol.nelec, nuc, ms, orbsym,
                   tol, float_format)

def from_scf(mf, filename, tol=TOL, float_format=DEFAULT_FLOAT_FORMAT,
             molpro_orbsym=MOLPRO_ORBSYM):
    '''Use the given SCF object to transfrom the 1-electron and 2-electron
    integrals then dump them to FCIDUMP.

    Kwargs:
        molpro_orbsym (bool): Whether to dump the orbsym in Molpro orbsym
            convention as documented in
            https://www.molpro.net/info/current/doc/manual/node36.html
    '''
    mol = mf.mol
    mo_coeff = mf.mo_coeff
    assert mo_coeff.dtype == numpy.double

    h1e = reduce(numpy.dot, (mo_coeff.T, mf.get_hcore(), mo_coeff))
    if mf._eri is None:
        if getattr(mf, 'exxdiv', None):  # PBC system
            eri = mf.with_df.ao2mo(mo_coeff)
        else:
            eri = ao2mo.full(mf.mol, mo_coeff)
    else:  # Handle cached integrals or customized systems
        eri = ao2mo.full(mf._eri, mo_coeff)
    orbsym = getattr(mo_coeff, 'orbsym', None)
    if molpro_orbsym and orbsym is not None:
        orbsym = [ORBSYM_MAP[mol.groupname][i] for i in orbsym]
    nuc = mf.energy_nuc()
    from_integrals(filename, h1e, eri, h1e.shape[0], mf.mol.nelec, nuc, 0, orbsym,
                   tol, float_format)


def read(filename, molpro_orbsym=MOLPRO_ORBSYM):
    '''Parse FCIDUMP.  Return a dictionary to hold the integrals and
    parameters with keys:  H1, H2, ECORE, NORB, NELEC, MS, ORBSYM, ISYM

    Kwargs:
        molpro_orbsym (bool): Whether the orbsym in the FCIDUMP file is in
            Molpro orbsym convention as documented in
            https://www.molpro.net/info/current/doc/manual/node36.html
            In return, orbsym is converted to pyscf symmetry convention
    '''
    print('Parsing %s' % filename)
    finp = open(filename, 'r')

    data = []
    for i in range(10):
        line = finp.readline().upper()
        data.append(line)
        if '&END' in line:
            break
    else:
        raise RuntimeError('Problematic FCIDUMP header')

    result = {}
    tokens = ','.join(data).replace('&FCI', '').replace('&END', '')
    tokens = tokens.replace(' ', '').replace('\n', '').replace(',,', ',')
    for token in re.split(',(?=[a-zA-Z])', tokens):
        key, val = token.split('=')
        if key in ('NORB', 'NELEC', 'MS2', 'ISYM'):
            result[key] = int(val.replace(',', ''))
        elif key in ('ORBSYM',):
            result[key] = [int(x) for x in val.replace(',', ' ').split()]
        else:
            result[key] = val

    # Convert to molpr orbsym convert_orbsym
    if 'ORBSYM' in result:
        if molpro_orbsym:
            # Guess which point group the orbsym belongs to. FCIDUMP does not
            # save the point group information, the guess might be wrong if
            # the high symmetry numbering of orbitals are not presented.
            orbsym = result['ORBSYM']
            if max(orbsym) > 4:
                result['ORBSYM'] = [ORBSYM_MAP['D2h'].index(i) for i in orbsym]
            elif max(orbsym) > 2:
                # Fortunately, without molecular orientation, B2 and B3 in D2
                # are not distinguishable
                result['ORBSYM'] = [ORBSYM_MAP['C2v'].index(i) for i in orbsym]
            elif max(orbsym) == 2:
                result['ORBSYM'] = [i-1 for i in orbsym]
            elif max(orbsym) == 1:
                result['ORBSYM'] = [0] * len(orbsym)
            else:
                raise RuntimeError('Unknown orbsym')
        elif max(result['ORBSYM']) >= 8:
            raise RuntimeError('Unknown orbsym convention')

    norb = result['NORB']
    norb_pair = norb * (norb+1) // 2
    h1e = numpy.zeros((norb,norb))
    h2e = numpy.zeros(norb_pair*(norb_pair+1)//2)
    dat = finp.readline().split()
    while dat:
        i, j, k, l = [int(x) for x in dat[1:5]]
        if k != 0:
            if i >= j:
                ij = i * (i-1) // 2 + j-1
            else:
                ij = j * (j-1) // 2 + i-1
            if k >= l:
                kl = k * (k-1) // 2 + l-1
            else:
                kl = l * (l-1) // 2 + k-1
            if ij >= kl:
                h2e[ij*(ij+1)//2+kl] = float(dat[0])
            else:
                h2e[kl*(kl+1)//2+ij] = float(dat[0])
        elif k == 0:
            if j != 0:
                h1e[i-1,j-1] = float(dat[0])
            else:
                result['ECORE'] = float(dat[0])
        dat = finp.readline().split()

    idx, idy = numpy.tril_indices(norb, -1)
    if numpy.linalg.norm(h1e[idy,idx]) == 0:
        h1e[idy,idx] = h1e[idx,idy]
    elif numpy.linalg.norm(h1e[idx,idy]) == 0:
        h1e[idx,idy] = h1e[idy,idx]
    result['H1'] = h1e
    result['H2'] = h2e
    finp.close()
    return result

def to_scf(filename, molpro_orbsym=MOLPRO_ORBSYM, mf=None, **kwargs):
    '''Use the Hamiltonians defined by FCIDUMP to build an SCF object'''
    ctx = read(filename, molpro_orbsym)
    mol = gto.M()
    mol.nelectron = ctx['NELEC']
    mol.spin = ctx['MS2']
    norb = mol.nao = ctx['NORB']
    if 'ECORE' in ctx:
        mol.energy_nuc = lambda *args: ctx['ECORE']
    mol.incore_anyway = True

    if 'ORBSYM' in ctx:
        mol.symmetry = True
        mol.groupname = 'N/A'
        orbsym = numpy.asarray(ctx['ORBSYM'])
        mol.irrep_id = list(set(orbsym))
        mol.irrep_name = [('IR%d' % ir) for ir in mol.irrep_id]
        so = numpy.eye(norb)
        mol.symm_orb = []
        for ir in mol.irrep_id:
            mol.symm_orb.append(so[:,orbsym==ir])

    if mf is None:
        mf = mol.RHF(**kwargs)
    else:
        mf.mol = mol
    h1 = ctx['H1']
    idx, idy = numpy.tril_indices(norb, -1)
    if h1[idx,idy].max() == 0:
        h1[idx,idy] = h1[idy,idx]
    else:
        h1[idy,idx] = h1[idx,idy]
    mf.get_hcore = lambda *args: h1
    mf.get_ovlp = lambda *args: numpy.eye(norb)
    mf._eri = ctx['H2']

    return mf

def scf_from_fcidump(mf, filename, molpro_orbsym=MOLPRO_ORBSYM):
    '''Update the SCF object with the quantities defined in FCIDUMP file'''
    return to_scf(filename, molpro_orbsym, mf)

scf.hf.SCF.from_fcidump = scf_from_fcidump

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='Convert chkfile to FCIDUMP')
    parser.add_argument('chkfile', help='pyscf chkfile')
    parser.add_argument('fcidump', help='FCIDUMP file')
    args = parser.parse_args()

    # fcidump.py chkfile output
    from_chkfile(args.fcidump, args.chkfile)


# pyscf/tools/mo_mapping.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from functools import reduce
import numpy
from pyscf import lib
from pyscf import gto
from pyscf.lib import logger
from pyscf import lo
from pyscf import __config__

BASE = getattr(__config__, 'BASE', 0)
MAP_TOL = getattr(__config__, 'mo_mapping_mo_map_tol', 0.5)
ORTH_METHOD = getattr(__config__, 'mo_mapping_mo_comps_orth_method', 'meta_lowdin')


def mo_map(mol1, mo1, mol2, mo2, base=BASE, tol=.5):
    '''Given two orbitals, based on their overlap <i|j>, search all
    orbital-pairs which have significant overlap.

    Returns:
        Two lists.  First list is the orbital-pair indices, second is the
        overlap value.
    '''
    s = gto.intor_cross('int1e_ovlp', mol1, mol2)
    s = reduce(numpy.dot, (mo1.T, s, mo2))
    idx = numpy.argwhere(abs(s) > tol) + base
    for i,j in idx:
        logger.info(mol1, '<mo-1|mo-2>  %d  %d  %12.8f',
                    i, j, s[i,j])
    return idx, s

def mo_1to1map(s):
    '''Given <i|j>, search for the 1-to-1 mapping between i and j.

    Returns:
        a list [j-close-to-i for i in <bra|]
    '''
    s1 = abs(s)
    like_input = []
    for i in range(s1.shape[0]):
        k = numpy.argmax(s1[i])
        like_input.append(k)
        s1[:,k] = 0
    return like_input

def mo_comps(aolabels_or_baslst, mol, mo_coeff, cart=False,
             orth_method=ORTH_METHOD):
    '''Given AO(s), show how the AO(s) are distributed in MOs.

    Args:
        aolabels_or_baslst : filter function or AO labels or AO index
            If it's a function,  the AO indices are the items for which the
            function return value is true.

    Kwargs:
        cart : bool
            whether the orbital coefficients are based on cartesian basis.
        orth_method : str
            The localization method to generated orthogonal AO upon which the AO
            contribution are computed.  It can be one of 'meta_lowdin',
            'lowdin' or 'nao'.

    Returns:
        A list of float to indicate the total contributions (normalized to 1) of
        localized AOs

    Examples:

    >>> from pyscf import gto, scf
    >>> from pyscf.tools import mo_mapping
    >>> mol = gto.M(atom='H 0 0 0; F 0 0 1', basis='6-31g')
    >>> mf = scf.RHF(mol).run()
    >>> comp = mo_mapping.mo_comps('F 2s', mol, mf.mo_coeff)
    >>> print('MO-id    F-2s components')
    >>> for i,c in enumerate(comp):
    ...     print('%-3d      %.10f' % (i, c))
    MO-id    components
    0        0.0000066344
    1        0.8796915532
    2        0.0590259826
    3        0.0000000000
    4        0.0000000000
    5        0.0435028851
    6        0.0155889103
    7        0.0000000000
    8        0.0000000000
    9        0.0000822361
    10       0.0021017982
    '''
    with lib.temporary_env(mol, cart=cart):
        assert(mo_coeff.shape[0] == mol.nao)
        s = mol.intor_symmetric('int1e_ovlp')
        lao = lo.orth.orth_ao(mol, orth_method, s=s)

        idx = gto.mole._aolabels2baslst(mol, aolabels_or_baslst)
        if len(idx) == 0:
            logger.warn(mol, 'Required orbitals are not found')
        mo1 = reduce(numpy.dot, (lao[:,idx].T, s, mo_coeff))
        s1 = numpy.einsum('ki,ki->i', mo1, mo1)
    return s1

del(BASE, MAP_TOL, ORTH_METHOD)


# pyscf/tools/molden.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

# MOLDEN format:
# http://www.cmbi.ru.nl/molden/molden_format.html

import sys
import re
import numpy
import pyscf
from pyscf import lib
from pyscf import gto
from pyscf.lib import logger
from pyscf import __config__

IGNORE_H = getattr(__config__, 'molden_ignore_h', True)


def orbital_coeff(mol, fout, mo_coeff, spin='Alpha', symm=None, ene=None,
                  occ=None, ignore_h=IGNORE_H):
    from pyscf.symm import label_orb_symm
    if mol.cart:
        # pyscf Cartesian GTOs are not normalized. This may not be consistent
        # with the requirements of molden format. Normalize Cartesian GTOs here
        norm = mol.intor('int1e_ovlp').diagonal() ** .5
        mo_coeff = numpy.einsum('i,ij->ij', norm, mo_coeff)

    if ignore_h:
        mol, mo_coeff = remove_high_l(mol, mo_coeff)

    aoidx = order_ao_index(mol)
    nmo = mo_coeff.shape[1]
    if symm is None:
        symm = ['A']*nmo
        if mol.symmetry:
            try:
                symm = label_orb_symm(mol, mol.irrep_name, mol.symm_orb,
                                      mo_coeff, tol=1e-5)
            except ValueError as e:
                logger.warn(mol, str(e))
    if ene is None:
        ene = numpy.arange(nmo)
    assert(spin == 'Alpha' or spin == 'Beta')
    if occ is None:
        occ = numpy.zeros(nmo)
        neleca, nelecb = mol.nelec
        if spin == 'Alpha':
            occ[:neleca] = 1
        else:
            occ[:nelecb] = 1

    if spin == 'Alpha':
        # Avoid duplicated [MO] session when dumping beta orbitals
        fout.write('[MO]\n')

    for imo in range(nmo):
        fout.write(' Sym= %s\n' % symm[imo])
        fout.write(' Ene= %15.10g\n' % ene[imo])
        fout.write(' Spin= %s\n' % spin)
        fout.write(' Occup= %10.5f\n' % occ[imo])
        for i,j in enumerate(aoidx):
            fout.write(' %3d    %18.14g\n' % (i+1, mo_coeff[j,imo]))

def from_mo(mol, filename, mo_coeff, spin='Alpha', symm=None, ene=None,
            occ=None, ignore_h=IGNORE_H):
    '''Dump the given MOs in Molden format'''
    with open(filename, 'w') as f:
        header(mol, f, ignore_h)
        orbital_coeff(mol, f, mo_coeff, spin, symm, ene, occ, ignore_h)


def from_scf(mf, filename, ignore_h=IGNORE_H):
    '''Dump the given SCF object in Molden format'''
    dump_scf(mf, filename, ignore_h)
def dump_scf(mf, filename, ignore_h=IGNORE_H):
    import pyscf.scf
    mol = mf.mol
    mo_coeff = mf.mo_coeff
    with open(filename, 'w') as f:
        header(mol, f, ignore_h)
        if isinstance(mf, pyscf.scf.uhf.UHF) or 'UHF' == mf.__class__.__name__:
            orbital_coeff(mol, f, mo_coeff[0], spin='Alpha',
                          ene=mf.mo_energy[0], occ=mf.mo_occ[0],
                          ignore_h=ignore_h)
            orbital_coeff(mol, f, mo_coeff[1], spin='Beta',
                          ene=mf.mo_energy[1], occ=mf.mo_occ[1],
                          ignore_h=ignore_h)
        else:
            orbital_coeff(mf.mol, f, mf.mo_coeff,
                          ene=mf.mo_energy, occ=mf.mo_occ, ignore_h=ignore_h)

def from_mcscf(mc, filename, ignore_h=IGNORE_H, cas_natorb=False):
    mol = mc.mol
    dm1 = mc.make_rdm1()
    if cas_natorb:
        mo_coeff, _, mo_energy = mc.canonicalize(sort=True, cas_natorb=cas_natorb)
    else:
        mo_coeff, mo_energy = mc.mo_coeff, mc.mo_energy

    mo_inv = numpy.dot(mc._scf.get_ovlp(), mo_coeff)
    occ = numpy.einsum('pi,pq,qi->i', mo_inv, dm1, mo_inv)
    with open(filename, 'w') as f:
        header(mol, f, ignore_h)
        orbital_coeff(mol, f, mo_coeff, ene=mo_energy, occ=occ, ignore_h=ignore_h)

def from_chkfile(filename, chkfile, key='scf/mo_coeff', ignore_h=IGNORE_H):
    import pyscf.scf
    with open(filename, 'w') as f:
        if key == 'scf/mo_coeff':
            mol, mf = pyscf.scf.chkfile.load_scf(chkfile)
            header(mol, f, ignore_h)
            ene = mf['mo_energy']
            occ = mf['mo_occ']
            mo = mf['mo_coeff']
        else:
            mol = pyscf.scf.chkfile.load_mol(chkfile)
            header(mol, f, ignore_h)
            dat = pyscf.scf.chkfile.load(chkfile, key.split('/')[0])
            if 'mo_energy' in dat:
                ene = dat['mo_energy']
            else:
                ene = None
            occ = dat['mo_occ']
            mo = dat['mo_coeff']

        if isinstance(ene, str) and ene == 'None':
            ene = None
        if isinstance(ene, str) and occ == 'None':
            occ = None
        if occ.ndim == 2:
            orbital_coeff(mol, f, mo[0], spin='Alpha', ene=ene[0], occ=occ[0],
                          ignore_h=ignore_h)
            orbital_coeff(mol, f, mo[1], spin='Beta', ene=ene[1], occ=occ[1],
                          ignore_h=ignore_h)
        else:
            orbital_coeff(mol, f, mo, ene=ene, occ=occ, ignore_h=ignore_h)


_SEC_RE = re.compile(r'\[[^]]+\]')

def _read_one_section(molden_fp):
    sec = [None]
    last_pos = 0
    while True:
        line = molden_fp.readline()
        if not line:
            break

        line = line.strip()
        if line == '' or line[0] == '#':  # comment or blank line
            continue

        mo = _SEC_RE.match(line)
        if mo:
            if sec[0] is None:
                sec[0] = line
            else:
                # Next section? rewind the fp pointer
                molden_fp.seek(last_pos)
                break
        else:
            sec.append(line)

        last_pos = molden_fp.tell()

    return sec

def _parse_natoms(lines, envs):
    envs['natm'] = natm = int(lines[1])
    return natm

def _parse_atoms(lines, envs):
    if 'ANG' in lines[0].upper():
        envs['unit'] = 1
    unit = envs['unit']

    envs['atoms'] = atoms = []
    for line in lines[1:]:
        dat = line.split()
        symb, atmid, chg = dat[:3]
        coord = numpy.array([float(x) for x in dat[3:]])*unit
        atoms.append((gto.mole._std_symbol(symb)+atmid, coord))

    if envs['natm'] is not None and envs['natm'] != len(atoms):
        sys.stderr.write('Number of atoms in section ATOMS does not equal to N_ATOMS\n')
    return atoms

def _parse_charge(lines, envs):
    mulliken_charges = [float(_d2e(x)) for x in lines[1:]]
    return mulliken_charges

def _parse_gto(lines, envs):
    mol = envs['mol']
    atoms = envs['atoms']
    basis = {}
    lines_iter = iter(lines)
    next(lines_iter)  # skip section header

# * Do not use iter() here. Python 2 and 3 are different in iter()
    def read_one_bas(lsym, nb, fac=1):
        fac = float(fac)
        if fac == float(0):
            fac = float(1)
        bas = [lib.param.ANGULARMAP[lsym.lower()],]
        for i in range(int(nb)):
            dat = _d2e(next(lines_iter)).split()
            bas.append((float(dat[0]), float(dat[1])*fac))
        return bas

# * Be careful with the atom sequence in [GTO] session, it does not correspond
# to the atom sequence in [Atoms] session.
    atom_seq = []

    for line in lines_iter:
        dat = line.split()
        if dat[0].isdigit():
            atom_seq.append(int(dat[0])-1)
            symb = atoms[int(dat[0])-1][0]
            basis[symb] = []

        elif dat[0].upper() in 'SPDFGHIJ':
            basis[symb].append(read_one_bas(*dat))

    mol.basis = envs['basis'] = basis
    mol.atom = [atoms[i] for i in atom_seq]
    return mol

def _parse_mo(lines, envs):
    mol = envs['mol']
    if not mol._built:
        try:
            mol.build(0, 0)
        except RuntimeError:
            mol.build(0, 0, spin=1)

    irrep_labels = []
    mo_energy = []
    spins = []
    mo_occ = []
    mo_coeff_prim = [] # primary data, will be reworked for missing values
    coeff_idx = []
    mo_id = 0
    for line in lines[1:]:
        line = line.upper()
        if 'SYM' in line:
            irrep_labels.append(line.split('=')[1].strip())
        elif 'ENE' in line:
            mo_energy.append(float(_d2e(line).split('=')[1].strip()))
            mo_id = len(mo_energy) - 1
        elif 'SPIN' in line:
            spins.append(line.split('=')[1].strip())
        elif 'OCC' in line:
            mo_occ.append(float(_d2e(line.split('=')[1].strip())))
        else:
            ao_id, c = line.split()[:2]
            coeff_idx.append([int(ao_id) - 1, mo_id])
            mo_coeff_prim.append(float(c))

    coeff_idx = numpy.array(coeff_idx)
    number_of_aos, number_of_mos = coeff_idx.max(axis=0) + 1
    mo_coeff = numpy.zeros([number_of_aos, number_of_mos])
    mo_coeff[coeff_idx[:,0], coeff_idx[:,1]] = mo_coeff_prim

    mo_energy = numpy.array(mo_energy)
    mo_occ = numpy.array(mo_occ)
    aoidx = numpy.argsort(order_ao_index(mol))
    mo_coeff = mo_coeff[aoidx]
    if mol.cart:
        # Cartesian GTOs are normalized in molden format but they are not in pyscf
        s = mol.intor('int1e_ovlp')
        mo_coeff = numpy.einsum('i,ij->ij', numpy.sqrt(1/s.diagonal()), mo_coeff)


    return mol, mo_energy, mo_coeff, mo_occ, irrep_labels, spins


def _parse_core(lines, envs):
    mol = envs['mol']
    atoms = envs['atoms']
    for line in lines[1:]:
        dat = line.split(':')
        if dat[0].strip().isdigit():
            atm_id = int(dat[0].strip()) - 1
            nelec_core = int(dat[1].strip())
            mol.ecp[atoms[atm_id][0]] = [nelec_core, []]

    if mol.ecp:
        sys.stderr.write('\nECP were dectected in the molden file.\n'
                         'Note Molden format does not support ECP data. '
                         'ECP information was lost when saving to molden format.\n\n')
    return mol.ecp

_SEC_PARSER = {'N_ATOMS'  : _parse_natoms,
               'ATOMS'    : _parse_atoms,
               'GTO'      : _parse_gto,
               'CHARGE'   : _parse_charge,
               'MO'       : _parse_mo,
               'CORE'     : _parse_core,
               'MOLDEN FORMAT' : lambda *args: None,}

_SEC_ORDER = ['N_ATOMS', 'ATOMS', 'GTO', 'CHARGE', 'MO', 'CORE', 'MOLDEN FORMAT']

def load(moldenfile, verbose=0):
    '''Extract mol and orbitals from molden file
    '''
    sec_kinds = {} # found sections and their lines are stored in this dic
    with open(moldenfile, 'r') as f:
        mol = gto.Mole()
        mol.cart = True
        tokens = {'natm'  : None,
                  'unit'  : lib.param.BOHR,
                  'mol'   : mol,
                  'atoms' : None,
                  'basis' : None,}

        while True:
            lines = _read_one_section(f)
            sec_title = lines[0]
            if sec_title is None:
                break

            sec_title = sec_title[1:sec_title.index(']')].upper()
            if sec_title in _SEC_PARSER:
                if sec_title not in sec_kinds:
                    sec_kinds.update({sec_title : [lines]})
                else:
                    sec_kinds[sec_title].append(lines)

            elif sec_title[:2] in ('5D', '7F', '9G'):
                mol.cart = False

            elif sec_title[:2] == '6D' or sec_title[:3] in ('10F', '15G'):
                mol.cart = True

            else:
                sys.stderr.write('Unknown section %s\n' % sec_title)

    for sec_kind in _SEC_ORDER:
        if sec_kind == 'MO' and 'MO' in sec_kinds:
            if len(sec_kinds['MO']) == 1:
                mol, mo_energy, mo_coeff, mo_occ, irrep_labels, spins = \
                        _parse_mo(sec_kinds['MO'][0], tokens)
                # If found only one MO section while 'B' appears in the spins
                # labels, the MOs so obtained are spin orbitals, with beta
                # orbitals at the second half of the mo_coeff matrix.
                if any(s[0] == 'B' for s in spins):
                    if mo_coeff.shape[0] == mo_coeff.shape[1]:
                        # general spin orbitals which allows to mix spin alpha
                        # and spin beta components in the same orbitals
                        raise NotImplementedError
                    else:
                        # Regular spin orbitals, alpha and beta do not mix
                        beta_idx = numpy.array([s[0] == 'B' for s in spins])
                        alpha_idx = ~beta_idx
                        mo_energy = mo_energy[alpha_idx], mo_energy[beta_idx]
                        mo_coeff = mo_coeff[:,alpha_idx], mo_coeff[:,beta_idx]
                        mo_occ = mo_occ[alpha_idx], mo_occ[beta_idx]
                        irrep_labels = numpy.array(irrep_labels)
                        irrep_labels = irrep_labels[alpha_idx], irrep_labels[beta_idx]
                        spins = numpy.array(spins)
                        spins = spins[alpha_idx], spins[beta_idx]

            elif len(sec_kinds['MO']) == 2:
                res_a = _parse_mo(sec_kinds['MO'][0], tokens)
                res_b = _parse_mo(sec_kinds['MO'][1], tokens)
                mo_energy, mo_coeff, mo_occ, irrep_labels, spins = \
                        list(zip(res_a[1:], res_b[1:]))
                mol = res_b[0]

        if sec_kind in sec_kinds:
            for n, content in enumerate(sec_kinds[sec_kind]):
                _SEC_PARSER[sec_kind](content, tokens)

    if isinstance(mo_occ, tuple):
        mol.spin = int(mo_occ[0].sum() - mo_occ[1].sum())

    return mol, mo_energy, mo_coeff, mo_occ, irrep_labels, spins

parse = read = load

def _d2e(token):
    return token.replace('D', 'e').replace('d', 'e')

def header(mol, fout, ignore_h=IGNORE_H):
    if ignore_h:
        mol = remove_high_l(mol)[0]
    fout.write('[Molden Format]\n')
    fout.write('made by pyscf v[%s]\n' % pyscf.__version__)
    fout.write('[Atoms] (AU)\n')
    for ia in range(mol.natm):
        symb = mol.atom_pure_symbol(ia)
        chg = mol.atom_charge(ia)
        fout.write('%s   %d   %d   ' % (symb, ia+1, chg))
        coord = mol.atom_coord(ia)
        fout.write('%18.14f   %18.14f   %18.14f\n' % tuple(coord))

    fout.write('[GTO]\n')
    for ia, (sh0, sh1, p0, p1) in enumerate(mol.offset_nr_by_atom()):
        fout.write('%d 0\n' %(ia+1))
        for ib in range(sh0, sh1):
            l = mol.bas_angular(ib)
            nprim = mol.bas_nprim(ib)
            nctr = mol.bas_nctr(ib)
            es = mol.bas_exp(ib)
            cs = mol.bas_ctr_coeff(ib)
            for ic in range(nctr):
                fout.write(' %s   %2d 1.00\n' % (lib.param.ANGULAR[l], nprim))
                for ip in range(nprim):
                    fout.write('    %18.14g  %18.14g\n' % (es[ip], cs[ip,ic]))
        fout.write('\n')

    if mol.cart:
        fout.write('[6d]\n[10f]\n[15g]\n')
    else:
        fout.write('[5d]\n[7f]\n[9g]\n')

    if mol.has_ecp():  # See https://github.com/zorkzou/Molden2AIM
        fout.write('[core]\n')
        for ia in range(mol.natm):
            nelec_ecp_core = mol.atom_nelec_core(ia)
            if nelec_ecp_core != 0:
                fout.write('%s : %d\n' % (ia+1, nelec_ecp_core))
    fout.write('\n')

def order_ao_index(mol):
    # reorder d,f,g fucntion to
    #  5D: D 0, D+1, D-1, D+2, D-2
    #  6D: xx, yy, zz, xy, xz, yz
    #
    #  7F: F 0, F+1, F-1, F+2, F-2, F+3, F-3
    # 10F: xxx, yyy, zzz, xyy, xxy, xxz, xzz, yzz, yyz, xyz
    #
    #  9G: G 0, G+1, G-1, G+2, G-2, G+3, G-3, G+4, G-4
    # 15G: xxxx yyyy zzzz xxxy xxxz yyyx yyyz zzzx zzzy xxyy xxzz yyzz xxyz yyxz zzxy
    idx = []
    off = 0
    if mol.cart:
        for ib in range(mol.nbas):
            l = mol.bas_angular(ib)
            for n in range(mol.bas_nctr(ib)):
                if l == 2:
                    idx.extend([off+0,off+3,off+5,off+1,off+2,off+4])
                elif l == 3:
                    idx.extend([off+0,off+6,off+9,off+3,off+1,
                                off+2,off+5,off+8,off+7,off+4])
                elif l == 4:
                    idx.extend([off+0 , off+10, off+14, off+1 , off+2 ,
                                off+6 , off+11, off+9 , off+13, off+3 ,
                                off+5 , off+12, off+4 , off+7 , off+8 ,])
                elif l > 4:
                    raise RuntimeError('l=5 is not supported')
                else:
                    idx.extend(range(off,off+(l+1)*(l+2)//2))
                off += (l+1)*(l+2)//2
    else:  # spherical orbitals
        for ib in range(mol.nbas):
            l = mol.bas_angular(ib)
            for n in range(mol.bas_nctr(ib)):
                if l == 2:
                    idx.extend([off+2,off+3,off+1,off+4,off+0])
                elif l == 3:
                    idx.extend([off+3,off+4,off+2,off+5,off+1,off+6,off+0])
                elif l == 4:
                    idx.extend([off+4,off+5,off+3,off+6,off+2,
                                off+7,off+1,off+8,off+0])
                elif l > 4:
                    raise RuntimeError('l=5 is not supported')
                else:
                    idx.extend(range(off,off+l*2+1))
                off += l * 2 + 1
    return idx

def remove_high_l(mol, mo_coeff=None):
    '''Remove high angular momentum (l >= 5) functions before dumping molden file.
    If molden function raised error message ``RuntimeError l=5 is not supported``,
    you can use this function to format orbitals.

    Note the formated orbitals may have normalization problem.  Some visualization
    tool will complain about the orbital normalization error.

    Examples:

    >>> mol1, orb1 = remove_high_l(mol, mf.mo_coeff)
    >>> molden.from_mo(mol1, outputfile, orb1)
    '''
    pmol = mol.copy()
    pmol.basis = {}
    for symb, bas in mol._basis.items():
        pmol.basis[symb] = [b for b in bas if b[0] <= 4]
    pmol.build(0, 0)
    if mo_coeff is None:
        return pmol, None
    else:
        p1 = 0
        idx = []
        for ib in range(mol.nbas):
            l = mol.bas_angular(ib)
            nc = mol.bas_nctr(ib)
            if mol.cart:
                nd = (l + 1) * (l + 2) // 2
            else:
                nd = l * 2 + 1
            p0, p1 = p1, p1 + nd * nc
            if l <= 4:
                idx.append(range(p0, p1))

        idx = numpy.hstack(idx)
        return pmol, mo_coeff[idx]



if __name__ == '__main__':
    from pyscf import scf
    import tempfile
    mol = gto.Mole()
    mol.verbose = 5
    mol.output = None#'out_gho'
    mol.atom = [['C', (0.,0.,0.)],
                ['H', ( 1, 1, 1)],
                ['H', (-1,-1, 1)],
                ['H', ( 1,-1,-1)],
                ['H', (-1, 1,-1)], ]
    mol.basis = {
        'C': 'sto-3g',
        'H': 'sto-3g'}
    mol.build(dump_input=False)
    m = scf.RHF(mol)
    m.scf()
    header(mol, mol.stdout)
    print(order_ao_index(mol))
    orbital_coeff(mol, mol.stdout, m.mo_coeff)

    ftmp = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
    from_mo(mol, ftmp.name, m.mo_coeff)

    print(parse(ftmp.name))


# pyscf/tools/ring.py

# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy

def make(nat, b=1.):
    r = b/2 / numpy.sin(numpy.pi/nat)
    atms = []
    for i in range(nat):
        theta = i * (2*numpy.pi/nat)
        atms.append((r*numpy.cos(theta), r*numpy.sin(theta), 0))
    return atms

if __name__ == '__main__':
    for c in make(6):
        print(c)



# pyscf/tools/wfn_format.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#         Jose Luis Casals Sainz <jluiscasalssainz@gmail.com>
#

'''
GAMESS WFN File format
'''

from pyscf import gto
from pyscf import lib
import numpy

# types
# 1 S
# 2 PX
# 3 PY
# 4 PZ
# 5 DXX
# 6 DYY
# 7 DZZ
# 8 DXY
# 9 DXZ
# 10 DYZ
# 11 FXXX
# 12 FYYY
# 13 FZZZ
# 14 FXXY
# 15 FXXZ
# 16 FYYZ
# 17 FXYY
# 18 FXZZ
# 19 FYZZ
# 20 FXYZ
# 21 GXXXX
# 22 GYYYY
# 23 GZZZZ
# 24 GXXXY
# 25 GXXXZ
# 26 GXYYY
# 27 GYYYZ
# 28 GXZZZ
# 29 GYZZZ
# 30 GXXYY
# 31 GXXZZ
# 32 GYYZZ
# 33 GXXYZ
# 34 GXYYZ
# 35 GXYZZ
# 36 HZZZZZ
# 37 HYZZZZ
# 38 HYYZZZ
# 39 HYYYZZ
# 40 HYYYYZ
# 41 HYYYYY
# 42 HXZZZZ
# 43 HXYZZZ
# 44 HXYYZZ
# 45 HXYYYZ
# 46 HXYYYY
# 47 HXXZZZ
# 48 HXXYZZ
# 49 HXXYYZ
# 50 HXXYYY
# 51 HXXXZZ
# 52 HXXXYZ
# 53 HXXXYY
# 54 HXXXXZ
# 55 HXXXXY
# 56 HXXXXX
TYPE_MAP = [
    [1],  # S
    [2, 3, 4],  # P
    [5, 8, 9, 6, 10, 7],  # D
    [11,14,15,17,20,18,12,16,19,13],  # F
    [21,24,25,30,33,31,26,34,35,28,22,27,32,29,23],  # G
    [56,55,54,53,52,51,50,49,48,47,46,45,44,43,42,41,40,39,38,37,36],  # H
]
def write_mo(fout, mol, mo_coeff, mo_energy=None, mo_occ=None):
    if mol.cart:
        raise NotImplementedError('Cartesian basis not available')

    #FIXME: Duplicated primitives may lead to problems.  x2c._uncontract_mol
    # is the workaround at the moment to remove duplicated primitives.
    from pyscf.x2c import x2c
    mol, ctr = x2c._uncontract_mol(mol, True, 0.)
    mo_coeff = numpy.dot(ctr, mo_coeff)

    nmo = mo_coeff.shape[1]
    mo_cart = []
    centers = []
    types = []
    exps = []
    p0 = 0
    for ib in range(mol.nbas):
        ia = mol.bas_atom(ib)
        l = mol.bas_angular(ib)
        es = mol.bas_exp(ib)
        c = mol._libcint_ctr_coeff(ib)
        np, nc = c.shape
        nd = nc*(2*l+1)
        mosub = mo_coeff[p0:p0+nd].reshape(-1,nc,nmo)
        c2s = gto.cart2sph(l)
        mosub = numpy.einsum('yki,cy,pk->pci', mosub, c2s, c)
        mo_cart.append(mosub.transpose(1,0,2).reshape(-1,nmo))

        for t in TYPE_MAP[l]:
            types.append([t]*np)
        ncart = mol.bas_len_cart(ib)
        exps.extend([es]*ncart)
        centers.extend([ia+1]*(np*ncart))
        p0 += nd
    mo_cart = numpy.vstack(mo_cart)
    centers = numpy.hstack(centers)
    types = numpy.hstack(types)
    exps = numpy.hstack(exps)
    nprim, nmo = mo_cart.shape

    fout.write('From PySCF\n')
    fout.write('GAUSSIAN            %3d MOL ORBITALS    %3d PRIMITIVES      %3d NUCLEI\n'
               % (mo_cart.shape[1], mo_cart.shape[0], mol.natm))
    for ia in range(mol.natm):
        x, y, z = mol.atom_coord(ia)
        fout.write('  %-4s %-4d (CENTRE%3d)  %11.8f %11.8f %11.8f  CHARGE = %.1f\n'
                   % (mol.atom_pure_symbol(ia), ia+1, ia+1, x, y, z,
                      mol.atom_charge(ia)))
    for i0, i1 in lib.prange(0, nprim, 20):
        fout.write('CENTRE ASSIGNMENTS  %s\n' % ''.join('%3d'%x for x in centers[i0:i1]))
    for i0, i1 in lib.prange(0, nprim, 20):
        fout.write('TYPE ASSIGNMENTS    %s\n' % ''.join('%3d'%x for x in types[i0:i1]))
    for i0, i1 in lib.prange(0, nprim, 5):
        fout.write('EXPONENTS  %s\n' % ' '.join('%13.7E'%x for x in exps[i0:i1]))

    for k in range(nmo):
        mo = mo_cart[:,k]
        if mo_energy is None and mo_occ is None:
            fout.write('CANMO %d\n' % (k+1))
        elif mo_energy is None:
            fout.write('MO  %-4d                  OCC NO = %12.8f ORB. ENERGY = 0.0000\n' %
                       (k+1, mo_occ[k]))
        else:
            fout.write('MO  %-4d                  OCC NO = %12.8f ORB. ENERGY = %12.8f\n' %
                       (k+1, mo_occ[k], mo_energy[k]))
        for i0, i1 in lib.prange(0, nprim, 5):
            fout.write(' %s\n' % ' '.join('%15.8E'%x for x in mo[i0:i1]))
    fout.write('END DATA\n')
    if mo_energy is None or mo_occ is None:
        fout.write('ALDET    ENERGY =        0.0000000000   VIRIAL(-V/T)  =   0.00000000\n')
    elif mo_energy is None and mo_occ is None:
        pass
    else :
        fout.write('RHF      ENERGY =        0.0000000000   VIRIAL(-V/T)  =   0.00000000\n')

def write_ci(fout, fcivec, norb, nelec, ncore=0):
    from pyscf import fci
    if isinstance(nelec, (int, numpy.number)):
        nelecb = nelec//2
        neleca = nelec - nelecb
    else:
        neleca, nelecb = nelec
    fout.write('NELACTIVE, NDETS, NORBCORE, NORBACTIVE\n')
    fout.write(' %5d %5d %5d %5d\n' % (neleca+nelecb, fcivec.size, ncore, norb))
    fout.write('COEFFICIENT/ OCCUPIED ACTIVE SPIN ORBITALS\n')

    nb = fci.cistring.num_strings(norb, nelecb)
    stringsa = fci.cistring.gen_strings4orblist(range(norb), neleca)
    stringsb = fci.cistring.gen_strings4orblist(range(norb), nelecb)
    def str2orbidx(string, ncore):
        bstring = bin(string)
        return [i+1+ncore for i,s in enumerate(bstring[::-1]) if s == '1']

    addrs = numpy.argsort(abs(fcivec.ravel()))
    for iaddr in reversed(addrs):
        addra, addrb = divmod(iaddr, nb)
        idxa = ['%3d' % x for x in str2orbidx(stringsa[addra], ncore)]
        idxb = ['%3d' % (-x) for x in str2orbidx(stringsb[addrb], ncore)]
        #TODO:add a cuttoff and a counter for ndets
        fout.write('%18.10E %s %s\n' % (fcivec[addra,addrb], ' '.join(idxa), ' '.join(idxb)))

if __name__ == '__main__':
    from pyscf import scf, mcscf, symm
    from pyscf.tools import molden
    mol = gto.M(atom='N 0 0 0; N 0 0 2.88972599',
                unit='B', basis='ccpvtz', verbose=4,
                symmetry=1, symmetry_subgroup='d2h')
    mf = scf.RHF(mol).run()
    coeff = mf.mo_coeff[:,mf.mo_occ>0]
    energy = mf.mo_energy[mf.mo_occ>0]
    occ = mf.mo_occ[mf.mo_occ>0]
    with open('n2_hf.wfn', 'w') as f2:
        write_mo(f2, mol, coeff, energy, occ)
#
    mc = mcscf.CASSCF(mf, 10, 10)
    mc.kernel()
    nmo = mc.ncore + mc.ncas
    nelecas = mc.nelecas[0] + mc.nelecas[1]
    casdm1, casdm2 = mc.fcisolver.make_rdm12(mc.ci, mc.ncas, mc.nelecas)
    rdm1, rdm2 = mcscf.addons._make_rdm12_on_mo(casdm1, casdm2, mc.ncore, mc.ncas, nmo)
    den_file = 'n2_cas.den'
    fspt = open(den_file,'w')
    fspt.write('CCIQA    ENERGY =      0.000000000000 THE VIRIAL(-V/T)=   2.00000000\n')
    fspt.write('La matriz D es:\n')
    for i in range(nmo):
        for j in range(nmo):
            fspt.write('%i %i %.16f\n' % ((i+1), (j+1), rdm1[i,j]))
    fspt.write('La matriz d es:\n')
    for i in range(nmo):
        for j in range(nmo):
            for k in range(nmo):
                for l in range(nmo):
                    if (abs(rdm2[i,j,k,l]) > 1e-12):
                        fspt.write('%i %i %i %i %.16f\n' % ((i+1), (j+1), (k+1), (l+1), rdm2[i,j,k,l]))
    fspt.close()
    orbsym = symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, mc.mo_coeff[:,:nmo])
    natocc, natorb = symm.eigh(-rdm1, orbsym)
    for i, k in enumerate(numpy.argmax(abs(natorb), axis=0)):
        if natorb[k,i] < 0:
            natorb[:,i] *= -1
    natorb = numpy.dot(mc.mo_coeff[:,:nmo], natorb)
    natocc = -natocc
    with open('n2_cas.det', 'w') as f3:
        write_ci(f3, mc.ci, mc.ncas, mc.nelecas, mc.ncore)
    with open('n2_ref_nat.mol', 'w') as f2:
        molden.header(mol, f2)
        molden.orbital_coeff(mol, f2, natorb, occ=natocc)
    with open('n2_ref.mol', 'w') as f2:
        molden.header(mol, f2)
        molden.orbital_coeff(mol, f2, mc.mo_coeff[:,:nmo], occ=mf.mo_occ[:nmo])
    with open('n2_cas.wfn', 'w') as f2:
        write_mo(f2, mol, natorb, mo_occ=natocc)
        write_mo(f2, mol, mc.mo_coeff[:,:nmo])
    with open('rdm_wfn.wfn', 'w') as f2:
        write_mo(f2, mol, mc.mo_coeff[:,:nmo], mo_occ=mf.mo_occ[:nmo])



# pyscf/x2c/__init__.py

# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pyscf.x2c.x2c import RHF, UHF
try:
    from pyscf.x2c.dft import RKS, UKS
except ImportError:
    pass


# pyscf/x2c/_response_functions.py

#!/usr/bin/env python
# Copyright 2022 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''
Generate X2C-SCF response functions
'''

from pyscf.x2c import x2c
from pyscf.x2c._response_functions import _gen_ghf_response

def _gen_x2chf_response(mf, mo_coeff=None, mo_occ=None,
                        with_j=True, hermi=0, max_memory=None):
    '''Generate a function to compute the product of X2C-HF response function
    and density matrices.
    '''
    return _gen_ghf_response(mf, mo_coeff, mo_occ, with_j, hermi, max_memory)

x2c.UHF.gen_response = _gen_x2chf_response
x2c.RHF.gen_response = _gen_x2chf_response


# pyscf/x2c/dft.py

#!/usr/bin/env python
# Copyright 2022 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''
X2C 2-component DFT methods
'''

from pyscf.x2c import x2c
from pyscf.dft import dks

class UKS(dks.KohnShamDFT, x2c.UHF):
    def __init__(self, mol, xc='LDA,VWN'):
        x2c.UHF.__init__(self, mol)
        dks.KohnShamDFT.__init__(self, xc)

    def dump_flags(self, verbose=None):
        x2c.UHF.dump_flags(self, verbose)
        dks.KohnShamDFT.dump_flags(self, verbose)
        return self

    def to_hf(self):
        '''Convert the input mean-field object to an X2C-HF object.

        Note this conversion only changes the class of the mean-field object.
        The total energy and wave-function are the same as them in the input
        mean-field object.
        '''
        mf = self.view(x2c.UHF)
        mf.converged = False
        return mf

X2C_UKS = UKS

class RKS(dks.KohnShamDFT, x2c.RHF):
    def __init__(self, mol, xc='LDA,VWN'):
        x2c.RHF.__init__(self, mol)
        dks.KohnShamDFT.__init__(self, xc)

    def dump_flags(self, verbose=None):
        x2c.RHF.dump_flags(self, verbose)
        dks.KohnShamDFT.dump_flags(self, verbose)
        return self

    def to_hf(self):
        '''Convert the input mean-field object to an X2C-HF object.

        Note this conversion only changes the class of the mean-field object.
        The total energy and wave-function are the same as them in the input
        mean-field object.
        '''
        mf = self.view(x2c.RHF)
        mf.converged = False
        return mf

X2C_RKS = RKS


# pyscf/x2c/newton_ah.py

#!/usr/bin/env python
# Copyright 2022 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''
Second order X2C-SCF solver
'''

import numpy
from pyscf.scf import hf
from pyscf.x2c import x2c
# To ensure .gen_response() methods are registered
from pyscf.x2c import _response_functions  # noqa
from pyscf.soscf import newton_ah

gen_g_hop_x2chf = newton_ah.gen_g_hop_dhf

def newton(mf):
    '''Co-iterative augmented hessian (CIAH) second order SCF solver

    Examples:

    >>> mol = gto.M(atom='H 0 0 0; H 0 0 1.1', basis='cc-pvdz')
    >>> mf = x2c.RHF(mol).newton()
    >>> mf.kernel()
    -1.0811707843774987
    '''
    assert isinstance(mf, x2c.SCF)

    if isinstance(mf, newton_ah._CIAH_SOSCF):
        return mf

    if isinstance(mf, x2c.RHF):
        class SecondOrderX2CHF(newton_ah._CIAH_SOSCF, mf.__class__):
            gen_g_hop = gen_g_hop_x2chf

            def update_rotate_matrix(self, dx, mo_occ, u0=1, mo_coeff=None):
                nmo = mo_occ.size
                nocc = numpy.count_nonzero(mo_occ)
                nvir = nmo - nocc
                dx = dx.reshape(nvir, nocc)
                dx_aa = dx[::2,::2]
                dr_aa = hf.unpack_uniq_var(dx_aa.ravel, mo_occ[::2])
                u = numpy.zeros((nmo, nmo), dtype=dr_aa.dtype)
                # Allows only the rotation within the up-up space and down-down space
                u[::2,::2] = u[1::2,1::2] = newton_ah.expmat(dr_aa)
                return numpy.dot(u0, u)

            def rotate_mo(self, mo_coeff, u, log=None):
                mo = numpy.dot(mo_coeff, u)
                return mo
    else:
        class SecondOrderX2CHF(newton_ah._CIAH_SOSCF, mf.__class__):
            gen_g_hop = gen_g_hop_x2chf

            def update_rotate_matrix(self, dx, mo_occ, u0=1, mo_coeff=None):
                dr = hf.unpack_uniq_var(dx, mo_occ)
                return numpy.dot(u0, newton_ah.expmat(dr))

            def rotate_mo(self, mo_coeff, u, log=None):
                mo = numpy.dot(mo_coeff, u)
                return mo

    return SecondOrderX2CHF(mf)


# pyscf/x2c/sfx2c1e.py

#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''
1-electron Spin-free X2C approximation
'''

from functools import reduce
import numpy
import scipy.linalg
from pyscf import lib
from pyscf import gto
from pyscf.lib import logger
from pyscf.scf import hf
from pyscf.scf import ghf
from pyscf.x2c import x2c
from pyscf.data import nist


def sfx2c1e(mf):
    '''Spin-free X2C.
    For the given SCF object, it updates the hcore constructor.  All integrals
    are computed in the real spherical GTO basis.

    Args:
        mf : an SCF object

    Returns:
        An SCF object

    Examples:

    >>> mol = gto.M(atom='H 0 0 0; F 0 0 1', basis='ccpvdz', verbose=0)
    >>> mf = scf.RHF(mol).sfx2c1e()
    >>> mf.scf()

    >>> import pyscf.x2c.sfx2c1e
    >>> mol.symmetry = 1
    >>> mol.build(0, 0)
    >>> mf = pyscf.x2c.sfx2c1e.sfx2c1e(scf.UHF(mol))
    >>> mf.scf()
    '''
    assert isinstance(mf, hf.SCF)

    if isinstance(mf, x2c._X2C_SCF):
        if mf.with_x2c is None:
            mf.with_x2c = SpinFreeX2CHelper(mf.mol)
            return mf
        elif not isinstance(mf.with_x2c, SpinFreeX2CHelper):
            # An object associated to x2c.SpinOrbitalX2CHelper
            raise NotImplementedError
        else:
            return mf

    mf_class = mf.__class__
    if mf_class.__doc__ is None:
        doc = ''
    else:
        doc = mf_class.__doc__

    class SFX2C1E_SCF(x2c._X2C_SCF, mf_class):
        __doc__ = doc + '''
        Attributes for spin-free X2C:
            with_x2c : X2C object
        '''
        def __init__(self, mol, *args, **kwargs):
            mf_class.__init__(self, mol, *args, **kwargs)
            self.with_x2c = SpinFreeX2CHelper(mf.mol)
            self._keys = self._keys.union(['with_x2c'])

        def get_hcore(self, mol=None):
            if self.with_x2c:
                hcore = self.with_x2c.get_hcore(mol)
                if isinstance(self, ghf.GHF):
                    hcore = scipy.linalg.block_diag(hcore, hcore)
                return hcore
            else:
                return mf_class.get_hcore(self, mol)

        def dump_flags(self, verbose=None):
            mf_class.dump_flags(self, verbose)
            if self.with_x2c:
                self.with_x2c.dump_flags(verbose)
            return self

        def reset(self, mol):
            self.with_x2c.reset(mol)
            return mf_class.reset(self, mol)

        def dip_moment(self, mol=None, dm=None, unit='Debye', verbose=logger.NOTE,
                       picture_change=True, **kwargs):
            r''' Dipole moment calculation with picture change correction

            Args:
                 mol: an instance of :class:`Mole`
                 dm : a 2D ndarrays density matrices

            Kwarg:
                picture_chang (bool) : Whether to compute the dipole moment with
                picture change correction.

            Return:
                A list: the dipole moment on x, y and z component
            '''
            if mol is None: mol = self.mol
            if dm is None: dm = self.make_rdm1()
            log = logger.new_logger(mol, verbose)

            if not (isinstance(dm, numpy.ndarray) and dm.ndim == 2):
                # UHF denisty matrices
                dm = dm[0] + dm[1]

            if isinstance(self, ghf.GHF):
                nao = mol.nao_nr()
                dm = dm[:nao,:nao] + dm[nao:,nao:]

            with mol.with_common_orig((0,0,0)):
                if picture_change:
                    xmol = self.with_x2c.get_xmol()[0]
                    nao = xmol.nao
                    prp = xmol.intor_symmetric('int1e_sprsp').reshape(3,4,nao,nao)[:,0]
                    ao_dip = self.with_x2c.picture_change(('int1e_r', prp))
                else:
                    ao_dip = mol.intor_symmetric('int1e_r')

            el_dip = numpy.einsum('xij,ji->x', ao_dip, dm).real

            charges = mol.atom_charges()
            coords  = mol.atom_coords()
            nucl_dip = numpy.einsum('i,ix->x', charges, coords)
            mol_dip = nucl_dip - el_dip

            if unit.upper() == 'DEBYE':
                mol_dip *= nist.AU2DEBYE
                log.note('Dipole moment(X, Y, Z, Debye): %8.5f, %8.5f, %8.5f', *mol_dip)
            else:
                log.note('Dipole moment(X, Y, Z, A.U.): %8.5f, %8.5f, %8.5f', *mol_dip)
            return mol_dip

    with_x2c = SpinFreeX2CHelper(mf.mol)
    return mf.view(SFX2C1E_SCF).add_keys(with_x2c=with_x2c)

sfx2c = sfx2c1e


class SpinFreeX2CHelper(x2c.X2CHelperMixin):
    '''1-component X2c (spin-free part only)
    '''
    def get_hcore(self, mol=None):
        '''1-component X2c Foldy-Wouthuysen (FW Hamiltonian  (spin-free part only)
        '''
        if mol is None: mol = self.mol
        if mol.has_ecp():
            raise NotImplementedError

        xmol, contr_coeff = self.get_xmol(mol)
        c = lib.param.LIGHT_SPEED
        assert('1E' in self.approx.upper())
        t = xmol.intor_symmetric('int1e_kin')
        v = xmol.intor_symmetric('int1e_nuc')
        s = xmol.intor_symmetric('int1e_ovlp')
        w = xmol.intor_symmetric('int1e_pnucp')
        if 'get_xmat' in self.__dict__:
            # If the get_xmat method is overwritten by user, build the X
            # matrix with the external get_xmat method
            x = self.get_xmat(xmol)
            h1 = x2c._get_hcore_fw(t, v, w, s, x, c)

        elif 'ATOM' in self.approx.upper():
            atom_slices = xmol.offset_nr_by_atom()
            nao = xmol.nao_nr()
            x = numpy.zeros((nao,nao))
            for ia in range(xmol.natm):
                ish0, ish1, p0, p1 = atom_slices[ia]
                shls_slice = (ish0, ish1, ish0, ish1)
                t1 = xmol.intor('int1e_kin', shls_slice=shls_slice)
                s1 = xmol.intor('int1e_ovlp', shls_slice=shls_slice)
                with xmol.with_rinv_at_nucleus(ia):
                    z = -xmol.atom_charge(ia)
                    v1 = z * xmol.intor('int1e_rinv', shls_slice=shls_slice)
                    w1 = z * xmol.intor('int1e_prinvp', shls_slice=shls_slice)
                x[p0:p1,p0:p1] = x2c._x2c1e_xmatrix(t1, v1, w1, s1, c)
            h1 = x2c._get_hcore_fw(t, v, w, s, x, c)

        else:
            h1 = x2c._x2c1e_get_hcore(t, v, w, s, c)

        if self.basis is not None:
            s22 = xmol.intor_symmetric('int1e_ovlp')
            s21 = gto.intor_cross('int1e_ovlp', xmol, mol)
            c = lib.cho_solve(s22, s21)
            h1 = reduce(numpy.dot, (c.T, h1, c))
        if self.xuncontract and contr_coeff is not None:
            h1 = reduce(numpy.dot, (contr_coeff.T, h1, contr_coeff))
        return h1

    def picture_change(self, even_operator=(None, None), odd_operator=None):
        '''Picture change for even_operator + odd_operator

        even_operator has two terms at diagonal blocks
        [ v  0 ]
        [ 0  w ]

        odd_operator has the term at off-diagonal blocks
        [ 0    p ]
        [ p^T  0 ]

        v, w, and p can be strings (integral name) or matrices.
        '''
        mol = self.mol
        xmol, c = self.get_xmol(mol)
        pc_mat = self._picture_change(xmol, even_operator, odd_operator)

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

        elif self.xuncontract:
            pass

        else:
            return pc_mat

        if pc_mat.ndim == 2:
            return lib.einsum('pi,pq,qj->ij', c, pc_mat, c)
        else:
            return lib.einsum('pi,xpq,qj->xij', c, pc_mat, c)

    def get_xmat(self, mol=None):
        if mol is None:
            xmol = self.get_xmol(mol)[0]
        else:
            xmol = mol
        c = lib.param.LIGHT_SPEED
        assert('1E' in self.approx.upper())

        if 'ATOM' in self.approx.upper():
            atom_slices = xmol.offset_nr_by_atom()
            nao = xmol.nao_nr()
            x = numpy.zeros((nao,nao))
            for ia in range(xmol.natm):
                ish0, ish1, p0, p1 = atom_slices[ia]
                shls_slice = (ish0, ish1, ish0, ish1)
                t1 = xmol.intor('int1e_kin', shls_slice=shls_slice)
                s1 = xmol.intor('int1e_ovlp', shls_slice=shls_slice)
                with xmol.with_rinv_at_nucleus(ia):
                    z = -xmol.atom_charge(ia)
                    v1 = z * xmol.intor('int1e_rinv', shls_slice=shls_slice)
                    w1 = z * xmol.intor('int1e_prinvp', shls_slice=shls_slice)
                x[p0:p1,p0:p1] = x2c._x2c1e_xmatrix(t1, v1, w1, s1, c)
        else:
            t = xmol.intor_symmetric('int1e_kin')
            v = xmol.intor_symmetric('int1e_nuc')
            s = xmol.intor_symmetric('int1e_ovlp')
            w = xmol.intor_symmetric('int1e_pnucp')
            x = x2c._x2c1e_xmatrix(t, v, w, s, c)
        return x

    def _get_rmat(self, x=None):
        '''The matrix (in AO basis) that changes metric from NESC metric to NR metric'''
        xmol = self.get_xmol()[0]
        if x is None:
            x = self.get_xmat(xmol)
        c = lib.param.LIGHT_SPEED
        s = xmol.intor_symmetric('int1e_ovlp')
        t = xmol.intor_symmetric('int1e_kin')
        s1 = s + reduce(numpy.dot, (x.conj().T, t, x)) * (.5/c**2)
        return x2c._get_r(s, s1)

    def hcore_deriv_generator(self, mol=None, deriv=1):
        from pyscf.x2c import sfx2c1e_grad
        from pyscf.x2c import sfx2c1e_hess
        if deriv == 1:
            return sfx2c1e_grad.hcore_grad_generator(self, mol)
        elif deriv == 2:
            return sfx2c1e_hess.hcore_hess_generator(self, mol)
        else:
            raise NotImplementedError

SpinFreeX2C = SpinFreeX2CHelper


if __name__ == '__main__':
    mol = gto.Mole()
    mol.build(
        verbose = 0,
        atom = [["O" , (0. , 0.     , 0.)],
                [1   , (0. , -0.757 , 0.587)],
                [1   , (0. , 0.757  , 0.587)] ],
        basis = 'ccpvdz-dk',
    )

    method = hf.RHF(mol)
    enr = method.kernel()
    print('E(NR) = %.12g' % enr)

    method = sfx2c1e(hf.RHF(mol))
    esfx2c = method.kernel()
    print('E(SFX2C1E) = %.12g' % esfx2c)
    method.with_x2c.basis = 'unc-ccpvqz-dk'
    print('E(SFX2C1E) = %.12g' % method.kernel())
    method.with_x2c.approx = 'atom1e'
    print('E(SFX2C1E) = %.12g' % method.kernel())



# pyscf/x2c/sfx2c1e_grad.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''
Analytical nuclear gradients for 1-electron spin-free x2c method

Ref.
JCP 135, 084114 (2011); DOI:10.1063/1.3624397
'''

from functools import reduce
import numpy
import scipy.linalg
from pyscf import lib
from pyscf import gto
from pyscf.x2c import x2c

def hcore_grad_generator(x2cobj, mol=None):
    '''nuclear gradients of 1-component X2c hcore Hamiltonian  (spin-free part only)
    '''
    if mol is None: mol = x2cobj.mol
    xmol, contr_coeff = x2cobj.get_xmol(mol)

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

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


def gen_sf_hfw(mol, approx='1E'):
    approx = approx.upper()
    c = lib.param.LIGHT_SPEED

    h0, s0 = _get_h0_s0(mol)
    e0, c0 = scipy.linalg.eigh(h0, s0)

    aoslices = mol.aoslice_by_atom()
    nao = mol.nao_nr()
    if 'ATOM' in approx:
        x0 = numpy.zeros((nao,nao))
        for ia in range(mol.natm):
            ish0, ish1, p0, p1 = aoslices[ia]
            shls_slice = (ish0, ish1, ish0, ish1)
            t1 = mol.intor('int1e_kin', shls_slice=shls_slice)
            s1 = mol.intor('int1e_ovlp', shls_slice=shls_slice)
            with mol.with_rinv_at_nucleus(ia):
                z = -mol.atom_charge(ia)
                v1 = z * mol.intor('int1e_rinv', shls_slice=shls_slice)
                w1 = z * mol.intor('int1e_prinvp', shls_slice=shls_slice)
            x0[p0:p1,p0:p1] = x2c._x2c1e_xmatrix(t1, v1, w1, s1, c)
    else:
        cl0 = c0[:nao,nao:]
        cs0 = c0[nao:,nao:]
        x0 = scipy.linalg.solve(cl0.T, cs0.T).T

    s_nesc0 = s0[:nao,:nao] + reduce(numpy.dot, (x0.T, s0[nao:,nao:], x0))
    R0 = x2c._get_r(s0[:nao,:nao], s_nesc0)
    c_fw0 = numpy.vstack((R0, numpy.dot(x0, R0)))
    h0_fw_half = numpy.dot(h0, c_fw0)

    get_h1_etc = _gen_first_order_quantities(mol, e0, c0, x0, approx)

    def hcore_deriv(ia):
        h1_ao, s1_ao, e1, c1, x1, s_nesc1, R1, c_fw1 = get_h1_etc(ia)
        hfw1 = lib.einsum('xpi,pj->xij', c_fw1, h0_fw_half)
        hfw1 = hfw1 + hfw1.transpose(0,2,1)
        hfw1+= lib.einsum('pi,xpq,qj->xij', c_fw0, h1_ao, c_fw0)
        return hfw1
    return hcore_deriv

def _get_h0_s0(mol):
    c = lib.param.LIGHT_SPEED
    s = mol.intor_symmetric('int1e_ovlp')
    t = mol.intor_symmetric('int1e_kin')
    v = mol.intor_symmetric('int1e_nuc')
    w = mol.intor_symmetric('int1e_pnucp')
    nao = s.shape[0]
    n2 = nao * 2
    h = numpy.zeros((n2,n2), dtype=v.dtype)
    m = numpy.zeros((n2,n2), dtype=v.dtype)
    h[:nao,:nao] = v
    h[:nao,nao:] = t
    h[nao:,:nao] = t
    h[nao:,nao:] = w * (.25/c**2) - t
    m[:nao,:nao] = s
    m[nao:,nao:] = t * (.5/c**2)
    return h, m

def _gen_h1_s1(mol):
    c = lib.param.LIGHT_SPEED
    s1 = mol.intor('int1e_ipovlp', comp=3)
    t1 = mol.intor('int1e_ipkin', comp=3)
    v1 = mol.intor('int1e_ipnuc', comp=3)
    w1 = mol.intor('int1e_ippnucp', comp=3)

    aoslices = mol.aoslice_by_atom()
    nao = s1.shape[1]
    n2 = nao * 2
    def get_h1_s1(ia):
        h1 = numpy.zeros((3,n2,n2), dtype=v1.dtype)
        m1 = numpy.zeros((3,n2,n2), dtype=v1.dtype)
        ish0, ish1, i0, i1 = aoslices[ia]
        with mol.with_rinv_origin(mol.atom_coord(ia)):
            z = mol.atom_charge(ia)
            rinv1   = -z*mol.intor('int1e_iprinv', comp=3)
            prinvp1 = -z*mol.intor('int1e_ipprinvp', comp=3)
        rinv1  [:,i0:i1,:] -= v1[:,i0:i1]
        prinvp1[:,i0:i1,:] -= w1[:,i0:i1]

        for i in range(3):
            s1cc = numpy.zeros((nao,nao))
            t1cc = numpy.zeros((nao,nao))
            s1cc[i0:i1,:] =-s1[i,i0:i1]
            s1cc[:,i0:i1]-= s1[i,i0:i1].T
            t1cc[i0:i1,:] =-t1[i,i0:i1]
            t1cc[:,i0:i1]-= t1[i,i0:i1].T
            v1cc = rinv1[i]   + rinv1[i].T
            w1cc = prinvp1[i] + prinvp1[i].T

            h1[i,:nao,:nao] = v1cc
            h1[i,:nao,nao:] = t1cc
            h1[i,nao:,:nao] = t1cc
            h1[i,nao:,nao:] = w1cc * (.25/c**2) - t1cc
            m1[i,:nao,:nao] = s1cc
            m1[i,nao:,nao:] = t1cc * (.5/c**2)
        return h1, m1
    return get_h1_s1

def _gen_first_order_quantities(mol, e0, c0, x0, approx='1E'):
    c = lib.param.LIGHT_SPEED
    nao = e0.size // 2
    n2 = nao * 2

    epq = e0[:,None] - e0
    degen_mask = abs(epq) < 1e-7
    epq[degen_mask] = 1e200

    cl0 = c0[:nao,nao:]
    # cs0 = c0[nao:,nao:]
    s0 = mol.intor('int1e_ovlp')
    t0 = mol.intor('int1e_kin')
    t0x0 = numpy.dot(t0, x0) * (.5/c**2)
    s_nesc0 = s0[:nao,:nao] + numpy.dot(x0.T, t0x0)

    w_s, v_s = scipy.linalg.eigh(s0)
    w_sqrt = numpy.sqrt(w_s)
    s_nesc0_vbas = reduce(numpy.dot, (v_s.T, s_nesc0, v_s))
    R0_mid = numpy.einsum('i,ij,j->ij', 1./w_sqrt, s_nesc0_vbas, 1./w_sqrt)
    wr0, vr0 = scipy.linalg.eigh(R0_mid)
    wr0_sqrt = numpy.sqrt(wr0)
    # R0 in v_s basis
    R0 = numpy.dot(vr0/wr0_sqrt, vr0.T)
    R0 *= w_sqrt
    R0 /= w_sqrt[:,None]
    # Transform R0 back
    R0 = reduce(numpy.dot, (v_s, R0, v_s.T))

    get_h1_s1 = _gen_h1_s1(mol)
    def get_first_order(ia):
        h1ao, s1ao = get_h1_s1(ia)
        h1mo = lib.einsum('pi,xpq,qj->xij', c0.conj(), h1ao, c0)
        s1mo = lib.einsum('pi,xpq,qj->xij', c0.conj(), s1ao, c0)

        if 'ATOM' in approx:
            e1 = c1_ao = x1 = None
            s_nesc1 = lib.einsum('pi,xpq,qj->xij', x0, s1ao[:,nao:,nao:], x0)
            s_nesc1+= s1ao[:,:nao,:nao]
        else:
            f1 = h1mo[:,:,nao:] - s1mo[:,:,nao:] * e0[nao:]
            c1 = f1 / -epq[:,nao:]
            e1 = f1[:,nao:]
            e1[:,~degen_mask[nao:,nao:]] = 0

            c1_ao = lib.einsum('pq,xqi->xpi', c0, c1)
            cl1 = c1_ao[:,:nao]
            cs1 = c1_ao[:,nao:]
            tmp = cs1 - lib.einsum('pq,xqi->xpi', x0, cl1)
            x1 = scipy.linalg.solve(cl0.T, tmp.reshape(-1,nao).T)
            x1 = x1.T.reshape(3,nao,nao)

            s_nesc1 = lib.einsum('xpi,pj->xij', x1, t0x0)
            s_nesc1 = s_nesc1 + s_nesc1.transpose(0,2,1)
            s_nesc1+= lib.einsum('pi,xpq,qj->xij', x0, s1ao[:,nao:,nao:], x0)
            s_nesc1+= s1ao[:,:nao,:nao]

        R1 = numpy.empty((3,nao,nao))
        c_fw1 = numpy.empty((3,n2,nao))
        for i in range(3):
            R1[i] = _get_r1((w_sqrt,v_s), s_nesc0_vbas,
                            s1ao[i,:nao,:nao], s_nesc1[i], (wr0_sqrt,vr0))
            c_fw1[i,:nao] = R1[i]
            c_fw1[i,nao:] = numpy.dot(x0, R1[i])
            if 'ATOM' not in approx:
                c_fw1[i,nao:] += numpy.dot(x1[i], R0)
        return h1ao, s1ao, e1, c1_ao, x1, s_nesc1, R1, c_fw1
    return get_first_order

def _get_r1(s0_roots, s_nesc0, s1, s_nesc1, r0_roots):
    # See JCP 135, 084114 (2011); DOI:10.1063/1.3624397, Eq (34)
    w_sqrt, v_s = s0_roots
    w_invsqrt = 1. / w_sqrt
    wr0_sqrt, vr0 = r0_roots
    wr0_invsqrt = 1. / wr0_sqrt

    s1 = reduce(numpy.dot, (v_s.T, s1, v_s))
    s_nesc1 = reduce(numpy.dot, (v_s.T, s_nesc1, v_s))

    s1_sqrt = s1 / (w_sqrt[:,None] + w_sqrt)
    s1_invsqrt = (numpy.einsum('i,ij,j->ij', w_invsqrt**2, s1, w_invsqrt**2)
                  / -(w_invsqrt[:,None] + w_invsqrt))
    R1_mid = numpy.dot(s1_invsqrt, s_nesc0) * w_invsqrt
    R1_mid = R1_mid + R1_mid.T
    R1_mid += numpy.einsum('i,ij,j->ij', w_invsqrt, s_nesc1, w_invsqrt)

    R1_mid = reduce(numpy.dot, (vr0.T, R1_mid, vr0))
    R1_mid /= -(wr0_invsqrt[:,None] + wr0_invsqrt)
    R1_mid = numpy.einsum('i,ij,j->ij', wr0_invsqrt**2, R1_mid, wr0_invsqrt**2)
    vr0_wr0_sqrt = vr0 * wr0_invsqrt
    vr0_s0_sqrt = vr0.T * w_sqrt
    vr0_s0_invsqrt = vr0.T * w_invsqrt

    R1  = reduce(numpy.dot, (vr0_s0_invsqrt.T, R1_mid, vr0_s0_sqrt))
    R1 += reduce(numpy.dot, (s1_invsqrt, vr0_wr0_sqrt, vr0_s0_sqrt))
    R1 += reduce(numpy.dot, (vr0_s0_invsqrt.T, vr0_wr0_sqrt.T, s1_sqrt))
    R1 = reduce(numpy.dot, (v_s, R1, v_s.T))
    return R1


if __name__ == '__main__':
    bak = lib.param.LIGHT_SPEED
    lib.param.LIGHT_SPEED = 10
    def get_h(mol):
        c = lib.param.LIGHT_SPEED
        t = mol.intor_symmetric('int1e_kin')
        v = mol.intor_symmetric('int1e_nuc')
        s = mol.intor_symmetric('int1e_ovlp')
        w = mol.intor_symmetric('int1e_pnucp')
        return x2c._x2c1e_get_hcore(t, v, w, s, c)

    mol = gto.M(
        verbose = 0,
        atom = [["O" , (0. , 0.     , 0.0001)],
                [1   , (0. , -0.757 , 0.587)],
                [1   , (0. , 0.757  , 0.587)]],
        basis = '3-21g',
    )
    h_1 = get_h(mol)

    mol = gto.M(
        verbose = 0,
        atom = [["O" , (0. , 0.     ,-0.0001)],
                [1   , (0. , -0.757 , 0.587)],
                [1   , (0. , 0.757  , 0.587)]],
        basis = '3-21g',
    )
    h_2 = get_h(mol)
    h_ref = (h_1 - h_2) / 0.0002 * lib.param.BOHR

    mol = gto.M(
        verbose = 0,
        atom = [["O" , (0. , 0.     , 0.   )],
                [1   , (0. , -0.757 , 0.587)],
                [1   , (0. , 0.757  , 0.587)]],
        basis = '3-21g',
    )
    hcore_deriv = gen_sf_hfw(mol)
    h1 = hcore_deriv(0)
    print(abs(h1[2]-h_ref).max())
    lib.param.LIGHT_SPEED = bak

    print(lib.finger(h1) - -1.4618392662849411)
    hcore_deriv = gen_sf_hfw(mol, approx='atom1e')
    h1 = hcore_deriv(0)
    print(lib.finger(h1) - -1.3596826558976405)


# pyscf/x2c/sfx2c1e_hess.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''
Analytical nuclear hessian for 1-electron spin-free x2c method

Ref.
JCP 135, 244104 (2011); DOI:10.1063/1.3667202
JCTC 8, 2617 (2012); DOI:10.1021/ct300127e
'''

from functools import reduce
import numpy
import scipy.linalg
from pyscf import lib
from pyscf import gto
from pyscf.x2c import x2c
from pyscf.x2c import sfx2c1e_grad

def hcore_hess_generator(x2cobj, mol=None):
    '''nuclear gradients of 1-component X2c hcore Hamiltonian  (spin-free part only)
    '''
    if mol is None: mol = x2cobj.mol
    xmol, contr_coeff = x2cobj.get_xmol(mol)

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

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


def gen_sf_hfw(mol, approx='1E'):
    approx = approx.upper()
    c = lib.param.LIGHT_SPEED

    h0, s0 = sfx2c1e_grad._get_h0_s0(mol)
    e0, c0 = scipy.linalg.eigh(h0, s0)
    c0[:,c0[1]<0] *= -1

    aoslices = mol.aoslice_by_atom()
    nao = mol.nao_nr()
    if 'ATOM' in approx:
        x0 = numpy.zeros((nao,nao))
        for ia in range(mol.natm):
            ish0, ish1, p0, p1 = aoslices[ia]
            shls_slice = (ish0, ish1, ish0, ish1)
            t1 = mol.intor('int1e_kin', shls_slice=shls_slice)
            s1 = mol.intor('int1e_ovlp', shls_slice=shls_slice)
            with mol.with_rinv_at_nucleus(ia):
                z = -mol.atom_charge(ia)
                v1 = z * mol.intor('int1e_rinv', shls_slice=shls_slice)
                w1 = z * mol.intor('int1e_prinvp', shls_slice=shls_slice)
            x0[p0:p1,p0:p1] = x2c._x2c1e_xmatrix(t1, v1, w1, s1, c)
    else:
        cl0 = c0[:nao,nao:]
        cs0 = c0[nao:,nao:]
        x0 = scipy.linalg.solve(cl0.T, cs0.T).T

    t0x0 = numpy.dot(s0[nao:,nao:], x0)
    s_nesc0 = s0[:nao,:nao] + numpy.dot(x0.T, t0x0)

    w_s, v_s = scipy.linalg.eigh(s0[:nao,:nao])
    w_sqrt = numpy.sqrt(w_s)
    s_nesc0_vbas = reduce(numpy.dot, (v_s.T, s_nesc0, v_s))
    R0_mid = numpy.einsum('i,ij,j->ij', 1./w_sqrt, s_nesc0_vbas, 1./w_sqrt)
    wr0, vr0 = scipy.linalg.eigh(R0_mid)
    wr0_sqrt = numpy.sqrt(wr0)
    # R0 in v_s basis
    R0 = numpy.dot(vr0/wr0_sqrt, vr0.T)
    R0 *= w_sqrt
    R0 /= w_sqrt[:,None]
    # Transform R0 back
    R0 = reduce(numpy.dot, (v_s, R0, v_s.T))
    R0 = x2c._get_r(s0[:nao,:nao], s_nesc0)
    c_fw0 = numpy.vstack((R0, numpy.dot(x0, R0)))
    h0_fw_half = numpy.dot(h0, c_fw0)

    epq = e0[:,None] - e0
    degen_mask = abs(epq) < 1e-7
    epq[degen_mask] = 1e200
    s2aa = mol.intor('int1e_ipipovlp', comp=9).reshape(3,3,nao,nao)
    t2aa = mol.intor('int1e_ipipkin', comp=9).reshape(3,3,nao,nao)
    v2aa = mol.intor('int1e_ipipnuc', comp=9).reshape(3,3,nao,nao)
    w2aa = mol.intor('int1e_ipippnucp', comp=9).reshape(3,3,nao,nao)
    s2ab = mol.intor('int1e_ipovlpip', comp=9).reshape(3,3,nao,nao)
    t2ab = mol.intor('int1e_ipkinip', comp=9).reshape(3,3,nao,nao)
    v2ab = mol.intor('int1e_ipnucip', comp=9).reshape(3,3,nao,nao)
    w2ab = mol.intor('int1e_ippnucpip', comp=9).reshape(3,3,nao,nao)
    n2 = nao * 2
    h2ao = numpy.zeros((3,3,n2,n2), dtype=v2aa.dtype)
    s2ao = numpy.zeros((3,3,n2,n2), dtype=v2aa.dtype)

    get_h1_etc = sfx2c1e_grad._gen_first_order_quantities(mol, e0, c0, x0, approx)

    def hcore_deriv(ia, ja):
        ish0, ish1, i0, i1 = aoslices[ia]
        jsh0, jsh1, j0, j1 = aoslices[ja]

        s2cc = numpy.zeros_like(s2aa)
        t2cc = numpy.zeros_like(s2aa)
        v2cc = numpy.zeros_like(s2aa)
        w2cc = numpy.zeros_like(s2aa)
        if ia == ja:
            with mol.with_rinv_origin(mol.atom_coord(ia)):
                z = mol.atom_charge(ia)
                rinv2aa = z*mol.intor('int1e_ipiprinv', comp=9).reshape(3,3,nao,nao)
                rinv2ab = z*mol.intor('int1e_iprinvip', comp=9).reshape(3,3,nao,nao)
                prinvp2aa = z*mol.intor('int1e_ipipprinvp', comp=9).reshape(3,3,nao,nao)
                prinvp2ab = z*mol.intor('int1e_ipprinvpip', comp=9).reshape(3,3,nao,nao)
            s2cc[:,:,i0:i1      ] = s2aa[:,:,i0:i1      ]
            s2cc[:,:,i0:i1,j0:j1]+= s2ab[:,:,i0:i1,j0:j1]
            t2cc[:,:,i0:i1      ] = t2aa[:,:,i0:i1      ]
            t2cc[:,:,i0:i1,j0:j1]+= t2ab[:,:,i0:i1,j0:j1]
            v2cc -= rinv2aa + rinv2ab
            v2cc[:,:,i0:i1      ]+= v2aa[:,:,i0:i1      ]
            v2cc[:,:,i0:i1,j0:j1]+= v2ab[:,:,i0:i1,j0:j1]
            v2cc[:,:,i0:i1      ]+= rinv2aa[:,:,i0:i1]
            v2cc[:,:,i0:i1      ]+= rinv2ab[:,:,i0:i1]
            v2cc[:,:,:    ,i0:i1]+= rinv2aa[:,:,i0:i1].transpose(0,1,3,2)
            v2cc[:,:,:    ,i0:i1]+= rinv2ab[:,:,:,i0:i1]
            w2cc -= prinvp2aa + prinvp2ab
            w2cc[:,:,i0:i1      ]+= w2aa[:,:,i0:i1      ]
            w2cc[:,:,i0:i1,j0:j1]+= w2ab[:,:,i0:i1,j0:j1]
            w2cc[:,:,i0:i1      ]+= prinvp2aa[:,:,i0:i1]
            w2cc[:,:,i0:i1      ]+= prinvp2ab[:,:,i0:i1]
            w2cc[:,:,:    ,i0:i1]+= prinvp2aa[:,:,i0:i1].transpose(0,1,3,2)
            w2cc[:,:,:    ,i0:i1]+= prinvp2ab[:,:,:,i0:i1]

        else:
            s2cc[:,:,i0:i1,j0:j1] = s2ab[:,:,i0:i1,j0:j1]
            t2cc[:,:,i0:i1,j0:j1] = t2ab[:,:,i0:i1,j0:j1]
            v2cc[:,:,i0:i1,j0:j1] = v2ab[:,:,i0:i1,j0:j1]
            w2cc[:,:,i0:i1,j0:j1] = w2ab[:,:,i0:i1,j0:j1]
            zi = mol.atom_charge(ia)
            zj = mol.atom_charge(ja)
            with mol.with_rinv_at_nucleus(ia):
                shls_slice = (jsh0, jsh1, 0, mol.nbas)
                rinv2aa = mol.intor('int1e_ipiprinv', comp=9, shls_slice=shls_slice)
                rinv2ab = mol.intor('int1e_iprinvip', comp=9, shls_slice=shls_slice)
                prinvp2aa = mol.intor('int1e_ipipprinvp', comp=9, shls_slice=shls_slice)
                prinvp2ab = mol.intor('int1e_ipprinvpip', comp=9, shls_slice=shls_slice)
                rinv2aa = zi * rinv2aa.reshape(3,3,j1-j0,nao)
                rinv2ab = zi * rinv2ab.reshape(3,3,j1-j0,nao)
                prinvp2aa = zi * prinvp2aa.reshape(3,3,j1-j0,nao)
                prinvp2ab = zi * prinvp2ab.reshape(3,3,j1-j0,nao)
                v2cc[:,:,j0:j1] += rinv2aa
                v2cc[:,:,j0:j1] += rinv2ab.transpose(1,0,2,3)
                w2cc[:,:,j0:j1] += prinvp2aa
                w2cc[:,:,j0:j1] += prinvp2ab.transpose(1,0,2,3)

            with mol.with_rinv_at_nucleus(ja):
                shls_slice = (ish0, ish1, 0, mol.nbas)
                rinv2aa = mol.intor('int1e_ipiprinv', comp=9, shls_slice=shls_slice)
                rinv2ab = mol.intor('int1e_iprinvip', comp=9, shls_slice=shls_slice)
                prinvp2aa = mol.intor('int1e_ipipprinvp', comp=9, shls_slice=shls_slice)
                prinvp2ab = mol.intor('int1e_ipprinvpip', comp=9, shls_slice=shls_slice)
                rinv2aa = zj * rinv2aa.reshape(3,3,i1-i0,nao)
                rinv2ab = zj * rinv2ab.reshape(3,3,i1-i0,nao)
                prinvp2aa = zj * prinvp2aa.reshape(3,3,i1-i0,nao)
                prinvp2ab = zj * prinvp2ab.reshape(3,3,i1-i0,nao)
                v2cc[:,:,i0:i1] += rinv2aa
                v2cc[:,:,i0:i1] += rinv2ab
                w2cc[:,:,i0:i1] += prinvp2aa
                w2cc[:,:,i0:i1] += prinvp2ab
        s2cc = s2cc + s2cc.transpose(0,1,3,2)
        t2cc = t2cc + t2cc.transpose(0,1,3,2)
        v2cc = v2cc + v2cc.transpose(0,1,3,2)
        w2cc = w2cc + w2cc.transpose(0,1,3,2)
        h2ao[:,:,:nao,:nao] = v2cc
        h2ao[:,:,:nao,nao:] = t2cc
        h2ao[:,:,nao:,:nao] = t2cc
        h2ao[:,:,nao:,nao:] = w2cc * (.25/c**2) - t2cc
        s2ao[:,:,:nao,:nao] = s2cc
        s2ao[:,:,nao:,nao:] = t2cc * (.5/c**2)

        h1i, s1i, e1i, c1i, x1i, s_nesc1i, R1i, c_fw1i = get_h1_etc(ia)
        h1j, s1j, e1j, c1j, x1j, s_nesc1j, R1j, c_fw1j = get_h1_etc(ja)
        if 'ATOM' not in approx:
            f2 = lib.einsum('xypq,qj->xypj', h2ao, c0[:,nao:])
            f2+= lib.einsum('xpq,yqj->xypj', h1i, c1j)
            f2+= lib.einsum('ypq,xqj->xypj', h1j, c1i)
            sc2 = lib.einsum('xypq,qj->xypj', s2ao, c0[:,nao:])
            sc2+= lib.einsum('xpq,yqj->xypj', s1i, c1j)
            sc2+= lib.einsum('ypq,xqj->xypj', s1j, c1i)
            f2-= sc2 * e0[nao:]
            sc1i = lib.einsum('xpq,qj->xpj', s1i, c0[:,nao:])
            sc1j = lib.einsum('xpq,qj->xpj', s1j, c0[:,nao:])
            sc1i+= lib.einsum('pq,xqj->xpj', s0, c1i)
            sc1j+= lib.einsum('pq,xqj->xpj', s0, c1j)
            f2-= lib.einsum('xpq,yqj->xypj', sc1i, e1j)
            f2-= lib.einsum('ypq,xqj->xypj', sc1j, e1i)

            c2 = lib.einsum('pi,xypj->xyij', c0.conj(), f2) / -epq[:,nao:]
            c2_ao = lib.einsum('pq,xyqi->xypi', c0, c2)
            cl2 = c2_ao[:,:,:nao]
            cs2 = c2_ao[:,:,nao:]

            tmp = cs2 - lib.einsum('pq,xyqi->xypi', x0, cl2)
            tmp-= lib.einsum('xpq,yqi->xypi', x1i, c1j[:,:nao])
            tmp-= lib.einsum('ypq,xqi->xypi', x1j, c1i[:,:nao])
            x2 = scipy.linalg.solve(cl0.T, tmp.reshape(-1,nao).T).T.reshape(3,3,nao,nao)

        hfw2 = numpy.empty((3,3,nao,nao))
        for i in range(3):
            for j in range(3):
                if 'ATOM' in approx:
                    s_nesc2  = reduce(numpy.dot, (x0.T, s2ao[i,j,nao:,nao:], x0))
                    s_nesc2 += s2ao[i,j,:nao,:nao]
                    R2 = _get_r2((w_sqrt,v_s), s_nesc0,
                                 s1i[i,:nao,:nao], s_nesc1i[i],
                                 s1j[j,:nao,:nao], s_nesc1j[j],
                                 s2ao[i,j,:nao,:nao], s_nesc2, (wr0_sqrt,vr0))
                    c_fw2 = numpy.vstack((R2, numpy.dot(x0, R2)))
                else:
                    s_nesc2  = numpy.dot(x2[i,j].T, t0x0)
                    s_nesc2 += reduce(numpy.dot, (x1i[i].T, s1j[j,nao:,nao:], x0))
                    s_nesc2 += reduce(numpy.dot, (x0.T, s1i[i,nao:,nao:], x1j[j]))
                    s_nesc2 += reduce(numpy.dot, (x1i[i].T, s0[nao:,nao:], x1j[j]))
                    s_nesc2  = s_nesc2 + s_nesc2.T
                    s_nesc2 += reduce(numpy.dot, (x0.T, s2ao[i,j,nao:,nao:], x0))
                    s_nesc2 += s2ao[i,j,:nao,:nao]
                    R2 = _get_r2((w_sqrt,v_s), s_nesc0,
                                 s1i[i,:nao,:nao], s_nesc1i[i],
                                 s1j[j,:nao,:nao], s_nesc1j[j],
                                 s2ao[i,j,:nao,:nao], s_nesc2, (wr0_sqrt,vr0))
                    c_fw_s = (numpy.dot(x0, R2) + numpy.dot(x1i[i], R1j[j]) +
                              numpy.dot(x1j[j], R1i[i]) + numpy.dot(x2[i,j], R0))
                    c_fw2 = numpy.vstack((R2, c_fw_s))
                tmp  = numpy.dot(c_fw2.T, h0_fw_half)
                tmp += reduce(numpy.dot, (c_fw1i[i].T, h1j[j], c_fw0))
                tmp += reduce(numpy.dot, (c_fw0.T, h1i[i], c_fw1j[j]))
                tmp += reduce(numpy.dot, (c_fw1i[i].T, h0, c_fw1j[j]))
                hfw2[i,j] = tmp + tmp.T
                hfw2[i,j]+= reduce(numpy.dot, (c_fw0.T, h2ao[i,j], c_fw0))
        return hfw2

    return hcore_deriv

def _get_r2(s0_roots, sa0, s1i, sa1i, s1j, sa1j, s2, sa2, r0_roots):
    w_sqrt, v_s = s0_roots
    w_invsqrt = 1. / w_sqrt
    wr0_sqrt, vr0 = r0_roots
    wr0_invsqrt = 1. / wr0_sqrt

    sa0  = lib.einsum('pi,pq,qj->ij', v_s, sa0 , v_s)
    s1i  = lib.einsum('pi,pq,qj->ij', v_s, s1i , v_s)
    s1j  = lib.einsum('pi,pq,qj->ij', v_s, s1j , v_s)
    s2   = lib.einsum('pi,pq,qj->ij', v_s, s2  , v_s)
    sa1i = lib.einsum('pi,pq,qj->ij', v_s, sa1i, v_s)
    sa1j = lib.einsum('pi,pq,qj->ij', v_s, sa1j, v_s)
    sa2  = lib.einsum('pi,pq,qj->ij', v_s, sa2 , v_s)

    s1i_sqrt = s1i / (w_sqrt[:,None] + w_sqrt)
    s1i_invsqrt = (numpy.einsum('i,ij,j->ij', w_invsqrt**2, s1i, w_invsqrt**2)
                   / -(w_invsqrt[:,None] + w_invsqrt))
    s1j_sqrt = s1j / (w_sqrt[:,None] + w_sqrt)
    s1j_invsqrt = (numpy.einsum('i,ij,j->ij', w_invsqrt**2, s1j, w_invsqrt**2)
                   / -(w_invsqrt[:,None] + w_invsqrt))

    tmp = numpy.dot(s1i_sqrt, s1j_sqrt)
    s2_sqrt = (s2 - tmp - tmp.T) / (w_sqrt[:,None] + w_sqrt)
    tmp = numpy.dot(s1i*w_invsqrt**2, s1j)
    tmp = s2 - tmp - tmp.T
    tmp = -numpy.einsum('i,ij,j->ij', w_invsqrt**2, tmp, w_invsqrt**2)
    tmp1 = numpy.dot(s1i_invsqrt, s1j_invsqrt)
    s2_invsqrt = (tmp - tmp1 - tmp1.T) / (w_invsqrt[:,None] + w_invsqrt)

    R1i_mid = lib.einsum('ip,pj,j->ij', s1i_invsqrt, sa0, w_invsqrt)
    R1i_mid = R1i_mid + R1i_mid.T
    R1i_mid+= numpy.einsum('i,ij,j->ij', w_invsqrt, sa1i, w_invsqrt)
    R1i_mid = tmpi = lib.einsum('pi,pq,qj->ij', vr0, R1i_mid, vr0)
    R1i_mid = (numpy.einsum('i,ij,j->ij', wr0_invsqrt**2, R1i_mid, wr0_invsqrt**2)
               / -(wr0_invsqrt[:,None] + wr0_invsqrt))

    R1j_mid = lib.einsum('ip,pj,j->ij', s1j_invsqrt, sa0, w_invsqrt)
    R1j_mid = R1j_mid + R1j_mid.T
    R1j_mid+= numpy.einsum('i,ij,j->ij', w_invsqrt, sa1j, w_invsqrt)
    R1j_mid = tmpj = lib.einsum('pi,pq,qj->ij', vr0, R1j_mid, vr0)
    R1j_mid = (numpy.einsum('i,ij,j->ij', wr0_invsqrt**2, R1j_mid, wr0_invsqrt**2)
               / -(wr0_invsqrt[:,None] + wr0_invsqrt))

# second derivative of (s_invsqrt * sa * s_invsqrt), 9 terms
    R2_mid = lib.einsum('ip,pj,j->ij', s2_invsqrt , sa0 , w_invsqrt)
    R2_mid+= lib.einsum('ip,pj,j->ij', s1i_invsqrt, sa1j, w_invsqrt)
    R2_mid+= lib.einsum('i,ip,pj->ij', w_invsqrt  , sa1i, s1j_invsqrt)
    R2_mid+= lib.einsum('ip,pq,qj->ij', s1i_invsqrt, sa0 , s1j_invsqrt)
    R2_mid  = R2_mid + R2_mid.T
    R2_mid+= numpy.einsum('i,ij,j->ij', w_invsqrt, sa2, w_invsqrt)
    R2_mid = lib.einsum('pi,pq,qj->ij', vr0, R2_mid, vr0)
    tmp = numpy.dot(tmpi*wr0_invsqrt**2, tmpj)
    tmp = R2_mid - tmp - tmp.T
    tmp = -numpy.einsum('i,ij,j->ij', wr0_invsqrt**2, tmp, wr0_invsqrt**2)
    tmp1 = numpy.dot(R1i_mid, R1j_mid)
    R2_mid = (tmp - tmp1 - tmp1.T) / (wr0_invsqrt[:,None] + wr0_invsqrt)

    R0_mid = numpy.dot(vr0*wr0_invsqrt, vr0.T)
    R1i_mid = reduce(numpy.dot, (vr0, R1i_mid, vr0.T))
    R1j_mid = reduce(numpy.dot, (vr0, R1j_mid, vr0.T))
    R2_mid = reduce(numpy.dot, (vr0, R2_mid, vr0.T))

    R2  = lib.einsum('ip,pj,j->ij' , s2_invsqrt , R0_mid , w_sqrt)
    R2 += lib.einsum('ip,pj,j->ij' , s1i_invsqrt, R1j_mid, w_sqrt)
    R2 += lib.einsum('ip,pq,qj->ij', s1i_invsqrt, R0_mid , s1j_sqrt)
    R2 += lib.einsum('ip,pj,j->ij' , s1j_invsqrt, R1i_mid, w_sqrt)
    R2 += numpy.einsum('i,ij,j->ij', w_invsqrt  , R2_mid , w_sqrt)
    R2 += lib.einsum('i,iq,qj->ij' , w_invsqrt  , R1i_mid, s1j_sqrt)
    R2 += lib.einsum('ip,pq,qj->ij', s1j_invsqrt, R0_mid , s1i_sqrt)
    R2 += lib.einsum('i,iq,qj->ij' , w_invsqrt  , R1j_mid, s1i_sqrt)
    R2 += lib.einsum('i,iq,qj->ij' , w_invsqrt  , R0_mid , s2_sqrt)
    R2 = reduce(numpy.dot, (v_s, R2, v_s.T))
    return R2


if __name__ == '__main__':
    bak = lib.param.LIGHT_SPEED
    lib.param.LIGHT_SPEED = 10

    mol = gto.M(
        verbose = 0,
        atom = [["O" , (0. , 0.     , 0.0001)],
                [1   , (0. , -0.757 , 0.587)],
                [1   , (0. , 0.757  , 0.587)]],
        basis = '3-21g',
    )
    h1_deriv_1 = sfx2c1e_grad.gen_sf_hfw(mol, approx='1E')

    mol = gto.M(
        verbose = 0,
        atom = [["O" , (0. , 0.     ,-0.0001)],
                [1   , (0. , -0.757 , 0.587)],
                [1   , (0. , 0.757  , 0.587)]],
        basis = '3-21g',
    )
    h1_deriv_2 = sfx2c1e_grad.gen_sf_hfw(mol, approx='1E')

    mol = gto.M(
        verbose = 0,
        atom = [["O" , (0. , 0.     , 0.   )],
                [1   , (0. , -0.757 , 0.587)],
                [1   , (0. , 0.757  , 0.587)]],
        basis = '3-21g',
    )
    h2_deriv = gen_sf_hfw(mol)

    h2 = h2_deriv(0,0)
    h2_ref = (h1_deriv_1(0)[2] - h1_deriv_2(0)[2]) / 0.0002 * lib.param.BOHR
    print(abs(h2[2,2]-h2_ref).max())
    print(lib.finger(h2) - 33.71188112440316)

    h2 = h2_deriv(1,0)
    h2_ref = (h1_deriv_1(1)[2] - h1_deriv_2(1)[2]) / 0.0002 * lib.param.BOHR
    print(abs(h2[2,2]-h2_ref).max())
    print(lib.finger(h2) - -23.609411428378138)
    lib.param.LIGHT_SPEED = bak



# pyscf/x2c/stability.py

#!/usr/bin/env python
# Copyright 2022 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''
Generate X2C-SCF response functions
'''

from pyscf.x2c import _response_functions  # noqa
from pyscf.scf.stability import dhf_stability

def x2chf_stability(mf, verbose=None, return_status=False):
    '''
    Stability analysis for X2C-HF/X2C-KS method.

    Args:
        mf : DHF or DKS object

    Kwargs:
        return_status: bool
            Whether to return `stable_i` and `stable_e`

    Returns:
        If return_status is False (default), the return value includes
        a new set of orbitals, which are more close to the stable condition.

        Else, another one boolean variable (indicating current status:
        stable or unstable) is returned.
    '''
    return dhf_stability(mf, verbose, return_status)


# pyscf/x2c/tdscf.py

#!/usr/bin/env python
# Copyright 2022 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from functools import reduce
import numpy
from pyscf import lib
from pyscf import dft
from pyscf import ao2mo
from pyscf.lib import logger
from pyscf.tdscf import ghf, gks
from pyscf.x2c import x2c
from pyscf.x2c import dft as x2c_dft
# To ensure .gen_response() methods are registered
from pyscf.x2c import _response_functions  # noqa
from pyscf import __config__

OUTPUT_THRESHOLD = getattr(__config__, 'tdscf_uhf_get_nto_threshold', 0.3)
REAL_EIG_THRESHOLD = getattr(__config__, 'tdscf_uhf_TDDFT_pick_eig_threshold', 1e-4)

# Low excitation filter to avoid numerical instability
POSTIVE_EIG_THRESHOLD = getattr(__config__, 'tdscf_uhf_TDDFT_positive_eig_threshold', 1e-3)

def gen_tda_operation(mf, fock_ao=None):
    '''A x
    '''
    return ghf.gen_tda_operation(mf, fock_ao, None)
gen_tda_hop = gen_tda_operation

def get_ab(mf, mo_energy=None, mo_coeff=None, mo_occ=None):
    r'''A and B matrices for TDDFT response function.

    A[i,a,j,b] = \delta_{ab}\delta_{ij}(E_a - E_i) + (ia||bj)
    B[i,a,j,b] = (ia||jb)
    '''
    if mo_energy is None: mo_energy = mf.mo_energy
    if mo_coeff is None: mo_coeff = mf.mo_coeff
    if mo_occ is None: mo_occ = mf.mo_occ
    assert(mo_coeff.dtype == numpy.double)

    mol = mf.mol
    nao, nmo = mo_coeff.shape
    n2c = nmo // 2
    occidx = n2c + numpy.where(mo_occ[n2c:] == 1)[0]
    viridx = n2c + numpy.where(mo_occ[n2c:] == 0)[0]
    orbv = mo_coeff[:,viridx]
    orbo = mo_coeff[:,occidx]
    nvir = orbv.shape[1]
    nocc = orbo.shape[1]
    nmo = nocc + nvir
    mo = numpy.hstack((orbo, orbv))

    e_ia = lib.direct_sum('a-i->ia', mo_energy[viridx], mo_energy[occidx])
    a = numpy.diag(e_ia.ravel()).reshape(nocc,nvir,nocc,nvir)
    b = numpy.zeros_like(a)

    def add_hf_(a, b, hyb=1):
        eri_mo = ao2mo.kernel(mol, [orbo, mo, mo, mo], intor='int2e_spinor')
        eri_mo = eri_mo.reshape(nocc,nmo,nmo,nmo)

        a += numpy.einsum('iabj->iajb', eri_mo[:nocc,nocc:,nocc:,:nocc])
        a -= numpy.einsum('ijba->iajb', eri_mo[:nocc,:nocc,nocc:,nocc:]) * hyb
        b += numpy.einsum('iajb->iajb', eri_mo[:nocc,nocc:,:nocc,nocc:])
        b -= numpy.einsum('jaib->iajb', eri_mo[:nocc,nocc:,:nocc,nocc:]) * hyb

    if isinstance(mf, dft.KohnShamDFT):
        ni = mf._numint
        ni.libxc.test_deriv_order(mf.xc, 2, raise_error=True)
        if getattr(mf, 'nlc', '') != '':
            raise NotImplementedError

        if not mf.collinear:
            raise NotImplementedError

        omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin)

        add_hf_(a, b, hyb)

        xctype = ni._xc_type(mf.xc)
        dm0 = mf.make_rdm1(mo_coeff, mo_occ)
        make_rho = ni._gen_rho_evaluator(mol, dm0, hermi=1)[0]
        mem_now = lib.current_memory()[0]
        max_memory = max(2000, mf.max_memory*.8-mem_now)

        if xctype == 'LDA':
            ao_deriv = 0
            for ao, mask, weight, coords \
                    in ni.block_loop(mol, mf.grids, nao, ao_deriv, max_memory):
                rho = make_rho(0, ao, mask, xctype)
                fxc = ni.eval_xc(mf.xc, rho, 1, deriv=2)[2]
                u_u, u_d, d_d = fxc[0].T

                aoa, aob = ao
                mo_oa = lib.einsum('rp,pi->ri', aoa, orbo)
                mo_va = lib.einsum('rp,pi->ri', aoa, orbv)
                mo_ob = lib.einsum('rp,pi->ri', aob, orbo)
                mo_vb = lib.einsum('rp,pi->ri', aob, orbv)
                rho_ov_a = numpy.einsum('ri,ra->ria', mo_oa.conj(), mo_va)
                rho_ov_b = numpy.einsum('ri,ra->ria', mo_ob.conj(), mo_vb)
                rho_vo_a = rho_ov_a.conj()
                rho_vo_b = rho_ov_b.conj()

                w_ov_a = numpy.einsum('ria,r->ria', u_u * rho_ov_a + u_d * rho_ov_b, weight)
                a += lib.einsum('ria,rjb->iajb', w_ov_a, rho_vo_a)
                b += lib.einsum('ria,rjb->iajb', w_ov_a, rho_ov_a)

                w_ov_b = numpy.einsum('ria,r->ria', u_d * rho_ov_a + d_d * rho_ov_b, weight)
                a += lib.einsum('ria,rjb->iajb', w_ov_b, rho_vo_b)
                b += lib.einsum('ria,rjb->iajb', w_ov_b, rho_ov_b)

        elif xctype == 'GGA':
            ao_deriv = 1
            for ao, mask, weight, coords \
                    in ni.block_loop(mol, mf.grids, nao, ao_deriv, max_memory):
                rho0 = make_rho(0, ao, mask, xctype)
                vxc, fxc = ni.eval_xc(mf.xc, rho0, 1, deriv=2)[1:3]
                r0, (mx0, my0, mz0) = rho0
                rho0a = (r0 + mz0) * .5
                rho0b = (r0 - mz0) * .5
                uu, ud, dd = vxc[1].T
                u_u, u_d, d_d = fxc[0].T
                u_uu, u_ud, u_dd, d_uu, d_ud, d_dd = fxc[1].T
                uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd = fxc[2].T

                aoa, aob = ao
                mo_oa = lib.einsum('xrp,pi->xri', aoa, orbo)
                mo_va = lib.einsum('xrp,pi->xri', aoa, orbv)
                mo_ob = lib.einsum('xrp,pi->xri', aob, orbo)
                mo_vb = lib.einsum('xrp,pi->xri', aob, orbv)
                rho_ov_a = numpy.einsum('xri,ra->xria', mo_oa.conj(), mo_va[0])
                rho_ov_b = numpy.einsum('xri,ra->xria', mo_ob.conj(), mo_vb[0])
                rho_ov_a[1:4] += numpy.einsum('ri,xra->xria', mo_oa[0].conj(), mo_va[1:4])
                rho_ov_b[1:4] += numpy.einsum('ri,xra->xria', mo_ob[0].conj(), mo_vb[1:4])
                rho_vo_a = rho_ov_a.conj()
                rho_vo_b = rho_vo_a.conj()
                # sigma1 ~ \nabla(\rho_\alpha+\rho_\beta) dot \nabla(|b><j|) z_{bj}
                a0a1 = numpy.einsum('xr,xria->ria', rho0a[1:4], rho_ov_a[1:4])
                a0b1 = numpy.einsum('xr,xria->ria', rho0a[1:4], rho_ov_b[1:4])
                b0a1 = numpy.einsum('xr,xria->ria', rho0b[1:4], rho_ov_a[1:4])
                b0b1 = numpy.einsum('xr,xria->ria', rho0b[1:4], rho_ov_b[1:4])

                # aaaa
                w_ov = numpy.empty_like(rho_ov_a)
                w_ov[0]  = numpy.einsum('r,ria->ria', u_u, rho_ov_a[0])
                w_ov[0] += numpy.einsum('r,ria->ria', 2*u_uu, a0a1)
                w_ov[0] += numpy.einsum('r,ria->ria',   u_ud, b0a1)
                f_ov_a  = numpy.einsum('r,ria->ria', 4*uu_uu, a0a1)
                f_ov_b  = numpy.einsum('r,ria->ria', 2*uu_ud, a0a1)
                f_ov_a += numpy.einsum('r,ria->ria', 2*uu_ud, b0a1)
                f_ov_b += numpy.einsum('r,ria->ria',   ud_ud, b0a1)
                f_ov_a += numpy.einsum('r,ria->ria', 2*u_uu, rho_ov_a[0])
                f_ov_b += numpy.einsum('r,ria->ria',   u_ud, rho_ov_a[0])
                w_ov[1:] = numpy.einsum('ria,xr->xria', f_ov_a, rho0a[1:4])
                w_ov[1:]+= numpy.einsum('ria,xr->xria', f_ov_b, rho0b[1:4])
                w_ov[1:]+= numpy.einsum('r,xria->xria', 2*uu, rho_ov_a[1:4])
                w_ov *= weight[:,None,None]
                a += lib.einsum('xria,xrjb->iajb', w_ov, rho_vo_a)
                b += lib.einsum('xria,xrjb->iajb', w_ov, rho_ov_a)

                # bbbb
                w_ov = numpy.empty_like(rho_ov_b)
                w_ov[0]  = numpy.einsum('r,ria->ria', d_d, rho_ov_b[0])
                w_ov[0] += numpy.einsum('r,ria->ria', 2*d_dd, b0b1)
                w_ov[0] += numpy.einsum('r,ria->ria',   d_ud, a0b1)
                f_ov_b  = numpy.einsum('r,ria->ria', 4*dd_dd, b0b1)
                f_ov_a  = numpy.einsum('r,ria->ria', 2*ud_dd, b0b1)
                f_ov_b += numpy.einsum('r,ria->ria', 2*ud_dd, a0b1)
                f_ov_a += numpy.einsum('r,ria->ria',   ud_ud, a0b1)
                f_ov_b += numpy.einsum('r,ria->ria', 2*d_dd, rho_ov_b[0])
                f_ov_a += numpy.einsum('r,ria->ria',   d_ud, rho_ov_b[0])
                w_ov[1:] = numpy.einsum('ria,xr->xria', f_ov_a, rho0a[1:4])
                w_ov[1:]+= numpy.einsum('ria,xr->xria', f_ov_b, rho0b[1:4])
                w_ov[1:]+= numpy.einsum('r,xria->xria', 2*dd, rho_ov_b[1:4])
                w_ov *= weight[:,None,None]
                a += lib.einsum('xria,xrjb->iajb', w_ov, rho_vo_b)
                b += lib.einsum('xria,xrjb->iajb', w_ov, rho_ov_b)

                # aabb and bbaa
                w_ov = numpy.empty_like(rho_ov_b)
                w_ov[0]  = numpy.einsum('r,ria->ria', u_d, rho_ov_b[0])
                w_ov[0] += numpy.einsum('r,ria->ria', 2*u_dd, b0b1)
                w_ov[0] += numpy.einsum('r,ria->ria',   u_ud, a0b1)
                f_ov_a  = numpy.einsum('r,ria->ria', 4*uu_dd, b0b1)
                f_ov_b  = numpy.einsum('r,ria->ria', 2*ud_dd, b0b1)
                f_ov_a += numpy.einsum('r,ria->ria', 2*uu_ud, a0b1)
                f_ov_b += numpy.einsum('r,ria->ria',   ud_ud, a0b1)
                f_ov_a += numpy.einsum('r,ria->ria', 2*d_uu, rho_ov_b[0])
                f_ov_b += numpy.einsum('r,ria->ria',   d_ud, rho_ov_b[0])
                w_ov[1:] = numpy.einsum('ria,xr->xria', f_ov_a, rho0a[1:4])
                w_ov[1:]+= numpy.einsum('ria,xr->xria', f_ov_b, rho0b[1:4])
                w_ov[1:]+= numpy.einsum('r,xria->xria', ud, rho_ov_b[1:4])
                w_ov *= weight[:,None,None]
                a_iajb = lib.einsum('xria,xrjb->iajb', w_ov, rho_vo_a)
                b_iajb = lib.einsum('xria,xrjb->iajb', w_ov, rho_ov_a)
                a += a_iajb
                a += a_iajb.conj().transpose(2,3,0,1)
                b += b_iajb * 2

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

    else:
        add_hf_(a, b)

    return a, b

def get_nto(tdobj, state=1, threshold=OUTPUT_THRESHOLD, verbose=None):
    raise NotImplementedError('get_nto')

def analyze(tdobj, verbose=None):
    raise NotImplementedError('analyze')

def _contract_multipole(tdobj, ints, hermi=True, xy=None):
    raise NotImplementedError


class TDMixin(ghf.TDMixin):

    @lib.with_doc(get_ab.__doc__)
    def get_ab(self, mf=None):
        if mf is None: mf = self._scf
        return get_ab(mf)

    analyze = analyze
    get_nto = get_nto
    _contract_multipole = _contract_multipole  # needed by transition dipoles

    def nuc_grad_method(self):
        raise NotImplementedError


class TDA(TDMixin, ghf.TDA):
    def gen_vind(self, mf):
        '''Compute Ax'''
        return gen_tda_hop(mf)

    def init_guess(self, mf, nstates=None, wfnsym=None):
        assert self.wfnsym is None
        return ghf.TDA.init_guess(self, mf, nstates, None)

    kernel = ghf.TDA.kernel


def gen_tdhf_operation(mf, fock_ao=None):
    '''Generate function to compute

    [ A  B][X]
    [-B -A][Y]
    '''
    return ghf.gen_tdhf_operation(mf, fock_ao, None)


class TDHF(TDMixin, ghf.TDHF):
    @lib.with_doc(gen_tdhf_operation.__doc__)
    def gen_vind(self, mf):
        return gen_tdhf_operation(mf)

    def init_guess(self, mf, nstates=None, wfnsym=None):
        assert self.wfnsym is None
        return ghf.TDHF.init_guess(self, mf, nstates, None)

    kernel = ghf.TDHF.kernel

TDDFT = TDHF

x2c.RHF.TDA  = x2c.UHF.TDA  = lib.class_as_method(TDA)
x2c.RHF.TDHF = x2c.UHF.TDHF = lib.class_as_method(TDHF)

x2c_dft.RKS.TDA   = x2c_dft.UKS.TDA   = lib.class_as_method(TDA)
x2c_dft.RKS.TDHF  = x2c_dft.UKS.TDHF  = None
x2c_dft.RKS.TDDFT = x2c_dft.UKS.TDDFT = lib.class_as_method(TDDFT)


# pyscf/x2c/x2c.py

#!/usr/bin/env python
# Copyright 2014-2022 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''
X2C 2-component HF methods
'''

from functools import reduce
import copy
import numpy
import scipy.linalg
from pyscf import lib
from pyscf.gto import mole
from pyscf.lib import logger
from pyscf.scf import hf, ghf, dhf
from pyscf.scf import _vhf
from pyscf.data import nist
from pyscf import __config__

LINEAR_DEP_THRESHOLD = 1e-9

class X2CHelperMixin(lib.StreamObject):
    '''2-component X2c (including spin-free and spin-dependent terms) in
    the j-adapted spinor basis.
    '''
    approx = getattr(__config__, 'x2c_X2C_approx', '1e')  # 'atom1e'
    xuncontract = getattr(__config__, 'x2c_X2C_xuncontract', True)
    basis = getattr(__config__, 'x2c_X2C_basis', None)
    def __init__(self, mol):
        self.mol = mol
        self.stdout = mol.stdout
        self.verbose = mol.verbose

    def dump_flags(self, verbose=None):
        log = logger.new_logger(self, verbose)
        log.info('\n')
        log.info('******** %s ********', self.__class__)
        log.info('approx = %s',    self.approx)
        log.info('xuncontract = %d', self.xuncontract)
        if self.basis is not None:
            log.info('basis for X matrix = %s', self.basis)
        return self

    def get_xmol(self, mol=None):
        if mol is None:
            mol = self.mol

        if self.basis is not None:
            xmol = copy.copy(mol)
            xmol.build(False, False, basis=self.basis)
            return xmol, None
        elif self.xuncontract:
            xmol, contr_coeff = _uncontract_mol(mol, self.xuncontract)
            return xmol, contr_coeff
        else:
            return mol, None

    def get_hcore(self, mol=None):
        '''2-component X2c Foldy-Wouthuysen (FW) Hamiltonian (including
        spin-free and spin-dependent terms) in the j-adapted spinor basis.
        '''
        if mol is None: mol = self.mol
        if mol.has_ecp():
            raise NotImplementedError

        xmol, contr_coeff_nr = self.get_xmol(mol)
        c = lib.param.LIGHT_SPEED
        assert('1E' in self.approx.upper())
        s = xmol.intor_symmetric('int1e_ovlp_spinor')
        t = xmol.intor_symmetric('int1e_spsp_spinor') * .5
        v = xmol.intor_symmetric('int1e_nuc_spinor')
        w = xmol.intor_symmetric('int1e_spnucsp_spinor')
        if 'get_xmat' in self.__dict__:
            # If the get_xmat method is overwritten by user, build the X
            # matrix with the external get_xmat method
            x = self.get_xmat(xmol)
            h1 = _get_hcore_fw(t, v, w, s, x, c)

        elif 'ATOM' in self.approx.upper():
            atom_slices = xmol.offset_2c_by_atom()
            n2c = xmol.nao_2c()
            x = numpy.zeros((n2c,n2c), dtype=numpy.complex128)
            for ia in range(xmol.natm):
                ish0, ish1, p0, p1 = atom_slices[ia]
                shls_slice = (ish0, ish1, ish0, ish1)
                s1 = xmol.intor('int1e_ovlp_spinor', shls_slice=shls_slice)
                t1 = xmol.intor('int1e_spsp_spinor', shls_slice=shls_slice) * .5
                with xmol.with_rinv_at_nucleus(ia):
                    z = -xmol.atom_charge(ia)
                    v1 = z*xmol.intor('int1e_rinv_spinor', shls_slice=shls_slice)
                    w1 = z*xmol.intor('int1e_sprinvsp_spinor', shls_slice=shls_slice)
                x[p0:p1,p0:p1] = _x2c1e_xmatrix(t1, v1, w1, s1, c)
            h1 = _get_hcore_fw(t, v, w, s, x, c)

        else:
            h1 = _x2c1e_get_hcore(t, v, w, s, c)

        if self.basis is not None:
            s22 = xmol.intor_symmetric('int1e_ovlp_spinor')
            s21 = mole.intor_cross('int1e_ovlp_spinor', xmol, mol)
            c = lib.cho_solve(s22, s21)
            h1 = reduce(numpy.dot, (c.T.conj(), h1, c))
        elif self.xuncontract:
            np, nc = contr_coeff_nr.shape
            contr_coeff = numpy.zeros((np*2,nc*2))
            contr_coeff[0::2,0::2] = contr_coeff_nr
            contr_coeff[1::2,1::2] = contr_coeff_nr
            h1 = reduce(numpy.dot, (contr_coeff.T.conj(), h1, contr_coeff))
        return h1

    def _picture_change(self, xmol, even_operator=(None, None), odd_operator=None):
        '''Picture change for even_operator + odd_operator

        even_operator has two terms at diagonal blocks
        [ v  0 ]
        [ 0  w ]

        odd_operator has the term at off-diagonal blocks
        [ 0    p ]
        [ p^T  0 ]

        v, w, and p can be strings (integral name) or matrices.
        '''
        c = lib.param.LIGHT_SPEED
        v_op, w_op = even_operator
        if isinstance(v_op, str):
            v_op = xmol.intor(v_op)
        if isinstance(w_op, str):
            w_op = xmol.intor(w_op)
            w_op *= (.5/c)**2
        if isinstance(odd_operator, str):
            odd_operator = xmol.intor(odd_operator) * (.5/c)

        if v_op is not None:
            shape = v_op.shape
        elif w_op is not None:
            shape = w_op.shape
        elif odd_operator is not None:
            shape = odd_operator.shape
        else:
            raise RuntimeError('No operators provided')

        x = self.get_xmat()
        r = self._get_rmat(x)
        def transform(mat):
            nao = mat.shape[-1] // 2
            xv = mat[:nao] + x.conj().T.dot(mat[nao:])
            h = xv[:,:nao] + xv[:,nao:].dot(x)
            return reduce(numpy.dot, (r.T.conj(), h, r))

        nao = shape[-1]
        dtype = numpy.result_type(v_op, w_op, odd_operator)

        if len(shape) == 2:
            mat = numpy.zeros((nao*2,nao*2), dtype)
            if v_op is not None:
                mat[:nao,:nao] = v_op
            if w_op is not None:
                mat[nao:,nao:] = w_op
            if odd_operator is not None:
                mat[:nao,nao:] = odd_operator
                mat[nao:,:nao] = odd_operator.conj().T
            pc_mat = transform(mat)

        else:
            assert len(shape) == 3
            mat = numpy.zeros((shape[0],nao*2,nao*2), dtype)
            if v_op is not None:
                mat[:,:nao,:nao] = v_op
            if w_op is not None:
                mat[:,nao:,nao:] = w_op
            if odd_operator is not None:
                mat[:,:nao,nao:] = odd_operator
                mat[:,nao:,:nao] = odd_operator.conj().transpose(0,2,1)
            pc_mat = numpy.asarray([transform(m) for m in mat])

        return pc_mat

    def picture_change(self, even_operator=(None, None), odd_operator=None):
        '''Picture change for even_operator + odd_operator

        even_operator has two terms at diagonal blocks
        [ v  0 ]
        [ 0  w ]

        odd_operator has the term at off-diagonal blocks
        [ 0    p ]
        [ p^T  0 ]

        v, w, and p can be strings (integral name) or matrices.
        '''
        mol = self.mol
        xmol, contr_coeff_nr = self.get_xmol(mol)
        pc_mat = self._picture_change(xmol, even_operator, odd_operator)

        if self.basis is not None:
            s22 = xmol.intor_symmetric('int1e_ovlp_spinor')
            s21 = mole.intor_cross('int1e_ovlp_spinor', xmol, mol)
            c = lib.cho_solve(s22, s21)

        elif self.xuncontract:
            np, nc = contr_coeff_nr.shape
            c = numpy.zeros((np*2,nc*2))
            c[0::2,0::2] = contr_coeff_nr
            c[1::2,1::2] = contr_coeff_nr

        else:
            return pc_mat

        if pc_mat.ndim == 2:
            return lib.einsum('pi,pq,qj->ij', c.conj(), pc_mat, c)
        else:
            return lib.einsum('pi,xpq,qj->xij', c.conj(), pc_mat, c)

    def get_xmat(self, mol=None):
        if mol is None:
            xmol = self.get_xmol(mol)[0]
        else:
            xmol = mol
        c = lib.param.LIGHT_SPEED
        assert('1E' in self.approx.upper())

        if 'ATOM' in self.approx.upper():
            atom_slices = xmol.offset_2c_by_atom()
            n2c = xmol.nao_2c()
            x = numpy.zeros((n2c,n2c), dtype=numpy.complex128)
            for ia in range(xmol.natm):
                ish0, ish1, p0, p1 = atom_slices[ia]
                shls_slice = (ish0, ish1, ish0, ish1)
                s1 = xmol.intor('int1e_ovlp_spinor', shls_slice=shls_slice)
                t1 = xmol.intor('int1e_spsp_spinor', shls_slice=shls_slice) * .5
                with xmol.with_rinv_at_nucleus(ia):
                    z = -xmol.atom_charge(ia)
                    v1 = z*xmol.intor('int1e_rinv_spinor', shls_slice=shls_slice)
                    w1 = z*xmol.intor('int1e_sprinvsp_spinor', shls_slice=shls_slice)
                x[p0:p1,p0:p1] = _x2c1e_xmatrix(t1, v1, w1, s1, c)
        else:
            s = xmol.intor_symmetric('int1e_ovlp_spinor')
            t = xmol.intor_symmetric('int1e_spsp_spinor') * .5
            v = xmol.intor_symmetric('int1e_nuc_spinor')
            w = xmol.intor_symmetric('int1e_spnucsp_spinor')
            x = _x2c1e_xmatrix(t, v, w, s, c)
        return x

    def _get_rmat(self, x=None):
        '''The matrix (in AO basis) that changes metric from NESC metric to NR metric'''
        xmol = self.get_xmol()[0]
        if x is None:
            x = self.get_xmat(xmol)
        c = lib.param.LIGHT_SPEED
        s = xmol.intor_symmetric('int1e_ovlp_spinor')
        t = xmol.intor_symmetric('int1e_spsp_spinor') * .5
        s1 = s + reduce(numpy.dot, (x.conj().T, t, x)) * (.5/c**2)
        return _get_r(s, s1)

    def reset(self, mol):
        '''Reset mol and clean up relevant attributes for scanner mode'''
        self.mol = mol
        return self

class SpinorX2CHelper(X2CHelperMixin):
    '''2-component X2c (including spin-free and spin-dependent terms) in
    the j-adapted spinor basis.
    '''
    pass

X2C = SpinorX2CHelper

class SpinOrbitalX2CHelper(X2CHelperMixin):
    '''2-component X2c (including spin-free and spin-dependent terms) in
    the Gaussian type spin-orbital basis (as the spin-orbital basis in GHF)
    '''
    def get_hcore(self, mol=None):
        if mol is None: mol = self.mol
        if mol.has_ecp():
            raise NotImplementedError

        xmol, contr_coeff = self.get_xmol(mol)
        c = lib.param.LIGHT_SPEED
        assert('1E' in self.approx.upper())

        t = _block_diag(xmol.intor_symmetric('int1e_kin'))
        v = _block_diag(xmol.intor_symmetric('int1e_nuc'))
        s = _block_diag(xmol.intor_symmetric('int1e_ovlp'))
        w = _sigma_dot(xmol.intor('int1e_spnucsp'))
        if 'get_xmat' in self.__dict__:
            # If the get_xmat method is overwritten by user, build the X
            # matrix with the external get_xmat method
            x = self.get_xmat(xmol)
            h1 = _get_hcore_fw(t, v, w, s, x, c)

        elif 'ATOM' in self.approx.upper():
            atom_slices = xmol.offset_nr_by_atom()
            # spin-orbital basis is twice the size of NR basis
            atom_slices[:,2:] *= 2
            nao = xmol.nao_nr() * 2
            x = numpy.zeros((nao,nao))
            for ia in range(xmol.natm):
                ish0, ish1, p0, p1 = atom_slices[ia]
                shls_slice = (ish0, ish1, ish0, ish1)
                t1 = _block_diag(xmol.intor('int1e_kin', shls_slice=shls_slice))
                s1 = _block_diag(xmol.intor('int1e_ovlp', shls_slice=shls_slice))
                with xmol.with_rinv_at_nucleus(ia):
                    z = -xmol.atom_charge(ia)
                    v1 = _block_diag(z * xmol.intor('int1e_rinv', shls_slice=shls_slice))
                    w1 = _sigma_dot(z * xmol.intor('int1e_sprinvsp', shls_slice=shls_slice))
                x[p0:p1,p0:p1] = _x2c1e_xmatrix(t1, v1, w1, s1, c)
            h1 = _get_hcore_fw(t, v, w, s, x, c)

        else:
            h1 = _x2c1e_get_hcore(t, v, w, s, c)

        if self.basis is not None:
            s22 = xmol.intor_symmetric('int1e_ovlp')
            s21 = mole.intor_cross('int1e_ovlp', xmol, mol)
            c = _block_diag(lib.cho_solve(s22, s21))
            h1 = reduce(lib.dot, (c.T, h1, c))
        if self.xuncontract and contr_coeff is not None:
            contr_coeff = _block_diag(contr_coeff)
            h1 = reduce(lib.dot, (contr_coeff.T, h1, contr_coeff))
        return h1

    @lib.with_doc(X2CHelperMixin.picture_change.__doc__)
    def picture_change(self, even_operator=(None, None), odd_operator=None):
        mol = self.mol
        xmol, c = self.get_xmol(mol)
        pc_mat = self._picture_change(xmol, even_operator, odd_operator)

        if self.basis is not None:
            s22 = xmol.intor_symmetric('int1e_ovlp')
            s21 = mole.intor_cross('int1e_ovlp', xmol, mol)
            c = lib.cho_solve(s22, s21)

        elif self.xuncontract:
            pass

        else:
            return pc_mat

        c = _block_diag(c)
        if pc_mat.ndim == 2:
            return lib.einsum('pi,pq,qj->ij', c, pc_mat, c)
        else:
            return lib.einsum('pi,xpq,qj->xij', c, pc_mat, c)

    def get_xmat(self, mol=None):
        if mol is None:
            xmol = self.get_xmol(mol)[0]
        else:
            xmol = mol
        c = lib.param.LIGHT_SPEED
        assert('1E' in self.approx.upper())

        if 'ATOM' in self.approx.upper():
            atom_slices = xmol.offset_nr_by_atom()
            # spin-orbital basis is twice the size of NR basis
            atom_slices[:,2:] *= 2
            nao = xmol.nao_nr() * 2
            x = numpy.zeros((nao,nao))
            for ia in range(xmol.natm):
                ish0, ish1, p0, p1 = atom_slices[ia]
                shls_slice = (ish0, ish1, ish0, ish1)
                t1 = _block_diag(xmol.intor('int1e_kin', shls_slice=shls_slice))
                s1 = _block_diag(xmol.intor('int1e_ovlp', shls_slice=shls_slice))
                with xmol.with_rinv_at_nucleus(ia):
                    z = -xmol.atom_charge(ia)
                    v1 = _block_diag(z * xmol.intor('int1e_rinv', shls_slice=shls_slice))
                    w1 = _sigma_dot(z * xmol.intor('int1e_sprinvsp', shls_slice=shls_slice))
                x[p0:p1,p0:p1] = _x2c1e_xmatrix(t1, v1, w1, s1, c)
        else:
            t = _block_diag(xmol.intor_symmetric('int1e_kin'))
            v = _block_diag(xmol.intor_symmetric('int1e_nuc'))
            s = _block_diag(xmol.intor_symmetric('int1e_ovlp'))
            w = _sigma_dot(xmol.intor('int1e_spnucsp'))
            x = _x2c1e_xmatrix(t, v, w, s, c)
        return x

    def _get_rmat(self, x=None):
        '''The matrix (in AO basis) that changes metric from NESC metric to NR metric'''
        xmol = self.get_xmol()[0]
        if x is None:
            x = self.get_xmat(xmol)
        c = lib.param.LIGHT_SPEED
        s = _block_diag(xmol.intor_symmetric('int1e_ovlp'))
        t = _block_diag(xmol.intor_symmetric('int1e_kin'))
        s1 = s + reduce(numpy.dot, (x.conj().T, t, x)) * (.5/c**2)
        return _get_r(s, s1)


def get_hcore(mol):
    '''2-component X2c hcore Hamiltonian (including spin-free and
    spin-dependent terms) in the j-adapted spinor basis.
    '''
    return SpinorX2CHelper(mol).get_hcore(mol)

def get_jk(mol, dm, hermi=1, mf_opt=None, with_j=True, with_k=True, omega=None):
    '''non-relativistic J/K matrices (without SSO,SOO etc) in the j-adapted
    spinor basis.
    '''
    n2c = dm.shape[0]
    dd = numpy.zeros((n2c*2,)*2, dtype=numpy.complex128)
    dd[:n2c,:n2c] = dm
    dhf._call_veff_llll(mol, dd, hermi, None)
    vj, vk = _vhf.rdirect_mapdm('int2e_spinor', 's8',
                                ('ji->s2kl', 'jk->s1il'), dm, 1,
                                mol._atm, mol._bas, mol._env, mf_opt)
    return dhf._jk_triu_(vj, vk, hermi)

def make_rdm1(mo_coeff, mo_occ, **kwargs):
    return numpy.dot(mo_coeff*mo_occ, mo_coeff.T.conj())

def init_guess_by_minao(mol):
    '''Initial guess in terms of the overlap to minimal basis.'''
    dm = hf.init_guess_by_minao(mol)
    return _proj_dmll(mol, dm, mol)

def init_guess_by_1e(mol):
    '''Initial guess from one electron system.'''
    mf = UHF(mol)
    return mf.init_guess_by_1e(mol)

def init_guess_by_atom(mol):
    '''Initial guess from atom calculation.'''
    dm = hf.init_guess_by_atom(mol)
    return _proj_dmll(mol, dm, mol)

def init_guess_by_chkfile(mol, chkfile_name, project=None):
    dm = dhf.init_guess_by_chkfile(mol, chkfile_name, project)
    n2c = dm.shape[0] // 2
    return dm[:n2c,:n2c].copy()

def get_init_guess(mol, key='minao'):
    if callable(key):
        return key(mol)
    elif key.lower() == '1e':
        return init_guess_by_1e(mol)
    elif key.lower() == 'atom':
        return init_guess_by_atom(mol)
    elif key.lower() == 'chkfile':
        raise RuntimeError('Call pyscf.scf.hf.init_guess_by_chkfile instead')
    else:
        return init_guess_by_minao(mol)


class SCF(hf.SCF):
    '''The full X2C problem (scaler + soc terms) in j-adapted spinor basis'''
    def __init__(self, mol):
        hf.SCF.__init__(self, mol)
        self.with_x2c = SpinorX2CHelper(mol)
        #self.with_x2c.xuncontract = False
        self._keys = self._keys.union(['with_x2c'])

    def build(self, mol=None):
        if self.verbose >= logger.WARN:
            self.check_sanity()
        if self.direct_scf:
            self.opt = self.init_direct_scf(mol)
        return self

    def dump_flags(self, verbose=None):
        hf.SCF.dump_flags(self, verbose)
        self.with_x2c.dump_flags(verbose)
        return self

    def init_guess_by_minao(self, mol=None):
        '''Initial guess in terms of the overlap to minimal basis.'''
        if mol is None: mol = self.mol
        return init_guess_by_minao(mol)

    def init_guess_by_atom(self, mol=None):
        if mol is None: mol = self.mol
        return init_guess_by_atom(mol)

    def init_guess_by_chkfile(self, chkfile=None, project=None):
        if chkfile is None: chkfile = self.chkfile
        return init_guess_by_chkfile(self.mol, chkfile, project=project)

    def _eigh(self, h, s):
        e, c = scipy.linalg.eigh(h, s)
        idx = numpy.argmax(abs(c.real), axis=0)
        c[:,c[idx,range(len(e))].real<0] *= -1
        return e, c

    def get_hcore(self, mol=None):
        if mol is None: mol = self.mol
        return self.with_x2c.get_hcore(mol)

    def get_ovlp(self, mol=None):
        if mol is None: mol = self.mol
        return mol.intor_symmetric('int1e_ovlp_spinor')

    def get_occ(self, mo_energy=None, mo_coeff=None):
        if mo_energy is None: mo_energy = self.mo_energy
        mol = self.mol
        mo_occ = numpy.zeros_like(mo_energy)
        mo_occ[:mol.nelectron] = 1
        if mol.nelectron < len(mo_energy):
            logger.info(self, 'nocc = %d  HOMO = %.12g  LUMO = %.12g',
                        mol.nelectron, mo_energy[mol.nelectron-1],
                        mo_energy[mol.nelectron])
        else:
            logger.info(self, 'nocc = %d  HOMO = %.12g  no LUMO',
                        mol.nelectron, mo_energy[mol.nelectron-1])
        logger.debug(self, '  mo_energy = %s', mo_energy)
        return mo_occ

    def make_rdm1(self, mo_coeff=None, mo_occ=None, **kwargs):
        if mo_coeff is None: mo_coeff = self.mo_coeff
        if mo_occ is None: mo_occ = self.mo_occ
        return make_rdm1(mo_coeff, mo_occ, **kwargs)

    def init_direct_scf(self, mol=None):
        if mol is None: mol = self.mol
        def set_vkscreen(opt, name):
            opt._this.contents.r_vkscreen = _vhf._fpointer(name)
        opt = _vhf.VHFOpt(mol, 'int2e_spinor', 'CVHFrkbllll_prescreen',
                          'CVHFrkbllll_direct_scf',
                          'CVHFrkbllll_direct_scf_dm')
        opt.direct_scf_tol = self.direct_scf_tol
        set_vkscreen(opt, 'CVHFrkbllll_vkscreen')
        return opt

    def get_jk(self, mol=None, dm=None, hermi=1, with_j=True, with_k=True,
               omega=None):
        if mol is None: mol = self.mol
        if dm is None: dm = self.make_rdm1()
        t0 = (logger.process_clock(), logger.perf_counter())
        if self.direct_scf and self.opt is None:
            self.opt = self.init_direct_scf(mol)
        vj, vk = get_jk(mol, dm, hermi, self.opt, with_j, with_k)
        logger.timer(self, 'vj and vk', *t0)
        return vj, vk

    def get_veff(self, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
        '''Dirac-Coulomb'''
        if mol is None: mol = self.mol
        if dm is None: dm = self.make_rdm1()
        if self.direct_scf:
            ddm = numpy.array(dm, copy=False) - numpy.array(dm_last, copy=False)
            vj, vk = self.get_jk(mol, ddm, hermi=hermi)
            return numpy.array(vhf_last, copy=False) + vj - vk
        else:
            vj, vk = self.get_jk(mol, dm, hermi=hermi)
            return vj - vk

    def analyze(self, verbose=None):
        if verbose is None: verbose = self.verbose
        return dhf.analyze(self, verbose)

    def dip_moment(self, mol=None, dm=None, unit='Debye', verbose=logger.NOTE,
                   picture_change=True, **kwargs):
        r''' Dipole moment calculation with picture change correction

        Args:
             mol: an instance of :class:`Mole`
             dm : a 2D ndarrays density matrices

        Kwarg:
            picture_change (bool) : Whether to compute the dipole moment with
            picture change correction.

        Return:
            A list: the dipole moment on x, y and z component
        '''
        if mol is None: mol = self.mol
        if dm is None: dm =self.make_rdm1()
        log = logger.new_logger(mol, verbose)

        if not (isinstance(dm, numpy.ndarray) and dm.ndim == 2):
            # UHF denisty matrices
            dm = dm[0] + dm[1]

        with mol.with_common_orig((0,0,0)):
            if picture_change:
                ao_dip = self.with_x2c.picture_change(('int1e_r_spinor',
                                                       'int1e_sprsp_spinor'))
            else:
                ao_dip = mol.intor_symmetric('int1e_r_spinor')

        el_dip = numpy.einsum('xij,ji->x', ao_dip, dm).real

        charges = mol.atom_charges()
        coords  = mol.atom_coords()
        nucl_dip = numpy.einsum('i,ix->x', charges, coords)
        mol_dip = nucl_dip - el_dip

        if unit.upper() == 'DEBYE':
            mol_dip *= nist.AU2DEBYE
            log.note('Dipole moment(X, Y, Z, Debye): %8.5f, %8.5f, %8.5f', *mol_dip)
        else:
            log.note('Dipole moment(X, Y, Z, A.U.): %8.5f, %8.5f, %8.5f', *mol_dip)
        return mol_dip

    def x2c1e(self):
        return self
    x2c = x2c1e

    def sfx2c1e(self):
        raise NotImplementedError

    def newton(self):
        from pyscf.x2c.newton import newton
        return newton(self)

    def stability(self, internal=None, external=None, verbose=None, return_status=False):
        '''
        X2C-HF/X2C-KS stability analysis.

        See also pyscf.scf.stability.rhf_stability function.

        Kwargs:
            return_status: bool
                Whether to return `stable_i` and `stable_e`

        Returns:
            If return_status is False (default), the return value includes
            two set of orbitals, which are more close to the stable condition.
            The first corresponds to the internal stability
            and the second corresponds to the external stability.

            Else, another two boolean variables (indicating current status:
            stable or unstable) are returned.
            The first corresponds to the internal stability
            and the second corresponds to the external stability.
        '''
        from pyscf.x2c.stability import x2chf_stability
        return x2chf_stability(self, verbose, return_status)

    def nuc_grad_method(self):
        raise NotImplementedError

X2C_SCF = SCF

class UHF(SCF):
    def to_ks(self, xc='HF'):
        '''Convert the input mean-field object to an X2C-KS object.

        Note this conversion only changes the class of the mean-field object.
        The total energy and wave-function are the same as them in the input
        mean-field object.
        '''
        from pyscf.x2c import dft
        mf = self.view(dft.UKS)
        mf.converged = False
        return mf

X2C_UHF = UHF

class RHF(SCF):
    def __init__(self, mol):
        super().__init__(mol)
        if dhf.zquatev is None:
            raise RuntimeError('zquatev library is required to perform Kramers-restricted X2C-RHF')

    def _eigh(self, h, s):
        return dhf.zquatev.solve_KR_FCSCE(self.mol, h, s)

    def to_ks(self, xc='HF'):
        '''Convert the input mean-field object to an X2C-KS object.

        Note this conversion only changes the class of the mean-field object.
        The total energy and wave-function are the same as them in the input
        mean-field object.
        '''
        from pyscf.x2c import dft
        mf = self.view(dft.RKS)
        mf.converged = False
        return mf

X2C_RHF = RHF

def x2c1e_ghf(mf):
    '''
    For the given *GHF* object, generate X2C-GSCF object in GHF spin-orbital
    basis. Note the orbital basis of X2C_GSCF is different to the X2C_RHF and
    X2C_UHF objects. X2C_RHF and X2C_UHF use spinor basis.

    Args:
        mf : an GHF/GKS object

    Returns:
        An GHF/GKS object

    Examples:

    >>> mol = pyscf.M(atom='H 0 0 0; F 0 0 1', basis='ccpvdz', verbose=0)
    >>> mf = scf.GHF(mol).x2c1e().run()
    '''
    assert isinstance(mf, ghf.GHF)

    if isinstance(mf, _X2C_SCF):
        if mf.with_x2c is None:
            mf.with_x2c = SpinOrbitalX2CHelper(mf.mol)
            return mf
        elif not isinstance(mf.with_x2c, SpinOrbitalX2CHelper):
            # An object associated to sfx2c1e.SpinFreeX2CHelper
            raise NotImplementedError
        else:
            return mf

    mf_class = mf.__class__
    if mf_class.__doc__ is None:
        doc = ''
    else:
        doc = mf_class.__doc__

    class X2C1E_GSCF(_X2C_SCF, mf_class):
        __doc__ = doc + '''
        Attributes for spin-orbital X2C:
            with_x2c : X2C object
        '''
        def __init__(self, mol, *args, **kwargs):
            mf_class.__init__(self, mol, *args, **kwargs)
            self.with_x2c = SpinOrbitalX2CHelper(mf.mol)
            self._keys = self._keys.union(['with_x2c'])

        def get_hcore(self, mol=None):
            if mol is None: mol = self.mol
            return self.with_x2c.get_hcore(mol)

        def dump_flags(self, verbose=None):
            mf_class.dump_flags(self, verbose)
            if self.with_x2c:
                self.with_x2c.dump_flags(verbose)
            return self

        def reset(self, mol):
            self.with_x2c.reset(mol)
            return mf_class.reset(self, mol)

        def dip_moment(self, mol=None, dm=None, unit='Debye', verbose=logger.NOTE,
                       picture_change=True, **kwargs):
            r''' Dipole moment calculation with picture change correction

            Args:
                 mol: an instance of :class:`Mole`
                 dm : a 2D ndarrays density matrices

            Kwarg:
                picture_change (bool) : Whether to compute the dipole moment with
                picture change correction.

            Return:
                A list: the dipole moment on x, y and z component
            '''
            if mol is None: mol = self.mol
            if dm is None: dm = self.make_rdm1()
            log = logger.new_logger(mol, verbose)

            nao = mol.nao
            dm = dm[:nao,:nao] + dm[nao:,nao:]

            charges = mol.atom_charges()
            coords  = mol.atom_coords()
            nucl_dip = numpy.einsum('i,ix->x', charges, coords)
            with mol.with_common_orig(nucl_dip):
                ao_dip = _block_diag(mol.intor_symmetric('int1e_r'))
                if picture_change:
                    xmol = self.with_x2c.get_xmol()[0]
                    nao = xmol.nao
                    prp = xmol.intor_symmetric('int1e_sprsp').reshape(3,4,nao,nao)[:,0]
                    prp = numpy.array([_block_diag(x) for x in prp])
                    ao_dip = self.with_x2c.picture_change((ao_dip, prp))

            mol_dip = -numpy.einsum('xij,ji->x', ao_dip, dm).real

            if unit.upper() == 'DEBYE':
                mol_dip *= nist.AU2DEBYE
                log.note('Dipole moment(X, Y, Z, Debye): %8.5f, %8.5f, %8.5f', *mol_dip)
            else:
                log.note('Dipole moment(X, Y, Z, A.U.): %8.5f, %8.5f, %8.5f', *mol_dip)
            return mol_dip

    with_x2c = SpinOrbitalX2CHelper(mf.mol)
    return mf.view(X2C1E_GSCF).add_keys(with_x2c=with_x2c)


def _uncontract_mol(mol, xuncontract=False, exp_drop=0.2):
    '''mol._basis + uncontracted steep functions'''
    pmol = copy.copy(mol)
    _bas = []
    _env = []
    ptr = len(pmol._env)
    contr_coeff = []

    def _to_full_contraction(mol, bas_idx):
        es = numpy.hstack([mol.bas_exp(ib) for ib in bas_idx])
        cs = scipy.linalg.block_diag(*[mol._libcint_ctr_coeff(ib) for ib in bas_idx])

        es, e_idx, rev_idx = numpy.unique(es.round(9), True, True)
        cs_new = numpy.zeros((es.size, cs.shape[1]))
        for i, j in enumerate(rev_idx):
            cs_new[j] += cs[i]
        return es[::-1], cs_new[::-1]

    aoslices = mol.aoslice_by_atom()
    for ia, (ib0, ib1) in enumerate(aoslices[:,:2]):
        if isinstance(xuncontract, str):
            uncontract_me = ((xuncontract == mol.atom_pure_symbol(ia)) or
                             (xuncontract == mol.atom_symbol(ia)))
        elif isinstance(xuncontract, (tuple, list)):
            uncontract_me = ((mol.atom_pure_symbol(ia) in xuncontract) or
                             (mol.atom_symbol(ia) in xuncontract) or
                             (ia in xuncontract))
        else:
            uncontract_me = xuncontract

        if not uncontract_me:
            p0, p1 = aoslices[ia]
            _bas.append(mol._bas[ib0:ib1])
            contr_coeff.append(numpy.eye(p1-p0))
            continue

        lmax = mol._bas[ib0:ib1,mole.ANG_OF].max()
        assert(all(mol._bas[ib0:ib1, mole.KAPPA_OF] == 0))
        # TODO: loop based on kappa
        for l in range(lmax+1):
            bas_idx = ib0 + numpy.where(mol._bas[ib0:ib1,mole.ANG_OF] == l)[0]
            if len(bas_idx) == 0:
                continue

# Some input basis may be the segmented basis from a general contracted set.
# This may lead to duplicated pGTOs. First contract all basis to remove
# duplicated primitive functions.
            mol_exps, b_coeff = _to_full_contraction(mol, bas_idx)

            if mol.cart:
                degen = (l + 1) * (l + 2) // 2
            else:
                degen = l * 2 + 1
            np, nc = b_coeff.shape

            # Uncontract all basis. Use pGTO basis for X
            bs = numpy.zeros((np, mole.BAS_SLOTS), dtype=numpy.int32)
            bs[:,mole.ATOM_OF] = ia
            bs[:,mole.ANG_OF ] = l
            bs[:,mole.NCTR_OF] = bs[:,mole.NPRIM_OF] = 1
            norm = mole.gto_norm(l, mol_exps)
            _env.append(mol_exps)
            _env.append(norm)
            bs[:,mole.PTR_EXP] = numpy.arange(ptr, ptr+np)
            bs[:,mole.PTR_COEFF] = numpy.arange(ptr+np, ptr+np*2)
            _bas.append(bs)
            ptr += np * 2

            c = b_coeff * 1/norm[:,None]
            c = scipy.linalg.block_diag(*([c,] * degen))
            c = c.reshape((degen, np, degen, nc))
            c = c.transpose(1,0,3,2).reshape(np*degen, nc*degen)
            contr_coeff.append(c)

    pmol._bas = numpy.asarray(numpy.vstack(_bas), dtype=numpy.int32)
    pmol._env = numpy.hstack([mol._env,] + _env)
    contr_coeff = scipy.linalg.block_diag(*contr_coeff)

    return pmol, contr_coeff


def _sqrt(a, tol=1e-14):
    e, v = numpy.linalg.eigh(a)
    idx = e > tol
    return numpy.dot(v[:,idx]*numpy.sqrt(e[idx]), v[:,idx].T.conj())

def _invsqrt(a, tol=1e-14):
    e, v = numpy.linalg.eigh(a)
    idx = e > tol
    return numpy.dot(v[:,idx]/numpy.sqrt(e[idx]), v[:,idx].T.conj())

def _get_hcore_fw(t, v, w, s, x, c):
    s1 = s + reduce(numpy.dot, (x.T.conj(), t, x)) * (.5/c**2)
    tx = numpy.dot(t, x)
    h1 =(v + tx + tx.T.conj() - numpy.dot(x.T.conj(), tx) +
         reduce(numpy.dot, (x.T.conj(), w, x)) * (.25/c**2))

    r = _get_r(s, s1)
    h1 = reduce(numpy.dot, (r.T.conj(), h1, r))
    return h1

def _get_r(s, snesc):
    # R^dag \tilde{S} R = S
    # R = S^{-1/2} [S^{-1/2}\tilde{S}S^{-1/2}]^{-1/2} S^{1/2}
    w, v = numpy.linalg.eigh(s)
    idx = w > 1e-14
    v = v[:,idx]
    w_sqrt = numpy.sqrt(w[idx])
    w_invsqrt = 1 / w_sqrt

    # eigenvectors of S as the new basis
    snesc = reduce(numpy.dot, (v.conj().T, snesc, v))
    r_mid = numpy.einsum('i,ij,j->ij', w_invsqrt, snesc, w_invsqrt)
    w1, v1 = numpy.linalg.eigh(r_mid)
    idx1 = w1 > 1e-14
    v1 = v1[:,idx1]
    r_mid = numpy.dot(v1/numpy.sqrt(w1[idx1]), v1.conj().T)
    r = numpy.einsum('i,ij,j->ij', w_invsqrt, r_mid, w_sqrt)
    # Back transform to AO basis
    r = reduce(numpy.dot, (v, r, v.conj().T))
    return r

def _x2c1e_xmatrix(t, v, w, s, c):
    nao = s.shape[0]
    n2 = nao * 2
    dtype = numpy.result_type(t, v, w, s)
    h = numpy.zeros((n2,n2), dtype=dtype)
    m = numpy.zeros((n2,n2), dtype=dtype)
    h[:nao,:nao] = v
    h[:nao,nao:] = t
    h[nao:,:nao] = t
    h[nao:,nao:] = w * (.25/c**2) - t
    m[:nao,:nao] = s
    m[nao:,nao:] = t * (.5/c**2)

    try:
        e, a = scipy.linalg.eigh(h, m)
        cl = a[:nao,nao:]
        cs = a[nao:,nao:]
        x = numpy.linalg.solve(cl.T, cs.T).T  # B = XA
    except scipy.linalg.LinAlgError:
        d, t = numpy.linalg.eigh(m)
        idx = d>LINEAR_DEP_THRESHOLD
        t = t[:,idx] / numpy.sqrt(d[idx])
        tht = reduce(numpy.dot, (t.T.conj(), h, t))
        e, a = numpy.linalg.eigh(tht)
        a = numpy.dot(t, a)
        idx = e > -c**2
        cl = a[:nao,idx]
        cs = a[nao:,idx]
        # X = B A^{-1} = B A^T S
        x = cs.dot(cl.conj().T).dot(m)
    return x

def _x2c1e_get_hcore(t, v, w, s, c):
    nao = s.shape[0]
    n2 = nao * 2
    dtype = numpy.result_type(t, v, w, s)
    h = numpy.zeros((n2,n2), dtype=dtype)
    m = numpy.zeros((n2,n2), dtype=dtype)
    h[:nao,:nao] = v
    h[:nao,nao:] = t
    h[nao:,:nao] = t
    h[nao:,nao:] = w * (.25/c**2) - t
    m[:nao,:nao] = s
    m[nao:,nao:] = t * (.5/c**2)

    try:
        e, a = scipy.linalg.eigh(h, m)
        cl = a[:nao,nao:]
        # cs = a[nao:,nao:]
        e = e[nao:]
    except scipy.linalg.LinAlgError:
        d, t = numpy.linalg.eigh(m)
        idx = d>LINEAR_DEP_THRESHOLD
        t = t[:,idx] / numpy.sqrt(d[idx])
        tht = reduce(numpy.dot, (t.T.conj(), h, t))
        e, a = numpy.linalg.eigh(tht)
        a = numpy.dot(t, a)
        idx = e > -c**2
        cl = a[:nao,idx]
        # cs = a[nao:,idx]
        e = e[idx]

# The so obtaied X seems not numerically stable.  We changed to the
# transformed matrix
# [1 1] [ V T ] [1 0]
# [0 1] [ T W ] [1 1]
#            h[:nao,:nao] = h[:nao,nao:] = h[nao:,:nao] = h[nao:,nao:] = w * (.25/c**2)
#            m[:nao,:nao] = m[:nao,nao:] = m[nao:,:nao] = m[nao:,nao:] = t * (.5/c**2)
#            h[:nao,:nao]+= v + t
#            h[nao:,nao:]-= t
#            m[:nao,:nao]+= s
#            e, a = scipy.linalg.eigh(h, m)
#            cl = a[:nao,nao:]
#            cs = a[nao:,nao:]
#            x = numpy.eye(nao) + numpy.linalg.solve(cl.T, cs.T).T  # B = XA
#            h1 = _get_hcore_fw(t, v, w, s, x, c)

# Taking A matrix as basis and rewrite the FW Hcore formula, to avoid inversing matrix
#   R^dag \tilde{S} R = S
#   R = S^{-1/2} [S^{-1/2}\tilde{S}S^{-1/2}]^{-1/2} S^{1/2}
# Using A matrix as basis, the representation of R is
#   R[A] = (A^+ S A)^{1/2} = (A^+ S A)^{-1/2} A^+ S A
# Construct h = R^+ h1 R in two steps, first in basis A matrix, then back
# transformed to AO basis
#   h  = (A^+)^{-1} R[A]^+ (A^+ h1 A) R[A] A^{-1}         (0)
# Using (A^+)^{-1} = \tilde{S} A, h can be transformed to
#   h  = \tilde{S} A R[A]^+ A^+ h1 A R[A] A^+ \tilde{S}   (1)
# Using R[A] = R[A]^{-1} A^+ S A,  Eq (0) turns to
#      = S A R[A]^{-1}^+ A^+ h1 A R[A]^{-1} A^+ S
#      = S A R[A]^{-1}^+ e R[A]^{-1} A^+ S                (2)

    w, u = numpy.linalg.eigh(reduce(numpy.dot, (cl.T.conj(), s, cl)))
    idx = w > 1e-14
    # Adopt (2) here becuase X is not appeared in Eq (2).
    # R[A] = u w^{1/2} u^+,  so R[A]^{-1} A^+ S in Eq (2) is
    r = reduce(numpy.dot, (u[:,idx]/numpy.sqrt(w[idx]), u[:,idx].T.conj(),
                           cl.T.conj(), s))
    h1 = reduce(numpy.dot, (r.T.conj()*e, r))
    return h1


def _proj_dmll(mol_nr, dm_nr, mol):
    from pyscf.scf import addons
    proj = addons.project_mo_nr2r(mol_nr, numpy.eye(mol_nr.nao_nr()), mol)
    # *.5 because alpha and beta are summed in project_mo_nr2r
    dm_ll = reduce(numpy.dot, (proj, dm_nr*.5, proj.T.conj()))
    dm_ll = (dm_ll + dhf.time_reversal_matrix(mol, dm_ll)) * .5
    return dm_ll

def _block_diag(mat):
    '''
    [A 0]
    [0 A]
    '''
    return scipy.linalg.block_diag(mat, mat)

def _sigma_dot(mat):
    '''sigma dot A x B + A dot B'''
    quaternion = numpy.vstack([1j * lib.PauliMatrices, numpy.eye(2)[None,:,:]])
    nao = mat.shape[-1] * 2
    return lib.einsum('sxy,spq->xpyq', quaternion, mat).reshape(nao, nao)


# A tag to label the derived SCF class
class _X2C_SCF:
    pass


if __name__ == '__main__':
    from pyscf.x2c import dft
    mol = mole.Mole()
    mol.build(
        verbose = 0,
        atom = [["O" , (0. , 0.     , 0.)],
                [1   , (0. , -0.757 , 0.587)],
                [1   , (0. , 0.757  , 0.587)] ],
        basis = 'ccpvdz-dk',
    )

    method = hf.RHF(mol)
    enr = method.kernel()
    print('E(NR) = %.12g' % enr)

    method = UHF(mol)
    ex2c = method.kernel()
    print('E(X2C1E) = %.12g' % ex2c)
    method.with_x2c.basis = {'O': 'unc-ccpvqz', 'H':'unc-ccpvdz'}
    print('E(X2C1E) = %.12g' % method.kernel())
    method.with_x2c.approx = 'atom1e'
    print('E(X2C1E) = %.12g' % method.kernel())

    method = dft.UKS(mol)
    ex2c = method.kernel()
    print('E(X2C1E) = %.12g' % ex2c)


# setup.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
from setuptools import setup, find_packages, Extension
from setuptools.command.build_ext import build_ext

CLASSIFIERS = [
'Development Status :: 5 - Production/Stable',
'Intended Audience :: Science/Research',
'Intended Audience :: Developers',
'License :: OSI Approved :: Apache Software License',
'Programming Language :: C',
'Programming Language :: Python',
'Programming Language :: Python :: 3.6',
'Programming Language :: Python :: 3.7',
'Programming Language :: Python :: 3.8',
'Programming Language :: Python :: 3.9',
'Topic :: Software Development',
'Topic :: Scientific/Engineering',
'Operating System :: POSIX',
'Operating System :: Unix',
'Operating System :: MacOS',
]

NAME             = 'pyscf'
MAINTAINER       = 'Qiming Sun'
MAINTAINER_EMAIL = 'osirpt.sun@gmail.com'
DESCRIPTION      = 'PySCF: Python-based Simulations of Chemistry Framework'
#LONG_DESCRIPTION = ''
URL              = 'http://www.pyscf.org'
DOWNLOAD_URL     = 'http://github.com/pyscf/pyscf'
LICENSE          = 'Apache License 2.0'
AUTHOR           = 'Qiming Sun'
AUTHOR_EMAIL     = 'osirpt.sun@gmail.com'
PLATFORMS        = ['Linux', 'Mac OS-X', 'Unix']
def get_version():
    topdir = os.path.abspath(os.path.join(__file__, '..'))
    with open(os.path.join(topdir, 'pyscf', '__init__.py'), 'r') as f:
        for line in f.readlines():
            if line.startswith('__version__'):
                delim = '"' if '"' in line else "'"
                return line.split(delim)[1]
    raise ValueError("Version string not found")
VERSION = get_version()

EXTRAS = {
    'geomopt': ['pyberny>=0.6.2', 'geometric>=0.9.7.2', 'pyscf-qsdopt'],
    'dftd3': ['pyscf-dftd3'],
    'dmrgscf': ['pyscf-dmrgscf'],
    'doci': ['pyscf-doci'],
    'icmpspt': ['pyscf-icmpspt'],
    'properties': ['pyscf-properties'],
    'semiempirical': ['pyscf-semiempirical'],
    'shciscf': ['pyscf-shciscf'],
    'cppe': ['cppe'],
    'pyqmc':['pyqmc'],
}
EXTRAS['all'] = [p for extras in EXTRAS.values() for p in extras]
# extras which should not be installed by "all" components
EXTRAS['cornell_shci'] = ['pyscf-cornell-shci']
EXTRAS['nao'] = ['pyscf-nao']
EXTRAS['fciqmcscf'] = ['pyscf-fciqmc']
EXTRAS['tblis'] = ['pyscf-tblis']

class CMakeBuildExt(build_ext):
    def run(self):
        extension = self.extensions[0]
        assert extension.name == 'pyscf_lib_placeholder'
        self.build_cmake(extension)

    def build_cmake(self, extension):
        self.announce('Configuring extensions', level=3)
        src_dir = os.path.abspath(os.path.join(__file__, '..', 'pyscf', 'lib'))
        cmd = ['cmake', f'-S{src_dir}', f'-B{self.build_temp}']
        configure_args = os.getenv('CMAKE_CONFIGURE_ARGS')
        if configure_args:
            cmd.extend(configure_args.split(' '))
        self.spawn(cmd)

        self.announce('Building binaries', level=3)
        # Do not use high level parallel compilation. OOM may be triggered
        # when compiling certain functionals in libxc.
        cmd = ['cmake', '--build', self.build_temp, '-j2']
        build_args = os.getenv('CMAKE_BUILD_ARGS')
        if build_args:
            cmd.extend(build_args.split(' '))
        if self.dry_run:
            self.announce(' '.join(cmd))
        else:
            self.spawn(cmd)

    # To remove the infix string like cpython-37m-x86_64-linux-gnu.so
    # Python ABI updates since 3.5
    # https://www.python.org/dev/peps/pep-3149/
    def get_ext_filename(self, ext_name):
        ext_path = ext_name.split('.')
        filename = build_ext.get_ext_filename(self, ext_name)
        name, ext_suffix = os.path.splitext(filename)
        return os.path.join(*ext_path) + ext_suffix

# Here to change the order of sub_commands to ['build_py', ..., 'build_ext']
# C extensions by build_ext are installed in source directory.
# build_py then copy all .so files into "build_ext.build_lib" directory.
# We have to ensure build_ext being executed earlier than build_py.
# A temporary workaround is to modifying the order of sub_commands in build class
from distutils.command.build import build
build.sub_commands = ([c for c in build.sub_commands if c[0] == 'build_ext'] +
                      [c for c in build.sub_commands if c[0] != 'build_ext'])

# scipy bugs
# https://github.com/scipy/scipy/issues/12533
_scipy_version = 'scipy!=1.5.0,!=1.5.1'
import sys
if sys.platform == 'darwin':
    # https://github.com/scipy/scipy/issues/15362
    if sys.version_info < (3, 8):
        _scipy_version = 'scipy<=1.1.0'
    else:
        # https://github.com/scipy/scipy/issues/16151
        print('scipy>1.1.0 may crash with segmentation fault when calling scipy.linalg.eigh')

setup(
    name=NAME,
    version=VERSION,
    description=DESCRIPTION,
    url=URL,
    download_url=DOWNLOAD_URL,
    license=LICENSE,
    classifiers=CLASSIFIERS,
    author=AUTHOR,
    author_email=AUTHOR_EMAIL,
    platforms=PLATFORMS,
    #package_dir={'pyscf': 'pyscf'},  # packages are under directory pyscf
    #include *.so *.dat files. They are now placed in MANIFEST.in
    #package_data={'': ['*.so', '*.dylib', '*.dll', '*.dat']},
    include_package_data=True,  # include everything in source control
    packages=find_packages(exclude=['*test*', '*examples*']),
    # The ext_modules placeholder is to ensure build_ext getting initialized
    ext_modules=[Extension('pyscf_lib_placeholder', [])],
    cmdclass={'build_ext': CMakeBuildExt},
    install_requires=['numpy>=1.13,!=1.16,!=1.17',
                      'scipy<=1.1.0' if sys.platform == "darwin" else 'scipy!=1.5.0,!=1.5.1',
                      'h5py>=2.7'],
    extras_require=EXTRAS,
)


# pyscf/dft/test/test_dks.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest
import numpy
from pyscf import gto
from pyscf import lib
from pyscf.dft import dks

mol = gto.Mole()
mol.atom = '''
O    0    0    0
H    0.   -0.757   0.587
H    0.   0.757    0.587'''
mol.charge = 1
mol.spin = None
mol.basis = 'uncsto3g'
mol.verbose = 7
mol.output = '/dev/null'
mol.build()

def tearDownModule():
    global mol
    mol.stdout.close()
    del mol

class KnownValues(unittest.TestCase):
    def test_noncol_dks_lda(self):
        with lib.temporary_env(lib.param, LIGHT_SPEED=12):
            mf = dks.UDKS(mol)
            mf.xc = 'lda,vwn'
            eks4 = mf.kernel()
        self.assertAlmostEqual(eks4, -80.3600775975093, 8)

    def test_x2c_uks_lda(self):
        mf = dks.UDKS(mol).x2c()
        mf.xc = 'lda,'
        eks4 = mf.kernel()
        self.assertAlmostEqual(eks4, -73.9972147814834, 8)

    def test_dks_lda_omega(self):
        mf = dks.UDKS(mol)
        mf.xc = 'lda + .2*HF'
        eks4 = mf.kernel()
        self.assertAlmostEqual(eks4, -75.71593214601225, 8)

        mf = dks.UDKS(mol)
        mf.xc = 'lda + .2*HF'
        mf.omega = .5
        eks4 = mf.kernel()
        self.assertAlmostEqual(eks4, -75.26598450889301, 8)

    def test_collinear_dks_lda(self):
        mf = dks.UDKS(mol)
        mf.xc = 'lda,vwn'
        mf.collinear = True
        eks4 = mf.kernel()
        self.assertAlmostEqual(eks4, -74.60528573582566, 8)

    def test_collinear_dks_gga(self):
        mf = dks.UDKS(mol)
        mf.xc = 'pbe'
        mf.collinear = True
        eks4 = mf.kernel()
        self.assertAlmostEqual(eks4, -75.07046613221685, 8)

    def test_collinear_x2c_uks_gga(self):
        mf = dks.UDKS(mol).x2c()
        mf.xc = 'pbe'
        mf.collinear = True
        eks4 = mf.kernel()
        self.assertAlmostEqual(eks4, -75.06735429926894, 8)


if __name__ == "__main__":
    print("Test DKS")
    unittest.main()


# pyscf/dft/test/test_grids.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest
import numpy
from pyscf import lib
from pyscf import gto
from pyscf import dft
from pyscf.dft import gen_grid
from pyscf.dft import radi

def setUpModule():
    global h2o
    h2o = gto.Mole()
    h2o.verbose = 5
    h2o.output = '/dev/null'
    h2o.atom = [
        ["O" , (0. , 0.     , 0.)],
        [1   , (0. , -0.757 , 0.587)],
        [1   , (0. , 0.757  , 0.587)] ]

    h2o.basis = {"H": '6-31g',
                 "O": '6-31g',}
    h2o.build()

def tearDownModule():
    global h2o
    h2o.stdout.close()
    del h2o


class KnownValues(unittest.TestCase):
    def test_gen_grid(self):
        grid = gen_grid.Grids(h2o)
        grid.prune = None
        grid.radi_method = radi.gauss_chebyshev
        grid.becke_scheme = gen_grid.original_becke
        grid.radii_adjust = radi.becke_atomic_radii_adjust
        grid.atomic_radii = radi.BRAGG_RADII
        grid.atom_grid = {"H": (10, 50), "O": (10, 50),}
        grid.build(with_non0tab=False)
        self.assertAlmostEqual(numpy.linalg.norm(grid.coords), 185.91245945279027, 9)
        self.assertAlmostEqual(numpy.linalg.norm(grid.weights), 1720.1317185648893, 8)

        grid.becke_scheme = gen_grid.stratmann
        grid.build(with_non0tab=False)
        self.assertAlmostEqual(numpy.linalg.norm(grid.weights), 1730.3692983091271, 8)

        grid.atom_grid = {"O": (10, 50),}
        grid.radii_adjust = None
        grid.becke_scheme = gen_grid.stratmann
        grid.kernel(with_non0tab=False)
        self.assertAlmostEqual(numpy.linalg.norm(grid.weights), 2559.0064040257907, 8)

        grid.atom_grid = (10, 11)
        grid.becke_scheme = gen_grid.original_becke
        grid.radii_adjust = None
        grid.build(with_non0tab=False)
        self.assertAlmostEqual(numpy.linalg.norm(grid.weights), 1712.3069450297105, 8)

    def test_radi(self):
        grid = gen_grid.Grids(h2o)
        grid.prune = None
        grid.radii_adjust = radi.becke_atomic_radii_adjust
        grid.atomic_radii = radi.COVALENT_RADII
        grid.radi_method = radi.mura_knowles
        grid.atom_grid = {"H": (10, 50), "O": (10, 50),}
        grid.build(with_non0tab=False)
        self.assertAlmostEqual(numpy.linalg.norm(grid.weights), 1804.5437331817291, 9)

        grid.radi_method = radi.delley
        grid.build(with_non0tab=False)
        self.assertAlmostEqual(numpy.linalg.norm(grid.weights), 1686.3482864673697, 9)

        grid.radi_method = radi.becke
        grid.build(with_non0tab=False)
        self.assertAlmostEqual(lib.fp(grid.weights), 2486249.209827192, 7)

    def test_prune(self):
        grid = gen_grid.Grids(h2o)
        grid.prune = gen_grid.sg1_prune
        grid.atom_grid = {"H": (10, 50), "O": (10, 50),}
        grid.build(with_non0tab=False)
        self.assertAlmostEqual(numpy.linalg.norm(grid.coords), 202.17732600266302, 9)
        self.assertAlmostEqual(numpy.linalg.norm(grid.weights), 442.54536463517167, 9)

        grid.prune = gen_grid.nwchem_prune
        grid.build(with_non0tab=False)
        self.assertAlmostEqual(numpy.linalg.norm(grid.coords), 149.55023044392638, 9)
        self.assertAlmostEqual(numpy.linalg.norm(grid.weights), 586.36841824004455, 9)

        z = 16
        rad, dr = radi.gauss_chebyshev(50)
        angs = gen_grid.sg1_prune(z, rad, 434, radii=radi.SG1RADII)
        self.assertAlmostEqual(lib.fp(angs), -291.0794420982329, 9)

        angs = gen_grid.nwchem_prune(z, rad, 434, radii=radi.BRAGG_RADII)
        self.assertAlmostEqual(lib.fp(angs), -180.12023039394498, 9)

        angs = gen_grid.nwchem_prune(z, rad, 26, radii=radi.BRAGG_RADII)
        self.assertTrue(numpy.all(angs==26))

    def test_gen_atomic_grids(self):
        grid = gen_grid.Grids(h2o)
        grid.prune = None
        grid.atom_grid = {"H": (10, 58), "O": (10, 50),}
        self.assertRaises(ValueError, grid.build)

    def test_make_mask(self):
        grid = gen_grid.Grids(h2o)
        grid.atom_grid = {"H": (10, 110), "O": (10, 110),}
        grid.build()
        coords = grid.coords*10.
        non0 = gen_grid.make_mask(h2o, coords)
        self.assertEqual(non0.sum(), 122)
        self.assertAlmostEqual(lib.fp(non0), 0.554275491306796, 9)

    def test_overwriting_grids_attribute(self):
        g = gen_grid.Grids(h2o).run()
        self.assertEqual(g.weights.size, 34310)

        g.atom_grid = {"H": (10, 110), "O": (10, 110),}
        self.assertTrue(g.weights is None)


if __name__ == "__main__":
    print("Test Grids")
    unittest.main()



# pyscf/dft/test/test_h2o.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest
import numpy
from pyscf import gto
from pyscf import lib
from pyscf import dft

def setUpModule():
    global h2o, h2osym, h2o_cation, h2osym_cation
    h2o = gto.Mole()
    h2o.verbose = 5
    h2o.output = '/dev/null'
    h2o.atom = [
        ["O" , (0. , 0.     , 0.)],
        [1   , (0. , -0.757 , 0.587)],
        [1   , (0. , 0.757  , 0.587)] ]

    h2o.basis = {"H": '6-31g', "O": '6-31g',}
    h2o.build()

    h2osym = gto.Mole()
    h2osym.verbose = 5
    h2osym.output = '/dev/null'
    h2osym.atom = [
        ["O" , (0. , 0.     , 0.)],
        [1   , (0. , -0.757 , 0.587)],
        [1   , (0. , 0.757  , 0.587)] ]

    h2osym.basis = {"H": '6-31g', "O": '6-31g',}
    h2osym.symmetry = 1
    h2osym.build()

    h2o_cation = gto.M(
        verbose = 5,
        output = '/dev/null',
        atom = [
        ["O" , (0. , 0.     , 0.)],
        [1   , (0. , -0.757 , 0.587)],
        [1   , (0. , 0.757  , 0.587)]],
        charge = 1,
        spin = 1,
        basis = '631g')

    h2osym_cation = gto.M(
        verbose = 5,
        output = '/dev/null',
        atom = [
        ["O" , (0. , 0.     , 0.)],
        [1   , (0. , -0.757 , 0.587)],
        [1   , (0. , 0.757  , 0.587)]],
        symmetry = True,
        charge = 1,
        spin = 1,
        basis = '631g')

def tearDownModule():
    global h2o, h2osym, h2o_cation, h2osym_cation
    h2o.stdout.close()
    h2osym.stdout.close()
    h2o_cation.stdout.close()
    h2osym_cation.stdout.close()
    del h2o, h2osym, h2o_cation, h2osym_cation


class KnownValues(unittest.TestCase):
    def test_nr_lda(self):
        method = dft.RKS(h2o)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        method.xc = 'lda, vwn_rpa'
        self.assertAlmostEqual(method.scf(), -76.01330948329084, 8)

    def test_nr_pw91pw91(self):
        method = dft.RKS(h2o)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        method.xc = 'pw91, pw91'
        # Small change from libxc3 to libxc4
        self.assertAlmostEqual(method.scf(), -76.355310330095563, 7)

    def test_nr_b88vwn(self):
        method = dft.RKS(h2o)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        method.xc = 'b88, vwn'
        self.assertAlmostEqual(method.scf(), -76.690247578608236, 8)

    def test_nr_xlyp(self):
        method = dft.RKS(h2o)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        method.xc = 'xlyp'
        self.assertAlmostEqual(method.scf(), -76.4174879445209, 8)

    def test_nr_b3lypg(self):
        method = dft.RKS(h2o)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        method.xc = 'b3lypg'
        self.assertAlmostEqual(method.scf(), -76.384928891413438, 8)
        g = method.nuc_grad_method().kernel()
        self.assertAlmostEqual(lib.fp(g), -0.035648772973075241, 6)

    def test_nr_b3lypg_direct(self):
        method = dft.RKS(h2o)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.radi_method = dft.radi.gauss_chebyshev
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        method.xc = 'b3lypg'
        method.max_memory = 0
        method.direct_scf = True
        self.assertAlmostEqual(method.scf(), -76.384928823070567, 8)
        method.direct_scf = False
        self.assertAlmostEqual(method.scf(), -76.384928823070567, 8)

    def test_nr_ub3lypg(self):
        method = dft.UKS(h2o)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        method.xc = 'b3lypg'
        self.assertAlmostEqual(method.scf(), -76.384928891413438, 8)
        g = method.nuc_grad_method().kernel()
        self.assertAlmostEqual(lib.fp(g), -0.035648777277847155, 6)

    def test_nr_uks_lsda(self):
        method = dft.UKS(h2osym_cation)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        self.assertAlmostEqual(method.scf(), -75.350995324984709, 8)

    def test_nr_uks_b3lypg(self):
        method = dft.UKS(h2osym_cation)
        method.xc = 'b3lypg'
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        self.assertAlmostEqual(method.scf(), -75.927304010489976, 8)

    def test_nr_uks_b3lypg_direct(self):
        method = dft.UKS(h2o)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.xc = 'b3lypg'
        method.max_memory = 0
        method.direct_scf = True
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        self.assertAlmostEqual(method.scf(), -76.384928891413438, 8)

    def test_nr_uks_b3lypg_cart(self):
        mol1 = h2o.copy()
        mol1.basis = 'ccpvdz'
        mol1.charge = 1
        mol1.spin = 1
        mol1.cart = True
        mol1.build(0, 0)
        method = dft.UKS(mol1)
        method.xc = 'b3lypg'
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        self.assertAlmostEqual(method.scf(), -75.968190479564399, 8)

    def test_nr_roks_lsda(self):
        method = dft.RKS(h2o_cation)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        self.assertAlmostEqual(method.scf(), -75.350333965173704, 8)

    def test_nr_roks_b3lypg(self):
        method = dft.ROKS(h2o_cation)
        method.xc = 'b3lypg'
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        self.assertAlmostEqual(method.scf(), -75.926526046608529, 8)
        g = method.nuc_grad_method().kernel()
        self.assertAlmostEqual(lib.fp(g), -0.10184251826412283, 6)

    def test_nr_roks_b3lypg_direct(self):
        method = dft.ROKS(h2o_cation)
        method.xc = 'b3lypg'
        method.max_memory = 0
        method.direct_scf = True
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        self.assertAlmostEqual(method.scf(), -75.926526046608529, 8)

    def test_nr_gks_lsda(self):
        method = dft.GKS(h2osym_cation)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        self.assertAlmostEqual(method.scf(), -75.350995324984709, 8)

    def test_nr_gks_b3lypg(self):
        method = dft.GKS(h2osym_cation)
        method.xc = 'b3lypg'
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        self.assertAlmostEqual(method.scf(), -75.902391377392391, 8)

    def test_nr_gks_b3lypg_direct(self):
        method = dft.GKS(h2o_cation)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.xc = 'b3lypg'
        method.max_memory = 0
        method.direct_scf = True
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        self.assertAlmostEqual(method.scf(), -75.902391377392391, 8)

#########
    def test_nr_symm_lda(self):
        method = dft.RKS(h2osym)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        method.xc = 'lda, vwn_rpa'
        self.assertAlmostEqual(method.scf(), -76.01330948329084, 8)

    def test_nr_symm_pw91pw91(self):
        method = dft.RKS(h2osym)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        method.xc = 'pw91, pw91'
        # Small change from libxc3 to libxc4
        self.assertAlmostEqual(method.scf(), -76.355310330095563, 7)

    def test_nr_symm_b88vwn(self):
        method = dft.RKS(h2osym)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        method.xc = 'b88, vwn'
        self.assertAlmostEqual(method.scf(), -76.690247578608236, 8)

    def test_nr_symm_b88vwn_df(self):
        method = dft.density_fit(dft.RKS(h2osym), 'weigend')
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        method.xc = 'b88, vwn'
        self.assertAlmostEqual(method.scf(), -76.690346887915879, 8)

    def test_nr_symm_xlyp(self):
        method = dft.RKS(h2osym)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        method.xc = 'xlyp'
        self.assertAlmostEqual(method.scf(), -76.4174879445209, 8)

    def test_nr_symm_b3lypg(self):
        method = dft.RKS(h2osym)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        method.xc = 'b3lypg'
        self.assertAlmostEqual(method.scf(), -76.384928891413438, 8)

    def test_nr_symm_b3lypg_direct(self):
        method = dft.RKS(h2osym)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.radi_method = dft.radi.gauss_chebyshev
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        method.xc = 'b3lypg'
        method.max_memory = 0
        method.direct_scf = True
        self.assertAlmostEqual(method.scf(), -76.384928823070567, 8)
        method.direct_scf = False
        self.assertAlmostEqual(method.scf(), -76.384928823070567, 8)

    def test_nr_symm_ub3lypg(self):
        method = dft.UKS(h2osym)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        method.xc = 'b3lypg'
        self.assertAlmostEqual(method.scf(), -76.384928891413438, 8)

    def test_nr_symm_uks_lsda(self):
        method = dft.UKS(h2osym_cation)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        self.assertAlmostEqual(method.scf(), -75.350995324984709, 8)

    def test_nr_symm_uks_b3lypg(self):
        method = dft.UKS(h2osym_cation)
        method.xc = 'b3lypg'
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        self.assertAlmostEqual(method.scf(), -75.927304010489976, 8)

    def test_nr_symm_uks_b3lypg_direct(self):
        method = dft.UKS(h2osym)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.xc = 'b3lypg'
        method.max_memory = 0
        method.direct_scf = True
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        self.assertAlmostEqual(method.scf(), -76.384928891413438, 8)

    def test_nr_symm_roks_lsda(self):
        method = dft.RKS(h2osym_cation)
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        self.assertAlmostEqual(method.scf(), -75.350333965173704, 8)

    def test_nr_symm_roks_b3lypg(self):
        method = dft.ROKS(h2osym_cation)
        method.xc = 'b3lypg'
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        self.assertAlmostEqual(method.scf(), -75.926526046608529, 8)

    def test_nr_symm_roks_b3lypg_direct(self):
        method = dft.ROKS(h2osym_cation)
        method.xc = 'b3lypg'
        method.max_memory = 0
        method.direct_scf = True
        method.grids.prune = dft.gen_grid.treutler_prune
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        self.assertAlmostEqual(method.scf(), -75.926526046608529, 8)

    def test_nr_mgga(self):
        method = dft.RKS(h2o)
        method.xc = 'm06l,m06l'
        method.grids.prune = None
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        self.assertAlmostEqual(method.scf(), -76.3772366, 6)

    def test_nr_rks_vv10(self):
        method = dft.RKS(h2o)
        dm = method.get_init_guess()
        method.xc = 'wB97M_V'
        method.nlc = 'vv10'
        method.grids.prune = None
        method.grids.atom_grid = {"H": (30, 86), "O": (30, 86),}
        method.nlcgrids.prune = None
        method.nlcgrids.atom_grid = {"H": (20, 50), "O": (20, 50),}
        method.dump_flags()
        vxc = method.get_veff(h2o, dm)
        self.assertAlmostEqual(lib.fp(vxc), 22.767504283729778, 8)

        method._eri = None
        method.max_memory = 0
        vxc = method.get_veff(h2o, dm)
        self.assertAlmostEqual(lib.fp(vxc), 22.767504283729778, 8)

    def test_nr_uks_vv10(self):
        method = dft.UKS(h2o)
        dm = method.get_init_guess()
        dm = (dm[0], dm[0])
        method.xc = 'wB97M_V'
        method.nlc = 'vv10'
        method.grids.prune = None
        method.grids.atom_grid = {"H": (30, 86), "O": (30, 86),}
        method.nlcgrids.prune = None
        method.nlcgrids.atom_grid = {"H": (20, 50), "O": (20, 50),}
        method.dump_flags()
        vxc = method.get_veff(h2o, dm)
        self.assertAlmostEqual(lib.fp(vxc[0]), 22.767504283729778, 8)
        self.assertAlmostEqual(lib.fp(vxc[1]), 22.767504283729778, 8)

        method._eri = None
        method.max_memory = 0
        vxc = method.get_veff(h2o, dm)
        self.assertAlmostEqual(lib.fp(vxc[0]), 22.767504283729778, 8)
        self.assertAlmostEqual(lib.fp(vxc[1]), 22.767504283729778, 8)

    def test_nr_rks_rsh(self):
        method = dft.RKS(h2o)
        dm = method.get_init_guess()
        method.xc = 'wB97M_V'
        vxc = method.get_veff(h2o, dm)
        self.assertAlmostEqual(lib.fp(vxc), 22.759558596896344, 8)

        method._eri = None
        method.max_memory = 0
        method.xc = 'wB97M_V'
        vxc = method.get_veff(h2o, dm, dm, vxc)
        self.assertAlmostEqual(lib.fp(vxc), 22.759558596896344, 8)

        method.xc = 'B97M_V'
        vxc = method.get_veff(h2o, dm)
        self.assertAlmostEqual(lib.fp(vxc), 23.058813088809824, 8)

        method._eri = None
        method.max_memory = 0
        method.xc = 'B97M_V'
        vxc = method.get_veff(h2o, dm, dm, vxc)
        self.assertAlmostEqual(lib.fp(vxc), 23.058813088809824, 8)

    def test_nr_rks_rsh_cart(self):
        mol1 = h2o.copy()
        mol1.basis = 'ccpvdz'
        mol1.cart = True
        mol1.build(0, 0)
        method = dft.RKS(mol1)
        method.xc = 'B97M_V'
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        self.assertAlmostEqual(method.kernel(), -76.44022393692919, 8)

    def test_nr_uks_rsh(self):
        method = dft.UKS(h2o)
        dm = method.get_init_guess()
        dm = (dm[0], dm[0])
        method.xc = 'wB97M_V'
        vxc = method.get_veff(h2o, dm)
        self.assertAlmostEqual(lib.fp(vxc[0]), 22.759558596896344, 8)
        self.assertAlmostEqual(lib.fp(vxc[1]), 22.759558596896344, 8)

        method._eri = None
        method.max_memory = 0
        method.xc = 'wB97M_V'
        vxc = method.get_veff(h2o, dm, dm, vxc)
        self.assertAlmostEqual(lib.fp(vxc[0]), 22.759558596896344, 8)
        self.assertAlmostEqual(lib.fp(vxc[1]), 22.759558596896344, 8)

        method.xc = 'B97M_V'
        vxc = method.get_veff(h2o, dm)
        self.assertAlmostEqual(lib.fp(vxc[0]), 23.058813088809824, 8)
        self.assertAlmostEqual(lib.fp(vxc[1]), 23.058813088809824, 8)

        method._eri = None
        method.max_memory = 0
        method.xc = 'B97M_V'
        vxc = method.get_veff(h2o, dm, dm, vxc)
        self.assertAlmostEqual(lib.fp(vxc[0]), 23.058813088809824, 8)
        self.assertAlmostEqual(lib.fp(vxc[1]), 23.058813088809824, 8)

    def test_nr_gks_rsh(self):
        method = dft.GKS(h2o)
        dm = method.get_init_guess()
        dm = dm + numpy.sin(dm)*.02j
        dm = dm + dm.conj().T
        method.xc = 'wB97M_V'
        vxc = method.get_veff(h2o, dm)
        self.assertAlmostEqual(lib.fp(vxc), 3.1818982731583274+0j, 8)

        method._eri = None
        method.max_memory = 0
        method.xc = 'wB97M_V'
        vxc = method.get_veff(h2o, dm, dm, vxc)
        self.assertAlmostEqual(lib.fp(vxc), 3.1818982731583274+0j, 8)

        method.xc = 'B97M_V'
        vxc = method.get_veff(h2o, dm)
        self.assertAlmostEqual(lib.fp(vxc), 2.0131447223203565+0j, 8)

        method._eri = None
        method.max_memory = 0
        method.xc = 'B97M_V'
        vxc = method.get_veff(h2o, dm, dm, vxc)
        self.assertAlmostEqual(lib.fp(vxc), 2.0131447223203565+0j, 8)

    def test_nr_rks_vv10_high_cost(self):
        method = dft.RKS(h2o)
        method.xc = 'wB97M_V'
        method.nlc = 'vv10'
        method.grids.prune = None
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        method.nlcgrids.prune = None
        method.nlcgrids.atom_grid = {"H": (40, 110), "O": (40, 110),}
        self.assertAlmostEqual(method.scf(), -76.352381513158718, 8)

    def test_nr_uks_vv10_high_cost(self):
        method = dft.UKS(h2o)
        method.xc = 'wB97M_V'
        method.nlc = 'vv10'
        method.grids.prune = None
        method.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        method.nlcgrids.prune = None
        method.nlcgrids.atom_grid = {"H": (40, 110), "O": (40, 110),}
        self.assertAlmostEqual(method.scf(), -76.352381513158718, 8)

    def test_camb3lyp_rsh_omega(self):
        mf = dft.RKS(h2o)
        mf.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        mf.run(xc='camb3lyp')
        self.assertAlmostEqual(mf.e_tot, -76.35549300028714, 9)

        mf1 = dft.RKS(h2o)
        mf1.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        mf1.run(xc='camb3lyp', omega=0.15)
        self.assertAlmostEqual(mf1.e_tot, -76.36649222362115, 9)

        mf2 = dft.RKS(h2o)
        mf2.xc='RSH(.15,0.65,-0.46) + 0.46*ITYH + .35*B88 + VWN5*0.19, LYP*0.81'
        mf2.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        mf2.kernel()
        self.assertAlmostEqual(mf1.e_tot, -76.36649222362115, 9)

    def test_reset(self):
        mf = dft.RKS(h2o).newton()
        mf.reset(h2osym)
        self.assertTrue(mf.mol is h2osym)
        self.assertTrue(mf.grids.mol is h2osym)
        self.assertTrue(mf.nlcgrids.mol is h2osym)

    def test_init_guess_by_vsap(self):
        dm = dft.RKS(h2o).get_init_guess(key='vsap')
        self.assertAlmostEqual(lib.fp(dm), 1.7285100188309719, 9)

        dm = dft.ROKS(h2osym).get_init_guess(key='vsap')
        self.assertEqual(dm.ndim, 3)
        self.assertAlmostEqual(lib.fp(dm), 1.9698972986009409, 9)

        dm = dft.UKS(h2osym).init_guess_by_vsap()
        self.assertEqual(dm.ndim, 3)
        self.assertAlmostEqual(lib.fp(dm), 1.9698972986009409, 9)

    def test_init(self):
        mol_r = h2o
        mol_u = gto.M(atom='Li', spin=1, verbose=0)
        mol_r1 = gto.M(atom='H', spin=1, verbose=0)
        sym_mol_r = h2osym
        sym_mol_u = gto.M(atom='Li', spin=1, symmetry=1, verbose=0)
        sym_mol_r1 = gto.M(atom='H', spin=1, symmetry=1, verbose=0)
        self.assertTrue(isinstance(dft.RKS(mol_r), dft.rks.RKS))
        self.assertTrue(isinstance(dft.RKS(mol_u), dft.roks.ROKS))
        self.assertTrue(isinstance(dft.UKS(mol_r), dft.uks.UKS))
        self.assertTrue(isinstance(dft.ROKS(mol_r), dft.roks.ROKS))
        self.assertTrue(isinstance(dft.GKS(mol_r), dft.gks.GKS))
        self.assertTrue(isinstance(dft.KS(mol_r), dft.rks.RKS))
        self.assertTrue(isinstance(dft.KS(mol_u), dft.uks.UKS))
        self.assertTrue(isinstance(dft.DKS(mol_u), dft.dks.UDKS))

        self.assertTrue(isinstance(mol_r.RKS(), dft.rks.RKS))
        self.assertTrue(isinstance(mol_u.RKS(), dft.roks.ROKS))
        self.assertTrue(isinstance(mol_r.UKS(), dft.uks.UKS))
        self.assertTrue(isinstance(mol_r.ROKS(), dft.roks.ROKS))
        self.assertTrue(isinstance(mol_r.GKS(), dft.gks.GKS))
        self.assertTrue(isinstance(mol_r.KS(), dft.rks.RKS))
        self.assertTrue(isinstance(mol_u.KS(), dft.uks.UKS))
        self.assertTrue(isinstance(mol_u.DKS(), dft.dks.UDKS))
        #TODO: self.assertTrue(isinstance(dft.X2C(mol_r), x2c.dft.UKS))

if __name__ == "__main__":
    print("Full Tests for H2O")
    unittest.main()


# pyscf/dft/test/test_he.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest
import numpy
from pyscf import gto
from pyscf import lib
from pyscf import dft
from pyscf import scf

def setUpModule():
    global mol, method, mol1
    mol = gto.Mole()
    mol.verbose = 7
    mol.output = '/dev/null'
    mol.atom = [[2, (0.,0.,0.)], ]
    mol.basis = {"He": 'cc-pvdz'}
    mol.build()
    method = dft.RKS(mol)

    mol1 = gto.Mole()
    mol1.verbose = 0
    mol1.output = None
    mol1.atom = 'He'
    mol1.basis = 'cc-pvdz'
    mol1.charge = 1
    mol1.spin = 1
    mol1.build()

def tearDownModule():
    global mol, method, mol1
    mol.stdout.close()
    del mol, method, mol1


class KnownValues(unittest.TestCase):
    def test_nr_lda(self):
        method.xc = 'lda, vwn_rpa'
        self.assertAlmostEqual(method.scf(), -2.8641551904776055, 9)

    def test_dks_lda(self):
        m = mol.DKS()
        self.assertAlmostEqual(m.kernel(), -2.8268242330361373, 9)

        m = mol.DKS().x2c()
        self.assertAlmostEqual(m.kernel(), -2.826788817256218, 9)

    def test_udks_lda(self):
        m = dft.dks.UDKS(mol)
        self.assertAlmostEqual(m.kernel(), -2.8268242330361373, 9)

        m = dft.dks.UDKS(mol).x2c()
        self.assertAlmostEqual(m.kernel(), -2.826788817256218, 9)

    def test_nr_pw91pw91(self):
        method.xc = 'pw91, pw91'
        self.assertAlmostEqual(method.scf(), -2.8914066724838849, 9)

    def test_nr_b88vwn(self):
        method.xc = 'b88, vwn'
        self.assertAlmostEqual(method.scf(), -2.9670729652962606, 9)

    def test_nr_xlyp(self):
        method.xc = 'xlyp'
        self.assertAlmostEqual(method.scf(), -2.9045738259332161, 9)

    def test_nr_b3lypg(self):
        method.xc = 'b3lypg'
        self.assertAlmostEqual(method.scf(), -2.9070540942168002, 9)

        m = mol.UKS()
        m.xc = 'b3lyp'
        self.assertAlmostEqual(m.scf(), -2.89992555753, 9)

    def test_nr_lda_1e(self):
        mf = dft.RKS(mol1).run()
        self.assertAlmostEqual(mf.e_tot, -1.936332393935281, 9)

    def test_nr_b3lypg_1e(self):
        mf = dft.ROKS(mol1).set(xc='b3lypg').run()
        self.assertAlmostEqual(mf.e_tot, -1.9931564410562266, 9)

    def test_xcfun_nr_blyp(self):
        m = mol.RKS()
        m._numint.libxc = dft.xcfun
        m.xc = 'b88,lyp'
        self.assertAlmostEqual(m.scf(), -2.8978518405, 9)

    def test_nr_m06l(self):
        m = mol.RKS()
        m.xc = 'm06l'
        self.assertAlmostEqual(m.scf(), -2.9039230673864243, 9)

        m = mol.UKS()
        m.xc = 'm06l'
        self.assertAlmostEqual(m.scf(), -2.9039230673864243, 9)

    def test_1e(self):
        mf = dft.RKS(gto.M(atom='H', spin=1)).run()
        self.assertTrue(isinstance(mf, dft.roks.ROKS))
        self.assertAlmostEqual(mf.e_tot, -0.43567023283650547)

        mf = dft.RKS(gto.M(atom='H', spin=1, symmetry=1)).run()
        self.assertTrue(isinstance(mf, dft.rks_symm.ROKS))
        self.assertAlmostEqual(mf.e_tot, -0.43567023283650547)

    def test_convert(self):
        rhf = scf.RHF(mol)
        uhf = scf.UHF(mol)
        ghf = scf.GHF(mol)
        rks = dft.RKS(mol)
        uks = dft.UKS(mol)
        gks = dft.GKS(mol)
        dhf = scf.DHF(mol)
        dks = dft.DKS(mol)
        udhf = scf.dhf.UDHF(mol)
        udks = dft.dks.UDKS(mol)

        self.assertTrue(isinstance(rhf.to_rhf(), scf.rhf.RHF))
        self.assertTrue(isinstance(rhf.to_uhf(), scf.uhf.UHF))
        self.assertTrue(isinstance(rhf.to_ghf(), scf.ghf.GHF))
        self.assertTrue(isinstance(rhf.to_rks(), dft.rks.RKS))
        self.assertTrue(isinstance(rhf.to_uks(), dft.uks.UKS))
        self.assertTrue(isinstance(rhf.to_gks(), dft.gks.GKS))

        self.assertTrue(isinstance(rks.to_rhf(), scf.rhf.RHF))
        self.assertTrue(isinstance(rks.to_uhf(), scf.uhf.UHF))
        self.assertTrue(isinstance(rks.to_ghf(), scf.ghf.GHF))
        self.assertTrue(isinstance(rks.to_rks('pbe'), dft.rks.RKS))
        self.assertTrue(isinstance(rks.to_uks('pbe'), dft.uks.UKS))
        self.assertTrue(isinstance(rks.to_gks('pbe'), dft.gks.GKS))

        self.assertTrue(isinstance(uhf.to_rhf(), scf.rhf.RHF))
        self.assertTrue(isinstance(uhf.to_uhf(), scf.uhf.UHF))
        self.assertTrue(isinstance(uhf.to_ghf(), scf.ghf.GHF))
        self.assertTrue(isinstance(uhf.to_rks(), dft.rks.RKS))
        self.assertTrue(isinstance(uhf.to_uks(), dft.uks.UKS))
        self.assertTrue(isinstance(uhf.to_gks(), dft.gks.GKS))

        self.assertTrue(isinstance(rks.to_rhf(), scf.rhf.RHF))
        self.assertTrue(isinstance(rks.to_uhf(), scf.uhf.UHF))
        self.assertTrue(isinstance(rks.to_ghf(), scf.ghf.GHF))
        self.assertTrue(isinstance(uks.to_rks('pbe'), dft.rks.RKS))
        self.assertTrue(isinstance(uks.to_uks('pbe'), dft.uks.UKS))
        self.assertTrue(isinstance(uks.to_gks('pbe'), dft.gks.GKS))

        #self.assertTrue(isinstance(ghf.to_rhf(), scf.rhf.RHF))
        #self.assertTrue(isinstance(ghf.to_uhf(), scf.uhf.UHF))
        self.assertTrue(isinstance(ghf.to_ghf(), scf.ghf.GHF))
        #self.assertTrue(isinstance(ghf.to_rks(), dft.rks.RKS))
        #self.assertTrue(isinstance(ghf.to_uks(), dft.uks.UKS))
        self.assertTrue(isinstance(ghf.to_gks(), dft.gks.GKS))

        #self.assertTrue(isinstance(gks.to_rhf(), scf.rhf.RHF))
        #self.assertTrue(isinstance(gks.to_uhf(), scf.uhf.UHF))
        self.assertTrue(isinstance(gks.to_ghf(), scf.ghf.GHF))
        #self.assertTrue(isinstance(gks.to_rks('pbe'), dft.rks.RKS))
        #self.assertTrue(isinstance(gks.to_uks('pbe'), dft.uks.UKS))
        self.assertTrue(isinstance(gks.to_gks('pbe'), dft.gks.GKS))

        self.assertRaises(RuntimeError, dhf.to_rhf)
        self.assertRaises(RuntimeError, dhf.to_uhf)
        self.assertRaises(RuntimeError, dhf.to_ghf)
        self.assertRaises(RuntimeError, dks.to_rks)
        self.assertRaises(RuntimeError, dks.to_uks)
        self.assertRaises(RuntimeError, dks.to_gks)

        if scf.dhf.zquatev is not None:
            self.assertTrue(isinstance(dhf.to_dhf(), scf.dhf.RDHF))
            self.assertTrue(isinstance(dhf.to_dks(), dft.dks.RDKS))
            self.assertTrue(isinstance(dks.to_dhf(), scf.dhf.RDHF))
            self.assertTrue(isinstance(dks.to_dks('pbe'), dft.dks.RDKS))
        self.assertTrue(isinstance(dhf.to_dhf(), scf.dhf.DHF))
        self.assertTrue(isinstance(dhf.to_dks(), dft.dks.DKS))
        self.assertTrue(isinstance(dks.to_dhf(), scf.dhf.DHF))
        self.assertTrue(isinstance(dks.to_dks('pbe'), dft.dks.DKS))
        self.assertTrue(isinstance(udhf.to_dhf(), scf.dhf.DHF))
        self.assertTrue(isinstance(udhf.to_dks(), dft.dks.DKS))
        self.assertTrue(isinstance(udks.to_dhf(), scf.dhf.DHF))
        self.assertTrue(isinstance(udks.to_dks('pbe'), dft.dks.DKS))

if __name__ == "__main__":
    print("Full Tests for He")
    unittest.main()


# pyscf/dft/test/test_libxc.py

#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import ctypes
import unittest
import numpy
from pyscf import gto, scf
from pyscf import dft
from pyscf import lib

def setUpModule():
    global mol, mf, ao, rho
    mol = gto.Mole()
    mol.verbose = 0
    mol.output = None
    mol.atom = 'h 0 0 0; h 1 .5 0; h 0 4 1; h 1 0 .2'
    mol.basis = 'aug-ccpvdz'
    mol.build()
    #dm = scf.RHF(mol).run(conv_tol=1e-14).make_rdm1()
    dm = numpy.load(os.path.realpath(os.path.join(__file__, '..', 'dm_h4.npy')))
    mf = dft.RKS(mol)
    mf.grids.atom_grid = {"H": (50, 110)}
    mf.prune = None
    mf.grids.build(with_non0tab=False)
    nao = mol.nao_nr()
    ao = dft.numint.eval_ao(mol, mf.grids.coords, deriv=1)
    rho = dft.numint.eval_rho(mol, ao, dm, xctype='GGA')

def tearDownModule():
    global mol, mf, ao, rho
    del mol, mf, ao, rho

class KnownValues(unittest.TestCase):
    def test_parse_xc(self):
        hyb, fn_facs = dft.libxc.parse_xc('.5*HF+.5*B3LYP,VWN*.5')
        self.assertAlmostEqual(hyb[0], .6, 12)
        self.assertEqual([x[0] for x in fn_facs], [1,106,131,7])
        self.assertTrue(numpy.allclose([x[1] for x in fn_facs],
                                       (0.04, 0.36, 0.405, 0.595)))
        hyb, fn_facs = dft.libxc.parse_xc('HF,')
        self.assertEqual(hyb[0], 1)
        self.assertEqual(fn_facs, [])

        hyb, fn_facs = dft.libxc.parse_xc('B88 - SLATER')
        self.assertEqual(fn_facs, [(106, 1), (1, -1)])
        hyb, fn_facs = dft.libxc.parse_xc('B88 -SLATER*.5')
        self.assertEqual(fn_facs, [(106, 1), (1, -0.5)])

        hyb, fn_facs = dft.libxc.parse_xc('0.5*B3LYP\n+0.25*B3LYP')
        self.assertTrue(numpy.allclose(hyb, [.15, 0, 0]))
        hyb = dft.libxc.hybrid_coeff('0.5*B3LYP+0.25*B3LYP')
        self.assertAlmostEqual(hyb, .15, 12)

        hyb, fn_facs = dft.libxc.parse_xc('0.6*CAM_B3LYP+0.4*B3P86')
        self.assertTrue(numpy.allclose(hyb, [.08, 0, 0]))
        self.assertTrue(numpy.allclose(fn_facs,
                                       [(433, 0.6), (1, 0.032), (106, 0.288), (132, 0.324), (7, 0.076)]))
        rsh = dft.libxc.rsh_coeff('0.6*CAM_B3LYP+0.4*B3P86')
        self.assertTrue(numpy.allclose(rsh, (0.33, 0.39, -0.196)))

        hyb, fn_facs = dft.libxc.parse_xc('0.4*B3P86+0.6*CAM_B3LYP')
        self.assertTrue(numpy.allclose(hyb, [.08, 0, 0]))
        self.assertTrue(numpy.allclose(fn_facs,
                                       [(1, 0.032), (106, 0.288), (132, 0.324), (7, 0.076), (433, 0.6)]))
        rsh = dft.libxc.rsh_coeff('0.4*B3P86+0.6*CAM_B3LYP')
        self.assertTrue(numpy.allclose(rsh, (0.33, 0.39, -0.196)))

        hyb, fn_facs = dft.libxc.parse_xc('0.5*SR-HF(0.3) + .8*HF + .22*LR_HF')
        self.assertEqual(hyb, [1.3, 1.02, 0.3])

        hyb, fn_facs = dft.libxc.parse_xc('0.5*SR-HF + .22*LR_HF(0.3) + .8*HF')
        self.assertEqual(hyb, [1.3, 1.02, 0.3])

        hyb, fn_facs = dft.libxc.parse_xc('0.5*SR-HF + .8*HF + .22*LR_HF(0.3)')
        self.assertEqual(hyb, [1.3, 1.02, 0.3])

        hyb, fn_facs = dft.libxc.parse_xc('0.5*RSH(2.04;0.56;0.3) + 0.5*BP86')
        self.assertEqual(hyb, [1.3, 1.02, 0.3])
        self.assertEqual(fn_facs, [(106, 0.5), (132, 0.5)])

        hyb, fn_facs = dft.libxc.parse_xc('0.5*RSH(.3, 2.04, 0.56) + 0.5*BP86')
        self.assertEqual(hyb, [1.3, 1.02, 0.3])
        self.assertEqual(fn_facs, [(106, 0.5), (132, 0.5)])

        self.assertRaises(ValueError, dft.libxc.parse_xc, 'SR_HF(0.3) + LR_HF(.5)')
        self.assertRaises(ValueError, dft.libxc.parse_xc, 'LR-HF(0.3) + SR-HF(.5)')

        hyb = dft.libxc.hybrid_coeff('M05')
        self.assertAlmostEqual(hyb, 0.28, 9)

        hyb, fn_facs = dft.libxc.parse_xc('APBE,')
        self.assertEqual(fn_facs, [(184, 1)])

        #hyb, fn_facs = dft.libxc.parse_xc('TF,')
        #self.assertEqual(fn_facs, [(50, 1)])

        ref = [(1, 1), (7, 1)]
        self.assertEqual(dft.libxc.parse_xc_name('LDA,VWN'), (1,7))
        self.assertEqual(dft.libxc.parse_xc(('LDA','VWN'))[1], ref)
        self.assertEqual(dft.libxc.parse_xc((1, 7))[1], ref)
        self.assertEqual(dft.libxc.parse_xc('1, 7')[1], ref)
        self.assertEqual(dft.libxc.parse_xc(7)[1], [(7,1)])

        self.assertEqual(dft.libxc.parse_xc('M11-L')[1], [(226,1),(75,1)])
        self.assertEqual(dft.libxc.parse_xc('M11L' )[1], [(226,1),(75,1)])
        self.assertEqual(dft.libxc.parse_xc('M11-L,M11L' )[1], [(226,1),(75,1)])
        self.assertEqual(dft.libxc.parse_xc('M11_L,M11-L')[1], [(226,1),(75,1)])
        self.assertEqual(dft.libxc.parse_xc('M11L,M11_L' )[1], [(226,1),(75,1)])

        self.assertEqual(dft.libxc.parse_xc('Xpbe,')[1], [(123,1)])
        self.assertEqual(dft.libxc.parse_xc('pbe,' )[1], [(101,1)])
        hyb, fn_facs = dft.libxc.parse_xc('PBE*.4+LDA')
        self.assertEqual(fn_facs, [(101, 0.4), (130, 0.4), (1, 1)])
        self.assertRaises(KeyError, dft.libxc.parse_xc, 'PBE+VWN')

        self.assertTrue (dft.libxc.is_meta_gga('m05'))
        self.assertFalse(dft.libxc.is_meta_gga('pbe0'))
        self.assertFalse(dft.libxc.is_meta_gga('tf,'))
        self.assertFalse(dft.libxc.is_meta_gga('vv10'))
        self.assertTrue (dft.libxc.is_gga('PBE0'))
        self.assertFalse(dft.libxc.is_gga('m05'))
        self.assertFalse(dft.libxc.is_gga('tf,'))
        self.assertTrue (dft.libxc.is_lda('tf,'))
        self.assertFalse(dft.libxc.is_lda('vv10'))
        self.assertTrue (dft.libxc.is_hybrid_xc('m05'))
        self.assertTrue (dft.libxc.is_hybrid_xc('pbe0,'))
        self.assertTrue (dft.libxc.is_hybrid_xc('m05,'))
        self.assertFalse(dft.libxc.is_hybrid_xc('vv10'))
        self.assertTrue (dft.libxc.is_hybrid_xc((402,'vv10')))
        self.assertTrue (dft.libxc.is_hybrid_xc(('402','vv10')))

    def test_libxc_cam_beta(self):
        rsh_tmp = (ctypes.c_double*3)()
        dft.libxc._itrf.LIBXC_rsh_coeff(1, rsh_tmp)
        beta = rsh_tmp[2]
        self.assertEqual(beta, 0)

        dft.libxc._itrf.LIBXC_rsh_coeff(433, rsh_tmp)
        dft.libxc._itrf.LIBXC_rsh_coeff(1, rsh_tmp)
        beta = rsh_tmp[2]
        self.assertEqual(beta, 0)

        dft.libxc._itrf.LIBXC_is_hybrid(1)
        dft.libxc._itrf.LIBXC_rsh_coeff(1, rsh_tmp)
        beta = rsh_tmp[2]
        self.assertEqual(beta, 0)

    def test_nlc_coeff(self):
        self.assertEqual(dft.libxc.nlc_coeff('0.5*vv10'), [5.9, 0.0093])
        self.assertEqual(dft.libxc.nlc_coeff('pbe__vv10'), [5.9, 0.0093])

    def test_lda(self):
        e,v,f,k = dft.libxc.eval_xc('lda,', rho[0][:3], deriv=3)
        self.assertAlmostEqual(lib.fp(e)   , -0.4720562542635522, 8)
        self.assertAlmostEqual(lib.fp(v[0]), -0.6294083390180697, 8)
        self.assertAlmostEqual(lib.fp(f[0]), -1.1414693830969338, 8)
        self.assertAlmostEqual(lib.fp(k[0]),  4.1402447248393921, 8)

        e,v,f,k = dft.libxc.eval_xc('lda,', [rho[0][:3]*.5]*2, spin=1, deriv=3)
        self.assertAlmostEqual(lib.fp(e)   , -0.4720562542635522, 8)
        self.assertAlmostEqual(lib.fp(v[0].T[0]), -0.6294083390180697, 8)
        self.assertAlmostEqual(lib.fp(v[0].T[1]), -0.6294083390180697, 8)
        self.assertAlmostEqual(lib.fp(f[0].T[0]), -1.1414693830969338*2, 8)
        self.assertAlmostEqual(lib.fp(f[0].T[2]), -1.1414693830969338*2, 8)
        self.assertAlmostEqual(lib.fp(k[0].T[0]),  4.1402447248393921*4, 7)
        self.assertAlmostEqual(lib.fp(k[0].T[3]),  4.1402447248393921*4, 7)

    def test_lyp(self):
        e,v,f = dft.libxc.eval_xc(',LYP', rho, deriv=2)[:3]
        self.assertAlmostEqual(numpy.dot(rho[0],e), -62.114576182676615, 8)
        self.assertAlmostEqual(numpy.dot(rho[0],v[0]),-81.771670866308455, 8)
        self.assertAlmostEqual(numpy.dot(rho[0],v[1]), 27.485383255125743, 8)
        self.assertAlmostEqual(numpy.dot(rho[0],f[0]), 186.823806251777, 2)
        self.assertAlmostEqual(numpy.dot(rho[0],f[1]), -3391.2428894571085, 6)
        self.assertAlmostEqual(numpy.dot(rho[0],f[2]), 0, 8)
        self.assertAlmostEqual(abs(f[2]).sum(), 0, 3)

    #def test_tpss(self):
    #    #FIXME: raised numerical error
    #    rho_a = numpy.array([[3.67808547e-08,-2.02358682e-08, 2.16729780e-07, 2.27036045e-07,-1.47795869e-07,-1.45668997e-09]]).T
    #    e, v = dft.libxc.eval_xc('tpss,', rho_a, spin=0, deriv=1)[:2]
    #    rho_b = numpy.array([[4.53272893e-06, 4.18968775e-06,-2.83034672e-06, 2.61832978e-06, 5.63360737e-06, 8.97541777e-07]]).T
    #    e, v = dft.libxc.eval_xc('tpss,', (rho_a, rho_b), spin=1, deriv=1)[:2]

    def test_define_xc(self):
        def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, verbose=None):
            # A fictitious XC functional to demonstrate the usage
            rho0, dx, dy, dz = rho[:4]
            gamma = (dx**2 + dy**2 + dz**2)
            exc = .01 * rho0**2 + .02 * (gamma+.001)**.5
            vrho = .01 * 2 * rho0
            vgamma = .02 * .5 * (gamma+.001)**(-.5)
            vlapl = None
            vtau = None
            vxc = (vrho, vgamma, vlapl, vtau)
            fxc = None  # 2nd order functional derivative
            kxc = None  # 3rd order functional derivative
            return exc, vxc, fxc, kxc

        mf = dft.RKS(mol)
        ni = dft.libxc.define_xc(mf._numint, eval_xc, 'GGA', hyb=0.2)
        numpy.random.seed(1)
        rho = numpy.random.random((4,10))
        exc, vxc = ni.eval_xc(None, rho, 0, deriv=1)[:2]
        self.assertAlmostEqual(lib.fp(exc), 0.0012441814416833327, 9)
        self.assertAlmostEqual(lib.fp(vxc[0]), 0.0065565189784811129, 9)
        self.assertAlmostEqual(lib.fp(vxc[1]), 0.0049270110162854116, 9)

        mf = mf.define_xc_('0.5*B3LYP+0.5*B3LYP')
        exc0, vxc0 = mf._numint.eval_xc(None, rho, 0, deriv=1)[:2]
        exc1, vxc1 = dft.libxc.eval_xc('0.5*B3LYP+0.5*B3LYP', rho, 0, deriv=1)[:2]
        self.assertAlmostEqual(abs(exc0-exc1).max(), 0, 9)
        self.assertAlmostEqual(abs(vxc0[0]-vxc1[0]).max(), 0, 9)
        self.assertAlmostEqual(abs(vxc0[1]-vxc1[1]).max(), 0, 9)

        self.assertRaises(ValueError, dft.libxc.define_xc, mf._numint, 0.1)

        ni = dft.libxc.define_xc(mf._numint, 'PBE')
        exc1, vxc1 = dft.libxc.eval_xc('PBE', rho, 0, deriv=1)[:2]
        exc0, vxc0 = ni.eval_xc(None, rho, 0, deriv=1)[:2]
        self.assertAlmostEqual(abs(exc0-exc1).max(), 0, 9)
        self.assertAlmostEqual(abs(vxc0[0]-vxc1[0]).max(), 0, 9)
        self.assertAlmostEqual(abs(vxc0[1]-vxc1[1]).max(), 0, 9)

    def test_m05x(self):
        rho =(numpy.array([1., 1., 0., 0., 0., 0.165 ]).reshape(-1,1),
              numpy.array([.8, 1., 0., 0., 0., 0.1050]).reshape(-1,1))
        #test_ref = numpy.array([-1.57876583, -2.12127045,-2.11264351,-0.00315462,
        #                         0.00000000, -0.00444560, 3.45640232, 4.4349756])  # libxc-4.3.4
        test_ref = numpy.array([-1.57730394, -2.12127045,-2.11297165,-0.00315462,
                                 0.00000000, -0.00446935, 3.45640232, 4.42563831])  # libxc-5.1.2
        exc, vxc, fxc, kxc = dft.libxc.eval_xc('1.38888888889*m05,', rho, 1, deriv=1)
        self.assertAlmostEqual(float(exc)*1.8, test_ref[0], 5)
        self.assertAlmostEqual(abs(vxc[0]-test_ref[1:3]).max(), 0, 6)
        self.assertAlmostEqual(abs(vxc[1]-test_ref[3:6]).max(), 0, 6)
        self.assertAlmostEqual(abs(vxc[3]-test_ref[6:8]).max(), 0, 5)

        exc, vxc, fxc, kxc = dft.libxc.eval_xc('1.38888888889*m05,', rho[0], 0, deriv=1)
        self.assertAlmostEqual(float(exc), -0.5746231988116002, 5)
        self.assertAlmostEqual(float(vxc[0]), -0.8806121005703862, 6)
        self.assertAlmostEqual(float(vxc[1]), -0.0032300155406846756, 7)
        self.assertAlmostEqual(float(vxc[3]), 0.4474953100487698, 5)

    def test_camb3lyp(self):
        rho = numpy.array([1., 1., 0.1, 0.1]).reshape(-1,1)
        exc, vxc, fxc, kxc = dft.libxc.eval_xc('camb3lyp', rho, 0, deriv=1)
        self.assertAlmostEqual(float(exc), -0.5752559666317147, 7)
        self.assertAlmostEqual(float(vxc[0]), -0.7709812578936763, 7)
        self.assertAlmostEqual(float(vxc[1]), -0.0029862221286189846, 7)

        self.assertEqual(dft.libxc.rsh_coeff('camb3lyp'), [0.33, 0.65, -0.46])

        rho = numpy.array([1., 1., 0.1, 0.1]).reshape(-1,1)
        exc, vxc, fxc, kxc = dft.libxc.eval_xc('RSH(0.5,0.65,-0.46) + 0.46*ITYH + .35*B88,', rho, 0, deriv=1)
        self.assertAlmostEqual(float(exc), -0.48916154057161476, 9)
        self.assertAlmostEqual(float(vxc[0]), -0.6761177630311709, 9)
        self.assertAlmostEqual(float(vxc[1]), -0.002949151742087167, 9)

    def test_ityh(self):
        rho = numpy.array([1., 1., 0.1, 0.1]).reshape(-1,1)
        exc, vxc, fxc, kxc = dft.libxc.eval_xc('ityh,', rho, 0, deriv=1)
        self.assertAlmostEqual(float(exc), -0.6359945579326314, 7)
        self.assertAlmostEqual(float(vxc[0]), -0.8712041561251518, 7)
        self.assertAlmostEqual(float(vxc[1]), -0.003911167644579979, 7)
        self.assertEqual(dft.libxc.rsh_coeff('ityh,'), [0.2, 0.0, 0.0])

    def test_deriv_order(self):
        self.assertTrue(dft.libxc.test_deriv_order('lda', 3, raise_error=False))
        self.assertTrue(dft.libxc.test_deriv_order('m05', 2, raise_error=False))
        self.assertTrue(dft.libxc.test_deriv_order('camb3lyp', 3, True))
        self.assertTrue(dft.libxc.test_deriv_order('pbe0', 3, True))
        self.assertRaises(KeyError, dft.libxc.test_deriv_order, 'OL2', 3, True)

    def test_xc_type(self):
        self.assertEqual(dft.libxc.xc_type(416), 'GGA')
        self.assertEqual(dft.libxc.xc_type('hf'), 'HF')
        self.assertEqual(dft.libxc.xc_type(',vwn'), 'LDA')
        self.assertEqual(dft.libxc.xc_type('lda+b3lyp'), 'GGA')
        self.assertEqual(dft.libxc.xc_type('wb97m_v'), 'MGGA')
        self.assertEqual(dft.libxc.xc_type('bp86'), 'GGA')

    def test_m06(self):
        rho = numpy.array([0.11939647, -0.18865577, -0.11633254, 0.01779666, -0.55475521, 0.07092032,
                           0.11850155, -0.19297934, -0.11581427, 0.01373251, -0.57534216, 0.06596468,])
        rho = rho.reshape(2,6,1)
        exc, vxc, fxc, kxc = dft.libxc.eval_xc(',m06', rho[0], 0, deriv=3)
        exc_ref = numpy.array([-0.02809518])
        vxc_ref = numpy.array([-0.0447466, -0.002186, 0.0154814])
        fxc_ref = numpy.array([-0.66060719, 0.35848677, -1.37754938,
                               -1.63288299, 0.21999345, 1.20222259])
        kxc_ref = numpy.array([-8.91475141e+00, 2.57071018e+01, -4.37413046e+01, 9.92064877e+01,
                               -1.47859948e+01, 1.51506522e+01, -4.83492461e+00, -1.75713898e+01,
                               -2.82029718e+01, 6.66847006e+01])
        self.assertAlmostEqual(exc[0] - exc_ref[0], 0, 7)
        self.assertAlmostEqual(abs(numpy.hstack([vxc[i] for i in [0,1,3]])-vxc_ref).max(), 0, 7)
        self.assertAlmostEqual(abs(numpy.hstack([fxc[i] for i in [0,1,2,4,6,9]])-fxc_ref).max(), 0, 7)
        self.assertAlmostEqual(abs(numpy.hstack([kxc[i] for i in [0,1,2,3,5,7,10,12,15,19]])-kxc_ref).max(), 0, 7)

        exc, vxc, fxc, kxc = dft.libxc.eval_xc(',m06', rho, 1, deriv=3)
        exc_ref = numpy.array([-0.03730162])
        vxc_ref = numpy.array([-0.06168318,-0.0610502,-0.00957516,0,-0.01988745, 0.03161512, 0.04391615])
        fxc_ref = numpy.array([ 0.62207817,-1.02032218, 0.59318092,
                               -0.00267104, 0, 0.2768053, 0.26785479, 0, 0.02509479,
                               -1.34211352, 0,-0.0399852, 0, 0,-1.28860549,
                               -1.69697623,-0.03422618,-2.1835459,
                               -0.22523468, 0.26116507, 0.26545614,-0.23488143,
                               1.52817397,-0.13694284, 0, 0,-0.13797703, 1.73484678])
        kxc_ref = numpy.array([ 1.49056834e+01, 4.93235407e+00, 5.13177936e+00, 6.83179846e+00,
                               -1.97455841e+01, 0,-8.12321586e-01,-1.51226035e+00, 0,-1.66826458e+00,-8.99591711e-01, 0, -9.19638526e+00,
                               1.90010442e+01, 0, 1.34796896e-01, 0, 0,-7.73620959e-01,-6.10018848e-01, 0, 2.77414513e-01, 0, 0, 6.16594005e+00,
                               2.67162133e+01, 0,-9.09592410e-01, 0, 0,-9.28025786e-01, 0, 0, 0, 3.76462242e+01,
                               2.74381367e-01,-2.08192677e+00, 5.25100873e-01, 6.10851890e-01,-2.02662181e+00,-3.07707014e+00,
                               -4.45125589e-01, 6.81329714e-01, 0, 0,-4.69145558e-01, -2.42540301e+00,-2.45259140e+00,-5.57504449e-01, 0, 0, 6.02633158e-01, 2.58871134e+00,
                               2.87554064e+00, 2.85563442e-01,-1.94628304e-01, -1.88482952e-01, 2.90513272e-01, 3.07132889e+00,
                               -1.08483575e+00, 1.57821239e+00, 0, 0, 1.59013101e+00, 1.61019571e+00, 0, 0, 0, 0, 1.62235586e+00,-3.60275312e+00,
                               -3.90972925e+01,-1.10619931e-02,-1.12015763e-02, 0, 0, 0,-1.11455329e-02,-1.12861703e-02, -4.81639494e+01,
                               7.24911656e+01, 4.35367189e-02, 4.40860771e-02, 1.00010975e+02])
        self.assertAlmostEqual(exc[0] - exc_ref[0], 0, 7)
        self.assertAlmostEqual(abs(numpy.hstack([vxc[i] for i in [0,1,3]])-vxc_ref).max(), 0, 7)
        self.assertAlmostEqual(abs(numpy.hstack([fxc[i] for i in [0,1,2,4,6,9]])-fxc_ref).max(), 0, 7)
        self.assertAlmostEqual(abs(numpy.hstack([kxc[i] for i in [0,1,2,3,5,7,10,12,15,19]])-kxc_ref).max(), 0, 6)


if __name__ == "__main__":
    print("Test libxc")
    unittest.main()


# pyscf/dft/test/test_numint.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest
import numpy
from pyscf import gto
from pyscf import dft
from pyscf import lib
from pyscf.dft import numint

def setUpModule():
    numint.SWITCH_SIZE = 0
    global mol, mf, h4, mf_h4, mol1, h2o, nao, he2, mf_he2
    mol = gto.Mole()
    mol.verbose = 0
    mol.atom = [('h', (0,0,i*3)) for i in range(12)]
    mol.basis = 'ccpvtz'
    mol.build()
    mf = dft.RKS(mol)
    mf.grids.atom_grid = {"H": (50, 110)}
    mf.prune = None
    mf.grids.build(with_non0tab=False)
    nao = mol.nao_nr()
    ao_loc = mol.ao_loc_nr()

    h4 = gto.Mole()
    h4.verbose = 0
    h4.atom = 'H 0 0 0; H 0 0 9; H 0 9 0; H 0 9 9'
    h4.basis = 'ccpvtz'
    h4.build()
    mf_h4 = dft.RKS(h4)
    mf_h4.grids.atom_grid = {"H": (50, 110)}
    mf_h4.grids.build(with_non0tab=True)

    mol1 = gto.Mole()
    mol1.verbose = 0
    mol1.atom = [('h', (0,0,i*3)) for i in range(4)]
    mol1.basis = 'ccpvtz'
    mol1.build()

    h2o = gto.Mole()
    h2o.verbose = 5
    h2o.output = '/dev/null'
    h2o.atom = [
        ["O" , (0. , 0.     , 0.)],
        [1   , (0. , -0.757 , 0.587)],
        [1   , (0. , 0.757  , 0.587)] ]

    h2o.basis = {"H": '6-31g', "O": '6-31g',}
    h2o.build()

    he2 = gto.Mole()
    he2.verbose = 0
    he2.atom = 'He 0.0 0.0 0.0; He 0.0 0.0 20.0'
    he2.basis = 'aug-pc-4'
    he2.build()
    mf_he2 = dft.RKS(he2)
    mf_he2.grids.level = 0
    mf_he2.grids.build(with_non0tab=True)

def tearDownModule():
    numint.SWITCH_SIZE = 800
    global mol, mf, h4, mf_h4, mol1, h2o
    h2o.stdout.close()
    del mol, mf, h4, mf_h4, mol1, h2o

class KnownValues(unittest.TestCase):
    def test_make_mask(self):
        non0 = dft.numint.make_mask(mol, mf.grids.coords)
        self.assertEqual(non0.sum(), 12364)
        self.assertAlmostEqual(lib.fp(non0), 8.087267296357203, 9)

    def test_dot_ao_dm(self):
        dm = mf_h4.get_init_guess(key='minao')
        ao_loc = h4.ao_loc_nr()
        ao = mf_h4._numint.eval_ao(h4, mf_h4.grids.coords).copy() + 0j
        nao = ao.shape[1]
        v1 = dft.numint._dot_ao_dm(h4, ao, dm, mf_h4.grids.non0tab, (0,h4.nbas), ao_loc)
        v2 = dft.numint._dot_ao_dm(h4, ao, dm, None, None, None)
        self.assertAlmostEqual(abs(v1-v2).max(), 0, 9)

        dm = mf_he2.get_init_guess(key='minao')
        ao_loc = he2.ao_loc_nr()
        ao = mf_he2._numint.eval_ao(he2, mf_he2.grids.coords).copy()
        v1 = dft.numint._dot_ao_dm(he2, ao, dm, mf_he2.grids.non0tab, (0,he2.nbas), ao_loc)
        v2 = dft.numint._dot_ao_dm(he2, ao, dm, None, None, None)
        self.assertAlmostEqual(abs(v1-v2).max(), 0, 9)

    def test_dot_ao_dm_high_cost(self):
        non0tab = mf._numint.make_mask(mol, mf.grids.coords)
        ao = dft.numint.eval_ao(mol, mf.grids.coords)
        numpy.random.seed(1)
        nao = ao.shape[1]
        ao_loc = mol.ao_loc_nr()
        dm = numpy.random.random((nao,nao))
        dm = dm + dm.T
        res0 = lib.dot(ao, dm)
        res1 = dft.numint._dot_ao_dm(mol, ao, dm, non0tab,
                                     shls_slice=(0,mol.nbas), ao_loc=ao_loc)
        self.assertTrue(numpy.allclose(res0, res1))

    def test_dot_ao_ao(self):
        dm = mf_h4.get_init_guess(key='minao')
        ao_loc = h4.ao_loc_nr()
        ao = mf_h4._numint.eval_ao(h4, mf_h4.grids.coords).copy() + 0j
        nao = h4.nao_nr()
        v1 = dft.numint._dot_ao_ao(h4, ao, ao, mf_h4.grids.non0tab, (0,h4.nbas), ao_loc)
        v2 = dft.numint._dot_ao_ao(h4, ao, ao, None, None, None)
        self.assertAlmostEqual(abs(v1-v2).max(), 0, 9)

        dm = mf_he2.get_init_guess(key='minao')
        ao_loc = he2.ao_loc_nr()
        ao = mf_he2._numint.eval_ao(he2, mf_he2.grids.coords).copy()
        nao = he2.nao_nr()
        v1 = dft.numint._dot_ao_ao(he2, ao, ao, mf_he2.grids.non0tab, (0,he2.nbas), ao_loc)
        v2 = dft.numint._dot_ao_ao(he2, ao, ao, None, None, None)
        self.assertAlmostEqual(abs(v1-v2).max(), 0, 9)

    def test_dot_ao_ao_high_cost(self):
        non0tab = mf.grids.make_mask(mol, mf.grids.coords)
        ao = dft.numint.eval_ao(mol, mf.grids.coords, deriv=1)
        nao = ao.shape[1]
        ao_loc = mol.ao_loc_nr()
        res0 = lib.dot(ao[0].T, ao[1])
        res1 = dft.numint._dot_ao_ao(mol, ao[0], ao[1], non0tab,
                                     shls_slice=(0,mol.nbas), ao_loc=ao_loc)
        self.assertTrue(numpy.allclose(res0, res1))

    def test_eval_rho(self):
        numpy.random.seed(10)
        ngrids = 500
        coords = numpy.random.random((ngrids,3))*20
        dm = numpy.random.random((nao,nao))
        dm = dm + dm.T
        ao = dft.numint.eval_ao(mol, coords, deriv=2)

        e, mo_coeff = numpy.linalg.eigh(dm)
        mo_occ = numpy.ones(nao)
        mo_occ[-2:] = -1
        dm = numpy.einsum('pi,i,qi->pq', mo_coeff, mo_occ, mo_coeff)

        rho0 = numpy.zeros((6,ngrids))
        rho0[0] = numpy.einsum('pi,ij,pj->p', ao[0], dm, ao[0].conj())
        rho0[1] = numpy.einsum('pi,ij,pj->p', ao[1], dm, ao[0].conj()) + numpy.einsum('pi,ij,pj->p', ao[0], dm, ao[1].conj())
        rho0[2] = numpy.einsum('pi,ij,pj->p', ao[2], dm, ao[0].conj()) + numpy.einsum('pi,ij,pj->p', ao[0], dm, ao[2].conj())
        rho0[3] = numpy.einsum('pi,ij,pj->p', ao[3], dm, ao[0].conj()) + numpy.einsum('pi,ij,pj->p', ao[0], dm, ao[3].conj())
        rho0[4]+= numpy.einsum('pi,ij,pj->p', ao[0], dm, ao[4].conj()) + numpy.einsum('pi,ij,pj->p', ao[4], dm, ao[0].conj())
        rho0[4]+= numpy.einsum('pi,ij,pj->p', ao[0], dm, ao[7].conj()) + numpy.einsum('pi,ij,pj->p', ao[7], dm, ao[0].conj())
        rho0[4]+= numpy.einsum('pi,ij,pj->p', ao[0], dm, ao[9].conj()) + numpy.einsum('pi,ij,pj->p', ao[9], dm, ao[0].conj())
        rho0[5]+= numpy.einsum('pi,ij,pj->p', ao[1], dm, ao[1].conj())
        rho0[5]+= numpy.einsum('pi,ij,pj->p', ao[2], dm, ao[2].conj())
        rho0[5]+= numpy.einsum('pi,ij,pj->p', ao[3], dm, ao[3].conj())
        rho0[4]+= rho0[5]*2
        rho0[5] *= .5

        ni = dft.numint.NumInt()
        rho1 = ni.eval_rho (mol, ao, dm, xctype='MGGA')
        rho2 = ni.eval_rho2(mol, ao, mo_coeff, mo_occ, xctype='MGGA')
        self.assertTrue(numpy.allclose(rho0, rho1))
        self.assertTrue(numpy.allclose(rho0, rho2))

    def test_eval_mat(self):
        numpy.random.seed(10)
        ngrids = 500
        coords = numpy.random.random((ngrids,3))*20
        rho = numpy.random.random((6,ngrids))
        vxc = numpy.random.random((4,ngrids))
        weight = numpy.random.random(ngrids)
        ao = dft.numint.eval_ao(mol, coords, deriv=2)

        mat0 = numpy.einsum('pi,p,pj->ij', ao[0].conj(), weight*vxc[0], ao[0])
        mat1 = dft.numint.eval_mat(mol, ao[0], weight, rho, vxc[0], xctype='LDA')
        self.assertTrue(numpy.allclose(mat0, mat1))
        # UKS
        mat2 = dft.numint.eval_mat(mol, ao[0], weight, rho, [vxc[0]]*2, xctype='LDA', spin=1)
        self.assertTrue(numpy.allclose(mat0, mat2))

        vrho, vsigma = vxc[:2]
        wv = weight * vsigma * 2
        mat0  = numpy.einsum('pi,p,pj->ij', ao[0].conj(), weight*vrho, ao[0])
        mat0 += numpy.einsum('pi,p,pj->ij', ao[0].conj(), rho[1]*wv, ao[1]) + numpy.einsum('pi,p,pj->ij', ao[1].conj(), rho[1]*wv, ao[0])
        mat0 += numpy.einsum('pi,p,pj->ij', ao[0].conj(), rho[2]*wv, ao[2]) + numpy.einsum('pi,p,pj->ij', ao[2].conj(), rho[2]*wv, ao[0])
        mat0 += numpy.einsum('pi,p,pj->ij', ao[0].conj(), rho[3]*wv, ao[3]) + numpy.einsum('pi,p,pj->ij', ao[3].conj(), rho[3]*wv, ao[0])
        mat1 = dft.numint.eval_mat(mol, ao, weight, rho, vxc[:4], xctype='GGA')
        self.assertTrue(numpy.allclose(mat0, mat1))
        # UKS
        ngrids = weight.size
        vxc_1 = [vxc[0], numpy.vstack((vxc[1], numpy.zeros(ngrids))).T]
        mat2 = dft.numint.eval_mat(mol, ao, weight, [rho[:4]]*2, vxc_1, xctype='GGA', spin=1)
        self.assertTrue(numpy.allclose(mat0, mat2))

        vrho, vsigma, _, vtau = vxc
        vxc = (vrho, vsigma, None, vtau)
        wv = weight * vtau * .5
        mat2  = numpy.einsum('pi,p,pj->ij', ao[1].conj(), wv, ao[1])
        mat2 += numpy.einsum('pi,p,pj->ij', ao[2].conj(), wv, ao[2])
        mat2 += numpy.einsum('pi,p,pj->ij', ao[3].conj(), wv, ao[3])
        mat0 += mat2
        mat1 = dft.numint.eval_mat(mol, ao, weight, rho, vxc, xctype='MGGA')
        self.assertTrue(numpy.allclose(mat0, mat1))
        # UKS
        ngrids = weight.size
        vxc_1 = [vxc[0],
                 numpy.vstack((vxc[1], numpy.zeros(ngrids))).T,
                 numpy.zeros((ngrids,2)),
                 numpy.vstack((vxc[3], numpy.zeros(ngrids))).T]
        mat2 = dft.numint.eval_mat(mol, ao, weight, [rho]*2, vxc_1, xctype='MGGA', spin=1)
        self.assertTrue(numpy.allclose(mat0, mat2))

    def test_rks_vxc(self):
        numpy.random.seed(10)
        nao = mol.nao_nr()
        dms = numpy.random.random((2,nao,nao))
        v = mf._numint.nr_vxc(mol, mf.grids, 'B88,', dms, spin=0, hermi=0)[2]
        self.assertAlmostEqual(lib.fp(v), -0.70124686853021512, 8)

        v = mf._numint.nr_vxc(mol, mf.grids, 'HF', dms, spin=0, hermi=0)[2]
        self.assertAlmostEqual(abs(v).max(), 0, 9)
        v = mf._numint.nr_vxc(mol, mf.grids, '', dms, spin=0, hermi=0)[2]
        self.assertAlmostEqual(abs(v).max(), 0, 9)

    def test_uks_vxc(self):
        numpy.random.seed(10)
        nao = h2o.nao_nr()
        dms = numpy.random.random((2,nao,nao))
        grids = dft.gen_grid.Grids(h2o)
        v = mf._numint.nr_vxc(h2o, grids, 'B88,', dms, spin=1)[2]
        self.assertAlmostEqual(lib.fp(v), -7.7508525240447348, 8)

        v = mf._numint.nr_vxc(h2o, grids, 'HF', dms, spin=1)[2]
        self.assertAlmostEqual(abs(v).max(), 0, 9)
        v = mf._numint.nr_vxc(h2o, grids, '', dms, spin=1)[2]
        self.assertAlmostEqual(abs(v).max(), 0, 9)

    def test_uks_vxc_high_cost(self):
        numpy.random.seed(10)
        nao = mol.nao_nr()
        dms = numpy.random.random((2,nao,nao))
        v = mf._numint.nr_vxc(mol, mf.grids, 'B88,', dms, spin=1)[2]
        self.assertAlmostEqual(lib.fp(v), -0.73803886056633594, 8)

        v = mf._numint.nr_vxc(mol, mf.grids, 'HF', dms, spin=1)[2]
        self.assertAlmostEqual(abs(v).max(), 0, 9)
        v = mf._numint.nr_vxc(mol, mf.grids, '', dms, spin=1)[2]
        self.assertAlmostEqual(abs(v).max(), 0, 9)

    def test_rks_fxc(self):
        numpy.random.seed(10)
        nao = mol1.nao_nr()
        dm0 = numpy.random.random((nao,nao))
        _, mo_coeff = numpy.linalg.eigh(dm0)
        mo_occ = numpy.ones(nao)
        mo_occ[-2:] = -1
        dm0 = numpy.einsum('pi,i,qi->pq', mo_coeff, mo_occ, mo_coeff)
        dms = numpy.random.random((2,nao,nao))
        ni = dft.numint.NumInt()
        v = ni.nr_fxc(mol1, mf.grids, 'B88,', dm0, dms, spin=0, hermi=0)
        self.assertAlmostEqual(lib.fp(v), -7.5671368618070343, 8)

        # test cache_kernel
        rvf = ni.cache_xc_kernel(mol1, mf.grids, 'B88,', mo_coeff, mo_occ, spin=0)
        v1 = dft.numint.nr_fxc(mol1, mf.grids, 'B88,', dm0, dms, spin=0, hermi=0,
                               rho0=rvf[0], vxc=rvf[1], fxc=rvf[2])
        self.assertAlmostEqual(abs(v-v1).max(), 0, 8)

        v = ni.nr_fxc(mol1, mf.grids, 'LDA,', dm0, dms[0], spin=0, hermi=0)
        self.assertAlmostEqual(lib.fp(v), -3.0019207112626876, 8)
        # test cache_kernel
        rvf = ni.cache_xc_kernel(mol1, mf.grids, 'LDA,', mo_coeff, mo_occ, spin=0)
        v1 = dft.numint.nr_fxc(mol1, mf.grids, 'LDA,', dm0, dms[0], spin=0, hermi=0,
                               rho0=rvf[0], vxc=rvf[1], fxc=rvf[2])
        self.assertAlmostEqual(abs(v-v1).max(), 0, 8)

        v = ni.nr_fxc(mol1, mf.grids, 'HF', dm0, dms, spin=0, hermi=0)
        self.assertAlmostEqual(abs(v).max(), 0, 9)
        v = ni.nr_fxc(mol1, mf.grids, '', dm0, dms, spin=0, hermi=0)
        self.assertAlmostEqual(abs(v).max(), 0, 9)

        v = dft.numint.nr_fxc(mol1, mf.grids, 'm06l,', dm0, dms)
        self.assertAlmostEqual(lib.fp(v), -11.007699914447517, 8)
        rvf = ni.cache_xc_kernel(mol1, mf.grids, 'm06l,', mo_coeff, mo_occ, spin=0)
        v1 = dft.numint.nr_fxc(mol1, mf.grids, 'm06l,', dm0, dms,
                               rho0=rvf[0], vxc=rvf[1], fxc=rvf[2])
        self.assertAlmostEqual(abs(v-v1).max(), 0, 8)

    def test_rks_fxc_st(self):
        numpy.random.seed(10)
        nao = mol1.nao_nr()
        dm0 = numpy.random.random((nao,nao))
        _, mo_coeff = numpy.linalg.eigh(dm0)
        mo_occ = numpy.ones(nao)
        mo_occ[-2:] = -1
        dm0 = numpy.einsum('pi,i,qi->pq', mo_coeff, mo_occ, mo_coeff)
        dms = numpy.random.random((2,nao,nao))
        ni = dft.numint.NumInt()

        rvf = ni.cache_xc_kernel(mol1, mf.grids, 'B88,', [mo_coeff,mo_coeff],
                                 [mo_occ*.5]*2, spin=1)
        v = dft.numint.nr_rks_fxc_st(ni, mol1, mf.grids, 'B88,', dm0, dms, singlet=True)
        self.assertAlmostEqual(lib.fp(v), -7.5671368618070343*2, 8)
        v1 = dft.numint.nr_rks_fxc_st(ni, mol1, mf.grids, 'B88,', dm0, dms, singlet=True,
                                      rho0=rvf[0], vxc=rvf[1], fxc=rvf[2])
        self.assertAlmostEqual(abs(v-v1).max(), 0, 8)

        v = dft.numint.nr_rks_fxc_st(ni, mol1, mf.grids, 'B88,', dm0, dms, singlet=False)
        self.assertAlmostEqual(lib.fp(v), -7.5671368618070343*2, 8)
        v1 = dft.numint.nr_rks_fxc_st(ni, mol1, mf.grids, 'B88,', dm0, dms, singlet=False,
                                      rho0=rvf[0], vxc=rvf[1], fxc=rvf[2])
        self.assertAlmostEqual(abs(v-v1).max(), 0, 8)

        rvf = ni.cache_xc_kernel(mol1, mf.grids, 'LDA,', [mo_coeff,mo_coeff],
                                 [mo_occ*.5]*2, spin=1)
        v = dft.numint.nr_rks_fxc_st(ni, mol1, mf.grids, 'LDA,', dm0, dms[0], singlet=True)
        self.assertAlmostEqual(lib.fp(v), -3.0019207112626876*2, 8)
        v1 = dft.numint.nr_rks_fxc_st(ni, mol1, mf.grids, 'LDA,', dm0, dms[0], singlet=True,
                                      rho0=rvf[0], vxc=rvf[1], fxc=rvf[2])
        self.assertAlmostEqual(abs(v-v1).max(), 0, 8)

        v = dft.numint.nr_rks_fxc_st(ni, mol1, mf.grids, 'LDA,', dm0, dms[0], singlet=False)
        self.assertAlmostEqual(lib.fp(v), -3.0019207112626876*2, 8)
        v1 = dft.numint.nr_rks_fxc_st(ni, mol1, mf.grids, 'LDA,', dm0, dms[0], singlet=False,
                                      rho0=rvf[0], vxc=rvf[1], fxc=rvf[2])
        self.assertAlmostEqual(abs(v-v1).max(), 0, 8)

        v = dft.numint.nr_rks_fxc_st(ni, mol1, mf.grids, 'm06l,', dm0, dms, singlet=True)
        self.assertAlmostEqual(lib.fp(v), -11.007699914447517*2, 8)
        rvf = ni.cache_xc_kernel(mol1, mf.grids, 'm06l,', [mo_coeff,mo_coeff],
                                 [mo_occ*.5]*2, spin=1)
        v1 = dft.numint.nr_rks_fxc_st(ni, mol1, mf.grids, 'm06l,', dm0, dms, singlet=True,
                                      rho0=rvf[0], vxc=rvf[1], fxc=rvf[2])
        self.assertAlmostEqual(abs(v-v1).max(), 0, 8)

        v = dft.numint.nr_rks_fxc_st(ni, mol1, mf.grids, 'm06l,', dm0, dms, singlet=False)
        self.assertAlmostEqual(lib.fp(v), -11.007699914447517*2, 8)
        rvf = ni.cache_xc_kernel(mol1, mf.grids, 'm06l,', [mo_coeff,mo_coeff],
                                 [mo_occ*.5]*2, spin=1)
        v1 = dft.numint.nr_rks_fxc_st(ni, mol1, mf.grids, 'm06l,', dm0, dms, singlet=False,
                                      rho0=rvf[0], vxc=rvf[1], fxc=rvf[2])
        self.assertAlmostEqual(abs(v-v1).max(), 0, 8)

    def test_uks_fxc(self):
        numpy.random.seed(10)
        nao = mol1.nao_nr()
        dm0 = numpy.random.random((2,nao,nao))
        e, mo_coeff = numpy.linalg.eigh(dm0)
        mo_occ = numpy.ones((2,nao))
        mo_occ[:,-2:] = -1
        dm0 = numpy.einsum('xpi,xi,xqi->xpq', mo_coeff, mo_occ, mo_coeff)
        dms = numpy.random.random((2,nao,nao))
        ni = dft.numint.NumInt()
        v = ni.nr_fxc(mol1, mf.grids, 'B88,', dm0, dms, spin=1)
        self.assertAlmostEqual(lib.fp(v), -10.316443204083185, 8)

        # test cache_kernel
        rvf = ni.cache_xc_kernel(mol1, mf.grids, 'B88,', mo_coeff, mo_occ, spin=1)
        v1 = dft.numint.nr_fxc(mol1, mf.grids, 'B88,', dm0, dms, hermi=0, spin=1,
                               rho0=rvf[0], vxc=rvf[1], fxc=rvf[2])
        self.assertAlmostEqual(abs(v-v1).max(), 0, 8)

        v = ni.nr_fxc(mol1, mf.grids, 'LDA,', dm0, dms[0], spin=1)
        self.assertAlmostEqual(lib.fp(v), -5.6474405864697967, 8)
        # test cache_kernel
        rvf = ni.cache_xc_kernel(mol1, mf.grids, 'LDA,', mo_coeff, mo_occ, spin=1)
        v1 = dft.numint.nr_fxc(mol1, mf.grids, 'LDA,', dm0, dms[0], hermi=0, spin=1,
                               rho0=rvf[0], vxc=rvf[1], fxc=rvf[2])
        self.assertAlmostEqual(abs(v-v1).max(), 0, 8)

        v = ni.nr_fxc(mol1, mf.grids, 'm06l', dm0, dms[0], spin=1)
        self.assertAlmostEqual(lib.fp(v), -7.08393530172142, 8)
        # test cache_kernel
        rvf = ni.cache_xc_kernel(mol1, mf.grids, 'm06l', mo_coeff, mo_occ, spin=1)
        v1 = dft.numint.nr_fxc(mol1, mf.grids, 'm06l', dm0, dms[0], hermi=0, spin=1,
                               rho0=rvf[0], vxc=rvf[1], fxc=rvf[2])
        self.assertAlmostEqual(abs(v-v1).max(), 0, 8)

    def test_vv10nlc(self):
        numpy.random.seed(10)
        rho = numpy.random.random((4,20))
        coords = (numpy.random.random((20,3))-.5)*3
        vvrho = numpy.random.random((4,60))
        vvweight = numpy.random.random(60)
        vvcoords = (numpy.random.random((60,3))-.5)*3
        nlc_pars = .8, .3
        v = dft.numint._vv10nlc(rho, coords, vvrho, vvweight, vvcoords, nlc_pars)
        self.assertAlmostEqual(lib.fp(v[0]), 0.15894647203764295, 9)
        self.assertAlmostEqual(lib.fp(v[1]), 0.20500922537924576, 9)

    def test_nr_uks_vxc_vv10(self):
        method = dft.UKS(h2o)
        dm = method.get_init_guess()
        dm = (dm[0], dm[0])
        grids = dft.gen_grid.Grids(h2o)
        grids.atom_grid = {'H': (20, 50), 'O': (20,50)}
        v = dft.numint.nr_vxc(h2o, grids, 'wB97M_V__vv10', dm, spin=1, hermi=0)[2]
        self.assertAlmostEqual(lib.fp(v), 0.02293399033256055, 8)

    def test_uks_gga_wv1(self):
        numpy.random.seed(1)
        rho0 = [numpy.random.random((4,5))]*2
        rho1 = [numpy.random.random((4,5))]*2
        weight = numpy.ones(5)

        exc, vxc, fxc, kxc = dft.libxc.eval_xc('b88,', rho0, 1, 0, 3)
        wva, wvb = dft.numint._uks_gga_wv1(rho0, rho1, vxc, fxc, weight)
        exc, vxc, fxc, kxc = dft.libxc.eval_xc('b88,', rho0[0]+rho0[1], 0, 0, 3)
        wv = dft.numint._rks_gga_wv1(rho0[0]+rho0[1], rho1[0]+rho1[1], vxc, fxc, weight)
        self.assertAlmostEqual(abs(1 - wv/wva).max(), 0, 12)
        self.assertAlmostEqual(abs(1 - wv/wvb).max(), 0, 12)

    def test_uks_gga_wv2(self):
        numpy.random.seed(1)
        rho0 = [numpy.random.random((4,5))]*2
        rho1 = [numpy.random.random((4,5))]*2
        weight = numpy.ones(5)

        exc, vxc, fxc, kxc = dft.libxc.eval_xc('b88,', rho0, 1, 0, 3)
        wva, wvb = dft.numint._uks_gga_wv2(rho0, rho1, fxc, kxc, weight)
        exc, vxc, fxc, kxc = dft.libxc.eval_xc('b88,', rho0[0]+rho0[1], 0, 0, 3)
        wv = dft.numint._rks_gga_wv2(rho0[0]+rho0[1], rho1[0]+rho1[1], fxc, kxc, weight)
        self.assertAlmostEqual(abs(1 - wv/wva).max(), 0, 12)
        self.assertAlmostEqual(abs(1 - wv/wvb).max(), 0, 12)

    def test_complex_dm(self):
        mf = dft.RKS(h2o)
        mf.xc = 'b3lyp'
        nao = h2o.nao
        numpy.random.seed(1)
        dm = (numpy.random.random((nao,nao)) +
              numpy.random.random((nao,nao))*1j)
        dm = dm + dm.conj().T
        v = mf.get_veff(h2o, dm)
        self.assertAlmostEqual(lib.fp(v), 30.543789621782576-0.23207622637751305j, 9)

    def test_rsh_omega(self):
        rho0 = numpy.array([1., 1., 0.1, 0.1]).reshape(-1,1)
        ni = dft.numint.NumInt()
        ni.omega = 0.4
        omega = 0.2
        exc, vxc, fxc, kxc = ni.eval_xc('ITYH,', rho0, 0, 0, 1, omega)
        self.assertAlmostEqual(float(exc), -0.6359945579326314, 7)
        self.assertAlmostEqual(float(vxc[0]), -0.8712041561251518, 7)
        self.assertAlmostEqual(float(vxc[1]), -0.003911167644579979, 7)

        exc, vxc, fxc, kxc = ni.eval_xc('ITYH,', rho0, 0, 0, 1)
        self.assertAlmostEqual(float(exc), -0.5406095865415561, 7)
        self.assertAlmostEqual(float(vxc[0]), -0.772123720263471, 7)
        self.assertAlmostEqual(float(vxc[1]), -0.00301639097170439, 7)

    def test_cache_xc_kernel(self):
        mf = dft.RKS(h2o)
        mf.grids.atom_grid = {"H": (50, 194), "O": (50, 194),}
        mf.run()

        mf.xc = 'WB97XD'
        mf.omega = 0.9
        rho, vxc, fxc = mf._numint.cache_xc_kernel(mf.mol, mf.grids, mf.xc, mf.mo_coeff, mf.mo_occ)
        self.assertAlmostEqual(rho[0].dot(mf.grids.weights), 10, 4)
        self.assertAlmostEqual(numpy.einsum('g,g,ig->', mf.grids.weights, rho[0], rho   ), 81.0427569, 4)
        self.assertAlmostEqual(numpy.einsum('g,g, g->', mf.grids.weights, rho[0], vxc[0]), -5.9013868, 4)
        self.assertAlmostEqual(numpy.einsum('g,g, g->', mf.grids.weights, rho[0], vxc[1]),  0.3874269, 4)
        self.assertAlmostEqual(numpy.einsum('g,g, g->', mf.grids.weights, rho[0], fxc[0]), -15.939767, 1)
        #self.assertAlmostEqual(numpy.einsum('g,g, g->', mf.grids.weights, rho[0], fxc[1]), 15889.0167, 2)

        #mf.xc = 'camb3lyp'
        #mf.omega = 0.9
        #rho1, vxc1, fxc1 = mf._numint.cache_xc_kernel(mf.mol, mf.grids, mf.xc, mf.mo_coeff, mf.mo_occ)

        #mf.xc = 'camb3lyp'
        #mf._numint.libxc = dft.xcfun
        #mf.omega = 0.9
        #rho2, vxc2, fxc2 = mf._numint.cache_xc_kernel(mf.mol, mf.grids, mf.xc, mf.mo_coeff, mf.mo_occ)

        #self.assertAlmostEqual(abs(rho1 - rho2).max(), 0, 4)
        #self.assertAlmostEqual(abs(vxc1[0] - vxc2[0]), 0, 4)
        #self.assertAlmostEqual(abs(vxc1[1] - vxc2[1]), 0, 4)
        #self.assertAlmostEqual(abs(fxc1[0] - fxc2[0]), 0, 4)
        #self.assertAlmostEqual(abs(fxc1[1] - fxc2[1]), 0, 0)
        #self.assertAlmostEqual(abs(fxc1[2] - fxc2[2]), 0, 0)


if __name__ == "__main__":
    print("Test numint")
    unittest.main()


# pyscf/dft/test/test_r_numint.py

#!/usr/bin/env python

import unittest
import numpy
from pyscf import gto
from pyscf import dft
from pyscf import lib
from pyscf.dft import r_numint

def setUpModule():
    global mol
    mol = gto.Mole()
    mol.verbose = 0
    mol.atom = [
        ["O" , (0. , 0.     , 0.)],
        [1   , (0. , -0.757 , 0.587)],
        [1   , (0. , 0.757  , 0.587)] ]
    mol.basis = '6-31g'
    mol.build()

def tearDownModule():
    global mol
    del mol

class KnownValues(unittest.TestCase):
    def test_eval_rho(self):
        n2c = mol.nao_2c()
        numpy.random.seed(10)
        ngrids = 100
        coords = numpy.random.random((ngrids,3))*20
        coords = coords[70:75]
        dm = numpy.random.random((n2c,n2c))
        dm = dm + dm.T.conj()
        aoLa, aoLb, aoSa, aoSb = r_numint.eval_ao(mol, coords, deriv=1)

        rho0a = numpy.einsum('pi,ij,pj->p', aoLa[0], dm, aoLa[0].conj())
        rho0b = numpy.einsum('pi,ij,pj->p', aoLb[0], dm, aoLb[0].conj())
        rho0 = rho0a + rho0b

        aoL = numpy.array([aoLa[0],aoLb[0]])
        m0 = numpy.einsum('api,ji,bpj,xab->xp', aoL.conj(), dm, aoL, lib.PauliMatrices)

        ni = r_numint.RNumInt()
        rho1, m1 = ni.eval_rho(mol, (aoLa[0], aoLb[0]), dm, xctype='LDA')
        self.assertAlmostEqual(abs(rho1.imag).max(), 0, 9)
        self.assertAlmostEqual(abs(rho0-rho1).max(), 0, 9)
        self.assertTrue(numpy.allclose(m0, m1))

#    def test_eval_mat(self):
#        numpy.random.seed(10)
#        ngrids = 100
#        coords = numpy.random.random((ngrids,3))*20
#        rho = numpy.random.random((ngrids))
#        m = numpy.random.random((3,ngrids)) * .05
#        vxc = [numpy.random.random((2,ngrids)).T, None, None, None]
#        weight = numpy.random.random(ngrids)
#        aoLa, aoLb, aoSa, aoSb = r_numint.eval_ao(mol, coords, deriv=1)
#
#        s = numpy.linalg.norm(m, axis=0)
#        m_pauli = numpy.einsum('xp,xij,p->pij', m, lib.PauliMatrices, 1./(s+1e-300))
#        aoL = numpy.array([aoLa[0],aoLb[0]])
#
#        mat0 = numpy.einsum('pi,p,pj->ij', aoLa[0].conj(), weight*vxc[0][:,0], aoLa[0])
#        mat0+= numpy.einsum('pi,p,pj->ij', aoLb[0].conj(), weight*vxc[0][:,0], aoLb[0])
#        mat0+= numpy.einsum('api,p,pab,bpj->ij', aoL.conj(), weight*vxc[0][:,1], m_pauli, aoL)
#        mat1 = r_numint.eval_mat(mol, (aoLa[0], aoLb[0]), weight, (rho, m), vxc, xctype='LDA')
#        self.assertTrue(numpy.allclose(mat0, mat1))

    def test_rsh_omega(self):
        rho0 = numpy.array([[1., 1., 0.1, 0.1],
                            [.1, .1, 0.01, .01]]).reshape(2, 4, 1)
        ni = r_numint.RNumInt()
        ni.omega = 0.4
        omega = 0.2
        exc, vxc, fxc, kxc = ni.eval_xc('ITYH,', rho0, deriv=1, omega=omega)
        self.assertAlmostEqual(float(exc), -0.6394181669577297, 7)
        #self.assertAlmostEqual(float(vxc[0][0,0]), -0.8688965017309331, 7)  # libxc-4.3.4
        #self.assertAlmostEqual(float(vxc[0][0,1]), -0.04641346660681983, 7)  # libxc-4.3.4
        self.assertAlmostEqual(float(vxc[0][0,0]), -0.868874616534556, 7)  # libxc-5.1.2
        self.assertAlmostEqual(float(vxc[0][0,1]), -0.0465674503111865, 7)  # libxc-5.1.2
        # vsigma of GGA may be problematic?
        #?self.assertAlmostEqual(float(vxc[1][0,0]), 0, 7)
        #?self.assertAlmostEqual(float(vxc[1][0,1]), 0, 7)

        exc, vxc, fxc, kxc = ni.eval_xc('ITYH,', rho0, deriv=1)
        self.assertAlmostEqual(float(exc), -0.5439673757289064, 7)
        #self.assertAlmostEqual(float(vxc[0][0,0]), -0.7699824959456474, 7)  # libxc-4.3.4
        #self.assertAlmostEqual(float(vxc[0][0,1]), -0.04529004028228567, 7)  # libxc-4.3.4
        self.assertAlmostEqual(float(vxc[0][0,0]), -0.7698921068411966, 7)  # libxc-5.1.2
        self.assertAlmostEqual(float(vxc[0][0,1]), -0.04592601580190408, 7)  # libxc-5.1.2
        # vsigma of GGA may be problematic?
        #?self.assertAlmostEqual(float(vxc[1][0,0]), 0, 7)
        #?self.assertAlmostEqual(float(vxc[1][0,1]), 0, 7)


    def test_vxc1(self):
        mol = gto.M(atom=[
            ["O" , (0. , 0.     , 0.)],
            [1   , (0. , -0.757 , 0.587)],
            [1   , (0. , 0.757  , 0.587)]],
            basis='6311g*')
        mf = dft.dks.DKS(mol)
        mf.grids.atom_grid = {"H": (30, 194), "O": (30, 194),}
        mf.grids.prune = None
        mf.grids.build()
        dm = mf.get_init_guess(key='minao')
        res = mf._numint.r_vxc(mol, mf.grids, mf.xc, dm, spin=0)
        self.assertAlmostEqual(res[1], -8.631807003163278, 11)

#    def test_fxc(self):

if __name__ == "__main__":
    print("Test r_numint")
    unittest.main()



# pyscf/dft/test/test_xcfun.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import unittest
import numpy
from pyscf import gto, scf
from pyscf import dft
from pyscf import lib

def setUpModule():
    global mol, mf, ao, rho, dm
    mol = gto.Mole()
    mol.verbose = 0
    mol.output = None
    mol.atom = 'h 0 0 0; h 1 .5 0; h 0 4 1; h 1 0 .2'
    mol.basis = 'aug-ccpvdz'
    mol.build()
    #dm = scf.RHF(mol).run(conv_tol=1e-14).make_rdm1()
    dm = numpy.load(os.path.realpath(os.path.join(__file__, '..', 'dm_h4.npy')))
    mf = dft.RKS(mol)
    mf.grids.atom_grid = {"H": (50, 110)}
    mf.prune = None
    mf.grids.build(with_non0tab=False)
    ao = dft.numint.eval_ao(mol, mf.grids.coords, deriv=1)
    rho = dft.numint.eval_rho(mol, ao, dm, xctype='GGA')

def tearDownModule():
    global mol, mf, ao, rho
    del mol, mf, ao, rho

def finger(a):
    w = numpy.cos(numpy.arange(a.size))
    return numpy.dot(w, a.ravel())

class KnownValues(unittest.TestCase):
    def test_parse_xc(self):
        hyb, fn_facs = dft.xcfun.parse_xc('.5*HF+.5*B3LYP,VWN*.5')
        self.assertAlmostEqual(hyb[0], .6, 12)
        self.assertEqual([x[0] for x in fn_facs], [0,6,16,3])
        self.assertTrue(numpy.allclose([x[1] for x in fn_facs],
                                       (0.04, 0.36, 0.405, 0.595)))
        hyb, fn_facs = dft.xcfun.parse_xc('HF,')
        self.assertEqual(hyb[0], 1)
        self.assertEqual(fn_facs, [])

        hyb, fn_facs = dft.libxc.parse_xc('B88 - SLATER')
        self.assertEqual(fn_facs, [(106, 1), (1, -1)])
        hyb, fn_facs = dft.libxc.parse_xc('B88 -SLATER*.5')
        self.assertEqual(fn_facs, [(106, 1), (1, -0.5)])

        hyb, fn_facs = dft.xcfun.parse_xc('0.5*B3LYP+0.25*B3LYP')
        self.assertTrue(numpy.allclose(hyb, [.15, 0, 0]))
        hyb = dft.libxc.hybrid_coeff('0.5*B3LYP+0.25*B3LYP')
        self.assertAlmostEqual(hyb, .15, 12)

        hyb, fn_facs = dft.xcfun.parse_xc('CAM_B3LYP')
        self.assertTrue(numpy.allclose(hyb, [0.19, 0.65, 0.33]))

        hyb, fn_facs = dft.xcfun.parse_xc('0.6*CAM_B3LYP+0.4*B3P86')
        self.assertTrue(numpy.allclose(hyb, [.08+0.19*.6, 0.65*.6, 0.33]))
        self.assertTrue(numpy.allclose(fn_facs,
                                       [(8, 0.276), (6, 0.498), (3, 0.19), (16, 0.486), (0, 0.032), (56, 0.324)]))
        rsh = dft.xcfun.rsh_coeff('0.6*CAM_B3LYP+0.4*B3P86')
        self.assertTrue(numpy.allclose(rsh, (0.33, 0.39, -0.196)))

        hyb, fn_facs = dft.xcfun.parse_xc('0.4*B3P86+0.6*CAM_B3LYP')
        self.assertTrue(numpy.allclose(hyb, [.08+0.19*.6, 0.65*.6, 0.33]))
        self.assertTrue(numpy.allclose(fn_facs,
                                       [(0, 0.032), (6, 0.498), (56, 0.324), (3, 0.19), (8, 0.276), (16, 0.486)]))
        rsh = dft.xcfun.rsh_coeff('0.4*B3P86+0.6*CAM_B3LYP')
        self.assertTrue(numpy.allclose(rsh, (0.33, 0.39, -0.196)))

        hyb, fn_facs = dft.xcfun.parse_xc('0.5*SR-HF(0.3) + .8*HF + .22*LR_HF')
        self.assertEqual(hyb, [1.3, 1.02, 0.3])

        hyb, fn_facs = dft.xcfun.parse_xc('0.5*SR-HF + .22*LR_HF(0.3) + .8*HF')
        self.assertEqual(hyb, [1.3, 1.02, 0.3])

        hyb, fn_facs = dft.xcfun.parse_xc('0.5*SR-HF + .8*HF + .22*LR_HF(0.3)')
        self.assertEqual(hyb, [1.3, 1.02, 0.3])

        hyb, fn_facs = dft.xcfun.parse_xc('0.5*RSH(2.04;0.56;0.3) + 0.5*BP86')
        self.assertEqual(hyb, [1.3, 1.02, 0.3])
        self.assertEqual(fn_facs, [(6, 0.5), (56, 0.5)])

        self.assertRaises(ValueError, dft.xcfun.parse_xc, 'SR_HF(0.3) + LR_HF(.5)')
        self.assertRaises(ValueError, dft.xcfun.parse_xc, 'LR-HF(0.3) + SR-HF(.5)')

        hyb = dft.xcfun.hybrid_coeff('M05')
        self.assertAlmostEqual(hyb, 0.28, 9)

        hyb, fn_facs = dft.xcfun.parse_xc('APBE,')
        self.assertEqual(fn_facs[0][0], 68)

        hyb, fn_facs = dft.xcfun.parse_xc('VWN,')
        self.assertEqual(fn_facs, [(3, 1)])

        hyb, fn_facs = dft.xcfun.parse_xc('TF,')
        self.assertEqual(fn_facs, [(24, 1)])

        ref = [(0, 1), (3, 1)]
        self.assertEqual(dft.xcfun.parse_xc_name('LDA,VWN'), (0,3))
        self.assertEqual(dft.xcfun.parse_xc(('LDA','VWN'))[1], ref)
        self.assertEqual(dft.xcfun.parse_xc((0, 3))[1], ref)
        self.assertEqual(dft.xcfun.parse_xc('0, 3')[1], ref)
        self.assertEqual(dft.xcfun.parse_xc(3)[1], [(3,1)])

        #self.assertEqual(dft.xcfun.parse_xc('M11-L')[1], [(226,1),(75,1)])
        #self.assertEqual(dft.xcfun.parse_xc('M11L' )[1], [(226,1),(75,1)])
        #self.assertEqual(dft.xcfun.parse_xc('M11-L,M11L' )[1], [(226,1),(75,1)])
        #self.assertEqual(dft.xcfun.parse_xc('M11_L,M11-L')[1], [(226,1),(75,1)])
        #self.assertEqual(dft.xcfun.parse_xc('M11L,M11_L' )[1], [(226,1),(75,1)])

        #self.assertEqual(dft.xcfun.parse_xc('Xpbe,')[1], [(123,1)])
        #self.assertEqual(dft.xcfun.parse_xc('pbe,' )[1], [(101,1)])
        hyb, fn_facs = dft.xcfun.parse_xc('PBE*.4+LDA')
        self.assertEqual(fn_facs, [(5, 0.4), (4, 0.4), (0, 1)])
        hyb, fn_facs = dft.xcfun.parse_xc('PBE*.4+VWN')
        self.assertEqual(fn_facs, [(5, 0.4), (4, 0.4), (3, 1)])

        self.assertTrue (dft.xcfun.is_meta_gga('m05'))
        self.assertFalse(dft.xcfun.is_meta_gga('pbe0'))
        self.assertFalse(dft.xcfun.is_meta_gga('tf,'))
        #self.assertFalse(dft.xcfun.is_meta_gga('vv10'))
        self.assertTrue (dft.xcfun.is_gga('PBE0'))
        self.assertFalse(dft.xcfun.is_gga('m05'))
        self.assertFalse(dft.xcfun.is_gga('tf,'))
        self.assertTrue (dft.xcfun.is_lda('tf,'))
        #self.assertFalse(dft.xcfun.is_lda('vv10'))
        self.assertTrue (dft.xcfun.is_hybrid_xc('m05'))
        self.assertTrue (dft.xcfun.is_hybrid_xc('pbe0,'))
        self.assertFalse(dft.xcfun.is_hybrid_xc('m05,'))
        #self.assertFalse(dft.xcfun.is_hybrid_xc('vv10'))
        self.assertTrue (dft.xcfun.is_hybrid_xc(('b3lyp', 4, 'vv10')))

    def test_nlc_coeff(self):
        #self.assertEqual(dft.xcfun.nlc_coeff('0.5*vv10'), [5.9, 0.0093])
        self.assertEqual(dft.xcfun.nlc_coeff('pbe__vv10'), [5.9, 0.0093])

    def test_lda(self):
        e,v,f,k = dft.xcfun.eval_xc('lda,', rho[0][:3], deriv=3)
        self.assertAlmostEqual(lib.finger(e)   , -0.4720562542635522, 8)
        self.assertAlmostEqual(lib.finger(v[0]), -0.6294083390180697, 8)
        self.assertAlmostEqual(lib.finger(f[0]), -1.1414693830969338, 8)
        self.assertAlmostEqual(lib.finger(k[0]),  4.1402447248393921, 8)

        e,v,f,k = dft.xcfun.eval_xc('lda,', [rho[0][:3]*.5]*2, spin=1, deriv=3)
        self.assertAlmostEqual(lib.finger(e)   , -0.4720562542635522, 8)
        self.assertAlmostEqual(lib.finger(v[0].T[0]), -0.6294083390180697, 8)
        self.assertAlmostEqual(lib.finger(v[0].T[1]), -0.6294083390180697, 8)
        self.assertAlmostEqual(lib.finger(f[0].T[0]), -1.1414693830969338*2, 8)
        self.assertAlmostEqual(lib.finger(f[0].T[2]), -1.1414693830969338*2, 8)
        self.assertAlmostEqual(lib.finger(k[0].T[0]),  4.1402447248393921*4, 7)
        self.assertAlmostEqual(lib.finger(k[0].T[3]),  4.1402447248393921*4, 7)

    def test_lyp(self):
        e,v,f = dft.xcfun.eval_xc(',LYP', rho, deriv=3)[:3]
        self.assertAlmostEqual(numpy.dot(rho[0],e), -62.114576182676615, 8)
        self.assertAlmostEqual(numpy.dot(rho[0],v[0]),-81.771670866308455, 8)
        self.assertAlmostEqual(numpy.dot(rho[0],v[1]), 27.485383255125743, 8)
        self.assertAlmostEqual(numpy.dot(rho[0],f[0]), 186.823806251777, 7)
        self.assertAlmostEqual(numpy.dot(rho[0],f[1]), -3391.2428894571085, 6)
        self.assertAlmostEqual(numpy.dot(rho[0],f[2]), 0, 9)

    #def test_tpss(self):
    #    #FIXME: raised numerical error
    #    rho_a = numpy.array([[3.67808547e-08,-2.02358682e-08, 2.16729780e-07, 2.27036045e-07,-1.47795869e-07,-1.45668997e-09]]).T
    #    e, v = dft.xcfun.eval_xc('tpss,', rho_a, spin=0, deriv=1)[:2]
    #    rho_b = numpy.array([[4.53272893e-06, 4.18968775e-06,-2.83034672e-06, 2.61832978e-06, 5.63360737e-06, 8.97541777e-07]]).T
    #    e, v = dft.xcfun.eval_xc('tpss,', (rho_a, rho_b), spin=1, deriv=1)[:2]

    def test_beckex(self):
        rho =(numpy.array([1.    , 1., 0., 0.]).reshape(-1,1),
              numpy.array([    .8, 1., 0., 0.]).reshape(-1,1))
        e,v,f = dft.xcfun.eval_xc('b88,', rho, spin=1, deriv=3)[:3]
        self.assertAlmostEqual(lib.finger(e)   ,-0.9061911523772116   , 9)
        self.assertAlmostEqual(lib.finger(v[0]),-1.8531364353196298   , 9)
        self.assertAlmostEqual(lib.finger(v[1]),-0.0018308066137967724, 9)
        self.assertAlmostEqual(lib.finger(f[0]),-0.21602284426026866  , 9)
        self.assertAlmostEqual(lib.finger(f[1]), 0.0072053520662545617, 9)
        self.assertAlmostEqual(lib.finger(f[2]), 0.0002275350850255538, 9)

    def test_m05x(self):
        rho =(numpy.array([1., 1., 0., 0., 0., 0.165 ]).reshape(-1,1),
              numpy.array([.8, 1., 0., 0., 0., 0.1050]).reshape(-1,1))
        test_ref = numpy.array([-1.57876583, -2.12127045,-2.11264351,-0.00315462,
                                 0.00000000, -0.00444560, 3.45640232, 4.4349756])
        exc, vxc, fxc, kxc = dft.xcfun.eval_xc('m05,', rho, 1, deriv=3)
        self.assertAlmostEqual(float(exc)*1.8, test_ref[0], 5)
        self.assertAlmostEqual(abs(vxc[0]-test_ref[1:3]).max(), 0, 6)
        self.assertAlmostEqual(abs(vxc[1]-test_ref[3:6]).max(), 0, 6)
        self.assertAlmostEqual(abs(vxc[3]-test_ref[6:8]).max(), 0, 5)

        exc, vxc, fxc, kxc = dft.xcfun.eval_xc('m05,', rho[0], 0, deriv=3)
        self.assertAlmostEqual(float(exc), -0.5746231988116002, 5)
        self.assertAlmostEqual(float(vxc[0]), -0.8806121005703862, 6)
        self.assertAlmostEqual(float(vxc[1]), -0.0032300155406846756, 7)
        self.assertAlmostEqual(float(vxc[3]), 0.4474953100487698, 5)

    def test_camb3lyp(self):
        rho = numpy.array([1., 1., 0.1, 0.1]).reshape(-1,1)
        exc, vxc, fxc, kxc = dft.xcfun.eval_xc('camb3lyp', rho, 0, deriv=1)
        self.assertAlmostEqual(float(exc), -0.5752559666317147, 5)
        self.assertAlmostEqual(float(vxc[0]), -0.7709812578936763, 5)
        self.assertAlmostEqual(float(vxc[1]), -0.0029862221286189846, 7)

        self.assertEqual(dft.xcfun.rsh_coeff('camb3lyp'), (0.33, 0.65, -0.46))

        rho = numpy.array([1., 1., 0.1, 0.1]).reshape(-1,1)
        exc, vxc, fxc, kxc = dft.xcfun.eval_xc('RSH(0.65;-0.46;0.5) + BECKECAMX', rho, 0, deriv=1)
        self.assertAlmostEqual(float(exc), -0.48916154057161476, 9)
        self.assertAlmostEqual(float(vxc[0]), -0.6761177630311709, 9)
        self.assertAlmostEqual(float(vxc[1]), -0.002949151742087167, 9)

    def test_define_xc(self):
        def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, verbose=None):
            exc = vxc = fxc = kxc = None
            return exc, vxc, fxc, kxc

        mf = dft.RKS(mol)
        ni = dft.xcfun.define_xc(mf._numint, eval_xc, 'GGA', hyb=0.2)
        ni = dft.xcfun.define_xc(mf._numint, 'b3lyp+vwn', 'GGA', hyb=0.2)
        self.assertRaises(ValueError, dft.xcfun.define_xc, mf._numint, 0.1)

    def test_vs_libxc_rks(self):
        ao = dft.numint.eval_ao(mol, mf.grids.coords[:200], deriv=2)
        rho = dft.numint.eval_rho(mol, ao, dm, xctype='MGGA')
        rhoa = rho[:,:200]
        def check(xc_code, deriv=3, e_place=9, v_place=9, f_place=9, k_place=9):
            exc0, vxc0, fxc0, kxc0 = dft.libxc.eval_xc(xc_code, rhoa, 0, deriv=deriv)
            exc1, vxc1, fxc1, kxc1 = dft.xcfun.eval_xc(xc_code, rhoa, 0, deriv=deriv)
            self.assertAlmostEqual(abs(exc0-exc1).max(), 0, e_place)
            if deriv > 0:
                for v0, v1 in zip(vxc0, vxc1):
                    if v0 is not None and v1 is not None:
                        self.assertAlmostEqual(abs(v0-v1).max(), 0, v_place)
            if deriv > 1:
                for f0, f1 in zip(fxc0, fxc1):
                    if f0 is not None and f1 is not None:
                        self.assertAlmostEqual(abs(f0-f1).max(), 0, f_place)
            if deriv > 2:
                for k0, k1 in zip(kxc0, kxc1):
                    if k0 is not None and k1 is not None:
                        self.assertAlmostEqual(abs(k0-k1).max(), 0, k_place)

        check('lda,')

        check('pw86,')
        check('pbe,', e_place=6, v_place=6, f_place=5, k_place=4)
        #?check('becke,')
        #?check('br,')
        #?check('LDAERF,')
        check('optx,')
        check('OPTXCORR,')
        check('RPBE,')
        check('TF,'  )
        check('PW91,' , e_place=6, v_place=4, f_place=2, k_place=-1)
        check('m05'  , deriv=1, e_place=6, v_place=6)
        #check('m05-2x,', deriv=1, e_place=6, v_place=6)
        check('m06'  , deriv=1, e_place=6, v_place=6)
        check('m06,'  , deriv=1, e_place=6, v_place=6)
        check('m062x,', deriv=1, e_place=6, v_place=6)
        check('m06l,' , deriv=1, e_place=6, v_place=6)
        check('scan,', deriv=2, e_place=8, v_place=7, f_place=-2)
        check('TPSS,' ,                               f_place=-3, k_place=-4)
        #check('REVTPSS,', deriv=1)  # xcfun crash
        check('APBE,')
        #check('BLOC,' , deriv=2)
        check('PBEINT,', e_place=7, v_place=6, f_place=5, k_place=4)

        check(',vwn3')
        check(',vwn5')
        check(',pbe'  , deriv=3)
        #?check(',br')
        #?check(',LDAERF')
        check(',lyp'  , deriv=3,                                  k_place=0)
        check(',SPBE' , deriv=3, e_place=1, v_place=1, f_place=0, k_place=-2)
        check(',PW91' , deriv=3, e_place=5, v_place=3, f_place=0, k_place=-2)
        check(',m052x', deriv=1)
        check(',m05'  , deriv=1)
        check(',m06'  , deriv=1)
        check(',m062x', deriv=1)
        check(',m06l' , deriv=1)
        check(',scan' , deriv=1, e_place=8, v_place=7)
        check(',TPSS' , deriv=1)
        check(',REVTPSS', deriv=1, e_place=2, v_place=1)
        check(',p86'    , deriv=3, e_place=5, v_place=5, f_place=3, k_place=-1)
        check(',APBE'   , deriv=3)
        check(',PBEINT' , deriv=3)
        check(',TPSSLOC', deriv=1)

        #?check('br')
        check('revpbe', deriv=3, e_place=6, v_place=6, f_place=5, k_place=4)
        check('b97'   , deriv=3, e_place=6, v_place=5, f_place=3, k_place=-5)
        #?check('b97_1')
        #?check('b97_2')
        check('SVWN')
        check('BLYP'  , deriv=3,                                  k_place=0)
        check('BP86'  , deriv=3, e_place=5, v_place=5, f_place=3, k_place=-1)
        check('OLYP'  , deriv=3,                                  k_place=0)
        check('KT1'   , deriv=3,                                  k_place=0)
        check('KT2'   , deriv=3,                                  k_place=-1)
        #?check('KT3')
        check('PBE0'   , deriv=3, e_place=6, v_place=6, f_place=5, k_place=-2)
        check('B3P86'  , deriv=3, e_place=5, v_place=5, f_place=3, k_place=-1)
        check('B3P86G' , deriv=3, e_place=5, v_place=5, f_place=3, k_place=-2)
        check('B3PW91' , deriv=3, e_place=5, v_place=3, f_place=0, k_place=-2)
        check('B3PW91G', deriv=3, e_place=2, v_place=2, f_place=0, k_place=-4)
        check('B3LYP'  , deriv=3,                                  k_place=0)
        check('B3LYP5' , deriv=3,                                  k_place=0)
        check('B3LYPG' , deriv=3,                                  k_place=-2)
        check('O3LYP'  , deriv=3,                                  k_place=-2)
        check('X3LYP'  , deriv=3, e_place=7, v_place=5, f_place=2, k_place=0)
        check('CAMB3LYP', deriv=1)
        check('B97_1'   , deriv=2, e_place=6, v_place=5, f_place=3)
        check('B97_2'   , deriv=2, e_place=6, v_place=5, f_place=3)
        check('TPSSH'   , deriv=1)

    def test_vs_libxc_uks(self):
        ao = dft.numint.eval_ao(mol, mf.grids.coords[:400], deriv=2)
        rho = dft.numint.eval_rho(mol, ao, dm, xctype='MGGA')
        rhoa = rho[:,:200]
        rhob = rhoa + rho[:,200:400]
        def check(xc_code, deriv=3, e_place=9, v_place=9, f_place=9, k_place=9):
            exc0, vxc0, fxc0, kxc0 = dft.libxc.eval_xc(xc_code, (rhoa, rhob), 1, deriv=deriv)
            exc1, vxc1, fxc1, kxc1 = dft.xcfun.eval_xc(xc_code, (rhoa, rhob), 1, deriv=deriv)
            self.assertAlmostEqual(abs(exc0-exc1).max(), 0, e_place)
            if deriv > 0:
                for v0, v1 in zip(vxc0, vxc1):
                    if v0 is not None and v1 is not None:
                        self.assertAlmostEqual(abs(v0-v1).max(), 0, v_place)
            if deriv > 1:
                for f0, f1 in zip(fxc0, fxc1):
                    if f0 is not None and f1 is not None:
                        self.assertAlmostEqual(abs(f0-f1).max(), 0, f_place)
            if deriv > 2 and kxc0 is not None:
                for k0, k1 in zip(kxc0, kxc1):
                    if k0 is not None and k1 is not None:
                        self.assertAlmostEqual(abs(k0-k1).max(), 0, k_place)

        check('lda,')

        check('pw86,')
        check('pbe,', e_place=6, v_place=6, f_place=5, k_place=4)
        #?check('becke,')
        #?check('br,')
        #?check('LDAERF,')
        check('optx,')
        check('OPTXCORR,')
        check('RPBE,')
        check('TF,'  , e_place=0, v_place=-1, f_place=-2, k_place=-2)
        check('PW91,' , e_place=6, v_place=4, f_place=2, k_place=-1)
        check('m05'  , deriv=1, e_place=6, v_place=6)
        #check('m052x,', deriv=1, e_place=6, v_place=6)
        check('m06'  , deriv=1, e_place=6, v_place=6)
        check('m06,'  , deriv=1, e_place=6, v_place=6)
        check('m062x,', deriv=1, e_place=6, v_place=6)
        check('m06l,' , deriv=1, e_place=6, v_place=6)
        check('scan,', deriv=2, e_place=8, v_place=7, f_place=-4)
        check('TPSS,' ,                               f_place=-4, k_place=-4)
        #check('REVTPSS,', deriv=1)  # xcfun crash
        check('APBE,')
        #check('BLOC,' , deriv=2)
        check('PBEINT,', e_place=7, v_place=6, f_place=5, k_place=4)

        check(',vwn3', e_place=2, v_place=1, f_place=1, k_place=0)
        check(',vwn5')
        check(',pbe'  , deriv=3,                        k_place=0)
        #?check(',br')
        #?check(',LDAERF')
        check(',lyp'  , deriv=3,                                  k_place=-1)
        check(',SPBE' , deriv=3, e_place=1, v_place=1, f_place=0, k_place=-1)
        check(',PW91' , deriv=3, e_place=5, v_place=3, f_place=2, k_place=-2)
        check(',m052x', deriv=1)
        check(',m05'  , deriv=1)
        check(',m06'  , deriv=1)
        check(',m062x', deriv=1)
        check(',m06l' , deriv=1)
        check(',scan' , deriv=1, e_place=8, v_place=7)
        check(',TPSS' , deriv=1)
        check(',REVTPSS', deriv=1, e_place=2, v_place=1)
        check(',p86'    , deriv=3, e_place=5, v_place=5, f_place=3, k_place=-2)
        check(',APBE'   , deriv=3,                                  k_place=-1)
        check(',PBEINT' , deriv=3,                                  k_place=-1)
        check(',TPSSLOC', deriv=1)

        #?check('br')
        check('revpbe', deriv=3, e_place=6, v_place=6, f_place=5, k_place=0)
        check('b97'   , deriv=3, e_place=6, v_place=5, f_place=3, k_place=-5)
        #?check('b97_1')
        #?check('b97_2')
        check('SVWN')
        check('BLYP'  , deriv=3,                                  k_place=-1)
        check('BP86'  , deriv=3, e_place=5, v_place=5, f_place=3, k_place=-3)
        check('OLYP'  , deriv=3,                                  k_place=-1)
        check('KT1'   , deriv=3,                                  k_place=-2)
        check('KT2'   , deriv=3,                                  k_place=-2)
        #?check('KT3')
        check('PBE0'   , deriv=3, e_place=6, v_place=6, f_place=5, k_place=-2)
        check('B3P86'  , deriv=3, e_place=5, v_place=5, f_place=3, k_place=-2)
        check('B3P86G' , deriv=3, e_place=3, v_place=2, f_place=2, k_place=-3)
        check('B3PW91' , deriv=3, e_place=5, v_place=4, f_place=2, k_place=-1)
        check('B3PW91G', deriv=3, e_place=2, v_place=2, f_place=2, k_place=-2)
        check('B3LYP'  , deriv=3,                                  k_place=-1)
        check('B3LYP5' , deriv=3,                                  k_place=-1)
        check('B3LYPG' , deriv=3, e_place=3, v_place=2, f_place=2, k_place=-2)
        check('O3LYP'  , deriv=3, e_place=3, v_place=2, f_place=1, k_place=-2)
        check('X3LYP'  , deriv=3, e_place=7, v_place=5, f_place=2, k_place=-1)
        check('CAMB3LYP', deriv=1, v_place=2)
        check('B97_1'   , deriv=3, e_place=6, v_place=5, f_place=3, k_place=-4)
        check('B97_2'   , deriv=3, e_place=6, v_place=5, f_place=3, k_place=-3)
        check('TPSSH'   , deriv=1)

    def test_m06(self):
        rho = numpy.array([0.11939647, -0.18865577, -0.11633254, 0.01779666, -0.55475521, 0.07092032,
                           0.11850155, -0.19297934, -0.11581427, 0.01373251, -0.57534216, 0.06596468,])
        rho = rho.reshape(2,6,1)
        exc, vxc, fxc, kxc = dft.xcfun.eval_xc(',m06', rho[0], 0, deriv=3)
        exc_ref = numpy.array([-0.02809518])
        vxc_ref = numpy.array([-0.0447466, -0.002186, 0.0154814])
        fxc_ref = numpy.array([-0.66060719, 0.35848677, -1.37754938,
                               -1.63288299, 0.21999345, 1.20222259])
        kxc_ref = numpy.array([-8.91475141e+00, 2.57071018e+01, -4.37413046e+01, 9.92064877e+01,
                               -1.47859948e+01, 1.51506522e+01, -4.83492461e+00, -1.75713898e+01,
                               -2.82029718e+01, 6.66847006e+01])
        self.assertAlmostEqual(exc[0] - exc_ref[0], 0, 7)
        self.assertAlmostEqual(abs(numpy.hstack([x for x in vxc if x is not None])-vxc_ref).max(), 0, 7)
        self.assertAlmostEqual(abs(numpy.hstack([x for x in fxc if x is not None])-fxc_ref).max(), 0, 7)
        self.assertAlmostEqual(abs(numpy.hstack([x for x in kxc if x is not None])-kxc_ref).max(), 0, 7)

        exc, vxc, fxc, kxc = dft.xcfun.eval_xc(',m06', rho, 1, deriv=3)
        exc_ref = numpy.array([-0.03730162])
        vxc_ref = numpy.array([-0.06168318,-0.0610502,-0.00957516,0,-0.01988745, 0.03161512, 0.04391615])
        fxc_ref = numpy.array([ 0.62207817,-1.02032218, 0.59318092,
                               -0.00267104, 0, 0.2768053, 0.26785479, 0, 0.02509479,
                               -1.34211352, 0,-0.0399852, 0, 0,-1.28860549,
                               -1.69697623,-0.03422618,-2.1835459,
                               -0.22523468, 0.26116507, 0.26545614,-0.23488143,
                               1.52817397,-0.13694284, 0, 0,-0.13797703, 1.73484678])
        kxc_ref = numpy.array([ 1.49056834e+01, 4.93235407e+00, 5.13177936e+00, 6.83179846e+00,
                               -1.97455841e+01, 0,-8.12321586e-01,-1.51226035e+00, 0,-1.66826458e+00,-8.99591711e-01, 0, -9.19638526e+00,
                               1.90010442e+01, 0, 1.34796896e-01, 0, 0,-7.73620959e-01,-6.10018848e-01, 0, 2.77414513e-01, 0, 0, 6.16594005e+00,
                               2.67162133e+01, 0,-9.09592410e-01, 0, 0,-9.28025786e-01, 0, 0, 0, 3.76462242e+01,
                               2.74381367e-01,-2.08192677e+00, 5.25100873e-01, 6.10851890e-01,-2.02662181e+00,-3.07707014e+00,
                               -4.45125589e-01, 6.81329714e-01, 0, 0,-4.69145558e-01, -2.42540301e+00,-2.45259140e+00,-5.57504449e-01, 0, 0, 6.02633158e-01, 2.58871134e+00,
                               2.87554064e+00, 2.85563442e-01,-1.94628304e-01, -1.88482952e-01, 2.90513272e-01, 3.07132889e+00,
                               -1.08483575e+00, 1.57821239e+00, 0, 0, 1.59013101e+00, 1.61019571e+00, 0, 0, 0, 0, 1.62235586e+00,-3.60275312e+00,
                               -3.90972925e+01,-1.10619931e-02,-1.12015763e-02, 0, 0, 0,-1.11455329e-02,-1.12861703e-02, -4.81639494e+01,
                               7.24911656e+01, 4.35367189e-02, 4.40860771e-02, 1.00010975e+02])
        self.assertAlmostEqual(exc[0] - exc_ref[0], 0, 7)
        self.assertAlmostEqual(abs(numpy.hstack([x for x in vxc if x is not None])-vxc_ref).max(), 0, 7)
        self.assertAlmostEqual(abs(numpy.hstack([x for x in fxc if x is not None])-fxc_ref).max(), 0, 7)
        self.assertAlmostEqual(abs(numpy.hstack([x for x in kxc if x is not None])-kxc_ref).max(), 0, 6)


if __name__ == "__main__":
    print("Test xcfun")
    unittest.main()


# pyscf/dft/xc/utils.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

def remove_dup(fn_facs):
    fn_ids = []
    facs = []
    n = 0
    for key, val in fn_facs:
        if key in fn_ids:
            facs[fn_ids.index(key)] += val
        else:
            fn_ids.append(key)
            facs.append(val)
            n += 1
    return list(zip(fn_ids, facs))

def format_xc_code(description):
    '''Format the description (removing white space) then convert the
    RSH(omega, alpha, beta) notation to the internal notation RSH(alpha; beta; omega)
    '''
    description = description.replace(' ', '').replace('\n', '').upper()
    if 'RSH' not in description:
        return description

    frags = description.split('RSH')
    out = [frags[0]]
    for frag in frags[1:]:
        rsh_key, rest = frag.split(')')
        if ',' in rsh_key:
            omega, alpha, beta = rsh_key[1:].split(',')
            frag = '(' + ';'.join((alpha, beta, omega)) + ')' + rest
        out.append(frag)

    return 'RSH'.join(out)


# pyscf/__all__.py

from . import lib, gto, scf, ao2mo
from .post_scf import *  # noqa: F403
from . import grad
from . import gw
from . import hessian
from . import lo
from . import qmmm
#from . import semiempirical
from . import sgx
from . import solvent
from . import tools
from . import x2c

from . import geomopt

#try:
#    from . import mrpt
#except ImportError:
#    pass
#try:
#    from . import prop
#except ImportError:
#    pass
#try:
#    from . import dftd3
#except ImportError:
#    pass
#try:
#    from . import dmrgscf
#except ImportError:
#    pass
#try:
#    from . import icmpspt
#except ImportError:
#    pass
#try:
#    from . import shciscf
#except ImportError:
#    pass
#try:
#    from . import cornell_shci
#except ImportError:
#    pass

from . import pbc
from .pbc import __all__
del __all__



# pyscf/__config__.py

import os, sys

#
# All parameters initialized before loading pyscf_conf.py will be overwritten
# by the dynamic importing procedure.
#

DEBUG = False

MAX_MEMORY = int(os.environ.get('PYSCF_MAX_MEMORY', 4000)) # MB
TMPDIR = os.environ.get('TMPDIR', '.')
TMPDIR = os.environ.get('PYSCF_TMPDIR', TMPDIR)

VERBOSE = 3  # default logger level (logger.NOTE)
UNIT = 'angstrom'

#
# Loading pyscf_conf.py and overwriting above parameters
#
for conf_file in (os.environ.get('PYSCF_CONFIG_FILE', None),
                  os.path.join(os.path.abspath('.'), '.pyscf_conf.py'),
                  os.path.join(os.environ.get('HOME', '.'), '.pyscf_conf.py')):
    if conf_file is not None and os.path.isfile(conf_file):
        break
else:
    conf_file = None

if conf_file is not None:
    if sys.version_info < (3,0):
        import imp
        imp.load_source('pyscf.__config__', conf_file)
        del(imp)
    else:
        from importlib import machinery
        machinery.SourceFileLoader('pyscf.__config__', conf_file).load_module()
        del(machinery)
del(os, sys)

#
# All parameters initialized after loading pyscf_conf.py will be kept in the
# program.
#


# pyscf/__init__.py

# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''
*****************************************************
PySCF Python-based simulations of chemistry framework
*****************************************************

How to use
----------
There are two ways to access the documentation: the docstrings come with
the code, and an online program reference, available from
http://www.sunqm.net/pyscf/index.html

We recommend the enhanced Python interpreter `IPython <http://ipython.org>`_
and the web-based Python IDE `Ipython notebook <http://ipython.org/notebook.html>`_
to try out the package::

    >>> from pyscf import gto, scf
    >>> mol = gto.M(atom='H 0 0 0; H 0 0 1.2', basis='cc-pvdz')
    >>> mol.apply(scf.RHF).run()
    converged SCF energy = -1.06111199785749
    -1.06111199786

'''

__version__ = '2.0.1'

import os
import sys
# Avoid too many threads being created in OMP loops.
# See issue https://github.com/pyscf/pyscf/issues/317
if 'OPENBLAS_NUM_THREADS' not in os.environ:
    os.environ['OPENBLAS_NUM_THREADS'] = '1'
if 'MKL_NUM_THREADS' not in os.environ:
    os.environ['MKL_NUM_THREADS'] = '1'

# Load modules which are developed as plugins of the namespace package
PYSCF_EXT_PATH = os.getenv('PYSCF_EXT_PATH')
if PYSCF_EXT_PATH:
    for p in PYSCF_EXT_PATH.split(':'):
        if os.path.isdir(p):
            submodules = os.listdir(p)
            if 'pyscf' in submodules:
                # PYSCF_EXT_PATH points to the root directory of a submodule
                __path__.append(os.path.join(p, 'pyscf'))
            else:
                # Load all modules in PYSCF_EXT_PATH if it's a folder that
                # contains all extended modules
                for f in submodules:
                    __path__.append(os.path.join(p, f, 'pyscf'))
                del f
        elif os.path.exists(p):
            # Load all moduels defined inside the file PYSCF_EXT_PATH
            with open(p, 'r') as f:
                __path__.extend(f.read().splitlines())
            del f
    del p
elif '/site-packages/' in __file__ or '/dist-packages/' in __file__:
    # If pyscf has been installed in the standard runtime path (typically
    # under site-packages), and plugins are installed with the pip editable
    # mode, load namespace plugins. In this case, it is likely all modules are
    # managed by pip/conda or their virtual environments. It is safe to search
    # namespace plugins according to the old style of PEP 420.
    __path__ = __import__('pkgutil').extend_path(__path__, __name__)
else:
    # We need a better way to load plugins if pyscf is imported by the
    # environment variable PYTHONPATH. Current treatment may mix installed
    # plugins (e.g.  through pip install) with the developing plugins which
    # were accidentally placed under PYTHONPATH. When PYTHONPATH mechanism is
    # taken, an explicit list of extended paths (using environment
    # PYSCF_EXT_PATH) is recommended.
    __path__ = __import__('pkgutil').extend_path(__path__, __name__)
    if not all('/site-packages/' in p for p in __path__[1:]):
        sys.stderr.write('pyscf plugins found in \n%s\n'
                         'When PYTHONPATH is set, it is recommended to load '
                         'these plugins through the environment variable '
                         'PYSCF_EXT_PATH' % '\n'.join(__path__[1:]))

from distutils.version import LooseVersion
import numpy
if LooseVersion(numpy.__version__) <= '1.8.0':
    raise SystemError("You're using an old version of Numpy (%s). "
                      "It is recommended to upgrade numpy to 1.8.0 or newer. \n"
                      "You still can use all features of PySCF with the old numpy by removing this warning msg. "
                      "Some modules (DFT, CC, MRPT) might be affected because of the bug in old numpy." %
                      numpy.__version__)
elif '1.16.2' <= LooseVersion(numpy.__version__) < '1.18':
    #sys.stderr.write('Numpy 1.16 has memory leak bug  '
    #                 'https://github.com/numpy/numpy/issues/13808\n'
    #                 'It is recommended to downgrade to numpy 1.15 or older\n')
    import ctypes
    from numpy.core import _internal
    def _get_void_ptr(arr):
        simple_arr = numpy.asarray(_internal._unsafe_first_element_pointer(arr))
        c_arr = (ctypes.c_char * 0).from_buffer(simple_arr)
        return ctypes.cast(ctypes.byref(c_arr), ctypes.c_void_p)
    # patch _get_void_ptr as a workaround to numpy issue #13808
    _internal._get_void_ptr = _get_void_ptr

from pyscf import __config__
from pyscf import lib
from pyscf import gto
from pyscf import scf
from pyscf import ao2mo

# Whether to enable debug mode. When this flag is set, some modules may run
# extra debug code.
DEBUG = __config__.DEBUG

def M(**kwargs):
    '''Main driver to create Molecule object (mol) or Material crystal object (cell)'''
    from pyscf import __all__
    if kwargs.get('a') is not None:  # a is crystal lattice parameter
        return __all__.pbc.gto.M(**kwargs)
    else:  # Molecule
        return gto.M(**kwargs)

del os, sys, LooseVersion


# pyscf/post_scf.py

from . import mp, ci, cc, mcscf, fci, tdscf
try:
    from . import doci
except ImportError:
    pass


# pyscf/dft/__init__.py

# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''
Density functional theory
=========================

Simple usage::

    >>> from pyscf import gto, dft
    >>> mol = gto.M(atom='N 0 0 0; N 0 0 1', basis='def2-tzvp')
    >>> mf = dft.RKS(mol)
    >>> mf.xc = 'pbe,pbe'
    >>> mf.run()
'''

try:
    from pyscf.dft import libxc
    XC = {**libxc.XC, **libxc.XC_ALIAS}
except (ImportError, OSError):
    XC = None
try:
    from pyscf.dft import xcfun
    if XC is None:
        XC = {**xcfun.XC, **xcfun.XC_ALIAS}
except (ImportError, OSError):
    pass
#from pyscf.dft import xc
from pyscf.dft import rks
from pyscf.dft import roks
from pyscf.dft import uks
from pyscf.dft import gks
from pyscf.dft import rks_symm
from pyscf.dft import uks_symm
from pyscf.dft import gks_symm
from pyscf.dft import dks
from pyscf.dft import gen_grid as grid
from pyscf.dft import radi
from pyscf.dft import numint
from pyscf.df import density_fit
from pyscf.dft.rks import KohnShamDFT
from pyscf.dft.gen_grid import sg1_prune, nwchem_prune, treutler_prune, \
        stratmann, original_becke, Grids
from pyscf.dft.radi import BRAGG_RADII, COVALENT_RADII, \
        delley, mura_knowles, gauss_chebyshev, treutler, treutler_ahlrichs, \
        treutler_atomic_radii_adjust, becke_atomic_radii_adjust


def KS(mol, xc='LDA,VWN'):
    if mol.spin == 0:
        return RKS(mol, xc)
    else:
        return UKS(mol, xc)
KS.__doc__ = '''
A wrap function to create DFT object (RKS or UKS).\n
''' + rks.RKS.__doc__
DFT = KS

def RKS(mol, xc='LDA,VWN'):
    if mol.spin == 0:
        if not mol.symmetry or mol.groupname == 'C1':
            return rks.RKS(mol, xc)
        else:
            return rks_symm.RKS(mol, xc)
    else:
        return ROKS(mol, xc)
RKS.__doc__ = rks.RKS.__doc__

def ROKS(mol, xc='LDA,VWN'):
    if not mol.symmetry or mol.groupname == 'C1':
        return roks.ROKS(mol, xc)
    else:
        return rks_symm.ROKS(mol, xc)
ROKS.__doc__ = roks.ROKS.__doc__

def UKS(mol, xc='LDA,VWN'):
    if not mol.symmetry or mol.groupname == 'C1':
        return uks.UKS(mol, xc)
    else:
        return uks_symm.UKS(mol, xc)
UKS.__doc__ = uks.UKS.__doc__

def GKS(mol, xc='LDA,VWN'):
    if not mol.symmetry or mol.groupname == 'C1':
        return gks.GKS(mol, xc)
    else:
        return gks_symm.GKS(mol, xc)
GKS.__doc__ = gks.GKS.__doc__

def DKS(mol, xc='LDA,VWN'):
    from pyscf.scf import dhf
    if dhf.zquatev and mol.spin == 0:
        return dks.RDKS(mol, xc=xc)
    else:
        return dks.UDKS(mol, xc=xc)

UDKS = dks.UDKS
RDKS = dks.RDKS


# pyscf/dft/dks.py

#!/usr/bin/env python
# Copyright 2014-2022 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
4-component Dirac-Kohn-Sham
'''


import copy
import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf.scf import dhf
from pyscf.dft import rks
from pyscf.dft import gks
from pyscf.dft import r_numint


@lib.with_doc(gks.get_veff.__doc__)
def get_veff(ks, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
    if ks.nlc != '':
        raise NotImplementedError(ks.nlc)
    return gks.get_veff(ks, mol, dm, dm_last, vhf_last, hermi)


def energy_elec(ks, dm=None, h1e=None, vhf=None):
    r'''Electronic part of DKS energy.

    Note this function has side effects which cause mf.scf_summary updated.

    Args:
        ks : an instance of DFT class

        dm : 2D ndarray
            one-partical density matrix
        h1e : 2D ndarray
            Core hamiltonian

    Returns:
        DKS electronic energy and the 2-electron contribution
    '''
    e1, e2 = rks.energy_elec(ks, dm, h1e, vhf)
    if not ks.with_ssss and ks.ssss_approx == 'Visscher':
        e2 += dhf._vischer_ssss_correction(ks, dm)
        ks.scf_summary['e2'] = e2
    return e1, e2


class KohnShamDFT(rks.KohnShamDFT):
    def __init__(self, xc='LDA,VWN'):
        rks.KohnShamDFT.__init__(self, xc)
        self._numint = r_numint.RNumInt()

    get_veff = gks.get_veff
    energy_elec = gks.energy_elec

    @property
    def collinear(self):
        return self._numint.collinear
    @collinear.setter
    def collinear(self, val):
        self._numint.collinear = val

    def to_rhf(self):
        raise RuntimeError

    def to_uhf(self):
        raise RuntimeError

    def to_ghf(self):
        raise RuntimeError

    def to_rks(self, xc=None):
        raise RuntimeError

    def to_uks(self, xc=None):
        raise RuntimeError

    def to_gks(self, xc=None):
        raise RuntimeError

    def to_dhf(self):
        '''Convert the input mean-field object to a DHF object.

        Note this conversion only changes the class of the mean-field object.
        The total energy and wave-function are the same as them in the input
        mean-field object.
        '''
        mf = self.view(dhf.DHF)
        mf.converged = False
        return mf

    def to_dks(self, xc=None):
        if xc is not None and xc != self.xc:
            mf = copy.copy(self)
            mf.xc = xc
            mf.converged = False
        return self


class DKS(KohnShamDFT, dhf.DHF):
    '''Kramers unrestricted Dirac-Kohn-Sham'''
    def __init__(self, mol, xc='LDA,VWN'):
        dhf.DHF.__init__(self, mol)
        KohnShamDFT.__init__(self, xc)

    def dump_flags(self, verbose=None):
        dhf.DHF.dump_flags(self, verbose)
        KohnShamDFT.dump_flags(self, verbose)
        return self

    def x2c1e(self):
        from pyscf.x2c import dft
        x2chf = dft.UKS(self.mol)
        x2c_keys = x2chf._keys
        x2chf.__dict__.update(self.__dict__)
        x2chf._keys = self._keys.union(x2c_keys)
        return x2chf
    x2c = x2c1e

UKS = UDKS = DKS

class RDKS(KohnShamDFT, dhf.RDHF):
    '''Kramers restricted Dirac-Kohn-Sham'''
    def __init__(self, mol, xc='LDA,VWN'):
        dhf.RDHF.__init__(self, mol)
        KohnShamDFT.__init__(self, xc)

    def x2c1e(self):
        from pyscf.x2c import dft
        x2chf = dft.RKS(self.mol)
        x2c_keys = x2chf._keys
        x2chf.__dict__.update(self.__dict__)
        x2chf._keys = self._keys.union(x2c_keys)
        return x2chf
    x2c = x2c1e

    def to_dhf(self):
        '''Convert the input mean-field object to a DHF object.

        Note this conversion only changes the class of the mean-field object.
        The total energy and wave-function are the same as them in the input
        mean-field object.
        '''
        mf = self.view(dhf.RDHF)
        mf.converged = False
        return mf

RKS = RDKS


# pyscf/dft/gen_grid.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Generate DFT grids and weights, based on the code provided by Gerald Knizia <>

Reference for Lebedev-Laikov grid:
  V. I. Lebedev, and D. N. Laikov "A quadrature formula for the sphere of the
  131st algebraic order of accuracy", Doklady Mathematics, 59, 477-481 (1999)
'''


import ctypes
import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf.dft import radi
from pyscf import gto
from pyscf.gto.eval_gto import BLKSIZE
from pyscf import __config__

libdft = lib.load_library('libdft')

# ~= (L+1)**2/3
LEBEDEV_ORDER = {
    0  : 1   ,
    3  : 6   ,
    5  : 14  ,
    7  : 26  ,
    9  : 38  ,
    11 : 50  ,
    13 : 74  ,
    15 : 86  ,
    17 : 110 ,
    19 : 146 ,
    21 : 170 ,
    23 : 194 ,
    25 : 230 ,
    27 : 266 ,
    29 : 302 ,
    31 : 350 ,
    35 : 434 ,
    41 : 590 ,
    47 : 770 ,
    53 : 974 ,
    59 : 1202,
    65 : 1454,
    71 : 1730,
    77 : 2030,
    83 : 2354,
    89 : 2702,
    95 : 3074,
    101: 3470,
    107: 3890,
    113: 4334,
    119: 4802,
    125: 5294,
    131: 5810
}
LEBEDEV_NGRID = numpy.asarray((
    1   , 6   , 14  , 26  , 38  , 50  , 74  , 86  , 110 , 146 ,
    170 , 194 , 230 , 266 , 302 , 350 , 434 , 590 , 770 , 974 ,
    1202, 1454, 1730, 2030, 2354, 2702, 3074, 3470, 3890, 4334,
    4802, 5294, 5810))

# SG0
# S. Chien and P. Gill,  J. Comput. Chem. 27 (2006) 730-739.


def sg1_prune(nuc, rads, n_ang, radii=radi.SG1RADII):
    '''SG1, CPL, 209, 506

    Args:
        nuc : int
            Nuclear charge.

        rads : 1D array
            Grid coordinates on radical axis.

        n_ang : int
            Max number of grids over angular part.

    Kwargs:
        radii : 1D array
            radii (in Bohr) for atoms in periodic table

    Returns:
        A list has the same length as rads. The list element is the number of
        grids over angular part for each radial grid.
    '''
# In SG1 the ang grids for the five regions
#            6  38 86  194 86
    leb_ngrid = numpy.array([6, 38, 86, 194, 86])
    alphas = numpy.array((
        (0.25  , 0.5, 1.0, 4.5),
        (0.1667, 0.5, 0.9, 3.5),
        (0.1   , 0.4, 0.8, 2.5)))
    r_atom = radii[nuc] + 1e-200
    if nuc <= 2:  # H, He
        place = ((rads/r_atom).reshape(-1,1) > alphas[0]).sum(axis=1)
    elif nuc <= 10:  # Li - Ne
        place = ((rads/r_atom).reshape(-1,1) > alphas[1]).sum(axis=1)
    else:
        place = ((rads/r_atom).reshape(-1,1) > alphas[2]).sum(axis=1)
    return leb_ngrid[place]

def nwchem_prune(nuc, rads, n_ang, radii=radi.BRAGG_RADII):
    '''NWChem

    Args:
        nuc : int
            Nuclear charge.

        rads : 1D array
            Grid coordinates on radical axis.

        n_ang : int
            Max number of grids over angular part.

    Kwargs:
        radii : 1D array
            radii (in Bohr) for atoms in periodic table

    Returns:
        A list has the same length as rads. The list element is the number of
        grids over angular part for each radial grid.
    '''
    alphas = numpy.array((
        (0.25  , 0.5, 1.0, 4.5),
        (0.1667, 0.5, 0.9, 3.5),
        (0.1   , 0.4, 0.8, 2.5)))
    leb_ngrid = LEBEDEV_NGRID[4:]  # [38, 50, 74, 86, ...]
    if n_ang < 50:
        return numpy.repeat(n_ang, len(rads))
    elif n_ang == 50:
        leb_l = numpy.array([1, 2, 2, 2, 1])
    else:
        idx = numpy.where(leb_ngrid==n_ang)[0][0]
        leb_l = numpy.array([1, 3, idx-1, idx, idx-1])

    r_atom = radii[nuc] + 1e-200
    if nuc <= 2:  # H, He
        place = ((rads/r_atom).reshape(-1,1) > alphas[0]).sum(axis=1)
    elif nuc <= 10:  # Li - Ne
        place = ((rads/r_atom).reshape(-1,1) > alphas[1]).sum(axis=1)
    else:
        place = ((rads/r_atom).reshape(-1,1) > alphas[2]).sum(axis=1)
    angs = leb_l[place]
    angs = leb_ngrid[angs]
    return angs

# Prune scheme JCP 102, 346 (1995); DOI:10.1063/1.469408
def treutler_prune(nuc, rads, n_ang, radii=None):
    '''Treutler-Ahlrichs

    Args:
        nuc : int
            Nuclear charge.

        rads : 1D array
            Grid coordinates on radical axis.

        n_ang : int
            Max number of grids over angular part.

    Returns:
        A list has the same length as rads. The list element is the number of
        grids over angular part for each radial grid.
    '''
    nr = len(rads)
    leb_ngrid = numpy.empty(nr, dtype=int)
    leb_ngrid[:nr//3] = 14 # l=5
    leb_ngrid[nr//3:nr//2] = 50 # l=11
    leb_ngrid[nr//2:] = n_ang
    return leb_ngrid



###########################################################
# Becke partitioning

# Stratmann, Scuseria, Frisch. CPL, 257, 213 (1996), eq.11
def stratmann(g):
    '''Stratmann, Scuseria, Frisch. CPL, 257, 213 (1996); DOI:10.1016/0009-2614(96)00600-8'''
    a = .64  # for eq. 14
    g = numpy.asarray(g)
    ma = g/a
    ma2 = ma * ma
    g1 = numpy.asarray((1/16.)*(ma*(35 + ma2*(-35 + ma2*(21 - 5 *ma2)))))
    g1[g<=-a] = -1
    g1[g>= a] =  1
    return g1

def original_becke(g):
    '''Becke, JCP 88, 2547 (1988); DOI:10.1063/1.454033'''
#    This funciton has been optimized in the C code VXCgen_grid
#    g = (3 - g**2) * g * .5
#    g = (3 - g**2) * g * .5
#    g = (3 - g**2) * g * .5
#    return g
    pass

def gen_atomic_grids(mol, atom_grid={}, radi_method=radi.gauss_chebyshev,
                     level=3, prune=nwchem_prune, **kwargs):
    '''Generate number of radial grids and angular grids for the given molecule.

    Returns:
        A dict, with the atom symbol for the dict key.  For each atom type,
        the dict value has two items: one is the meshgrid coordinates wrt the
        atom center; the second is the volume of that grid.
    '''
    if isinstance(atom_grid, (list, tuple)):
        atom_grid = dict([(mol.atom_symbol(ia), atom_grid)
                          for ia in range(mol.natm)])
    atom_grids_tab = {}
    for ia in range(mol.natm):
        symb = mol.atom_symbol(ia)

        if symb not in atom_grids_tab:
            chg = gto.charge(symb)
            if symb in atom_grid:
                n_rad, n_ang = atom_grid[symb]
                if n_ang not in LEBEDEV_NGRID:
                    if n_ang in LEBEDEV_ORDER:
                        logger.warn(mol, 'n_ang %d for atom %d %s is not '
                                    'the supported Lebedev angular grids. '
                                    'Set n_ang to %d', n_ang, ia, symb,
                                    LEBEDEV_ORDER[n_ang])
                        n_ang = LEBEDEV_ORDER[n_ang]
                    else:
                        raise ValueError('Unsupported angular grids %d' % n_ang)
            else:
                n_rad = _default_rad(chg, level)
                n_ang = _default_ang(chg, level)
            rad, dr = radi_method(n_rad, chg, ia, **kwargs)

            rad_weight = 4*numpy.pi * rad**2 * dr

            if callable(prune):
                angs = prune(chg, rad, n_ang)
            else:
                angs = [n_ang] * n_rad
            logger.debug(mol, 'atom %s rad-grids = %d, ang-grids = %s',
                         symb, n_rad, angs)

            angs = numpy.array(angs)
            coords = []
            vol = []
            for n in sorted(set(angs)):
                grid = numpy.empty((n,4))
                libdft.MakeAngularGrid(grid.ctypes.data_as(ctypes.c_void_p),
                                       ctypes.c_int(n))
                idx = numpy.where(angs==n)[0]
                for i0, i1 in prange(0, len(idx), 12):  # 12 radi-grids as a group
                    coords.append(numpy.einsum('i,jk->jik',rad[idx[i0:i1]],
                                               grid[:,:3]).reshape(-1,3))
                    vol.append(numpy.einsum('i,j->ji', rad_weight[idx[i0:i1]],
                                            grid[:,3]).ravel())
            atom_grids_tab[symb] = (numpy.vstack(coords), numpy.hstack(vol))
    return atom_grids_tab


def get_partition(mol, atom_grids_tab,
                  radii_adjust=None, atomic_radii=radi.BRAGG_RADII,
                  becke_scheme=original_becke, concat=True):
    '''Generate the mesh grid coordinates and weights for DFT numerical integration.
    We can change radii_adjust, becke_scheme functions to generate different meshgrid.

    Kwargs:
        concat: bool
            Whether to concatenate grids and weights in return

    Returns:
        grid_coord and grid_weight arrays.  grid_coord array has shape (N,3);
        weight 1D array has N elements.
    '''
    if callable(radii_adjust) and atomic_radii is not None:
        f_radii_adjust = radii_adjust(mol, atomic_radii)
    else:
        f_radii_adjust = None
    atm_coords = numpy.asarray(mol.atom_coords() , order='C')
    atm_dist = gto.inter_distance(mol)
    if (becke_scheme is original_becke and
        (radii_adjust is radi.treutler_atomic_radii_adjust or
         radii_adjust is radi.becke_atomic_radii_adjust or
         f_radii_adjust is None)):
        if f_radii_adjust is None:
            p_radii_table = lib.c_null_ptr()
        else:
            f_radii_table = numpy.asarray([f_radii_adjust(i, j, 0)
                                           for i in range(mol.natm)
                                           for j in range(mol.natm)])
            p_radii_table = f_radii_table.ctypes.data_as(ctypes.c_void_p)

        def gen_grid_partition(coords):
            coords = numpy.asarray(coords, order='F')
            ngrids = coords.shape[0]
            pbecke = numpy.empty((mol.natm,ngrids))
            libdft.VXCgen_grid(pbecke.ctypes.data_as(ctypes.c_void_p),
                               coords.ctypes.data_as(ctypes.c_void_p),
                               atm_coords.ctypes.data_as(ctypes.c_void_p),
                               p_radii_table,
                               ctypes.c_int(mol.natm), ctypes.c_int(ngrids))
            return pbecke
    else:
        def gen_grid_partition(coords):
            ngrids = coords.shape[0]
            grid_dist = numpy.empty((mol.natm,ngrids))
            for ia in range(mol.natm):
                dc = coords - atm_coords[ia]
                grid_dist[ia] = numpy.sqrt(numpy.einsum('ij,ij->i',dc,dc))
            pbecke = numpy.ones((mol.natm,ngrids))
            for i in range(mol.natm):
                for j in range(i):
                    g = 1/atm_dist[i,j] * (grid_dist[i]-grid_dist[j])
                    if f_radii_adjust is not None:
                        g = f_radii_adjust(i, j, g)
                    g = becke_scheme(g)
                    pbecke[i] *= .5 * (1-g)
                    pbecke[j] *= .5 * (1+g)
            return pbecke

    coords_all = []
    weights_all = []
    for ia in range(mol.natm):
        coords, vol = atom_grids_tab[mol.atom_symbol(ia)]
        coords = coords + atm_coords[ia]
        pbecke = gen_grid_partition(coords)
        weights = vol * pbecke[ia] * (1./pbecke.sum(axis=0))
        coords_all.append(coords)
        weights_all.append(weights)

    if concat:
        coords_all = numpy.vstack(coords_all)
        weights_all = numpy.hstack(weights_all)
    return coords_all, weights_all
gen_partition = get_partition

def make_mask(mol, coords, relativity=0, shls_slice=None, verbose=None):
    '''Mask to indicate whether a shell is zero on grid

    Args:
        mol : an instance of :class:`Mole`

        coords : 2D array, shape (N,3)
            The coordinates of grids.

    Kwargs:
        relativity : bool
            No effects.
        shls_slice : 2-element list
            (shl_start, shl_end).
            If given, only part of AOs (shl_start <= shell_id < shl_end) are
            evaluated.  By default, all shells defined in mol will be evaluated.
        verbose : int or object of :class:`Logger`
            No effects.

    Returns:
        2D mask array of shape (N,nbas), where N is the number of grids, nbas
        is the number of shells.
    '''
    coords = numpy.asarray(coords, order='F')
    ngrids = len(coords)
    if shls_slice is None:
        shls_slice = (0, mol.nbas)
    nbas = shls_slice[1] - shls_slice[0]

    non0tab = numpy.empty(((ngrids+BLKSIZE-1)//BLKSIZE, nbas),
                          dtype=numpy.uint8)
    libdft.VXCnr_ao_screen(non0tab.ctypes.data_as(ctypes.c_void_p),
                           coords.ctypes.data_as(ctypes.c_void_p),
                           ctypes.c_int(ngrids),
                           mol._atm.ctypes.data_as(ctypes.c_void_p),
                           ctypes.c_int(mol.natm),
                           mol._bas[shls_slice[0]:].ctypes.data_as(ctypes.c_void_p),
                           ctypes.c_int(nbas),
                           mol._env.ctypes.data_as(ctypes.c_void_p))
    return non0tab



class Grids(lib.StreamObject):
    '''DFT mesh grids

    Attributes for Grids:
        level : int
            To control the number of radial and angular grids. Large number
            leads to large mesh grids. The default level 3 corresponds to
            (50,302) for H, He;
            (75,302) for second row;
            (80~105,434) for rest.

            Grids settings at other levels can be found in
            pyscf.dft.gen_grid.RAD_GRIDS and pyscf.dft.gen_grid.ANG_ORDER

        atomic_radii : 1D array
            | radi.BRAGG_RADII  (default)
            | radi.COVALENT_RADII
            | None : to switch off atomic radii adjustment

        radii_adjust : function(mol, atomic_radii) => (function(atom_id, atom_id, g) => array_like_g)
            Function to adjust atomic radii, can be one of
            | radi.treutler_atomic_radii_adjust
            | radi.becke_atomic_radii_adjust
            | None : to switch off atomic radii adjustment

        radi_method : function(n) => (rad_grids, rad_weights)
            scheme for radial grids, can be one of
            | radi.treutler  (default)
            | radi.delley
            | radi.mura_knowles
            | radi.gauss_chebyshev

        becke_scheme : function(v) => array_like_v
            weight partition function, can be one of
            | gen_grid.original_becke  (default)
            | gen_grid.stratmann

        prune : function(nuc, rad_grids, n_ang) => list_n_ang_for_each_rad_grid
            scheme to reduce number of grids, can be one of
            | gen_grid.nwchem_prune  (default)
            | gen_grid.sg1_prune
            | gen_grid.treutler_prune
            | None : to switch off grid pruning

        symmetry : bool
            whether to symmetrize mesh grids (TODO)

        atom_grid : dict
            Set (radial, angular) grids for particular atoms.
            Eg, grids.atom_grid = {'H': (20,110)} will generate 20 radial
            grids and 110 angular grids for H atom.

        Examples:

        >>> mol = gto.M(atom='H 0 0 0; H 0 0 1.1')
        >>> grids = dft.gen_grid.Grids(mol)
        >>> grids.level = 4
        >>> grids.build()
        '''
    def __init__(self, mol):
        import sys
        self.mol = mol
        self.stdout = mol.stdout
        self.verbose = mol.verbose
        self.symmetry = mol.symmetry
        self.atom_grid = {}
        self.non0tab = None

        cur_mod = sys.modules[__name__]

        def _load_conf(mod, name, default):
            var = getattr(__config__, name, None)
            if var is None:
                return default
            elif callable(var):
                return var
            elif mod is None:
                return cur_mod[var]
            else:
                return getattr(mod, var)

        self.atomic_radii = _load_conf(radi, 'dft_gen_grid_Grids_atomic_radii',
                                       radi.BRAGG_RADII)
        #self.atomic_radii = radi.COVALENT_RADII

        self.radii_adjust = _load_conf(radi, 'dft_gen_grid_Grids_radii_adjust',
                                       radi.treutler_atomic_radii_adjust)
        #self.radii_adjust = radi.becke_atomic_radii_adjust
        #self.radii_adjust = None # to switch off atomic radii adjustment

        self.radi_method = _load_conf(radi, 'dft_gen_grid_Grids_radi_method',
                                      radi.treutler)
        #self.radi_method = radi.gauss_chebyshev
        #self.radi_method = radi.mura_knowles
        #self.radi_method = radi.delley

        self.becke_scheme = _load_conf(None, 'dft_gen_grid_Grids_becke_scheme',
                                       original_becke)
        #self.becke_scheme = stratmann

        self.prune = _load_conf(None, 'dft_gen_grid_Grids_prune', nwchem_prune)

        self.level = getattr(__config__, 'dft_gen_grid_Grids_level', 3)

##################################################
# don't modify the following attributes, they are not input options
        self.coords  = None
        self.weights = None
        self._keys = set(self.__dict__.keys())

    @property
    def size(self):
        return getattr(self.weights, 'size', 0)

    def __setattr__(self, key, val):
        if key in ('atom_grid', 'atomic_radii', 'radii_adjust', 'radi_method',
                   'becke_scheme', 'prune', 'level'):
            self.reset()
        super(Grids, self).__setattr__(key, val)

    def dump_flags(self, verbose=None):
        logger.info(self, 'radial grids: %s', self.radi_method.__doc__)
        logger.info(self, 'becke partition: %s', self.becke_scheme.__doc__)
        logger.info(self, 'pruning grids: %s', self.prune)
        logger.info(self, 'grids dens level: %d', self.level)
        logger.info(self, 'symmetrized grids: %s', self.symmetry)
        if self.radii_adjust is not None:
            logger.info(self, 'atomic radii adjust function: %s',
                        self.radii_adjust)
            logger.debug2(self, 'atomic_radii : %s', self.atomic_radii)
        if self.atom_grid:
            logger.info(self, 'User specified grid scheme %s', str(self.atom_grid))
        return self

    def build(self, mol=None, with_non0tab=False, **kwargs):
        if mol is None: mol = self.mol
        if self.verbose >= logger.WARN:
            self.check_sanity()
        atom_grids_tab = self.gen_atomic_grids(mol, self.atom_grid,
                                               self.radi_method,
                                               self.level, self.prune, **kwargs)
        self.coords, self.weights = \
                self.get_partition(mol, atom_grids_tab,
                                   self.radii_adjust, self.atomic_radii,
                                   self.becke_scheme)
        if with_non0tab:
            self.non0tab = self.make_mask(mol, self.coords)
        else:
            self.non0tab = None
        logger.info(self, 'tot grids = %d', len(self.weights))
        return self

    def kernel(self, mol=None, with_non0tab=False):
        self.dump_flags()
        return self.build(mol, with_non0tab)

    def reset(self, mol=None):
        '''Reset mol and clean up relevant attributes for scanner mode'''
        if mol is not None:
            self.mol = mol
        self.coords = None
        self.weights = None
        self.non0tab = None
        return self

    @lib.with_doc(gen_atomic_grids.__doc__)
    def gen_atomic_grids(self, mol, atom_grid=None, radi_method=None,
                         level=None, prune=None, **kwargs):
        if atom_grid is None: atom_grid = self.atom_grid
        if radi_method is None: radi_method = self.radi_method
        if level is None: level = self.level
        if prune is None: prune = self.prune
        return gen_atomic_grids(mol, atom_grid, self.radi_method, level, prune, **kwargs)

    @lib.with_doc(get_partition.__doc__)
    def get_partition(self, mol, atom_grids_tab=None,
                      radii_adjust=None, atomic_radii=radi.BRAGG_RADII,
                      becke_scheme=original_becke, concat=True):
        if atom_grids_tab is None:
            atom_grids_tab = self.gen_atomic_grids(mol)
        return get_partition(mol, atom_grids_tab, radii_adjust, atomic_radii,
                             becke_scheme, concat=concat)

    gen_partition = get_partition

    @lib.with_doc(make_mask.__doc__)
    def make_mask(self, mol=None, coords=None, relativity=0, shls_slice=None,
                  verbose=None):
        if mol is None: mol = self.mol
        if coords is None: coords = self.coords
        return make_mask(mol, coords, relativity, shls_slice, verbose)


def _default_rad(nuc, level=3):
    '''Number of radial grids '''
    tab   = numpy.array( (2 , 10, 18, 36, 54, 86, 118))
    period = (nuc > tab).sum()
    return RAD_GRIDS[level,period]
#                Period    1   2   3   4   5   6   7        # level
RAD_GRIDS = numpy.array((( 10, 15, 20, 30, 35, 40, 50),     # 0
                         ( 30, 40, 50, 60, 65, 70, 75),     # 1
                         ( 40, 60, 65, 75, 80, 85, 90),     # 2
                         ( 50, 75, 80, 90, 95,100,105),     # 3
                         ( 60, 90, 95,105,110,115,120),     # 4
                         ( 70,105,110,120,125,130,135),     # 5
                         ( 80,120,125,135,140,145,150),     # 6
                         ( 90,135,140,150,155,160,165),     # 7
                         (100,150,155,165,170,175,180),     # 8
                         (200,200,200,200,200,200,200),))   # 9

def _default_ang(nuc, level=3):
    '''Order of angular grids. See LEBEDEV_ORDER for the mapping of
    the order and the number of angular grids'''
    tab   = numpy.array( (2 , 10, 18, 36, 54, 86, 118))
    period = (nuc > tab).sum()
    return LEBEDEV_ORDER[ANG_ORDER[level,period]]
#               Period    1   2   3   4   5   6   7         # level
ANG_ORDER = numpy.array(((11, 15, 17, 17, 17, 17, 17 ),     # 0
                         (17, 23, 23, 23, 23, 23, 23 ),     # 1
                         (23, 29, 29, 29, 29, 29, 29 ),     # 2
                         (29, 29, 35, 35, 35, 35, 35 ),     # 3
                         (35, 41, 41, 41, 41, 41, 41 ),     # 4
                         (41, 47, 47, 47, 47, 47, 47 ),     # 5
                         (47, 53, 53, 53, 53, 53, 53 ),     # 6
                         (53, 59, 59, 59, 59, 59, 59 ),     # 7
                         (59, 59, 59, 59, 59, 59, 59 ),     # 8
                         (65, 65, 65, 65, 65, 65, 65 ),))   # 9

def prange(start, end, step):
    for i in range(start, end, step):
        yield i, min(i+step, end)




if __name__ == '__main__':
    h2o = gto.Mole()
    h2o.verbose = 0
    h2o.output = None#"out_h2o"
    h2o.atom = [
        ['O' , (0. , 0.     , 0.)],
        ['H' , (0. , -0.757 , 0.587)],
        ['H' , (0. , 0.757  , 0.587)] ]
    h2o.build()
    g = Grids(h2o)
    g.build()
    print(g.coords.shape)



# pyscf/dft/gen_libxc_param.py

#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys
ldpath = os.path.abspath(os.path.join(__file__, '..', '..', 'lib', 'deps', 'lib'))
if ldpath not in os.environ['LD_LIBRARY_PATH']:
    sys.stderr.write(f'Set\n\tLD_LIBRARY_PATH={ldpath}:$LD_LIBRARY_PATH\n'
                     'and rerun the script\n')
    exit()

pypath = os.path.join(__file__, '..', '..', 'lib', 'build', 'deps', 'src', 'libxc')
sys.path.insert(0, os.path.abspath(pypath))
import pylibxc

for xcname in pylibxc.util.xc_available_functional_names():
    f = pylibxc.LibXCFunctional(xcname, 1)
    f_id = f.get_number()
    ref = f.get_references()
    key = f"'{xcname.upper()}'"
    print(f"{key:<31s}: {f_id:<3d}, # {ref[0]}")
    for r in ref[1:]:
        print(f"                                     # {r}")


# pyscf/dft/gen_xcfun_param.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pyscf import dft
import ctypes

libdft = dft.xcfun._itrf
libdft.xc_enumerate_parameters.restype = ctypes.c_char_p
libdft.XCFUN_xc_type.restype = ctypes.c_int
libdft.xc_describe_short.restype = ctypes.c_char_p
libdft.xc_describe_short.argtype = [ctypes.c_char_p]
libdft.xc_describe_long.restype = ctypes.c_char_p
libdft.xc_describe_long.argtype = [ctypes.c_char_p]
lda_ids = []
gga_ids = []
mgga_ids = []
mlgga_ids = []
xc_codes = {}
for i in range(68):
    name = libdft.xc_enumerate_parameters(ctypes.c_int(i))
    sdescr = libdft.xc_describe_short(name)
    #ldescr = libdft.xc_describe_long(ctypes.c_int(i))
    if sdescr is not None:
        #print("'%s' : %d,  # %s" % (name, i, sdescr.replace('\n', '. ')))
        print('%-16s: %2d,  # %s' % ("'%s'"%name, i, sdescr.replace('\n', '. ')))
        xc_codes[name] = i

    fntype = libdft.XCFUN_xc_type(ctypes.c_int(i))
    if fntype == 0:
        lda_ids.append(i)
    elif fntype == 1:
        gga_ids.append(i)
    elif fntype == 2:
        mgga_ids.append(i)

alias = {
    'SLATER': 'SLATERX',
    'LDA'   : 'SLATERX',
    'VWN'   : 'VWN5C',
    'VWN5'  : 'VWN5C',
    'B88'   : 'BECKEX',
    'LYP'   : 'LYPC',
}
for k, v in alias.items():
    print('%-16s: %2d,  # %s' % ("'%s'"%k, xc_codes[v], v))
print('LDA_IDS = %s' % lda_ids)
print('GGA_IDS = %s' % gga_ids)
print('MGGA_IDS = %s' % mgga_ids)

#define XC_D0 0
#define XC_D1 1
#define XC_D2 2
#define XC_D3 3
#define XC_D4 4

#define XC_D00 0
#define XC_D10 1
#define XC_D01 2
#define XC_D20 3
#define XC_D11 4
#define XC_D02 5
#define XC_D30 6
#define XC_D21 7
#define XC_D12 8
#define XC_D03 9
#define XC_D40 10
#define XC_D31 11
#define XC_D22 12
#define XC_D13 13
#define XC_D04 14

#define XC_D00000 0
#define XC_D10000 1
#define XC_D01000 2
#define XC_D00100 3
#define XC_D00010 4
#define XC_D00001 5
#define XC_D20000 6
#define XC_D11000 7
#define XC_D10100 8
#define XC_D10010 9
#define XC_D10001 10
#define XC_D02000 11
#define XC_D01100 12
#define XC_D01010 13
#define XC_D01001 14
#define XC_D00200 15
#define XC_D00110 16
#define XC_D00101 17
#define XC_D00020 18
#define XC_D00011 19
#define XC_D00002 20
#define XC_D30000 21
#define XC_D21000 22
#define XC_D20100 23
#define XC_D20010 24
#define XC_D20001 25
#define XC_D12000 26
#define XC_D11100 27
#define XC_D11010 28
#define XC_D11001 29
#define XC_D10200 30
#define XC_D10110 31
#define XC_D10101 32
#define XC_D10020 33
#define XC_D10011 34
#define XC_D10002 35
#define XC_D03000 36
#define XC_D02100 37
#define XC_D02010 38
#define XC_D02001 39
#define XC_D01200 40
#define XC_D01110 41
#define XC_D01101 42
#define XC_D01020 43
#define XC_D01011 44
#define XC_D01002 45
#define XC_D00300 46
#define XC_D00210 47
#define XC_D00201 48
#define XC_D00120 49
#define XC_D00111 50
#define XC_D00102 51
#define XC_D00030 52
#define XC_D00021 53
#define XC_D00012 54
#define XC_D00003 55
#define XC_D40000 56
#define XC_D31000 57
#define XC_D30100 58
#define XC_D30010 59
#define XC_D30001 60
#define XC_D22000 61
#define XC_D21100 62
#define XC_D21010 63
#define XC_D21001 64
#define XC_D20200 65
#define XC_D20110 66
#define XC_D20101 67
#define XC_D20020 68
#define XC_D20011 69
#define XC_D20002 70
#define XC_D13000 71
#define XC_D12100 72
#define XC_D12010 73
#define XC_D12001 74
#define XC_D11200 75
#define XC_D11110 76
#define XC_D11101 77
#define XC_D11020 78
#define XC_D11011 79
#define XC_D11002 80
#define XC_D10300 81
#define XC_D10210 82
#define XC_D10201 83
#define XC_D10120 84
#define XC_D10111 85
#define XC_D10102 86
#define XC_D10030 87
#define XC_D10021 88
#define XC_D10012 89
#define XC_D10003 90
#define XC_D04000 91
#define XC_D03100 92
#define XC_D03010 93
#define XC_D03001 94
#define XC_D02200 95
#define XC_D02110 96
#define XC_D02101 97
#define XC_D02020 98
#define XC_D02011 99
#define XC_D02002 100
#define XC_D01300 101
#define XC_D01210 102
#define XC_D01201 103
#define XC_D01120 104
#define XC_D01111 105
#define XC_D01102 106
#define XC_D01030 107
#define XC_D01021 108
#define XC_D01012 109
#define XC_D01003 110
#define XC_D00400 111
#define XC_D00310 112
#define XC_D00301 113
#define XC_D00220 114
#define XC_D00211 115
#define XC_D00202 116
#define XC_D00130 117
#define XC_D00121 118
#define XC_D00112 119
#define XC_D00103 120
#define XC_D00040 121
#define XC_D00031 122
#define XC_D00022 123
#define XC_D00013 124
#define XC_D00004 125


# pyscf/dft/gks.py

#!/usr/bin/env python
# Copyright 2014-2022 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Generalized Kohn-Sham
'''


import numpy
import scipy.linalg
from pyscf import lib
from pyscf.lib import logger
from pyscf.scf import ghf
from pyscf.dft import rks


def get_veff(ks, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
    '''Coulomb + XC functional

    .. note::
        This function will change the ks object.

    Args:
        ks : an instance of :class:`RKS`
            XC functional are controlled by ks.xc attribute.  Attribute
            ks.grids might be initialized.
        dm : ndarray or list of ndarrays
            A density matrix or a list of density matrices

    Kwargs:
        dm_last : ndarray or a list of ndarrays or 0
            The density matrix baseline.  If not 0, this function computes the
            increment of HF potential w.r.t. the reference HF potential matrix.
        vhf_last : ndarray or a list of ndarrays or 0
            The reference Vxc potential matrix.
        hermi : int
            Whether J, K matrix is hermitian

            | 0 : no hermitian or symmetric
            | 1 : hermitian
            | 2 : anti-hermitian

    Returns:
        matrix Veff = J + Vxc.  Veff can be a list matrices, if the input
        dm is a list of density matrices.
    '''
    if mol is None: mol = ks.mol
    if dm is None: dm = ks.make_rdm1()
    ks.initialize_grids(mol, dm)

    t0 = (logger.process_clock(), logger.perf_counter())

    ground_state = isinstance(dm, numpy.ndarray) and dm.ndim == 2

    if hermi == 2:  # because rho = 0
        n, exc, vxc = 0, 0, 0
    else:
        max_memory = ks.max_memory - lib.current_memory()[0]
        ni = ks._numint
        n, exc, vxc = ni.get_vxc(mol, ks.grids, ks.xc, dm,
                                 hermi=hermi, max_memory=max_memory)
        if ks.nlc != '':
            assert('VV10' in ks.nlc.upper())
            _, enlc, vnlc = ni.get_vxc(mol, ks.nlcgrids, ks.xc+'__'+ks.nlc, dm,
                                       hermi=hermi, max_memory=max_memory)
            exc += enlc
            vxc += vnlc
        logger.debug(ks, 'nelec by numeric integration = %s', n)
        t0 = logger.timer(ks, 'vxc', *t0)

    #enabling range-separated hybrids
    omega, alpha, hyb = ni.rsh_and_hybrid_coeff(ks.xc, spin=mol.spin)

    if abs(hyb) < 1e-10 and abs(alpha) < 1e-10:
        vk = None
        if (ks._eri is None and ks.direct_scf and
            getattr(vhf_last, 'vj', None) is not None):
            ddm = numpy.asarray(dm) - numpy.asarray(dm_last)
            vj = ks.get_j(mol, ddm, hermi)
            vj += vhf_last.vj
        else:
            vj = ks.get_j(mol, dm, hermi)
        vxc += vj
    else:
        if (ks._eri is None and ks.direct_scf and
            getattr(vhf_last, 'vk', None) is not None):
            ddm = numpy.asarray(dm) - numpy.asarray(dm_last)
            vj, vk = ks.get_jk(mol, ddm, hermi)
            vk *= hyb
            if abs(omega) > 1e-10:
                vklr = ks.get_k(mol, ddm, hermi, omega=omega)
                vklr *= (alpha - hyb)
                vk += vklr
            vj += vhf_last.vj
            vk += vhf_last.vk
        else:
            vj, vk = ks.get_jk(mol, dm, hermi)
            vk *= hyb
            if abs(omega) > 1e-10:
                vklr = ks.get_k(mol, dm, hermi, omega=omega)
                vklr *= (alpha - hyb)
                vk += vklr
        vxc += vj - vk

        if ground_state:
            exc -= numpy.einsum('ij,ji', dm, vk).real * .5
    if ground_state:
        ecoul = numpy.einsum('ij,ji', dm, vj).real * .5
    else:
        ecoul = None

    vxc = lib.tag_array(vxc, ecoul=ecoul, exc=exc, vj=vj, vk=vk)
    return vxc

energy_elec = rks.energy_elec


class GKS(rks.KohnShamDFT, ghf.GHF):
    '''Generalized Kohn-Sham'''
    def __init__(self, mol, xc='LDA,VWN'):
        from pyscf.dft.numint2c import NumInt2C
        ghf.GHF.__init__(self, mol)
        rks.KohnShamDFT.__init__(self, xc)
        self._numint = NumInt2C()

    def dump_flags(self, verbose=None):
        ghf.GHF.dump_flags(self, verbose)
        rks.KohnShamDFT.dump_flags(self, verbose)
        return self

    get_veff = get_veff
    energy_elec = energy_elec

    @property
    def collinear(self):
        return self._numint.collinear
    @collinear.setter
    def collinear(self, val):
        self._numint.collinear = val

    def nuc_grad_method(self):
        raise NotImplementedError


if __name__ == '__main__':
    from pyscf import gto
    mol = gto.Mole()
    mol.verbose = 3
    mol.atom = 'H 0 0 0; H 0 0 1; O .5 .6 .2'
    mol.basis = 'ccpvdz'
    mol.build()

    mf = GKS(mol)
    mf.xc = 'b3lyp'
    mf.kernel()

    dm = mf.init_guess_by_1e(mol)
    dm = dm + 0j
    nao = mol.nao_nr()
    numpy.random.seed(12)
    dm[:nao,nao:] = numpy.random.random((nao,nao)) * .1j
    dm[nao:,:nao] = dm[:nao,nao:].T.conj()
    mf.kernel(dm)
    mf.canonicalize(mf.mo_coeff, mf.mo_occ)
    mf.analyze()
    print(mf.spin_square())
    print(mf.e_tot - -76.2760115704274)


# pyscf/dft/gks_symm.py

#!/usr/bin/env python
# Copyright 2014-2022 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Generalized Kohn-Sham
'''

from pyscf.lib import logger
from pyscf.scf import ghf_symm
from pyscf.dft import gks
from pyscf.dft import rks


class GKS(ghf_symm.GHF, rks.KohnShamDFT):
    ''' Restricted Kohn-Sham '''
    def __init__(self, mol, xc='LDA,VWN'):
        from pyscf.dft.numint2c import NumInt2C
        ghf_symm.GHF.__init__(self, mol)
        rks.KohnShamDFT.__init__(self, xc)
        self._numint = NumInt2C()

    def dump_flags(self, verbose=None):
        ghf_symm.GHF.dump_flags(self, verbose)
        rks.KohnShamDFT.dump_flags(self, verbose)
        return self

    get_veff = gks.get_veff
    energy_elec = rks.energy_elec

    def nuc_grad_method(self):
        raise NotImplementedError


if __name__ == '__main__':
    import numpy
    from pyscf import gto
    mol = gto.Mole()
    mol.verbose = 3
    mol.atom = 'H 0 0 0; H 0 0 1; O .5 .6 .2'
    mol.symmetry = True
    mol.basis = 'ccpvdz'
    mol.build()

    mf = GKS(mol)
    mf.xc = 'b3lyp'
    mf.kernel()

    dm = mf.init_guess_by_1e(mol)
    dm = dm + 0j
    nao = mol.nao_nr()
    numpy.random.seed(12)
    dm[:nao,nao:] = numpy.random.random((nao,nao)) * .1j
    dm[nao:,:nao] = dm[:nao,nao:].T.conj()
    mf.kernel(dm)
    mf.canonicalize(mf.mo_coeff, mf.mo_occ)
    mf.analyze()
    print(mf.spin_square())
    print(mf.e_tot - -76.2760114849027)


# pyscf/dft/libxc.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Authors: Qiming Sun <osirpt.sun@gmail.com>
#          Susi Lehtola <susi.lehtola@gmail.com>

'''
XC functional, the interface to libxc
(http://www.tddft.org/programs/octopus/wiki/index.php/Libxc)
'''

import sys
import warnings
import copy
import ctypes
import math
import numpy
from pyscf import lib
from pyscf.dft.xc.utils import remove_dup, format_xc_code
from pyscf import __config__

_itrf = lib.load_library('libxc_itrf')
_itrf.LIBXC_is_lda.restype = ctypes.c_int
_itrf.LIBXC_is_gga.restype = ctypes.c_int
_itrf.LIBXC_is_meta_gga.restype = ctypes.c_int
_itrf.LIBXC_needs_laplacian.restype = ctypes.c_int
_itrf.LIBXC_needs_laplacian.argtypes = [ctypes.c_int]
_itrf.LIBXC_is_hybrid.restype = ctypes.c_int
_itrf.LIBXC_is_cam_rsh.restype = ctypes.c_int
_itrf.LIBXC_max_deriv_order.restype = ctypes.c_int
_itrf.LIBXC_number_of_functionals.restype = ctypes.c_int
_itrf.LIBXC_functional_numbers.argtypes = (numpy.ctypeslib.ndpointer(dtype=numpy.intc, ndim=1, flags=("W", "C", "A")), )
_itrf.LIBXC_functional_name.argtypes = [ctypes.c_int]
_itrf.LIBXC_functional_name.restype = ctypes.c_char_p
_itrf.LIBXC_hybrid_coeff.argtypes = [ctypes.c_int]
_itrf.LIBXC_hybrid_coeff.restype = ctypes.c_double
_itrf.LIBXC_nlc_coeff.argtypes = [ctypes.c_int,ctypes.POINTER(ctypes.c_double)]
_itrf.LIBXC_rsh_coeff.argtypes = [ctypes.c_int,ctypes.POINTER(ctypes.c_double)]

_itrf.LIBXC_version.restype = ctypes.c_char_p
_itrf.LIBXC_reference.restype = ctypes.c_char_p
_itrf.LIBXC_reference_doi.restype = ctypes.c_char_p
_itrf.LIBXC_xc_reference.argtypes = [ctypes.c_int, (ctypes.c_char_p * 8)]

def libxc_version():
    '''Returns the version of libxc'''
    return _itrf.LIBXC_version().decode("UTF-8")

def libxc_reference():
    '''Returns the reference to libxc'''
    return _itrf.LIBXC_reference().decode("UTF-8")

def libxc_reference_doi():
    '''Returns the reference to libxc'''
    return _itrf.LIBXC_reference_doi().decode("UTF-8")

__version__ = libxc_version()
__reference__ = libxc_reference()
__reference_doi__ = libxc_reference_doi()

# Runtime detection of available functionals
dynamic_func = getattr(__config__, 'dft_libxc_dynamic', False)

if dynamic_func:
    def available_libxc_functionals():
        # Number of functionals is
        nfunc = _itrf.LIBXC_number_of_functionals()
        # Get functional numbers
        numbers = numpy.zeros(nfunc, dtype=numpy.intc)
        _itrf.LIBXC_functional_numbers(numbers)
        # Returned array
        return {_itrf.LIBXC_functional_name(x).decode("UTF-8").upper() : x for x in numbers}

    XC = XC_CODES = available_libxc_functionals()
    PROBLEMATIC_XC = dict([])

else:
    # XC dict is generated by
    #import pylibxc
    #for xcname in pylibxc.util.xc_available_functional_names():
    #    f = pylibxc.LibXCFunctional(xcname, 1)
    #    f_id = f.get_number()
    #    ref = f.get_references()
    #    key = f"'{xcname.upper()}'"
    #    print(f"{key:<31s}: {f_id:<3d}, # {ref[0]}")
    XC = XC_CODES = {
        'LDA_C_1D_CSC'                 : 18 , # M. Casula, S. Sorella, and G. Senatore, Phys. Rev. B 74, 245427 (2006)
        'LDA_C_1D_LOOS'                : 26 , # P.-F. Loos, J. Chem. Phys. 138, 064108 (2013)
        'LDA_C_2D_AMGB'                : 15 , # C. Attaccalite, S. Moroni, P. Gori-Giorgi, and G. B. Bachelet, Phys. Rev. Lett. 88, 256601 (2002)
        'LDA_C_2D_PRM'                 : 16 , # S. Pittalis, E. Rasanen, and M. A. L. Marques, Phys. Rev. B 78, 195322 (2008)
        'LDA_C_BR78'                   : 552, # G. B. Jr. and S. M. Rothstein, J. Chem. Phys. 69, 1177 (1978)
        'LDA_C_CHACHIYO'               : 287, # T. Chachiyo, J. Chem. Phys. 145, 021101 (2016)
        'LDA_C_CHACHIYO_MOD'           : 307, # T. Chachiyo and H. Chachiyo, Comput. Theor. Chem. 1172, 112669 (2020)
        'LDA_C_GK72'                   : 578, # R. G. Gordon and Y. S. Kim, J. Chem. Phys. 56, 3122 (1972)
        'LDA_C_GL'                     : 5  , # O. Gunnarsson and B. I. Lundqvist, Phys. Rev. B 13, 4274 (1976)
        'LDA_C_GOMBAS'                 : 24 , # P. Gombas, Fortschr. Phys. 13, 137 (1965)
        'LDA_C_HL'                     : 4  , # L. Hedin and B. I. Lundqvist, J. Phys. C: Solid State Phys. 4, 2064 (1971)
        'LDA_C_KARASIEV'               : 579, # V. V. Karasiev, J. Chem. Phys. 145, 157101 (2016)
        'LDA_C_KARASIEV_MOD'           : 308, # T. Chachiyo and H. Chachiyo, Comput. Theor. Chem. 1172, 112669 (2020)
        'LDA_C_LP96'                   : 289, # S. Liu and R. G. Parr, Phys. Rev. A 53, 2211 (1996)
        'LDA_C_MCWEENY'                : 551, # R. McWeeny, in The New World of Quantum Chemistry, edited by B. Pullman and R. Parr (Reidel, Boston, 1976) pp. 3--31
        'LDA_C_ML1'                    : 22 , # E. I. Proynov and D. R. Salahub, Phys. Rev. B 49, 7874 (1994)
        'LDA_C_ML2'                    : 23 , # E. I. Proynov and D. R. Salahub, Phys. Rev. B 49, 7874 (1994)
        'LDA_C_OB_PW'                  : 14 , # G. Ortiz and P. Ballone, Phys. Rev. B 50, 1391 (1994)
        'LDA_C_OB_PZ'                  : 11 , # G. Ortiz and P. Ballone, Phys. Rev. B 50, 1391 (1994)
        'LDA_C_OW'                     : 574, # P. A. Stewart and P. M. W. Gill, J. Chem. Soc., Faraday Trans. 91, 4337 (1995)
        'LDA_C_OW_LYP'                 : 573, # P. A. Stewart and P. M. W. Gill, J. Chem. Soc., Faraday Trans. 91, 4337 (1995)
        'LDA_C_PK09'                   : 554, # E. Proynov and J. Kong, Phys. Rev. A 79, 014103 (2009)
        'LDA_C_PMGB06'                 : 590, # S. Paziani, S. Moroni, P. Gori-Giorgi, and G. B. Bachelet, Phys. Rev. B 73, 155111 (2006)
        'LDA_C_PW'                     : 12 , # J. P. Perdew and Y. Wang, Phys. Rev. B 45, 13244 (1992)
        'LDA_C_PW_MOD'                 : 13 , # J. P. Perdew and Y. Wang, Phys. Rev. B 45, 13244 (1992), added extra digits to some constants as in the PBE routine (http://dft.rutgers.edu/pubs/PBE.asc)
        'LDA_C_PW_RPA'                 : 25 , # J. P. Perdew and Y. Wang, Phys. Rev. B 45, 13244 (1992)
        'LDA_C_PZ'                     : 9  , # J. P. Perdew and A. Zunger, Phys. Rev. B 23, 5048 (1981)
        'LDA_C_PZ_MOD'                 : 10 , # J. P. Perdew and A. Zunger, Phys. Rev. B 23, 5048 (1981), modified to improve the matching between the low- and high-rs parts
        'LDA_C_RC04'                   : 27 , # S. Ragot and P. Cortona, J. Chem. Phys. 121, 7671 (2004)
        'LDA_C_RPA'                    : 3  , # M. Gell-Mann and K. A. Brueckner, Phys. Rev. 106, 364 (1957)
        'LDA_C_RPW92'                  : 684, # M. Ruggeri, P. L. Rios, and A. Alavi, Phys. Rev. B 98, 161105 (2018)
        'LDA_C_UPW92'                  : 683, # M. Ruggeri, P. L. Rios, and A. Alavi, Phys. Rev. B 98, 161105 (2018)
        'LDA_C_VBH'                    : 17 , # U. von Barth and L. Hedin, J. Phys. C: Solid State Phys. 5, 1629 (1972)
        'LDA_C_VWN'                    : 7  , # S. H. Vosko, L. Wilk, and M. Nusair, Can. J. Phys. 58, 1200 (1980)
        'LDA_C_VWN_1'                  : 28 , # S. H. Vosko, L. Wilk, and M. Nusair, Can. J. Phys. 58, 1200 (1980)
        'LDA_C_VWN_2'                  : 29 , # S. H. Vosko, L. Wilk, and M. Nusair, Can. J. Phys. 58, 1200 (1980)
        'LDA_C_VWN_3'                  : 30 , # S. H. Vosko, L. Wilk, and M. Nusair, Can. J. Phys. 58, 1200 (1980)
        'LDA_C_VWN_4'                  : 31 , # S. H. Vosko, L. Wilk, and M. Nusair, Can. J. Phys. 58, 1200 (1980)
        'LDA_C_VWN_RPA'                : 8  , # S. H. Vosko, L. Wilk, and M. Nusair, Can. J. Phys. 58, 1200 (1980)
        'LDA_C_WIGNER'                 : 2  , # E. Wigner, Trans. Faraday Soc. 34, 678 (1938)
        'LDA_C_XALPHA'                 : 6  , # J. C. Slater, Phys. Rev. 81, 385 (1951)
        'LDA_K_GDS08_WORKER'           : 100001, # L. M. Ghiringhelli and L. Delle Site, Phys. Rev. B 77, 073104 (2008)
        'LDA_K_LP'                     : 51 , # C. Lee and R. G. Parr, Phys. Rev. A 35, 2377 (1987)
        'LDA_K_LP96'                   : 580, # S. Liu and R. G. Parr, Phys. Rev. A 53, 2211 (1996)
        'LDA_K_TF'                     : 50 , # L. H. Thomas, Math. Proc. Cambridge Philos. Soc. 23, 542 (1927)
        'LDA_K_ZLP'                    : 550, # P. Fuentealba and O. Reyes, Chem. Phys. Lett. 232, 31 (1995)
        'LDA_X'                        : 1  , # P. A. M. Dirac, Math. Proc. Cambridge Philos. Soc. 26, 376 (1930)
        'LDA_X_1D_EXPONENTIAL'         : 600, # N. Helbig, J. I. Fuks, M. Casula, M. J. Verstraete, M. A. L. Marques, I. V. Tokatly, and A. Rubio, Phys. Rev. A 83, 032503 (2011)
        'LDA_X_1D_SOFT'                : 21 , # N. Helbig, J. I. Fuks, M. Casula, M. J. Verstraete, M. A. L. Marques, I. V. Tokatly, and A. Rubio, Phys. Rev. A 83, 032503 (2011)
        'LDA_X_2D'                     : 19 , # P. A. M. Dirac, Math. Proc. Cambridge Philos. Soc. 26, 376 (1930)
        'LDA_X_ERF'                    : 546, # P. M. W. Gill, R. D. Adamson, and J. A. Pople, Mol. Phys. 88, 1005 (1996)
        'LDA_X_RAE'                    : 549, # A. Rae, Chem. Phys. Lett. 18, 574 (1973)
        'LDA_X_REL'                    : 532, # A. K. Rajagopal, J. Phys. C: Solid State Phys. 11, L943 (1978)
        'LDA_X_SLOC'                   : 692, # K. Finzel and A. I. Baranov, Int. J. Quantum Chem. 117, 40 (2017)
        'LDA_XC_1D_EHWLRG_1'           : 536, # M. T. Entwistle, M. J. P. Hodgson, J. Wetherell, B. Longstaff, J. D. Ramsden, and R. W. Godby, Phys. Rev. B 94, 205134 (2016)
        'LDA_XC_1D_EHWLRG_2'           : 537, # M. T. Entwistle, M. J. P. Hodgson, J. Wetherell, B. Longstaff, J. D. Ramsden, and R. W. Godby, Phys. Rev. B 94, 205134 (2016)
        'LDA_XC_1D_EHWLRG_3'           : 538, # M. T. Entwistle, M. J. P. Hodgson, J. Wetherell, B. Longstaff, J. D. Ramsden, and R. W. Godby, Phys. Rev. B 94, 205134 (2016)
        'LDA_XC_BN05'                  : 588, # R. Baer and D. Neuhauser, Phys. Rev. Lett. 94, 043002 (2005)
        'LDA_XC_GDSMFB'                : 577, # S. Groth, T. Dornheim, T. Sjostrom, F. D. Malone, W. M. C. Foulkes, and M. Bonitz, Phys. Rev. Lett. 119, 135001 (2017)
        'LDA_XC_KSDT'                  : 259, # V. V. Karasiev, T. Sjostrom, J. Dufty, and S. B. Trickey, Phys. Rev. Lett. 112, 076403 (2014)
        'LDA_XC_LP_A'                  : 547, # C. Lee and R. G. Parr, Phys. Rev. A 42, 193 (1990)
        'LDA_XC_LP_B'                  : 548, # C. Lee and R. G. Parr, Phys. Rev. A 42, 193 (1990)
        'LDA_XC_TETER93'               : 20 , # S. Goedecker, M. Teter, and J. Hutter, Phys. Rev. B 54, 1703 (1996)
        'LDA_XC_TIH'                   : 599, # D. J. Tozer, V. E. Ingamells, and N. C. Handy, J. Chem. Phys. 105, 9200 (1996)
        'LDA_XC_ZLP'                   : 43 , # Q. Zhao, M. Levy, and R. G. Parr, Phys. Rev. A 47, 918 (1993)
        'HYB_LDA_XC_CAM_LDA0'          : 178, # M. A. Mosquera, C. H. Borca, M. A. Ratner, and G. C. Schatz, J. Phys. Chem. A 120, 1605 (2016)
        'HYB_LDA_XC_LDA0'              : 177, # P. Rinke, A. Schleife, E. Kioupakis, A. Janotti, C. Rodl, F. Bechstedt, M. Scheffler, and C. G. Van de Walle, Phys. Rev. Lett. 108, 126404 (2012)
        'GGA_C_ACGGA'                  : 39 , # A. Cancio, G. P. Chen, B. T. Krull, and K. Burke, J. Chem. Phys. 149, 084116 (2018)
        'GGA_C_ACGGAP'                 : 176, # A. Cancio, G. P. Chen, B. T. Krull, and K. Burke, J. Chem. Phys. 149, 084116 (2018)
        'GGA_C_AM05'                   : 135, # R. Armiento and A. E. Mattsson, Phys. Rev. B 72, 085108 (2005)
        'GGA_C_APBE'                   : 186, # L. A. Constantin, E. Fabiano, S. Laricchia, and F. Della Sala, Phys. Rev. Lett. 106, 186406 (2011)
        'GGA_C_BMK'                    : 280, # A. D. Boese and J. M. L. Martin, J. Chem. Phys. 121, 3405 (2004)
        'GGA_C_CCDF'                   : 313, # J. T. Margraf, C. Kunkel, and K. Reuter, J. Chem. Phys. 150, 244116 (2019)
        'GGA_C_CHACHIYO'               : 309, # T. Chachiyo and H. Chachiyo, Comput. Theor. Chem. 1172, 112669 (2020)
        'GGA_C_CS1'                    : 565, # N. C. Handy and A. J. Cohen, J. Chem. Phys. 116, 5411 (2002)
        'GGA_C_FT97'                   : 88 , # M. Filatov and W. Thiel, Int. J. Quantum Chem. 62, 603 (1997)
        'GGA_C_GAM'                    : 33 , # H. S. Yu, W. Zhang, P. Verma, X. He, and D. G. Truhlar, Phys. Chem. Chem. Phys. 17, 12146 (2015)
        'GGA_C_GAPC'                   : 555, # E. Fabiano, P. E. Trevisanutto, A. Terentjevs, and L. A. Constantin, J. Chem. Theory Comput. 10, 2016 (2014)
        'GGA_C_GAPLOC'                 : 556, # E. Fabiano, P. E. Trevisanutto, A. Terentjevs, and L. A. Constantin, J. Chem. Theory Comput. 10, 2016 (2014)
        'GGA_C_HCTH_A'                 : 97 , # F. A. Hamprecht, A. J. Cohen, D. J. Tozer, and N. C. Handy, J. Chem. Phys. 109, 6264 (1998)
        'GGA_C_HYB_TAU_HCTH'           : 283, # A. D. Boese and N. C. Handy, J. Chem. Phys. 116, 9559 (2002)
        'GGA_C_LM'                     : 137, # D. C. Langreth and M. J. Mehl, Phys. Rev. Lett. 47, 446 (1981)
        'GGA_C_LYP'                    : 131, # C. Lee, W. Yang, and R. G. Parr, Phys. Rev. B 37, 785 (1988)
        'GGA_C_MGGAC'                  : 712, # B. Patra, S. Jana, L. A. Constantin, and P. Samal, Phys. Rev. B 100, 155140 (2019)
        'GGA_C_N12'                    : 80 , # R. Peverati and D. G. Truhlar, J. Chem. Theory Comput. 8, 2310 (2012)
        'GGA_C_N12_SX'                 : 79 , # R. Peverati and D. G. Truhlar, Phys. Chem. Chem. Phys. 14, 16187 (2012)
        'GGA_C_OP_B88'                 : 87 , # T. Tsuneda, T. Suzumura, and K. Hirao, J. Chem. Phys. 110, 10664 (1999)
        'GGA_C_OP_G96'                 : 85 , # T. Tsuneda, T. Suzumura, and K. Hirao, J. Chem. Phys. 110, 10664 (1999)
        'GGA_C_OP_PBE'                 : 86 , # T. Tsuneda, T. Suzumura, and K. Hirao, J. Chem. Phys. 110, 10664 (1999)
        'GGA_C_OP_PW91'                : 262, # T. Tsuneda, T. Suzumura, and K. Hirao, J. Chem. Phys. 110, 10664 (1999)
        'GGA_C_OP_XALPHA'              : 84 , # T. Tsuneda, T. Suzumura, and K. Hirao, J. Chem. Phys. 110, 10664 (1999)
        'GGA_C_OPTC'                   : 200, # A. J. Cohen and N. C. Handy, Mol. Phys. 99, 607 (2001)
        'GGA_C_P86'                    : 132, # J. P. Perdew, Phys. Rev. B 33, 8822 (1986)
        'GGA_C_P86_FT'                 : 217, # J. P. Perdew, Phys. Rev. B 33, 8822 (1986)
        'GGA_C_P86VWN'                 : 252, # J. P. Perdew, Phys. Rev. B 33, 8822 (1986)
        'GGA_C_P86VWN_FT'              : 253, # J. P. Perdew, Phys. Rev. B 33, 8822 (1986)
        'GGA_C_PBE'                    : 130, # J. P. Perdew, K. Burke, and M. Ernzerhof, Phys. Rev. Lett. 77, 3865 (1996)
        'GGA_C_PBE_JRGX'               : 138, # L. S. Pedroza, A. J. R. da Silva, and K. Capelle, Phys. Rev. B 79, 201106 (2009)
        'GGA_C_PBE_MOL'                : 272, # J. M. del Campo, J. L. Gazquez, S. B. Trickey, and A. Vela, J. Chem. Phys. 136, 104108 (2012)
        'GGA_C_PBE_SOL'                : 133, # J. P. Perdew, A. Ruzsinszky, G. I. Csonka, O. A. Vydrov, G. E. Scuseria, L. A. Constantin, X. Zhou, and K. Burke, Phys. Rev. Lett. 100, 136406 (2008)
        'GGA_C_PBE_VWN'                : 216, # E. Kraisler, G. Makov, and I. Kelson, Phys. Rev. A 82, 042516 (2010)
        'GGA_C_PBEFE'                  : 258, # R. Sarmiento-Perez, S. Botti, and M. A. L. Marques, J. Chem. Theory Comput. 11, 3844 (2015)
        'GGA_C_PBEINT'                 : 62 , # E. Fabiano, L. A. Constantin, and F. Della Sala, Phys. Rev. B 82, 113104 (2010)
        'GGA_C_PBELOC'                 : 246, # L. A. Constantin, E. Fabiano, and F. Della Sala, Phys. Rev. B 86, 035130 (2012)
        'GGA_C_PW91'                   : 134, # J. P. Perdew, in Proceedings of the 75. WE-Heraeus-Seminar and 21st Annual International Symposium on Electronic Structure of Solids, edited by P. Ziesche and H. Eschrig (Akademie Verlag, Berlin, 1991) p. 11
        'GGA_C_Q2D'                    : 47 , # L. Chiodo, L. A. Constantin, E. Fabiano, and F. Della Sala, Phys. Rev. Lett. 108, 126402 (2012)
        'GGA_C_REGTPSS'                : 83 , # J. P. Perdew, A. Ruzsinszky, G. I. Csonka, L. A. Constantin, and J. Sun, Phys. Rev. Lett. 103, 026403 (2009)
        'GGA_C_REVTCA'                 : 99 , # V. Tognetti, P. Cortona, and C. Adamo, Chem. Phys. Lett. 460, 536 (2008)
        'GGA_C_RGE2'                   : 143, # A. Ruzsinszky, G. I. Csonka, and G. E. Scuseria, J. Chem. Theory Comput. 5, 763 (2009)
        'GGA_C_SCAN_E0'                : 553, # J. Sun, A. Ruzsinszky, and J. P. Perdew, Phys. Rev. Lett. 115, 036402 (2015)
        'GGA_C_SG4'                    : 534, # L. A. Constantin, A. Terentjevs, F. Della Sala, P. Cortona, and E. Fabiano, Phys. Rev. B 93, 045126 (2016)
        'GGA_C_SOGGA11'                : 152, # R. Peverati, Y. Zhao, and D. G. Truhlar, J. Phys. Chem. Lett. 2, 1991 (2011)
        'GGA_C_SOGGA11_X'              : 159, # R. Peverati and D. G. Truhlar, J. Chem. Phys. 135, 191102 (2011)
        'GGA_C_SPBE'                   : 89 , # M. Swart, M. Sola, and F. M. Bickelhaupt, J. Chem. Phys. 131, 094103 (2009)
        'GGA_C_TAU_HCTH'               : 281, # A. D. Boese and N. C. Handy, J. Chem. Phys. 116, 9559 (2002)
        'GGA_C_TCA'                    : 100, # V. Tognetti, P. Cortona, and C. Adamo, J. Chem. Phys. 128, 034101 (2008)
        'GGA_C_TM_LYP'                 : 559, # A. J. Thakkar and S. P. McCarthy, J. Chem. Phys. 131, 134109 (2009)
        'GGA_C_TM_PBE'                 : 560, # A. J. Thakkar and S. P. McCarthy, J. Chem. Phys. 131, 134109 (2009)
        'GGA_C_W94'                    : 561, # L. C. Wilson, Chem. Phys. 181, 337 (1994)
        'GGA_C_WI'                     : 148, # L. C. Wilson and S. Ivanov, Int. J. Quantum Chem. 69, 523 (1998)
        'GGA_C_WI0'                    : 153, # L. C. Wilson and S. Ivanov, Int. J. Quantum Chem. 69, 523 (1998)
        'GGA_C_WL'                     : 147, # L. C. Wilson and M. Levy, Phys. Rev. B 41, 12930 (1990)
        'GGA_C_XPBE'                   : 136, # X. Xu and W. A. Goddard, J. Chem. Phys. 121, 4068 (2004)
        'GGA_C_ZPBEINT'                : 61 , # L. A. Constantin, E. Fabiano, and F. Della Sala, Phys. Rev. B 84, 233103 (2011)
        'GGA_C_ZPBESOL'                : 63 , # L. A. Constantin, E. Fabiano, and F. Della Sala, Phys. Rev. B 84, 233103 (2011)
        'GGA_C_ZVPBEINT'               : 557, # L. A. Constantin, E. Fabiano, and F. D. Sala, J. Chem. Phys. 137, 194105 (2012)
        'GGA_C_ZVPBELOC'               : 606, # E. Fabiano, L. A. Constantin, P. Cortona, and F. Della Sala, J. Chem. Theory Comput. 11, 122 (2015)
        'GGA_C_ZVPBESOL'               : 558, # L. A. Constantin, E. Fabiano, and F. D. Sala, J. Chem. Phys. 137, 194105 (2012)
        'GGA_K_ABSP1'                  : 506, # P. K. Acharya, L. J. Bartolotti, S. B. Sears, and R. G. Parr, Proc. Natl. Acad. Sci. U. S. A. 77, 6978 (1980)
        'GGA_K_ABSP2'                  : 507, # P. K. Acharya, L. J. Bartolotti, S. B. Sears, and R. G. Parr, Proc. Natl. Acad. Sci. U. S. A. 77, 6978 (1980)
        'GGA_K_ABSP3'                  : 277, # P. K. Acharya, L. J. Bartolotti, S. B. Sears, and R. G. Parr, Proc. Natl. Acad. Sci. U. S. A. 77, 6978 (1980)
        'GGA_K_ABSP4'                  : 278, # P. K. Acharya, L. J. Bartolotti, S. B. Sears, and R. G. Parr, Proc. Natl. Acad. Sci. U. S. A. 77, 6978 (1980)
        'GGA_K_APBE'                   : 185, # L. A. Constantin, E. Fabiano, S. Laricchia, and F. Della Sala, Phys. Rev. Lett. 106, 186406 (2011)
        'GGA_K_APBEINT'                : 54 , # S. Laricchia, E. Fabiano, L. A. Constantin, and F. Della Sala, J. Chem. Theory Comput. 7, 2439 (2011)
        'GGA_K_BALTIN'                 : 504, # R. Baltin, Z. Naturforsch. A 27, 1176 (1972)
        'GGA_K_DK'                     : 516, # A. E. DePristo and J. D. Kress, Phys. Rev. A 35, 438 (1987)
        'GGA_K_ERNZERHOF'              : 520, # M. Ernzerhof, J. Mol. Struct.: THEOCHEM 501--502, 59 (2000)
        'GGA_K_EXP4'                   : 597, # V. V. Karasiev, S. B. Trickey, and F. E. Harris, J. Comput.-Aided Mater. Des. 13, 111 (2006)
        'GGA_K_FR_B88'                 : 514, # P. Fuentealba and O. Reyes, Chem. Phys. Lett. 232, 31 (1995)
        'GGA_K_FR_PW86'                : 515, # P. Fuentealba and O. Reyes, Chem. Phys. Lett. 232, 31 (1995)
        'GGA_K_GDS08'                  : 591, # L. M. Ghiringhelli and L. Delle Site, Phys. Rev. B 77, 073104 (2008)
        'GGA_K_GE2'                    : 501, # A. S. Kompaneets and E. S. Pavlovskii, Zh. Eksp. Teor. Fiz. 31, 427 (1956), [J. Exp. Theor. Phys. 4, 328 (1957)]
        'GGA_K_GHDS10'                 : 592, # L. M. Ghiringhelli, I. P. Hamilton, and L. D. Site, J. Chem. Phys. 132, 014106 (2010)
        'GGA_K_GHDS10R'                : 593, # S. B. Trickey, V. V. Karasiev, and A. Vela, Phys. Rev. B 84, 075146 (2011)
        'GGA_K_GOLDEN'                 : 502, # S. Golden, Phys. Rev. 105, 604 (1957)
        'GGA_K_GP85'                   : 510, # S. K. Ghosh and R. G. Parr, J. Chem. Phys. 82, 3307 (1985)
        'GGA_K_GR'                     : 508, # J. L. Gazquez and J. Robles, J. Chem. Phys. 76, 1467 (1982)
        'GGA_K_LC94'                   : 521, # A. Lembarki and H. Chermette, Phys. Rev. A 50, 5328 (1994)
        'GGA_K_LGAP'                   : 620, # L. A. Constantin, E. Fabiano, S. Smiga, and F. Della Sala, Phys. Rev. B 95, 115153 (2017)
        'GGA_K_LGAP_GE'                : 633, # L. A. Constantin, E. Fabiano, S. Smiga, and F. Della Sala, Phys. Rev. B 95, 115153 (2017)
        'GGA_K_LIEB'                   : 505, # E. H. Lieb, Rev. Mod. Phys. 53, 603 (1981)
        'GGA_K_LKT'                    : 613, # K. Luo, V. V. Karasiev, and S. B. Trickey, Phys. Rev. B 98, 041111 (2018)
        'GGA_K_LLP'                    : 522, # H. Lee, C. Lee, and R. G. Parr, Phys. Rev. A 44, 768 (1991)
        'GGA_K_LUDENA'                 : 509, # E. V. Ludena, in Cond. Matt. Theor., Vol. 1, edited by F. B. Malik (Plenum, New York, 1986) p. 183
        'GGA_K_MEYER'                  : 57 , # A. Meyer, G. C. Wang, and W. H. Young, Z. Naturforsch. A 31, 898 (1976)
        'GGA_K_OL1'                    : 512, # H. Ou-Yang and M. Levy, Int. J. Quantum Chem. 40, 379 (1991)
        'GGA_K_OL2'                    : 513, # H. Ou-Yang and M. Levy, Int. J. Quantum Chem. 40, 379 (1991)
        'GGA_K_PBE2'                   : 616, # V. V. Karasiev, S. B. Trickey, and F. E. Harris, J. Comput.-Aided Mater. Des. 13, 111 (2006)
        'GGA_K_PBE3'                   : 595, # V. V. Karasiev, S. B. Trickey, and F. E. Harris, J. Comput.-Aided Mater. Des. 13, 111 (2006)
        'GGA_K_PBE4'                   : 596, # V. V. Karasiev, S. B. Trickey, and F. E. Harris, J. Comput.-Aided Mater. Des. 13, 111 (2006)
        'GGA_K_PEARSON'                : 511, # D. J. Lacks and R. G. Gordon, J. Chem. Phys. 100, 4446 (1994)
        'GGA_K_PERDEW'                 : 517, # J. P. Perdew, Phys. Lett. A 165, 79 (1992)
        'GGA_K_PG1'                    : 219, # L. A. Constantin, E. Fabiano, and F. Della Sala, J. Phys. Chem. Lett. 9, 4385 (2018), pMID: 30019904
        'GGA_K_RATIONAL_P'             : 218, # J. Lehtomaki and O. Lopez-Acevedo, Phys. Rev. B 100, 165111 (2019)
        'GGA_K_REVAPBE'                : 55 , # L. A. Constantin, E. Fabiano, S. Laricchia, and F. Della Sala, Phys. Rev. Lett. 106, 186406 (2011)
        'GGA_K_REVAPBEINT'             : 53 , # S. Laricchia, E. Fabiano, L. A. Constantin, and F. Della Sala, J. Chem. Theory Comput. 7, 2439 (2011)
        'GGA_K_TFVW'                   : 52 , # C. F. von Weizsacker, Z. Phys. 96, 431 (1935)
        'GGA_K_TFVW_OPT'               : 635, # L. A. Espinosa Leal, A. Karpenko, M. A. Caro, and O. Lopez-Acevedo, Phys. Chem. Chem. Phys. 17, 31463 (2015)
        'GGA_K_THAKKAR'                : 523, # A. J. Thakkar, Phys. Rev. A 46, 6920 (1992)
        'GGA_K_TKVLN'                  : 594, # S. B. Trickey, V. V. Karasiev, and A. Vela, Phys. Rev. B 84, 075146 (2011)
        'GGA_K_TW1'                    : 187, # F. Tran and T. A. Wesolowski, Int. J. Quantum Chem. 89, 441 (2002)
        'GGA_K_TW2'                    : 188, # F. Tran and T. A. Wesolowski, Int. J. Quantum Chem. 89, 441 (2002)
        'GGA_K_TW3'                    : 189, # F. Tran and T. A. Wesolowski, Int. J. Quantum Chem. 89, 441 (2002)
        'GGA_K_TW4'                    : 190, # F. Tran and T. A. Wesolowski, Int. J. Quantum Chem. 89, 441 (2002)
        'GGA_K_VJKS'                   : 519, # L. Vitos, B. Johansson, J. Kollar, and H. L. Skriver, Phys. Rev. A 61, 052511 (2000)
        'GGA_K_VSK'                    : 518, # L. Vitos, H. L. Skriver, and J. Kollar, Phys. Rev. B 57, 12611 (1998)
        'GGA_K_VT84F'                  : 619, # V. V. Karasiev, D. Chakraborty, O. A. Shukruto, and S. B. Trickey, Phys. Rev. B 88, 161108 (2013)
        'GGA_K_VW'                     : 500, # C. F. von Weizsacker, Z. Phys. 96, 431 (1935)
        'GGA_K_YT65'                   : 503, # K. Yonei and Y. Tomishima, J. Phys. Soc. Jpn. 20, 1051 (1965)
        'GGA_X_2D_B86'                 : 128, # J. G. Vilhena, E. Rasanen, M. A. L. Marques, and S. Pittalis, J. Chem. Theory Comput. 10, 1837 (2014)
        'GGA_X_2D_B86_MGC'             : 124, # S. Pittalis, E. Rasanen, J. G. Vilhena, and M. A. L. Marques, Phys. Rev. A 79, 012503 (2009)
        'GGA_X_2D_B88'                 : 127, # J. G. Vilhena, E. Rasanen, M. A. L. Marques, and S. Pittalis, J. Chem. Theory Comput. 10, 1837 (2014)
        'GGA_X_2D_PBE'                 : 129, # J. G. Vilhena, E. Rasanen, M. A. L. Marques, and S. Pittalis, J. Chem. Theory Comput. 10, 1837 (2014)
        'GGA_X_AIRY'                   : 192, # L. A. Constantin, A. Ruzsinszky, and J. P. Perdew, Phys. Rev. B 80, 035125 (2009)
        'GGA_X_AK13'                   : 56 , # R. Armiento and S. Kummel, Phys. Rev. Lett. 111, 036402 (2013)
        'GGA_X_AM05'                   : 120, # R. Armiento and A. E. Mattsson, Phys. Rev. B 72, 085108 (2005)
        'GGA_X_APBE'                   : 184, # L. A. Constantin, E. Fabiano, S. Laricchia, and F. Della Sala, Phys. Rev. Lett. 106, 186406 (2011)
        'GGA_X_B86'                    : 103, # A. D. Becke, J. Chem. Phys. 84, 4524 (1986)
        'GGA_X_B86_MGC'                : 105, # A. D. Becke, J. Chem. Phys. 84, 4524 (1986)
        'GGA_X_B86_R'                  : 41 , # I. Hamada, Phys. Rev. B 89, 121103 (2014)
        'GGA_X_B88'                    : 106, # A. D. Becke, Phys. Rev. A 38, 3098 (1988)
        'GGA_X_B88_6311G'              : 179, # J. M. Ugalde, C. Sarasola, and M. Aguado, J. Phys. B: At., Mol. Opt. Phys. 27, 423 (1994)
        'GGA_X_B88M'                   : 570, # E. Proynov, H. Chermette, and D. R. Salahub, J. Chem. Phys. 113, 10013 (2000)
        'GGA_X_BAYESIAN'               : 125, # J. J. Mortensen, K. Kaasbjerg, S. L. Frederiksen, J. K. Norskov, J. P. Sethna, and K. W. Jacobsen, Phys. Rev. Lett. 95, 216401 (2005)
        'GGA_X_BCGP'                   : 38 , # K. Burke, A. Cancio, T. Gould, and S. Pittalis, ArXiv e-prints (2014), arXiv:1409.4834 [cond-mat.mtrl-sci]
        'GGA_X_BEEFVDW'                : 285, # J. Wellendorff, K. T. Lundgaard, A. Mogelhoj, V. Petzold, D. D. Landis, J. K. Norskov, T. Bligaard, and K. W. Jacobsen, Phys. Rev. B 85, 235149 (2012)
        'GGA_X_BPCCAC'                 : 98 , # E. Bremond, D. Pilard, I. Ciofini, H. Chermette, C. Adamo, and P. Cortona, Theor. Chem. Acc. 131, 1184 (2012)
        'GGA_X_C09X'                   : 158, # V. R. Cooper, Phys. Rev. B 81, 161104 (2010)
        'GGA_X_CAP'                    : 270, # J. Carmona-Espindola, J. L. Gazquez, A. Vela, and S. B. Trickey, J. Chem. Phys. 142, 054105 (2015)
        'GGA_X_CHACHIYO'               : 298, # T. Chachiyo and H. Chachiyo, Molecules 25, 3485 (2020)
        'GGA_X_DK87_R1'                : 111, # A. E. DePristo and J. D. Kress, J. Chem. Phys. 86, 1425 (1987)
        'GGA_X_DK87_R2'                : 112, # A. E. DePristo and J. D. Kress, J. Chem. Phys. 86, 1425 (1987)
        'GGA_X_EB88'                   : 271, # P. Elliott and K. Burke, Can. J. Chem. 87, 1485 (2009)
        'GGA_X_ECMV92'                 : 215, # E. Engel, J. A. Chevary, L. D. Macdonald, and S. H. Vosko, Z. Phys. D: At., Mol. Clusters 23, 7 (1992)
        'GGA_X_EV93'                   : 35 , # E. Engel and S. H. Vosko, Phys. Rev. B 47, 13164 (1993)
        'GGA_X_FD_LB94'                : 604, # A. P. Gaiduk and V. N. Staroverov, Phys. Rev. A 83, 012509 (2011)
        'GGA_X_FD_REVLB94'             : 605, # A. P. Gaiduk and V. N. Staroverov, Phys. Rev. A 83, 012509 (2011)
        'GGA_X_FT97_A'                 : 114, # M. Filatov and W. Thiel, Mol. Phys. 91, 847 (1997)
        'GGA_X_FT97_B'                 : 115, # M. Filatov and W. Thiel, Mol. Phys. 91, 847 (1997)
        'GGA_X_G96'                    : 107, # P. M. W. Gill, Mol. Phys. 89, 433 (1996)
        'GGA_X_GAM'                    : 32 , # H. S. Yu, W. Zhang, P. Verma, X. He, and D. G. Truhlar, Phys. Chem. Chem. Phys. 17, 12146 (2015)
        'GGA_X_GG99'                   : 535, # A. T. Gilbert and P. M. Gill, Chem. Phys. Lett. 312, 511 (1999)
        'GGA_X_HCTH_A'                 : 34 , # F. A. Hamprecht, A. J. Cohen, D. J. Tozer, and N. C. Handy, J. Chem. Phys. 109, 6264 (1998)
        'GGA_X_HERMAN'                 : 104, # F. Herman, J. P. V. Dyke, and I. B. Ortenburger, Phys. Rev. Lett. 22, 807 (1969)
        'GGA_X_HJS_B88'                : 527, # T. M. Henderson, B. G. Janesko, and G. E. Scuseria, J. Chem. Phys. 128, 194105 (2008)
        'GGA_X_HJS_B88_V2'             : 46 , # E. Weintraub, T. M. Henderson, and G. E. Scuseria, J. Chem. Theory Comput. 5, 754 (2009)
        'GGA_X_HJS_B97X'               : 528, # T. M. Henderson, B. G. Janesko, and G. E. Scuseria, J. Chem. Phys. 128, 194105 (2008)
        'GGA_X_HJS_PBE'                : 525, # T. M. Henderson, B. G. Janesko, and G. E. Scuseria, J. Chem. Phys. 128, 194105 (2008)
        'GGA_X_HJS_PBE_SOL'            : 526, # T. M. Henderson, B. G. Janesko, and G. E. Scuseria, J. Chem. Phys. 128, 194105 (2008)
        'GGA_X_HTBS'                   : 191, # P. Haas, F. Tran, P. Blaha, and K. Schwarz, Phys. Rev. B 83, 205117 (2011)
        'GGA_X_ITYH'                   : 529, # H. Iikura, T. Tsuneda, T. Yanai, and K. Hirao, J. Chem. Phys. 115, 3540 (2001)
        'GGA_X_ITYH_OPTX'              : 622, # N. C. Handy and A. J. Cohen, Mol. Phys. 99, 403 (2001)
        'GGA_X_ITYH_PBE'               : 623, # J. P. Perdew, K. Burke, and M. Ernzerhof, Phys. Rev. Lett. 77, 3865 (1996)
        'GGA_X_KGG99'                  : 544, # A. T. Gilbert and P. M. Gill, Chem. Phys. Lett. 312, 511 (1999)
        'GGA_X_KT1'                    : 145, # T. W. Keal and D. J. Tozer, J. Chem. Phys. 119, 3015 (2003)
        'GGA_X_LAG'                    : 193, # L. Vitos, B. Johansson, J. Kollar, and H. L. Skriver, Phys. Rev. B 62, 10046 (2000)
        'GGA_X_LAMBDA_CH_N'            : 44 , # M. M. Odashima, K. Capelle, and S. B. Trickey, J. Chem. Theory Comput. 5, 798 (2009)
        'GGA_X_LAMBDA_LO_N'            : 45 , # M. M. Odashima, K. Capelle, and S. B. Trickey, J. Chem. Theory Comput. 5, 798 (2009)
        'GGA_X_LAMBDA_OC2_N'           : 40 , # M. M. Odashima, K. Capelle, and S. B. Trickey, J. Chem. Theory Comput. 5, 798 (2009)
        'GGA_X_LB'                     : 160, # R. van Leeuwen and E. J. Baerends, Phys. Rev. A 49, 2421 (1994)
        'GGA_X_LBM'                    : 182, # P. R. T. Schipper, O. V. Gritsenko, S. J. A. van Gisbergen, and E. J. Baerends, J. Chem. Phys. 112, 1344 (2000)
        'GGA_X_LG93'                   : 113, # D. J. Lacks and R. G. Gordon, Phys. Rev. A 47, 4681 (1993)
        'GGA_X_LSPBE'                  : 168, # J. C. Pacheco-Kato, J. M. del Campo, J. L. Gazquez, S. Trickey, and A. Vela, Chem. Phys. Lett. 651, 268 (2016)
        'GGA_X_LSRPBE'                 : 169, # J. C. Pacheco-Kato, J. M. del Campo, J. L. Gazquez, S. Trickey, and A. Vela, Chem. Phys. Lett. 651, 268 (2016)
        'GGA_X_LV_RPW86'               : 58 , # K. Berland and P. Hyldgaard, Phys. Rev. B 89, 035412 (2014)
        'GGA_X_MB88'                   : 149, # V. Tognetti and C. Adamo, J. Phys. Chem. A 113, 14415 (2009)
        'GGA_X_MPBE'                   : 122, # C. Adamo and V. Barone, J. Chem. Phys. 116, 5933 (2002)
        'GGA_X_MPW91'                  : 119, # C. Adamo and V. Barone, J. Chem. Phys. 108, 664 (1998)
        'GGA_X_N12'                    : 82 , # R. Peverati and D. G. Truhlar, J. Chem. Theory Comput. 8, 2310 (2012)
        'GGA_X_NCAP'                   : 180, # J. Carmona-Espindola, J. L. Gazquez, A. Vela, and S. B. Trickey, J. Chem. Theory Comput. 15, 303 (2019)
        'GGA_X_OL2'                    : 183, # P. Fuentealba and O. Reyes, Chem. Phys. Lett. 232, 31 (1995)
        'GGA_X_OPTB86B_VDW'            : 171, # J. Klimes, D. R. Bowler, and A. Michaelides, Phys. Rev. B 83, 195131 (2011)
        'GGA_X_OPTB88_VDW'             : 139, # J. Klimes, D. R. Bowler, and A. Michaelides, J. Phys.: Condens. Matter 22, 022201 (2010)
        'GGA_X_OPTPBE_VDW'             : 141, # J. Klimes, D. R. Bowler, and A. Michaelides, J. Phys.: Condens. Matter 22, 022201 (2010)
        'GGA_X_OPTX'                   : 110, # N. C. Handy and A. J. Cohen, Mol. Phys. 99, 403 (2001)
        'GGA_X_PBE'                    : 101, # J. P. Perdew, K. Burke, and M. Ernzerhof, Phys. Rev. Lett. 77, 3865 (1996)
        'GGA_X_PBE_JSJR'               : 126, # L. S. Pedroza, A. J. R. da Silva, and K. Capelle, Phys. Rev. B 79, 201106 (2009)
        'GGA_X_PBE_MOL'                : 49 , # J. M. del Campo, J. L. Gazquez, S. B. Trickey, and A. Vela, J. Chem. Phys. 136, 104108 (2012)
        'GGA_X_PBE_R'                  : 102, # Y. Zhang and W. Yang, Phys. Rev. Lett. 80, 890 (1998)
        'GGA_X_PBE_SOL'                : 116, # J. P. Perdew, A. Ruzsinszky, G. I. Csonka, O. A. Vydrov, G. E. Scuseria, L. A. Constantin, X. Zhou, and K. Burke, Phys. Rev. Lett. 100, 136406 (2008)
        'GGA_X_PBE_TCA'                : 59 , # V. Tognetti, P. Cortona, and C. Adamo, Chem. Phys. Lett. 460, 536 (2008)
        'GGA_X_PBEA'                   : 121, # G. K. H. Madsen, Phys. Rev. B 75, 195108 (2007)
        'GGA_X_PBEFE'                  : 265, # R. Sarmiento-Perez, S. Botti, and M. A. L. Marques, J. Chem. Theory Comput. 11, 3844 (2015)
        'GGA_X_PBEINT'                 : 60 , # E. Fabiano, L. A. Constantin, and F. Della Sala, Phys. Rev. B 82, 113104 (2010)
        'GGA_X_PBEK1_VDW'              : 140, # J. Klimes, D. R. Bowler, and A. Michaelides, J. Phys.: Condens. Matter 22, 022201 (2010)
        'GGA_X_PBEPOW'                 : 539, # Eric Bremond, J. Chem. Phys. 145, 244102 (2016)
        'GGA_X_PBETRANS'               : 291, # Eric Bremond, I. Ciofini, and C. Adamo, Mol. Phys. 114, 1059 (2016)
        'GGA_X_PW86'                   : 108, # J. P. Perdew and W. Yue, Phys. Rev. B 33, 8800 (1986)
        'GGA_X_PW91'                   : 109, # J. P. Perdew, in Proceedings of the 75. WE-Heraeus-Seminar and 21st Annual International Symposium on Electronic Structure of Solids, edited by P. Ziesche and H. Eschrig (Akademie Verlag, Berlin, 1991) p. 11
        'GGA_X_PW91_MOD'               : 316, # J. P. Perdew, in Proceedings of the 75. WE-Heraeus-Seminar and 21st Annual International Symposium on Electronic Structure of Solids, edited by P. Ziesche and H. Eschrig (Akademie Verlag, Berlin, 1991) p. 11
        'GGA_X_Q2D'                    : 48 , # L. Chiodo, L. A. Constantin, E. Fabiano, and F. Della Sala, Phys. Rev. Lett. 108, 126402 (2012)
        'GGA_X_REVSSB_D'               : 312, # M. Swart, M. Sola, and F. M. Bickelhaupt, J. Comput. Chem. 32, 1117 (2011)
        'GGA_X_RGE2'                   : 142, # A. Ruzsinszky, G. I. Csonka, and G. E. Scuseria, J. Chem. Theory Comput. 5, 763 (2009)
        'GGA_X_RPBE'                   : 117, # B. Hammer, L. B. Hansen, and J. K. Norskov, Phys. Rev. B 59, 7413 (1999)
        'GGA_X_RPW86'                  : 144, # E. D. Murray, K. Lee, and D. C. Langreth, J. Chem. Theory Comput. 5, 2754 (2009)
        'GGA_X_S12G'                   : 495, # M. Swart, Chem. Phys. Lett. 580, 166 (2013)
        'GGA_X_SFAT'                   : 530, # A. Savin and H.-J. Flad, Int. J. Quantum Chem. 56, 327 (1995)
        'GGA_X_SFAT_PBE'               : 601, # A. Savin and H.-J. Flad, Int. J. Quantum Chem. 56, 327 (1995)
        'GGA_X_SG4'                    : 533, # L. A. Constantin, A. Terentjevs, F. Della Sala, P. Cortona, and E. Fabiano, Phys. Rev. B 93, 045126 (2016)
        'GGA_X_SOGGA'                  : 150, # Y. Zhao and D. G. Truhlar, J. Chem. Phys. 128, 184109 (2008)
        'GGA_X_SOGGA11'                : 151, # R. Peverati, Y. Zhao, and D. G. Truhlar, J. Phys. Chem. Lett. 2, 1991 (2011)
        'GGA_X_SSB'                    : 91 , # M. Swart, M. Sola, and F. M. Bickelhaupt, J. Chem. Phys. 131, 094103 (2009)
        'GGA_X_SSB_D'                  : 92 , # M. Swart, M. Sola, and F. M. Bickelhaupt, J. Chem. Phys. 131, 094103 (2009)
        'GGA_X_SSB_SW'                 : 90 , # M. Swart, M. Sola, and F. M. Bickelhaupt, J. Comput. Methods Sci. Eng. 9, 69 (2009)
        'GGA_X_VMT84_GE'               : 68 , # A. Vela, J. C. Pacheco-Kato, J. L. Gazquez, J. M. del Campo, and S. B. Trickey, J. Chem. Phys. 136, 144115 (2012)
        'GGA_X_VMT84_PBE'              : 69 , # A. Vela, J. C. Pacheco-Kato, J. L. Gazquez, J. M. del Campo, and S. B. Trickey, J. Chem. Phys. 136, 144115 (2012)
        'GGA_X_VMT_GE'                 : 70 , # A. Vela, V. Medel, and S. B. Trickey, J. Chem. Phys. 130, 244103 (2009)
        'GGA_X_VMT_PBE'                : 71 , # A. Vela, V. Medel, and S. B. Trickey, J. Chem. Phys. 130, 244103 (2009)
        'GGA_X_WC'                     : 118, # Z. Wu and R. E. Cohen, Phys. Rev. B 73, 235116 (2006)
        'GGA_X_WPBEH'                  : 524, # J. Heyd, G. E. Scuseria, and M. Ernzerhof, J. Chem. Phys. 118, 8207 (2003)
        'GGA_X_XPBE'                   : 123, # X. Xu and W. A. Goddard, J. Chem. Phys. 121, 4068 (2004)
        'GGA_XC_B97_D'                 : 170, # S. Grimme, J. Comput. Chem. 27, 1787 (2006)
        'GGA_XC_B97_GGA1'              : 96 , # A. J. Cohen and N. C. Handy, Chem. Phys. Lett. 316, 160 (2000)
        'GGA_XC_BEEFVDW'               : 286, # J. Wellendorff, K. T. Lundgaard, A. Mogelhoj, V. Petzold, D. D. Landis, J. K. Norskov, T. Bligaard, and K. W. Jacobsen, Phys. Rev. B 85, 235149 (2012)
        'GGA_XC_EDF1'                  : 165, # R. D. Adamson, P. M. W. Gill, and J. A. Pople, Chem. Phys. Lett. 284, 6 (1998)
        'GGA_XC_HCTH_120'              : 162, # A. D. Boese, N. L. Doltsinis, N. C. Handy, and M. Sprik, J. Chem. Phys. 112, 1670 (2000)
        'GGA_XC_HCTH_147'              : 163, # A. D. Boese, N. L. Doltsinis, N. C. Handy, and M. Sprik, J. Chem. Phys. 112, 1670 (2000)
        'GGA_XC_HCTH_407'              : 164, # A. D. Boese and N. C. Handy, J. Chem. Phys. 114, 5497 (2001)
        'GGA_XC_HCTH_407P'             : 93 , # A. D. Boese, A. Chandra, J. M. L. Martin, and D. Marx, J. Chem. Phys. 119, 5965 (2003)
        'GGA_XC_HCTH_93'               : 161, # F. A. Hamprecht, A. J. Cohen, D. J. Tozer, and N. C. Handy, J. Chem. Phys. 109, 6264 (1998)
        'GGA_XC_HCTH_P14'              : 95 , # G. Menconi, P. J. Wilson, and D. J. Tozer, J. Chem. Phys. 114, 3958 (2001)
        'GGA_XC_HCTH_P76'              : 94 , # G. Menconi, P. J. Wilson, and D. J. Tozer, J. Chem. Phys. 114, 3958 (2001)
        'GGA_XC_HLE16'                 : 545, # P. Verma and D. G. Truhlar, J. Phys. Chem. Lett. 8, 380 (2017)
        'GGA_XC_KT1'                   : 167, # T. W. Keal and D. J. Tozer, J. Chem. Phys. 119, 3015 (2003)
        'GGA_XC_KT2'                   : 146, # T. W. Keal and D. J. Tozer, J. Chem. Phys. 119, 3015 (2003)
        'GGA_XC_KT3'                   : 587, # T. W. Keal and D. J. Tozer, J. Chem. Phys. 121, 5654 (2004)
        'GGA_XC_LB07'                  : 589, # E. Livshits and R. Baer, Phys. Chem. Chem. Phys. 9, 2932 (2007)
        'GGA_XC_MOHLYP'                : 194, # N. E. Schultz, Y. Zhao, and D. G. Truhlar, J. Phys. Chem. A 109, 11127 (2005)
        'GGA_XC_MOHLYP2'               : 195, # J. Zheng, Y. Zhao, and D. G. Truhlar, J. Chem. Theory Comput. 5, 808 (2009)
        'GGA_XC_MPWLYP1W'              : 174, # E. E. Dahlke and D. G. Truhlar, J. Phys. Chem. B 109, 15677 (2005)
        'GGA_XC_NCAP'                  : 181, # J. Carmona-Espindola, J. L. Gazquez, A. Vela, and S. B. Trickey, J. Chem. Theory Comput. 15, 303 (2019)
        'GGA_XC_OBLYP_D'               : 67 , # L. Goerigk and S. Grimme, J. Chem. Theory Comput. 6, 107 (2010)
        'GGA_XC_OPBE_D'                : 65 , # L. Goerigk and S. Grimme, J. Chem. Theory Comput. 6, 107 (2010)
        'GGA_XC_OPWLYP_D'              : 66 , # L. Goerigk and S. Grimme, J. Chem. Theory Comput. 6, 107 (2010)
        'GGA_XC_PBE1W'                 : 173, # E. E. Dahlke and D. G. Truhlar, J. Phys. Chem. B 109, 15677 (2005)
        'GGA_XC_PBELYP1W'              : 175, # E. E. Dahlke and D. G. Truhlar, J. Phys. Chem. B 109, 15677 (2005)
        'GGA_XC_TH1'                   : 154, # D. J. Tozer and N. C. Handy, J. Chem. Phys. 108, 2545 (1998)
        'GGA_XC_TH2'                   : 155, # D. J. Tozer and N. C. Handy, J. Phys. Chem. A 102, 3162 (1998)
        'GGA_XC_TH3'                   : 156, # N. C. Handy and D. J. Tozer, Mol. Phys. 94, 707 (1998)
        'GGA_XC_TH4'                   : 157, # N. C. Handy and D. J. Tozer, Mol. Phys. 94, 707 (1998)
        'GGA_XC_TH_FC'                 : 197, # D. J. Tozer, N. C. Handy, and W. H. Green, Chem. Phys. Lett. 273, 183 (1997)
        'GGA_XC_TH_FCFO'               : 198, # D. J. Tozer, N. C. Handy, and W. H. Green, Chem. Phys. Lett. 273, 183 (1997)
        'GGA_XC_TH_FCO'                : 199, # D. J. Tozer, N. C. Handy, and W. H. Green, Chem. Phys. Lett. 273, 183 (1997)
        'GGA_XC_TH_FL'                 : 196, # D. J. Tozer, N. C. Handy, and W. H. Green, Chem. Phys. Lett. 273, 183 (1997)
        'GGA_XC_VV10'                  : 255, # O. A. Vydrov and T. Van Voorhis, J. Chem. Phys. 133, 244103 (2010)
        'GGA_XC_XLYP'                  : 166, # X. Xu and W. A. Goddard, Proc. Natl. Acad. Sci. U. S. A. 101, 2673 (2004)
        #'HYB_GGA_X_LC2GAU'             : 710, # J.-W. Song, M. A. Watson, and K. Hirao, J. Chem. Phys. 131, 144108 (2009)
        #'HYB_GGA_X_LCGAU'              : 708, # J.-W. Song, S. Tokura, T. Sato, M. A. Watson, and K. Hirao, J. Chem. Phys. 127, 154109 (2007)
        #'HYB_GGA_X_LCGAU_CORE'         : 709, # J.-W. Song, M. A. Watson, A. Nakata, and K. Hirao, J. Chem. Phys. 129, 184113 (2008)
        'HYB_GGA_X_N12_SX'             : 81 , # R. Peverati and D. G. Truhlar, Phys. Chem. Chem. Phys. 14, 16187 (2012)
        'HYB_GGA_X_S12H'               : 496, # M. Swart, Chem. Phys. Lett. 580, 166 (2013)
        'HYB_GGA_X_SOGGA11_X'          : 426, # R. Peverati and D. G. Truhlar, J. Chem. Phys. 135, 191102 (2011)
        'HYB_GGA_XC_APBE0'             : 607, # E. Fabiano, L. A. Constantin, P. Cortona, and F. Della Sala, J. Chem. Theory Comput. 11, 122 (2015)
        'HYB_GGA_XC_APF'               : 409, # A. Austin, G. A. Petersson, M. J. Frisch, F. J. Dobek, G. Scalmani, and K. Throssell, J. Chem. Theory Comput. 8, 4989 (2012)
        'HYB_GGA_XC_B1LYP'             : 416, # C. Adamo and V. Barone, Chem. Phys. Lett. 274, 242 (1997)
        'HYB_GGA_XC_B1PW91'            : 417, # C. Adamo and V. Barone, Chem. Phys. Lett. 274, 242 (1997)
        'HYB_GGA_XC_B1WC'              : 412, # D. I. Bilc, R. Orlando, R. Shaltaf, G.-M. Rignanese, J. Iniguez, and P. Ghosez, Phys. Rev. B 77, 165107 (2008)
        #'HYB_GGA_XC_B2PLYP'            : 713, # S. Grimme, J. Chem. Phys. 124, 034108 (2006)
        'HYB_GGA_XC_B3LYP'             : 402, # P. J. Stephens, F. J. Devlin, C. F. Chabalowski, and M. J. Frisch, J. Phys. Chem. 98, 11623 (1994)
        'HYB_GGA_XC_B3LYP5'            : 475, # P. J. Stephens, F. J. Devlin, C. F. Chabalowski, and M. J. Frisch, J. Phys. Chem. 98, 11623 (1994)
        'HYB_GGA_XC_B3LYP_MCM1'        : 461, # M. T. Caldeira and R. Custodio, J. Mol. Model. 25, 62 (2019)
        'HYB_GGA_XC_B3LYP_MCM2'        : 462, # M. T. Caldeira and R. Custodio, J. Mol. Model. 25, 62 (2019)
        'HYB_GGA_XC_B3LYPS'            : 459, # M. Reiher, O. Salomon, and B. A. Hess, Theor. Chem. Acc. 107, 48 (2001)
        'HYB_GGA_XC_B3P86'             : 403, # Defined through Gaussian implementation
        #'HYB_GGA_XC_B3P86_NWCHEM'      : 315, # Defined through NWChem implementation
        'HYB_GGA_XC_B3PW91'            : 401, # A. D. Becke, J. Chem. Phys. 98, 5648 (1993)
        'HYB_GGA_XC_B5050LYP'          : 572, # Y. Shao, M. Head-Gordon, and A. I. Krylov, J. Chem. Phys. 118, 4807 (2003)
        'HYB_GGA_XC_B97'               : 407, # A. D. Becke, J. Chem. Phys. 107, 8554 (1997)
        'HYB_GGA_XC_B97_1'             : 408, # F. A. Hamprecht, A. J. Cohen, D. J. Tozer, and N. C. Handy, J. Chem. Phys. 109, 6264 (1998)
        'HYB_GGA_XC_B97_1P'            : 266, # A. J. Cohen and N. C. Handy, Chem. Phys. Lett. 316, 160 (2000)
        'HYB_GGA_XC_B97_2'             : 410, # P. J. Wilson, T. J. Bradley, and D. J. Tozer, J. Chem. Phys. 115, 9233 (2001)
        'HYB_GGA_XC_B97_3'             : 414, # T. W. Keal and D. J. Tozer, J. Chem. Phys. 123, 121103 (2005)
        'HYB_GGA_XC_B97_K'             : 413, # A. D. Boese and J. M. L. Martin, J. Chem. Phys. 121, 3405 (2004)
        'HYB_GGA_XC_BHANDH'            : 435, # A. D. Becke, J. Chem. Phys. 98, 1372 (1993)
        'HYB_GGA_XC_BHANDHLYP'         : 436, # A. D. Becke, J. Chem. Phys. 98, 1372 (1993)
        'HYB_GGA_XC_BLYP35'            : 499, # M. Renz, K. Theilacker, C. Lambert, and M. Kaupp, J. Am. Chem. Soc. 131, 16292 (2009)
        'HYB_GGA_XC_CAM_B3LYP'         : 433, # T. Yanai, D. P. Tew, and N. C. Handy, Chem. Phys. Lett. 393, 51 (2004)
        'HYB_GGA_XC_CAM_O3LYP'         : 395, # M. P. Bircher and U. Rothlisberger, J. Chem. Theory Comput. 14, 3184 (2018)
        'HYB_GGA_XC_CAM_PBEH'          : 681, # W. Chen, G. Miceli, G.-M. Rignanese, and A. Pasquarello, Phys. Rev. Mater. 2, 073803 (2018)
        'HYB_GGA_XC_CAM_QTP_00'        : 490, # P. Verma and R. J. Bartlett, J. Chem. Phys. 140, 18A534 (2014)
        'HYB_GGA_XC_CAM_QTP_01'        : 482, # Y. Jin and R. J. Bartlett, J. Chem. Phys. 145, 034107 (2016)
        'HYB_GGA_XC_CAM_QTP_02'        : 491, # R. L. A. Haiduke and R. J. Bartlett, J. Chem. Phys. 148, 184106 (2018)
        'HYB_GGA_XC_CAMH_B3LYP'        : 614, # Y. Shao, Y. Mei, D. Sundholm, and V. R. I. Kaila, J. Chem. Theory Comput. 16, 587 (2020), https://doi.org/10.1021/acs.jctc.9b00823
        'HYB_GGA_XC_CAMY_B3LYP'        : 470, # M. Seth and T. Ziegler, J. Chem. Theory Comput. 8, 901 (2012)
        'HYB_GGA_XC_CAMY_BLYP'         : 455, # Y. Akinaga and S. Ten-no, Chem. Phys. Lett. 462, 348 (2008)
        'HYB_GGA_XC_CAMY_PBEH'         : 682, # W. Chen, G. Miceli, G.-M. Rignanese, and A. Pasquarello, Phys. Rev. Mater. 2, 073803 (2018)
        'HYB_GGA_XC_CAP0'              : 477, # J. Carmona-Espindola, J. L. Gazquez, A. Vela, and S. B. Trickey, Theor. Chem. Acc. 135, 120 (2016)
        'HYB_GGA_XC_EDF2'              : 476, # C. Y. Lin, M. W. George, and P. M. W. Gill, Aust. J. Chem. 57, 365 (2004)
        'HYB_GGA_XC_HAPBE'             : 608, # E. Fabiano, L. A. Constantin, P. Cortona, and F. Della Sala, J. Chem. Theory Comput. 11, 122 (2015)
        'HYB_GGA_XC_HFLYP'             : 314, # C. Lee, W. Yang, and R. G. Parr, Phys. Rev. B 37, 785 (1988)
        #'HYB_GGA_XC_HISS'              : 717, # T. M. Henderson, A. F. Izmaylov, G. E. Scuseria, and A. Savin, J. Chem. Phys. 127, 221103 (2007)
        'HYB_GGA_XC_HJS_B88'           : 431, # T. M. Henderson, B. G. Janesko, and G. E. Scuseria, J. Chem. Phys. 128, 194105 (2008)
        'HYB_GGA_XC_HJS_B97X'          : 432, # T. M. Henderson, B. G. Janesko, and G. E. Scuseria, J. Chem. Phys. 128, 194105 (2008)
        'HYB_GGA_XC_HJS_PBE'           : 429, # T. M. Henderson, B. G. Janesko, and G. E. Scuseria, J. Chem. Phys. 128, 194105 (2008)
        'HYB_GGA_XC_HJS_PBE_SOL'       : 430, # T. M. Henderson, B. G. Janesko, and G. E. Scuseria, J. Chem. Phys. 128, 194105 (2008)
        'HYB_GGA_XC_HPBEINT'           : 472, # E. Fabiano, L. A. Constantin, and F. Della Sala, Int. J. Quantum Chem. 113, 673 (2013)
        'HYB_GGA_XC_HSE03'             : 427, # J. Heyd, G. E. Scuseria, and M. Ernzerhof, J. Chem. Phys. 118, 8207 (2003)
        'HYB_GGA_XC_HSE06'             : 428, # J. Heyd, G. E. Scuseria, and M. Ernzerhof, J. Chem. Phys. 118, 8207 (2003)
        'HYB_GGA_XC_HSE12'             : 479, # J. E. Moussa, P. A. Schultz, and J. R. Chelikowsky, J. Chem. Phys. 136, 204117 (2012)
        'HYB_GGA_XC_HSE12S'            : 480, # J. E. Moussa, P. A. Schultz, and J. R. Chelikowsky, J. Chem. Phys. 136, 204117 (2012)
        'HYB_GGA_XC_HSE_SOL'           : 481, # L. Schimka, J. Harl, and G. Kresse, J. Chem. Phys. 134, 024116 (2011)
        'HYB_GGA_XC_KMLYP'             : 485, # J. K. Kang and C. B. Musgrave, J. Chem. Phys. 115, 11040 (2001)
        'HYB_GGA_XC_LC_BLYP'           : 400, # L. N. Anderson, M. B. Oviedo, and B. M. Wong, J. Chem. Theory Comput. 13, 1656 (2017)
        'HYB_GGA_XC_LC_BOP'            : 636, # J.-W. Song, T. Hirosawa, T. Tsuneda, and K. Hirao, J. Chem. Phys. 126, 154105 (2007)
        'HYB_GGA_XC_LC_PBEOP'          : 637, # Y. Tawada, T. Tsuneda, S. Yanagisawa, T. Yanai, and K. Hirao, J. Chem. Phys. 120, 8425 (2004)
        'HYB_GGA_XC_LC_QTP'            : 492, # R. L. A. Haiduke and R. J. Bartlett, J. Chem. Phys. 148, 184106 (2018)
        'HYB_GGA_XC_LC_VV10'           : 469, # O. A. Vydrov and T. Van Voorhis, J. Chem. Phys. 133, 244103 (2010)
        'HYB_GGA_XC_LC_WPBE'           : 478, # O. A. Vydrov and G. E. Scuseria, J. Chem. Phys. 125, 234109 (2006)
        'HYB_GGA_XC_LC_WPBE08_WHS'     : 488, # E. Weintraub, T. M. Henderson, and G. E. Scuseria, J. Chem. Theory Comput. 5, 754 (2009)
        'HYB_GGA_XC_LC_WPBE_WHS'       : 486, # E. Weintraub, T. M. Henderson, and G. E. Scuseria, J. Chem. Theory Comput. 5, 754 (2009)
        'HYB_GGA_XC_LC_WPBEH_WHS'      : 487, # E. Weintraub, T. M. Henderson, and G. E. Scuseria, J. Chem. Theory Comput. 5, 754 (2009)
        'HYB_GGA_XC_LC_WPBESOL_WHS'    : 489, # E. Weintraub, T. M. Henderson, and G. E. Scuseria, J. Chem. Theory Comput. 5, 754 (2009)
        'HYB_GGA_XC_LCY_BLYP'          : 468, # Y. Akinaga and S. Ten-no, Chem. Phys. Lett. 462, 348 (2008)
        'HYB_GGA_XC_LCY_PBE'           : 467, # M. Seth and T. Ziegler, J. Chem. Theory Comput. 8, 901 (2012)
        'HYB_GGA_XC_LRC_WPBE'          : 473, # M. A. Rohrdanz, K. M. Martins, and J. M. Herbert, J. Chem. Phys. 130, 054112 (2009)
        'HYB_GGA_XC_LRC_WPBEH'         : 465, # M. A. Rohrdanz, K. M. Martins, and J. M. Herbert, J. Chem. Phys. 130, 054112 (2009)
        'HYB_GGA_XC_MB3LYP_RC04'       : 437, # V. Tognetti, P. Cortona, and C. Adamo, Chem. Phys. Lett. 439, 381 (2007)
        'HYB_GGA_XC_MPW1K'             : 405, # B. J. Lynch, P. L. Fast, M. Harris, and D. G. Truhlar, J. Phys. Chem. A 104, 4811 (2000)
        'HYB_GGA_XC_MPW1LYP'           : 483, # C. Adamo and V. Barone, J. Chem. Phys. 108, 664 (1998)
        'HYB_GGA_XC_MPW1PBE'           : 484, # C. Adamo and V. Barone, J. Chem. Phys. 108, 664 (1998)
        'HYB_GGA_XC_MPW1PW'            : 418, # C. Adamo and V. Barone, J. Chem. Phys. 108, 664 (1998)
        'HYB_GGA_XC_MPW3LYP'           : 419, # Y. Zhao and D. G. Truhlar, J. Phys. Chem. A 108, 6908 (2004)
        'HYB_GGA_XC_MPW3PW'            : 415, # C. Adamo and V. Barone, J. Chem. Phys. 108, 664 (1998)
        'HYB_GGA_XC_MPWLYP1M'          : 453, # N. E. Schultz, Y. Zhao, and D. G. Truhlar, J. Phys. Chem. A 109, 11127 (2005)
        'HYB_GGA_XC_O3LYP'             : 404, # A. J. Cohen and N. C. Handy, Mol. Phys. 99, 607 (2001)
        'HYB_GGA_XC_PBE0_13'           : 456, # P. Cortona, J. Chem. Phys. 136, 086101 (2012)
        'HYB_GGA_XC_PBE50'             : 290, # Y. A. Bernard, Y. Shao, and A. I. Krylov, J. Chem. Phys. 136, 204103 (2012)
        'HYB_GGA_XC_PBE_MOL0'          : 273, # J. M. del Campo, J. L. Gazquez, S. B. Trickey, and A. Vela, J. Chem. Phys. 136, 104108 (2012)
        'HYB_GGA_XC_PBE_MOLB0'         : 276, # J. M. del Campo, J. L. Gazquez, S. B. Trickey, and A. Vela, J. Chem. Phys. 136, 104108 (2012)
        'HYB_GGA_XC_PBE_SOL0'          : 274, # J. M. del Campo, J. L. Gazquez, S. B. Trickey, and A. Vela, J. Chem. Phys. 136, 104108 (2012)
        'HYB_GGA_XC_PBEB0'             : 275, # J. M. del Campo, J. L. Gazquez, S. B. Trickey, and A. Vela, J. Chem. Phys. 136, 104108 (2012)
        'HYB_GGA_XC_PBEH'              : 406, # C. Adamo and V. Barone, J. Chem. Phys. 110, 6158 (1999)
        'HYB_GGA_XC_QTP17'             : 460, # Y. Jin and R. J. Bartlett, J. Chem. Phys. 149, 064111 (2018)
        'HYB_GGA_XC_RCAM_B3LYP'        : 610, # A. J. Cohen, P. Mori-Sanchez, and W. Yang, J. Chem. Phys. 126, 191109 (2007)
        'HYB_GGA_XC_REVB3LYP'          : 454, # L. Lu, H. Hu, H. Hou, and B. Wang, Comput. Theor. Chem. 1015, 64 (2013)
        'HYB_GGA_XC_SB98_1A'           : 420, # H. L. Schmider and A. D. Becke, J. Chem. Phys. 108, 9624 (1998)
        'HYB_GGA_XC_SB98_1B'           : 421, # H. L. Schmider and A. D. Becke, J. Chem. Phys. 108, 9624 (1998)
        'HYB_GGA_XC_SB98_1C'           : 422, # H. L. Schmider and A. D. Becke, J. Chem. Phys. 108, 9624 (1998)
        'HYB_GGA_XC_SB98_2A'           : 423, # H. L. Schmider and A. D. Becke, J. Chem. Phys. 108, 9624 (1998)
        'HYB_GGA_XC_SB98_2B'           : 424, # H. L. Schmider and A. D. Becke, J. Chem. Phys. 108, 9624 (1998)
        'HYB_GGA_XC_SB98_2C'           : 425, # H. L. Schmider and A. D. Becke, J. Chem. Phys. 108, 9624 (1998)
        #'HYB_GGA_XC_SRC1_BLYP'         : 714, # N. A. Besley, M. J. G. Peach, and D. J. Tozer, Phys. Chem. Chem. Phys. 11, 10350 (2009)
        #'HYB_GGA_XC_SRC2_BLYP'         : 715, # N. A. Besley, M. J. G. Peach, and D. J. Tozer, Phys. Chem. Chem. Phys. 11, 10350 (2009)
        'HYB_GGA_XC_TUNED_CAM_B3LYP'   : 434, # K. Okuno, Y. Shigeta, R. Kishi, H. Miyasaka, and M. Nakano, J. Photochem. Photobiol., A 235, 29 (2012)
        'HYB_GGA_XC_WB97'              : 463, # J.-D. Chai and M. Head-Gordon, J. Chem. Phys. 128, 084106 (2008)
        'HYB_GGA_XC_WB97X'             : 464, # J.-D. Chai and M. Head-Gordon, J. Chem. Phys. 128, 084106 (2008)
        'HYB_GGA_XC_WB97X_D'           : 471, # J.-D. Chai and M. Head-Gordon, Phys. Chem. Chem. Phys. 10, 6615 (2008)
        'HYB_GGA_XC_WB97X_D3'          : 399, # Y.-S. Lin, G.-D. Li, S.-P. Mao, and J.-D. Chai, J. Chem. Theory Comput. 9, 263 (2013)
        'HYB_GGA_XC_WB97X_V'           : 466, # N. Mardirossian and M. Head-Gordon, Phys. Chem. Chem. Phys. 16, 9904 (2014)
        'HYB_GGA_XC_WC04'              : 611, # K. W. Wiitala, T. R. Hoye, and C. J. Cramer, J. Chem. Theory Comput. 2, 1085 (2006)
        'HYB_GGA_XC_WHPBE0'            : 615, # Y. Shao, Y. Mei, D. Sundholm, and V. R. I. Kaila, J. Chem. Theory Comput. 16, 587 (2020), https://doi.org/10.1021/acs.jctc.9b00823
        'HYB_GGA_XC_WP04'              : 612, # K. W. Wiitala, T. R. Hoye, and C. J. Cramer, J. Chem. Theory Comput. 2, 1085 (2006)
        'HYB_GGA_XC_X3LYP'             : 411, # X. Xu and W. A. Goddard, Proc. Natl. Acad. Sci. U. S. A. 101, 2673 (2004)
        'MGGA_C_B88'                   : 571, # A. D. Becke, J. Chem. Phys. 88, 1053 (1988)
        'MGGA_C_B94'                   : 397, # A. D. Becke, Int. J. Quantum Chem. 52, 625 (1994)
        'MGGA_C_BC95'                  : 240, # A. D. Becke, J. Chem. Phys. 104, 1040 (1996)
        'MGGA_C_CS'                    : 72 , # R. Colle and O. Salvetti, Theor. Chim. Acta 37, 329 (1975)
        'MGGA_C_DLDF'                  : 37 , # K. Pernal, R. Podeszwa, K. Patkowski, and K. Szalewicz, Phys. Rev. Lett. 103, 263201 (2009)
        'MGGA_C_HLTAPW'                : 699, # S. Lehtola and M. A. L. Marques, Meta-local density functionals: a new rung on jacob's ladder, (2020), arXiv:2006.16835 [physics.chem-ph]
        'MGGA_C_KCIS'                  : 562, # J. Rey and A. Savin, Int. J. Quantum Chem. 69, 581 (1998)
        'MGGA_C_KCISK'                 : 638, # J. Rey and A. Savin, Int. J. Quantum Chem. 69, 581 (1998)
        'MGGA_C_M05'                   : 237, # Y. Zhao, N. E. Schultz, and D. G. Truhlar, J. Chem. Phys. 123, 161103 (2005)
        'MGGA_C_M05_2X'                : 238, # Y. Zhao, N. E. Schultz, and D. G. Truhlar, J. Chem. Theory Comput. 2, 364 (2006)
        'MGGA_C_M06'                   : 235, # Y. Zhao and D. G. Truhlar, Theor. Chem. Acc. 120, 215 (2008)
        'MGGA_C_M06_2X'                : 236, # Y. Zhao and D. G. Truhlar, Theor. Chem. Acc. 120, 215 (2008)
        'MGGA_C_M06_HF'                : 234, # Y. Zhao and D. G. Truhlar, J. Phys. Chem. A 110, 13126 (2006)
        'MGGA_C_M06_L'                 : 233, # Y. Zhao and D. G. Truhlar, J. Chem. Phys. 125, 194101 (2006)
        'MGGA_C_M06_SX'                : 311, # Y. Wang, P. Verma, L. Zhang, Y. Li, Z. Liu, D. G. Truhlar, and X. He, Proc. Natl. Acad. Sci. U. S. A. 117, 2294 (2020), https://www.pnas.org/content/117/5/2294.full.pdf
        'MGGA_C_M08_HX'                : 78 , # Y. Zhao and D. G. Truhlar, J. Chem. Theory Comput. 4, 1849 (2008)
        'MGGA_C_M08_SO'                : 77 , # Y. Zhao and D. G. Truhlar, J. Chem. Theory Comput. 4, 1849 (2008)
        'MGGA_C_M11'                   : 76 , # R. Peverati and D. G. Truhlar, J. Phys. Chem. Lett. 2, 2810 (2011)
        'MGGA_C_M11_L'                 : 75 , # R. Peverati and D. G. Truhlar, J. Phys. Chem. Lett. 3, 117 (2012)
        'MGGA_C_MN12_L'                : 74 , # R. Peverati and D. G. Truhlar, Phys. Chem. Chem. Phys. 14, 13171 (2012)
        'MGGA_C_MN12_SX'               : 73 , # R. Peverati and D. G. Truhlar, Phys. Chem. Chem. Phys. 14, 16187 (2012)
        'MGGA_C_MN15'                  : 269, # H. S. Yu, X. He, S. L. Li, and D. G. Truhlar, Chem. Sci. 7, 5032 (2016)
        'MGGA_C_MN15_L'                : 261, # H. S. Yu, X. He, and D. G. Truhlar, J. Chem. Theory Comput. 12, 1280 (2016)
        'MGGA_C_PKZB'                  : 239, # J. P. Perdew, S. Kurth, A. Zupan, and P. Blaha, Phys. Rev. Lett. 82, 2544 (1999)
        'MGGA_C_R2SCAN'                : 498, # J. W. Furness, A. D. Kaplan, J. Ning, J. P. Perdew, and J. Sun, J. Phys. Chem. Lett. 11, 8208 (2020)
        'MGGA_C_R2SCANL'               : 719, # D. Mejia-Rodriguez and S. B. Trickey, Phys. Rev. B 102, 121109 (2020)
        'MGGA_C_REVM06'                : 306, # Y. Wang, P. Verma, X. Jin, D. G. Truhlar, and X. He, Proc. Natl. Acad. Sci. U. S. A. 115, 10257 (2018)
        'MGGA_C_REVM06_L'              : 294, # Y. Wang, X. Jin, H. S. Yu, D. G. Truhlar, and X. He, Proc. Natl. Acad. Sci. U. S. A. 114, 8487 (2017)
        'MGGA_C_REVM11'                : 172, # P. Verma, Y. Wang, S. Ghosh, X. He, and D. G. Truhlar, J. Phys. Chem. A 123, 2966 (2019)
        'MGGA_C_REVSCAN'               : 582, # P. D. Mezei, G. I. Csonka, and M. Kallay, J. Chem. Theory Comput. 14, 2469 (2018)
        'MGGA_C_REVSCAN_VV10'          : 585, # P. D. Mezei, G. I. Csonka, and M. Kallay, J. Chem. Theory Comput. 14, 2469 (2018)
        'MGGA_C_REVTM'                 : 694, # S. Jana, K. Sharma, and P. Samal, J. Phys. Chem. A 123, 6356 (2019)
        'MGGA_C_REVTPSS'               : 241, # J. P. Perdew, A. Ruzsinszky, G. I. Csonka, L. A. Constantin, and J. Sun, Phys. Rev. Lett. 103, 026403 (2009)
        'MGGA_C_RSCAN'                 : 494, # A. P. Bartok and J. R. Yates, J. Chem. Phys. 150, 161101 (2019)
        'MGGA_C_SCAN'                  : 267, # J. Sun, A. Ruzsinszky, and J. P. Perdew, Phys. Rev. Lett. 115, 036402 (2015)
        'MGGA_C_SCAN_RVV10'            : 292, # H. Peng, Z.-H. Yang, J. P. Perdew, and J. Sun, Phys. Rev. X 6, 041005 (2016)
        'MGGA_C_SCAN_VV10'             : 584, # J. G. Brandenburg, J. E. Bates, J. Sun, and J. P. Perdew, Phys. Rev. B 94, 115144 (2016)
        'MGGA_C_SCANL'                 : 702, # D. Mejia-Rodriguez and S. B. Trickey, Phys. Rev. A 96, 052512 (2017)
        'MGGA_C_SCANL_RVV10'           : 703, # D. Mejia-Rodriguez and S. B. Trickey, Phys. Rev. A 96, 052512 (2017)
        'MGGA_C_SCANL_VV10'            : 704, # D. Mejia-Rodriguez and S. B. Trickey, Phys. Rev. A 96, 052512 (2017)
        'MGGA_C_TM'                    : 251, # J. Tao and Y. Mo, Phys. Rev. Lett. 117, 073001 (2016)
        'MGGA_C_TPSS'                  : 231, # J. Tao, J. P. Perdew, V. N. Staroverov, and G. E. Scuseria, Phys. Rev. Lett. 91, 146401 (2003)
        'MGGA_C_TPSSLOC'               : 247, # L. A. Constantin, E. Fabiano, and F. Della Sala, Phys. Rev. B 86, 035130 (2012)
        'MGGA_C_VSXC'                  : 232, # T. V. Voorhis and G. E. Scuseria, J. Chem. Phys. 109, 400 (1998)
        'MGGA_K_CSK1'                  : 629, # A. C. Cancio, D. Stewart, and A. Kuna, J. Chem. Phys. 144, 084107 (2016)
        'MGGA_K_CSK4'                  : 630, # A. C. Cancio, D. Stewart, and A. Kuna, J. Chem. Phys. 144, 084107 (2016)
        'MGGA_K_CSK_LOC1'              : 631, # A. C. Cancio, D. Stewart, and A. Kuna, J. Chem. Phys. 144, 084107 (2016)
        'MGGA_K_CSK_LOC4'              : 632, # A. C. Cancio, D. Stewart, and A. Kuna, J. Chem. Phys. 144, 084107 (2016)
        'MGGA_K_GEA2'                  : 627, # A. S. Kompaneets and E. S. Pavlovskii, Zh. Eksp. Teor. Fiz. 31, 427 (1956), [J. Exp. Theor. Phys. 4, 328 (1957)]
        'MGGA_K_GEA4'                  : 628, # C. H. Hodges, Can. J. Phys. 51, 1428 (1973)
        'MGGA_K_L04'                   : 617, # S. Laricchia, L. A. Constantin, E. Fabiano, and F. Della Sala, J. Chem. Theory Comput. 10, 164 (2014)
        'MGGA_K_L06'                   : 618, # S. Laricchia, L. A. Constantin, E. Fabiano, and F. Della Sala, J. Chem. Theory Comput. 10, 164 (2014)
        'MGGA_K_PC07'                  : 543, # J. P. Perdew and L. A. Constantin, Phys. Rev. B 75, 155109 (2007)
        'MGGA_K_PC07_OPT'              : 634, # D. Mejia-Rodriguez and S. B. Trickey, Phys. Rev. A 96, 052512 (2017)
        'MGGA_K_PGSL025'               : 220, # L. A. Constantin, E. Fabiano, and F. Della Sala, J. Phys. Chem. Lett. 9, 4385 (2018), pMID: 30019904
        'MGGA_K_RDA'                   : 621, # V. V. Karasiev, R. S. Jones, S. B. Trickey, and F. E. Harris, Phys. Rev. B 80, 245120 (2009)
        'MGGA_X_2D_JS17'               : 609, # S. Jana and P. Samal, J. Phys. Chem. A 121, 4804 (2017)
        'MGGA_X_2D_PRHG07'             : 210, # S. Pittalis, E. Rasanen, N. Helbig, and E. K. U. Gross, Phys. Rev. B 76, 235314 (2007)
        'MGGA_X_2D_PRHG07_PRP10'       : 211, # S. Pittalis, E. Rasanen, N. Helbig, and E. K. U. Gross, Phys. Rev. B 76, 235314 (2007)
        'MGGA_X_B00'                   : 284, # A. D. Becke, J. Chem. Phys. 112, 4020 (2000)
        'MGGA_X_BJ06'                  : 207, # A. D. Becke and E. R. Johnson, J. Chem. Phys. 124, 221101 (2006)
        'MGGA_X_BLOC'                  : 244, # L. A. Constantin, E. Fabiano, and F. Della Sala, J. Chem. Theory Comput. 9, 2256 (2013)
        'MGGA_X_BR89'                  : 206, # A. D. Becke and M. R. Roussel, Phys. Rev. A 39, 3761 (1989)
        'MGGA_X_BR89_1'                : 214, # A. D. Becke and M. R. Roussel, Phys. Rev. A 39, 3761 (1989)
        'MGGA_X_BR89_EXPLICIT'         : 586, # A. D. Becke and M. R. Roussel, Phys. Rev. A 39, 3761 (1989)
        'MGGA_X_BR89_EXPLICIT_1'       : 602, # A. D. Becke and M. R. Roussel, Phys. Rev. A 39, 3761 (1989)
        'MGGA_X_EDMGGA'                : 686, # J. Tao, J. Chem. Phys. 115, 3519 (2001)
        'MGGA_X_GDME_0'                : 689, # R. M. Koehl, G. K. Odom, and G. E. Scuseria, Mol. Phys. 87, 835 (1996)
        'MGGA_X_GDME_KOS'              : 690, # R. M. Koehl, G. K. Odom, and G. E. Scuseria, Mol. Phys. 87, 835 (1996)
        'MGGA_X_GDME_NV'               : 687, # J. W. Negele and D. Vautherin, Phys. Rev. C 5, 1472 (1972)
        'MGGA_X_GDME_VT'               : 691, # R. M. Koehl, G. K. Odom, and G. E. Scuseria, Mol. Phys. 87, 835 (1996)
        'MGGA_X_GVT4'                  : 204, # T. V. Voorhis and G. E. Scuseria, J. Chem. Phys. 109, 400 (1998)
        'MGGA_X_GX'                    : 575, # P.-F. Loos, J. Chem. Phys. 146, 114108 (2017)
        'MGGA_X_HLTA'                  : 698, # S. Lehtola and M. A. L. Marques, Meta-local density functionals: a new rung on jacob's ladder, (2020), arXiv:2006.16835 [physics.chem-ph]
        'MGGA_X_JK'                    : 256, # P. Jemmer and P. J. Knowles, Phys. Rev. A 51, 3571 (1995)
        'MGGA_X_LTA'                   : 201, # M. Ernzerhof and G. E. Scuseria, J. Chem. Phys. 111, 911 (1999)
        'MGGA_X_M06_L'                 : 203, # Y. Zhao and D. G. Truhlar, J. Chem. Phys. 125, 194101 (2006)
        'MGGA_X_M11_L'                 : 226, # R. Peverati and D. G. Truhlar, J. Phys. Chem. Lett. 3, 117 (2012)
        'MGGA_X_MBEEF'                 : 249, # J. Wellendorff, K. T. Lundgaard, K. W. Jacobsen, and T. Bligaard, J. Chem. Phys. 140, 144107 (2014)
        'MGGA_X_MBEEFVDW'              : 250, # K. T. Lundgaard, J. Wellendorff, J. Voss, K. W. Jacobsen, and T. Bligaard, Phys. Rev. B 93, 235162 (2016)
        'MGGA_X_MBR'                   : 716, # A. Patra, S. Jana, H. Myneni, and P. Samal, Phys. Chem. Chem. Phys. 21, 19639 (2019)
        'MGGA_X_MBRXC_BG'              : 696, # B. Patra, S. Jana, L. A. Constantin, and P. Samal, Phys. Rev. B 100, 045147 (2019)
        'MGGA_X_MBRXH_BG'              : 697, # B. Patra, S. Jana, L. A. Constantin, and P. Samal, Phys. Rev. B 100, 045147 (2019)
        'MGGA_X_MGGAC'                 : 711, # B. Patra, S. Jana, L. A. Constantin, and P. Samal, Phys. Rev. B 100, 155140 (2019)
        'MGGA_X_MK00'                  : 230, # F. R. Manby and P. J. Knowles, J. Chem. Phys. 112, 7002 (2000)
        'MGGA_X_MK00B'                 : 243, # F. R. Manby and P. J. Knowles, J. Chem. Phys. 112, 7002 (2000)
        'MGGA_X_MN12_L'                : 227, # R. Peverati and D. G. Truhlar, Phys. Chem. Chem. Phys. 14, 13171 (2012)
        'MGGA_X_MN15_L'                : 260, # H. S. Yu, X. He, and D. G. Truhlar, J. Chem. Theory Comput. 12, 1280 (2016)
        'MGGA_X_MODTPSS'               : 245, # J. P. Perdew, A. Ruzsinszky, J. Tao, G. I. Csonka, and G. E. Scuseria, Phys. Rev. A 76, 042506 (2007)
        'MGGA_X_MS0'                   : 221, # J. Sun, B. Xiao, and A. Ruzsinszky, J. Chem. Phys. 137, 051101 (2012)
        'MGGA_X_MS1'                   : 222, # J. Sun, R. Haunschild, B. Xiao, I. W. Bulik, G. E. Scuseria, and J. P. Perdew, J. Chem. Phys. 138, 044113 (2013)
        'MGGA_X_MS2'                   : 223, # J. Sun, R. Haunschild, B. Xiao, I. W. Bulik, G. E. Scuseria, and J. P. Perdew, J. Chem. Phys. 138, 044113 (2013)
        'MGGA_X_MS2_REV'               : 228, # J. Sun, R. Haunschild, B. Xiao, I. W. Bulik, G. E. Scuseria, and J. P. Perdew, J. Chem. Phys. 138, 044113 (2013)
        'MGGA_X_MS2B'                  : 300, # J. W. Furness and J. Sun, Phys. Rev. B 99, 041119 (2019)
        'MGGA_X_MS2BS'                 : 301, # J. W. Furness and J. Sun, ArXiv e-prints (2018), arXiv:1805.11707v1 [physics.chem-ph]
        'MGGA_X_MVS'                   : 257, # J. Sun, J. P. Perdew, and A. Ruzsinszky, Proc. Natl. Acad. Sci. U. S. A. 112, 685 (2015)
        'MGGA_X_MVSB'                  : 302, # J. W. Furness and J. Sun, ArXiv e-prints (2018), arXiv:1805.11707v1 [physics.chem-ph]
        'MGGA_X_MVSBS'                 : 303, # J. W. Furness and J. Sun, ArXiv e-prints (2018), arXiv:1805.11707v1 [physics.chem-ph]
        'MGGA_X_PBE_GX'                : 576, # P.-F. Loos, J. Chem. Phys. 146, 114108 (2017)
        'MGGA_X_PKZB'                  : 213, # J. P. Perdew, S. Kurth, A. Zupan, and P. Blaha, Phys. Rev. Lett. 82, 2544 (1999)
        'MGGA_X_R2SCAN'                : 497, # J. W. Furness, A. D. Kaplan, J. Ning, J. P. Perdew, and J. Sun, J. Phys. Chem. Lett. 11, 8208 (2020)
        'MGGA_X_R2SCANL'               : 718, # D. Mejia-Rodriguez and S. B. Trickey, Phys. Rev. B 102, 121109 (2020)
        'MGGA_X_REGTPSS'               : 603, # A. Ruzsinszky, J. Sun, B. Xiao, and G. I. Csonka, J. Chem. Theory Comput. 8, 2078 (2012)
        'MGGA_X_REVM06_L'              : 293, # Y. Wang, X. Jin, H. S. Yu, D. G. Truhlar, and X. He, Proc. Natl. Acad. Sci. U. S. A. 114, 8487 (2017)
        'MGGA_X_REVSCAN'               : 581, # P. D. Mezei, G. I. Csonka, and M. Kallay, J. Chem. Theory Comput. 14, 2469 (2018)
        'MGGA_X_REVSCANL'              : 701, # D. Mejia-Rodriguez and S. B. Trickey, Phys. Rev. A 96, 052512 (2017)
        'MGGA_X_REVTM'                 : 693, # S. Jana, K. Sharma, and P. Samal, J. Phys. Chem. A 123, 6356 (2019)
        'MGGA_X_REVTPSS'               : 212, # J. P. Perdew, A. Ruzsinszky, G. I. Csonka, L. A. Constantin, and J. Sun, Phys. Rev. Lett. 103, 026403 (2009)
        'MGGA_X_RLDA'                  : 688, # X. Campi and A. Bouyssy, Phys. Lett. B 73, 263 (1978)
        'MGGA_X_RPP09'                 : 209, # E. Rasanen, S. Pittalis, and C. R. Proetto, J. Chem. Phys. 132, 044112 (2010)
        'MGGA_X_RSCAN'                 : 493, # A. P. Bartok and J. R. Yates, J. Chem. Phys. 150, 161101 (2019)
        'MGGA_X_RTPSS'                 : 299, # A. J. Garza, A. T. Bell, and M. Head-Gordon, J. Chem. Theory Comput. 14, 3083 (2018)
        'MGGA_X_SA_TPSS'               : 542, # L. A. Constantin, E. Fabiano, J. M. Pitarke, and F. Della Sala, Phys. Rev. B 93, 115127 (2016)
        'MGGA_X_SCAN'                  : 263, # J. Sun, A. Ruzsinszky, and J. P. Perdew, Phys. Rev. Lett. 115, 036402 (2015)
        'MGGA_X_SCANL'                 : 700, # D. Mejia-Rodriguez and S. B. Trickey, Phys. Rev. A 96, 052512 (2017)
        'MGGA_X_TASK'                  : 707, # T. Aschebrock and S. Kummel, Phys. Rev. Res. 1, 033082 (2019)
        'MGGA_X_TAU_HCTH'              : 205, # A. D. Boese and N. C. Handy, J. Chem. Phys. 116, 9559 (2002)
        'MGGA_X_TB09'                  : 208, # F. Tran and P. Blaha, Phys. Rev. Lett. 102, 226401 (2009)
        'MGGA_X_TH'                    : 225, # T. Tsuneda and K. Hirao, Phys. Rev. B 62, 15527 (2000)
        'MGGA_X_TLDA'                  : 685, # F. G. Eich and M. Hellgren, J. Chem. Phys. 141, 224107 (2014)
        'MGGA_X_TM'                    : 540, # J. Tao and Y. Mo, Phys. Rev. Lett. 117, 073001 (2016)
        'MGGA_X_TPSS'                  : 202, # J. Tao, J. P. Perdew, V. N. Staroverov, and G. E. Scuseria, Phys. Rev. Lett. 91, 146401 (2003)
        'MGGA_X_VT84'                  : 541, # J. M. del Campo, J. L. Gazquez, S. Trickey, and A. Vela, Chem. Phys. Lett. 543, 179 (2012)
        'MGGA_XC_B97M_V'               : 254, # N. Mardirossian and M. Head-Gordon, J. Chem. Phys. 142, 074111 (2015)
        'MGGA_XC_CC06'                 : 229, # A. C. Cancio and M. Y. Chou, Phys. Rev. B 74, 081202 (2006)
        'MGGA_XC_HLE17'                : 288, # P. Verma and D. G. Truhlar, J. Phys. Chem. C 121, 7144 (2017)
        'MGGA_XC_LP90'                 : 564, # C. Lee and R. G. Parr, Phys. Rev. A 42, 193 (1990)
        'MGGA_XC_OTPSS_D'              : 64 , # L. Goerigk and S. Grimme, J. Chem. Theory Comput. 6, 107 (2010)
        'MGGA_XC_TPSSLYP1W'            : 242, # E. E. Dahlke and D. G. Truhlar, J. Phys. Chem. B 109, 15677 (2005)
        'MGGA_XC_ZLP'                  : 42 , # Q. Zhao, M. Levy, and R. G. Parr, Phys. Rev. A 47, 918 (1993)
        'HYB_MGGA_X_BMK'               : 279, # A. D. Boese and J. M. L. Martin, J. Chem. Phys. 121, 3405 (2004)
        'HYB_MGGA_X_DLDF'              : 36 , # K. Pernal, R. Podeszwa, K. Patkowski, and K. Szalewicz, Phys. Rev. Lett. 103, 263201 (2009)
        'HYB_MGGA_X_JS18'              : 705, # S. Jana and P. Samal, Phys. Chem. Chem. Phys. 20, 8999 (2018)
        'HYB_MGGA_X_M05'               : 438, # Y. Zhao, N. E. Schultz, and D. G. Truhlar, J. Chem. Phys. 123, 161103 (2005)
        'HYB_MGGA_X_M05_2X'            : 439, # Y. Zhao, N. E. Schultz, and D. G. Truhlar, J. Chem. Theory Comput. 2, 364 (2006)
        'HYB_MGGA_X_M06'               : 449, # Y. Zhao and D. G. Truhlar, Theor. Chem. Acc. 120, 215 (2008)
        'HYB_MGGA_X_M06_2X'            : 450, # Y. Zhao and D. G. Truhlar, Theor. Chem. Acc. 120, 215 (2008)
        'HYB_MGGA_X_M06_HF'            : 444, # Y. Zhao and D. G. Truhlar, J. Phys. Chem. A 110, 13126 (2006)
        'HYB_MGGA_X_M06_SX'            : 310, # Y. Wang, P. Verma, L. Zhang, Y. Li, Z. Liu, D. G. Truhlar, and X. He, Proc. Natl. Acad. Sci. U. S. A. 117, 2294 (2020), https://www.pnas.org/content/117/5/2294.full.pdf
        'HYB_MGGA_X_M08_HX'            : 295, # Y. Zhao and D. G. Truhlar, J. Chem. Theory Comput. 4, 1849 (2008)
        'HYB_MGGA_X_M08_SO'            : 296, # Y. Zhao and D. G. Truhlar, J. Chem. Theory Comput. 4, 1849 (2008)
        'HYB_MGGA_X_M11'               : 297, # R. Peverati and D. G. Truhlar, J. Phys. Chem. Lett. 2, 2810 (2011)
        'HYB_MGGA_X_MN12_SX'           : 248, # R. Peverati and D. G. Truhlar, Phys. Chem. Chem. Phys. 14, 16187 (2012)
        'HYB_MGGA_X_MN15'              : 268, # H. S. Yu, X. He, S. L. Li, and D. G. Truhlar, Chem. Sci. 7, 5032 (2016)
        'HYB_MGGA_X_MS2H'              : 224, # J. Sun, R. Haunschild, B. Xiao, I. W. Bulik, G. E. Scuseria, and J. P. Perdew, J. Chem. Phys. 138, 044113 (2013)
        'HYB_MGGA_X_MVSH'              : 474, # J. Sun, J. P. Perdew, and A. Ruzsinszky, Proc. Natl. Acad. Sci. U. S. A. 112, 685 (2015)
        'HYB_MGGA_X_PJS18'             : 706, # B. Patra, S. Jana, and P. Samal, Phys. Chem. Chem. Phys. 20, 8991 (2018)
        'HYB_MGGA_X_REVM06'            : 305, # Y. Wang, P. Verma, X. Jin, D. G. Truhlar, and X. He, Proc. Natl. Acad. Sci. U. S. A. 115, 10257 (2018)
        'HYB_MGGA_X_REVM11'            : 304, # P. Verma, Y. Wang, S. Ghosh, X. He, and D. G. Truhlar, J. Phys. Chem. A 123, 2966 (2019)
        'HYB_MGGA_X_REVSCAN0'          : 583, # P. D. Mezei, G. I. Csonka, and M. Kallay, J. Chem. Theory Comput. 14, 2469 (2018)
        'HYB_MGGA_X_SCAN0'             : 264, # K. Hui and J.-D. Chai, J. Chem. Phys. 144, 044114 (2016)
        'HYB_MGGA_X_TAU_HCTH'          : 282, # A. D. Boese and N. C. Handy, J. Chem. Phys. 116, 9559 (2002)
        'HYB_MGGA_XC_B0KCIS'           : 563, # J. Toulouse, A. Savin, and C. Adamo, J. Chem. Phys. 117, 10465 (2002)
        'HYB_MGGA_XC_B86B95'           : 441, # A. D. Becke, J. Chem. Phys. 104, 1040 (1996)
        'HYB_MGGA_XC_B88B95'           : 440, # A. D. Becke, J. Chem. Phys. 104, 1040 (1996)
        'HYB_MGGA_XC_B94_HYB'          : 398, # A. D. Becke, Int. J. Quantum Chem. 52, 625 (1994)
        'HYB_MGGA_XC_B98'              : 598, # A. D. Becke, J. Chem. Phys. 109, 2092 (1998)
        'HYB_MGGA_XC_BB1K'             : 443, # Y. Zhao, B. J. Lynch, and D. G. Truhlar, J. Phys. Chem. A 108, 2715 (2004)
        'HYB_MGGA_XC_EDMGGAH'          : 695, # J. Tao, J. Chem. Phys. 116, 2335 (2002)
        'HYB_MGGA_XC_MPW1B95'          : 445, # Y. Zhao and D. G. Truhlar, J. Phys. Chem. A 108, 6908 (2004)
        'HYB_MGGA_XC_MPW1KCIS'         : 566, # Y. Zhao, N. Gonzalez-Garcia, and D. G. Truhlar, J. Phys. Chem. A 109, 2012 (2005)
        'HYB_MGGA_XC_MPWB1K'           : 446, # Y. Zhao and D. G. Truhlar, J. Phys. Chem. A 108, 6908 (2004)
        'HYB_MGGA_XC_MPWKCIS1K'        : 567, # Y. Zhao, N. Gonzalez-Garcia, and D. G. Truhlar, J. Phys. Chem. A 109, 2012 (2005)
        'HYB_MGGA_XC_PBE1KCIS'         : 568, # Y. Zhao and D. G. Truhlar, J. Chem. Theory Comput. 1, 415 (2005)
        'HYB_MGGA_XC_PW6B95'           : 451, # Y. Zhao and D. G. Truhlar, J. Phys. Chem. A 109, 5656 (2005)
        'HYB_MGGA_XC_PW86B95'          : 442, # A. D. Becke, J. Chem. Phys. 104, 1040 (1996)
        'HYB_MGGA_XC_PWB6K'            : 452, # Y. Zhao and D. G. Truhlar, J. Phys. Chem. A 109, 5656 (2005)
        'HYB_MGGA_XC_REVTPSSH'         : 458, # G. I. Csonka, J. P. Perdew, and A. Ruzsinszky, J. Chem. Theory Comput. 6, 3688 (2010)
        'HYB_MGGA_XC_TPSS0'            : 396, # S. Grimme, J. Phys. Chem. A 109, 3067 (2005)
        'HYB_MGGA_XC_TPSS1KCIS'        : 569, # Y. Zhao, B. J. Lynch, and D. G. Truhlar, Phys. Chem. Chem. Phys. 7, 43 (2005)
        'HYB_MGGA_XC_TPSSH'            : 457, # V. N. Staroverov, G. E. Scuseria, J. Tao, and J. P. Perdew, J. Chem. Phys. 119, 12129 (2003)
        'HYB_MGGA_XC_WB97M_V'          : 531, # N. Mardirossian and M. Head-Gordon, J. Chem. Phys. 144, 214110 (2016)
        'HYB_MGGA_XC_X1B95'            : 447, # Y. Zhao and D. G. Truhlar, J. Phys. Chem. A 108, 6908 (2004)
        'HYB_MGGA_XC_XB1K'             : 448, # Y. Zhao and D. G. Truhlar, J. Phys. Chem. A 108, 6908 (2004)
    }

    #PROBLEMATIC_XC = dict([(XC_CODES[x], x) for x in
    #                       ('GGA_C_SPBE', 'MGGA_X_REVTPSS')])
    PROBLEMATIC_XC = dict([])

def _xc_key_without_underscore(xc_keys):
    new_xc = []
    for key, xc_id in xc_keys.items():
        for delimeter in ('_XC_', '_X_', '_C_', '_K_'):
            if delimeter in key:
                key0, key1 = key.split(delimeter)
                new_key1 = key1.replace('_', '').replace('-', '')
                if key1 != new_key1:
                    new_xc.append((key0+delimeter+new_key1, xc_id))
                break
    return new_xc
XC_CODES.update(_xc_key_without_underscore(XC_CODES))
del(_xc_key_without_underscore)

#
# alias
#
XC_CODES.update({
    'GGA_C_BCGP'    : 'GGA_C_ACGGA',
    'LDA'           : 1 ,
    'SLATER'        : 1 ,
    'VWN3'          : 8,
    'VWNRPA'        : 8,
    'VWN5'          : 7,
    'B88'           : 106,
    'PBE0'          : 406,
    'PBE1PBE'       : 406,
    'OPTXCORR'      : '0.7344536875999693*SLATER - 0.6984752285760186*OPTX,',
    'B3LYP'         : 'B3LYP5',  # VWN5 version
    'B3LYP5'        : '.2*HF + .08*SLATER + .72*B88, .81*LYP + .19*VWN',
    'B3LYPG'        : 402,  # VWN3, used by Gaussian
    'B3P86'         : 'B3P865',  # VWN5 version
    'B3P865'        : '.2*HF + .08*SLATER + .72*B88, .81*P86 + .19*VWN',
    # FIXME: Check if Gaussian takes a different form for B3P86
    #'B3P86G'        : 403,  # VWN3, used by Gaussian
    'B3P86G'        : '.2*HF + .08*SLATER + .72*B88, .81*P86 + .19*VWN3',
    'B3PW91'        : 'B3PW915',
    'B3PW915'       : '.2*HF + .08*SLATER + .72*B88, .81*PW91 + .19*VWN',
    #'B3PW91G'       : '.2*HF + .08*SLATER + .72*B88, .81*PW91 + .19*VWN3',
    'B3PW91G'       : 401,
    #'O3LYP5'        : '.1161*HF + .9262*SLATER + .8133*OPTXCORR, .81*LYP + .19*VWN5',
    #'O3LYPG'        : '.1161*HF + .9262*SLATER + .8133*OPTXCORR, .81*LYP + .19*VWN3',
    'O3LYP'         : 404, # in libxc == '.1161*HF + 0.071006917*SLATER + .8133*OPTX, .81*LYP + .19*VWN5', may be erroreous
    'MPW3PW'        : 'MPW3PW5',  # VWN5 version
    'MPW3PW5'       : '.2*HF + .08*SLATER + .72*MPW91, .81*PW91 + .19*VWN',
    'MPW3PWG'       : 415,  # VWN3, used by Gaussian
    'MPW3LYP'       : 'MPW3LYP5',  # VWN5 version
    'MPW3LYP5'      : '.218*HF + .073*SLATER + .709*MPW91, .871*LYP + .129*VWN',
    'MPW3LYPG'      : 419,  # VWN3, used by Gaussian
    'REVB3LYP'      : 'REVB3LYP5',  # VWN5 version
    'REVB3LYP5'     : '.2*HF + .13*SLATER + .67*B88, .84*LYP + .16*VWN',
    'REVB3LYPG'     : 454,  # VWN3, used by Gaussian
    'X3LYP'         : 'X3LYP5',  # VWN5 version
    'X3LYP5'        : '.218*HF + .073*SLATER + .542385*B88 + .166615*PW91, .871*LYP + .129*VWN',
    'X3LYPG'        : 411,  # VWN3, used by Gaussian
    'CAMB3LYP'      : 'HYB_GGA_XC_CAM_B3LYP',
    'CAMYBLYP'      : 'HYB_GGA_XC_CAMY_BLYP',
    'CAMYB3LYP'     : 'HYB_GGA_XC_CAMY_B3LYP',
    'B5050LYP'      : '.5*HF + .08*SLATER + .42*B88, .81*LYP + .19*VWN',
    'MPW1LYP'       : '.25*HF + .75*MPW91, LYP',
    'MPW1PBE'       : '.25*HF + .75*MPW91, PBE',
    'PBE50'         : '.5*HF + .5*PBE, PBE',
    'REVPBE0'       : '.25*HF + .75*PBE_R, PBE',
    'B1B95'         : 440,
    'TPSS0'         : '.25*HF + .75*TPSS, TPSS',
})  # noqa: E501

XC_KEYS = set(XC_CODES.keys())

# Some XC functionals have conventional name, like M06-L means M06-L for X
# functional and M06-L for C functional, PBE mean PBE-X plus PBE-C. If the
# conventional name was placed in the XC_CODES, it may lead to recursive
# reference when parsing the xc description.  These names (as exceptions of
# XC_CODES) are listed in XC_ALIAS below and they should be treated as a
# shortcut for XC functional.
XC_ALIAS = {
    # Conventional name : name in XC_CODES
    'BLYP'              : 'B88,LYP',
    'BP86'              : 'B88,P86',
    'PW91'              : 'PW91,PW91',
    'PBE'               : 'PBE,PBE',
    'REVPBE'            : 'PBE_R,PBE',
    'PBESOL'            : 'PBE_SOL,PBE_SOL',
    'PKZB'              : 'PKZB,PKZB',
    'TPSS'              : 'TPSS,TPSS',
    'REVTPSS'           : 'REVTPSS,REVTPSS',
    'SCAN'              : 'SCAN,SCAN',
    'RSCAN'             : 'RSCAN,RSCAN',
    'R2SCAN'            : 'R2SCAN,R2SCAN',
    'SCANL'             : 'SCANL,SCANL',
    'R2SCANL'           : 'R2SCANL,R2SCANL',
    'SOGGA'             : 'SOGGA,PBE',
    'BLOC'              : 'BLOC,TPSSLOC',
    'OLYP'              : 'OPTX,LYP',
    'OPBE'              : 'OPTX,PBE',
    'RPBE'              : 'RPBE,PBE',
    'BPBE'              : 'B88,PBE',
    'MPW91'             : 'MPW91,PW91',
    'HFLYP'             : 'HF,LYP',
    'HFPW92'            : 'HF,PW_MOD',
    'SPW92'             : 'SLATER,PW_MOD',
    'SVWN'              : 'SLATER,VWN',
    'MS0'               : 'MS0,REGTPSS',
    'MS1'               : 'MS1,REGTPSS',
    'MS2'               : 'MS2,REGTPSS',
    'MS2H'              : 'MS2H,REGTPSS',
    'MVS'               : 'MVS,REGTPSS',
    'MVSH'              : 'MVSH,REGTPSS',
    'SOGGA11'           : 'SOGGA11,SOGGA11',
    'SOGGA11_X'         : 'SOGGA11_X,SOGGA11_X',
    'KT1'               : 'KT1,VWN',
    'KT2'               : 'GGA_XC_KT2',
    'KT3'               : 'GGA_XC_KT3',
    'DLDF'              : 'DLDF,DLDF',
    'GAM'               : 'GAM,GAM',
    'M06_L'             : 'M06_L,M06_L',
    'M06_SX'            : 'M06_SX,M06_SX',
    'M11_L'             : 'M11_L,M11_L',
    'MN12_L'            : 'MN12_L,MN12_L',
    'MN15_L'            : 'MN15_L,MN15_L',
    'N12'               : 'N12,N12',
    'N12_SX'            : 'N12_SX,N12_SX',
    'MN12_SX'           : 'MN12_SX,MN12_SX',
    'MN15'              : 'MN15,MN15',
    'MBEEF'             : 'MBEEF,PBE_SOL',
    'SCAN0'             : 'SCAN0,SCAN',
    'PBEOP'             : 'PBE,OP_PBE',
    'BOP'               : 'B88,OP_B88',
    # new in libxc-4.2.3
    'REVSCAN'           : 'MGGA_X_REVSCAN,MGGA_C_REVSCAN',
    'REVSCAN_VV10'      : 'MGGA_X_REVSCAN,MGGA_C_REVSCAN_VV10',
    'SCAN_VV10'         : 'MGGA_X_SCAN,MGGA_C_SCAN_VV10',
    'SCAN_RVV10'        : 'MGGA_X_SCAN,MGGA_C_SCAN_RVV10',
    'M05'               : 'HYB_MGGA_X_M05,MGGA_C_M05',
    'M06'               : 'HYB_MGGA_X_M06,MGGA_C_M06',
    'M05_2X'            : 'HYB_MGGA_X_M05_2X,MGGA_C_M05_2X',
    'M06_2X'            : 'HYB_MGGA_X_M06_2X,MGGA_C_M06_2X',
    # extra aliases
    'SOGGA11X'          : 'SOGGA11_X',
    'M06L'              : 'M06_L',
    'M11L'              : 'M11_L',
    'MN12L'             : 'MN12_L',
    'MN15L'             : 'MN15_L',
    'N12SX'             : 'N12_SX',
    'MN12SX'            : 'MN12_SX',
    'M052X'             : 'M05_2X',
    'M062X'             : 'M06_2X',
}  # noqa: E122
XC_ALIAS.update([(key.replace('-',''), XC_ALIAS[key])
                 for key in XC_ALIAS if '-' in key])

VV10_XC = set(('B97M_V', 'WB97M_V', 'WB97X_V', 'VV10', 'LC_VV10',
               'REVSCAN_VV10',
               'SCAN_VV10', 'SCAN_RVV10', 'SCANL_VV10', 'SCANL_RVV10'))
VV10_XC = VV10_XC.union(set([x.replace('_', '') for x in VV10_XC]))

def xc_reference(xc_code):
    '''Returns the reference to the individual XC functional'''
    hyb, fn_facs = parse_xc(xc_code)
    refs = []
    c_refs = (ctypes.c_char_p * 8)()
    for xid, fac in fn_facs:
        _itrf.LIBXC_xc_reference(xid, c_refs)
        for ref in c_refs:
            if ref:
                refs.append(ref.decode("UTF-8"))
    return refs

def xc_type(xc_code):
    if xc_code is None:
        return None
    elif isinstance(xc_code, str):
        if is_nlc(xc_code):
            return 'NLC'
        hyb, fn_facs = parse_xc(xc_code)
    else:
        fn_facs = [(xc_code, 1)]  # mimic fn_facs
    if not fn_facs:
        return 'HF'
    elif all(_itrf.LIBXC_is_lda(ctypes.c_int(xid)) for xid, fac in fn_facs):
        return 'LDA'
    elif any(_itrf.LIBXC_is_meta_gga(ctypes.c_int(xid)) for xid, fac in fn_facs):
        return 'MGGA'
    else:
        # any(_itrf.LIBXC_is_gga(ctypes.c_int(xid)) for xid, fac in fn_facs)
        # include hybrid_xc
        return 'GGA'

def is_lda(xc_code):
    return xc_type(xc_code) == 'LDA'

def is_hybrid_xc(xc_code):
    if xc_code is None:
        return False
    elif isinstance(xc_code, str):
        if xc_code.isdigit():
            return _itrf.LIBXC_is_hybrid(ctypes.c_int(int(xc_code)))
        else:
            if 'HF' in xc_code:
                return True
            if hybrid_coeff(xc_code) != 0:
                return True
            if rsh_coeff(xc_code) != [0, 0, 0]:
                return True
            return False
    elif isinstance(xc_code, int):
        return _itrf.LIBXC_is_hybrid(ctypes.c_int(xc_code))
    else:
        return any((is_hybrid_xc(x) for x in xc_code))

def is_meta_gga(xc_code):
    return xc_type(xc_code) == 'MGGA'

def is_gga(xc_code):
    return xc_type(xc_code) == 'GGA'

def needs_laplacian(xc_code):
    return _itrf.LIBXC_needs_laplacian(xc_code) != 0

def is_nlc(xc_code):
    return '__VV10' in xc_code.upper()

def max_deriv_order(xc_code):
    hyb, fn_facs = parse_xc(xc_code)
    if fn_facs:
        return min(_itrf.LIBXC_max_deriv_order(ctypes.c_int(xid)) for xid, fac in fn_facs)
    else:
        return 3

def test_deriv_order(xc_code, deriv, raise_error=False):
    support = deriv <= max_deriv_order(xc_code)
    if not support and raise_error:
        from pyscf.dft import xcfun
        msg = ('libxc library does not support derivative order %d for  %s' %
               (deriv, xc_code))
        try:
            if xcfun.test_deriv_order(xc_code, deriv, raise_error=False):
                msg += ('''
    This functional derivative is supported in the xcfun library.
    The following code can be used to change the libxc library to xcfun library:

        from pyscf.dft import xcfun
        mf._numint.libxc = xcfun
''')
            raise NotImplementedError(msg)
        except KeyError as e:
            sys.stderr.write('\n'+msg+'\n')
            sys.stderr.write('%s not found in xcfun library\n\n' % xc_code)
            raise e
    return support

def hybrid_coeff(xc_code, spin=0):
    '''Support recursively defining hybrid functional
    '''
    hyb, fn_facs = parse_xc(xc_code)
    for xid, fac in fn_facs:
        hyb[0] += fac * _itrf.LIBXC_hybrid_coeff(ctypes.c_int(xid))
    return hyb[0]

def nlc_coeff(xc_code):
    '''Get NLC coefficients
    '''
    nlc_code = None
    if isinstance(xc_code, str) and '__VV10' in xc_code.upper():
        xc_code, nlc_code = xc_code.upper().split('__', 1)

    hyb, fn_facs = parse_xc(xc_code)
    nlc_pars = [0, 0]
    nlc_tmp = (ctypes.c_double*2)()
    for xid, fac in fn_facs:
        _itrf.LIBXC_nlc_coeff(xid, nlc_tmp)
        nlc_pars[0] += nlc_tmp[0]
        nlc_pars[1] += nlc_tmp[1]

    if nlc_pars[0] == 0 and nlc_pars[1] == 0:
        if nlc_code is not None:
            # Use VV10 NLC parameters by default for the general case
            _itrf.LIBXC_nlc_coeff(XC_CODES['GGA_XC_' + nlc_code], nlc_tmp)
            nlc_pars[0] += nlc_tmp[0]
            nlc_pars[1] += nlc_tmp[1]
        else:
            raise NotImplementedError(
                '%s does not have NLC part. Available functionals are %s' %
                (xc_code, ', '.join(VV10_XC.keys())))
    return nlc_pars

def rsh_coeff(xc_code):
    '''Range-separated parameter and HF exchange components: omega, alpha, beta

    Exc_RSH = c_LR * LR_HFX + c_SR * SR_HFX + (1-c_SR) * Ex_SR + (1-c_LR) * Ex_LR + Ec
            = alpha * HFX   + beta * SR_HFX + (1-c_SR) * Ex_SR + (1-c_LR) * Ex_LR + Ec
            = alpha * LR_HFX + hyb * SR_HFX + (1-c_SR) * Ex_SR + (1-c_LR) * Ex_LR + Ec

    SR_HFX = < pi | e^{-omega r_{12}}/r_{12} | iq >
    LR_HFX = < pi | (1-e^{-omega r_{12}})/r_{12} | iq >
    alpha = c_LR
    beta = c_SR - c_LR = hyb - alpha
    '''
    if xc_code is None:
        return 0, 0, 0

    check_omega = True
    if isinstance(xc_code, str) and ',' in xc_code:
        # Parse only X part for the RSH coefficients.  This is to handle
        # exceptions for C functionals such as M11.
        xc_code = format_xc_code(xc_code)
        xc_code = xc_code.split(',')[0] + ','
        if 'SR_HF' in xc_code or 'LR_HF' in xc_code or 'RSH(' in xc_code:
            check_omega = False

    hyb, fn_facs = parse_xc(xc_code)

    hyb, alpha, omega = hyb
    beta = hyb - alpha
    rsh_pars = [omega, alpha, beta]
    rsh_tmp = (ctypes.c_double*3)()
    _itrf.LIBXC_rsh_coeff(433, rsh_tmp)
    for xid, fac in fn_facs:
        _itrf.LIBXC_rsh_coeff(xid, rsh_tmp)
        if rsh_pars[0] == 0:
            rsh_pars[0] = rsh_tmp[0]
        elif check_omega:
            # Check functional is actually a CAM functional
            if rsh_tmp[0] != 0 and not _itrf.LIBXC_is_cam_rsh(ctypes.c_int(xid)):
                raise KeyError('Libxc functional %i employs a range separation '
                               'kernel that is not supported in PySCF' % xid)
            # Check omega
            if (rsh_tmp[0] != 0 and rsh_pars[0] != rsh_tmp[0]):
                raise ValueError('Different values of omega found for RSH functionals')
        rsh_pars[1] += rsh_tmp[1] * fac
        rsh_pars[2] += rsh_tmp[2] * fac
    return rsh_pars

def parse_xc_name(xc_name='LDA,VWN'):
    '''Convert the XC functional name to libxc library internal ID.
    '''
    fn_facs = parse_xc(xc_name)[1]
    return fn_facs[0][0], fn_facs[1][0]

def parse_xc(description):
    r'''Rules to input functional description:

    * The given functional description must be a one-line string.
    * The functional description is case-insensitive.
    * The functional description string has two parts, separated by ",".  The
      first part describes the exchange functional, the second is the correlation
      functional.

      - If "," was not in string, the entire string is considered as a
        compound XC functional (including both X and C functionals, such as b3lyp).
      - To input only X functional (without C functional), leave the second
        part blank. E.g. description='slater,' means pure LDA functional.
      - To neglect X functional (just apply C functional), leave the first
        part blank. E.g. description=',vwn' means pure VWN functional.
      - If compound XC functional is specified, no matter whehter it is in the
        X part (the string in front of comma) or the C part (the string behind
        comma), both X and C functionals of the compound XC functional will be
        used.

    * The functional name can be placed in arbitrary order.  Two name needs to
      be separated by operators "+" or "-".  Blank spaces are ignored.
      NOTE the parser only reads operators "+" "-" "*".  / is not in support.
    * A functional name can have at most one factor.  If the factor is not
      given, it is set to 1.  Compound functional can be scaled as a unit. For
      example '0.5*b3lyp' is equivalent to
      'HF*0.1 + .04*LDA + .36*B88, .405*LYP + .095*VWN'
    * String "HF" stands for exact exchange (HF K matrix).  Putting "HF" in
      correlation functional part is the same to putting "HF" in exchange
      part.
    * String "RSH" means range-separated operator. Its format is
      RSH(omega, alpha, beta).  Another way to input RSH is to use keywords
      SR_HF and LR_HF: "SR_HF(0.1) * alpha_plus_beta" and "LR_HF(0.1) *
      alpha" where the number in parenthesis is the value of omega.
    * Be careful with the libxc convention on GGA functional, in which the LDA
      contribution has been included.

    Args:
        xc_code : str
            A string to describe the linear combination of different XC functionals.
            The X and C functional are separated by comma like '.8*LDA+.2*B86,VWN'.
            If "HF" was appeared in the string, it stands for the exact exchange.
        rho : ndarray
            Shape of ((*,N)) for electron density (and derivatives) if spin = 0;
            Shape of ((*,N),(*,N)) for alpha/beta electron density (and derivatives) if spin > 0;
            where N is number of grids.
            rho (*,N) are ordered as (den,grad_x,grad_y,grad_z,laplacian,tau)
            where grad_x = d/dx den, laplacian = \nabla^2 den, tau = 1/2(\nabla f)^2
            In spin unrestricted case,
            rho is ((den_u,grad_xu,grad_yu,grad_zu,laplacian_u,tau_u)
                    (den_d,grad_xd,grad_yd,grad_zd,laplacian_d,tau_d))

    Kwargs:
        spin : int
            spin polarized if spin > 0
        relativity : int
            No effects.
        verbose : int or object of :class:`Logger`
            No effects.

    Returns:
        ex, vxc, fxc, kxc

        where

        * vxc = (vrho, vsigma, vlapl, vtau) for restricted case

        * vxc for unrestricted case
          | vrho[:,2]   = (u, d)
          | vsigma[:,3] = (uu, ud, dd)
          | vlapl[:,2]  = (u, d)
          | vtau[:,2]   = (u, d)

        * fxc for restricted case:
          (v2rho2, v2rhosigma, v2sigma2, v2lapl2, vtau2, v2rholapl, v2rhotau, v2lapltau, v2sigmalapl, v2sigmatau)

        * fxc for unrestricted case:
          | v2rho2[:,3]     = (u_u, u_d, d_d)
          | v2rhosigma[:,6] = (u_uu, u_ud, u_dd, d_uu, d_ud, d_dd)
          | v2sigma2[:,6]   = (uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd)
          | v2lapl2[:,3]
          | v2tau2[:,3]     = (u_u, u_d, d_d)
          | v2rholapl[:,4]
          | v2rhotau[:,4]   = (u_u, u_d, d_u, d_d)
          | v2lapltau[:,4]
          | v2sigmalapl[:,6]
          | v2sigmatau[:,6] = (uu_u, uu_d, ud_u, ud_d, dd_u, dd_d)

        * kxc for restricted case:
          (v3rho3, v3rho2sigma, v3rhosigma2, v3sigma3,
           v3rho2lapl, v3rho2tau,
           v3rhosigmalapl, v3rhosigmatau,
           v3rholapl2, v3rholapltau, v3rhotau2,
           v3sigma2lapl, v3sigma2tau,
           v3sigmalapl2, v3sigmalapltau, v3sigmatau2,
           v3lapl3, v3lapl2tau, v3lapltau2, v3tau3)

        * kxc for unrestricted case:
          | v3rho3[:,4]         = (u_u_u, u_u_d, u_d_d, d_d_d)
          | v3rho2sigma[:,9]    = (u_u_uu, u_u_ud, u_u_dd, u_d_uu, u_d_ud, u_d_dd, d_d_uu, d_d_ud, d_d_dd)
          | v3rhosigma2[:,12]   = (u_uu_uu, u_uu_ud, u_uu_dd, u_ud_ud, u_ud_dd, u_dd_dd, d_uu_uu, d_uu_ud, d_uu_dd, d_ud_ud, d_ud_dd, d_dd_dd)
          | v3sigma3[:,10]      = (uu_uu_uu, uu_uu_ud, uu_uu_dd, uu_ud_ud, uu_ud_dd, uu_dd_dd, ud_ud_ud, ud_ud_dd, ud_dd_dd, dd_dd_dd)
          | v3rho2lapl[:,6]
          | v3rho2tau[:,6]      = (u_u_u, u_u_d, u_d_u, u_d_d, d_d_u, d_d_d)
          | v3rhosigmalapl[:,12]
          | v3rhosigmatau[:,12] = (u_uu_u, u_uu_d, u_ud_u, u_ud_d, u_dd_u, u_dd_d,
                                   d_uu_u, d_uu_d, d_ud_u, d_ud_d, d_dd_u, d_dd_d)
          | v3rholapl2[:,6]
          | v3rholapltau[:,8]
          | v3rhotau2[:,6]      = (u_u_u, u_u_d, u_d_d, d_u_u, d_u_d, d_d_d)
          | v3sigma2lapl[:,12]
          | v3sigma2tau[:,12]   = (uu_uu_u, uu_uu_d, uu_ud_u, uu_ud_d, uu_dd_u, uu_dd_d,
                                   ud_ud_u, ud_ud_d, ud_dd_u, ud_dd_d, dd_dd_u, dd_dd_d)
          | v3sigmalapl2[:,9]
          | v3sigmalapltau[:,12]
          | v3sigmatau2[:,9]    = (uu_u_u, uu_u_d, uu_d_d, ud_u_u, ud_u_d, ud_d_d, dd_u_u, dd_u_d, dd_d_d)
          | v3lapl3[:,4]
          | v3lapl2tau[:,6]
          | v3lapltau2[:,6]
          | v3tau3[:,4]         = (u_u_u, u_u_d, u_d_d, d_d_d)

        see also libxc_itrf.c
    '''  # noqa: E501
    hyb = [0, 0, 0]  # hybrid, alpha, omega (== SR_HF, LR_HF, omega)
    if description is None:
        return hyb, []
    elif isinstance(description, int):
        return hyb, [(description, 1.)]
    elif not isinstance(description, str): #isinstance(description, (tuple,list)):
        return parse_xc('%s,%s' % tuple(description))

    def assign_omega(omega, hyb_or_sr, lr=0):
        if hyb[2] == omega or omega == 0:
            hyb[0] += hyb_or_sr
            hyb[1] += lr
        elif hyb[2] == 0:
            hyb[0] += hyb_or_sr
            hyb[1] += lr
            hyb[2] = omega
        else:
            raise ValueError('Different values of omega found for RSH functionals')
    fn_facs = []

    def parse_token(token, ftype, search_xc_alias=False):
        if token:
            if token[0] == '-':
                sign = -1
                token = token[1:]
            else:
                sign = 1
            if '*' in token:
                fac, key = token.split('*')
                if fac[0].isalpha():
                    fac, key = key, fac
                fac = sign * float(fac)
            else:
                fac, key = sign, token

            if key[:3] == 'RSH':
                # RSH(alpha; beta; omega): Range-separated-hybrid functional
                # See also utils.format_xc_code
                alpha, beta, omega = [float(x) for x in key[4:-1].split(';')]
                assign_omega(omega, fac*(alpha+beta), fac*alpha)
            elif key == 'HF':
                hyb[0] += fac
                hyb[1] += fac  # also add to LR_HF
            elif 'SR_HF' in key:
                if '(' in key:
                    omega = float(key.split('(')[1].split(')')[0])
                    assign_omega(omega, fac, 0)
                else:  # Assuming this omega the same to the existing omega
                    hyb[0] += fac
            elif 'LR_HF' in key:
                if '(' in key:
                    omega = float(key.split('(')[1].split(')')[0])
                    assign_omega(omega, 0, fac)
                else:
                    hyb[1] += fac  # == alpha
            elif key.isdigit():
                fn_facs.append((int(key), fac))
            else:
                if search_xc_alias and key in XC_ALIAS:
                    x_id = XC_ALIAS[key]
                elif key in XC_CODES:
                    x_id = XC_CODES[key]
                else:
                    possible_xc_for = fpossible_dic[ftype]
                    possible_xc = XC_KEYS.intersection(possible_xc_for(key))
                    if possible_xc:
                        if len(possible_xc) > 1:
                            sys.stderr.write('Possible xc_code %s matches %s. '
                                             % (list(possible_xc), key))
                            for x_id in possible_xc:  # Prefer X functional
                                if '_X_' in x_id:
                                    break
                            else:
                                x_id = possible_xc.pop()
                            sys.stderr.write('XC parser takes %s\n' % x_id)
                            sys.stderr.write('You can add prefix to %s for a '
                                             'specific functional (e.g. X_%s, '
                                             'HYB_MGGA_X_%s)\n'
                                             % (key, key, key))
                        else:
                            x_id = possible_xc.pop()
                        x_id = XC_CODES[x_id]
                    else:
                        raise KeyError('Unknown %s functional  %s' % (ftype, key))
                if isinstance(x_id, str):
                    hyb1, fn_facs1 = parse_xc(x_id)
# Recursively scale the composed functional, to support e.g. '0.5*b3lyp'
                    if hyb1[0] != 0 or hyb1[1] != 0:
                        assign_omega(hyb1[2], hyb1[0]*fac, hyb1[1]*fac)
                    fn_facs.extend([(xid, c*fac) for xid, c in fn_facs1])
                elif x_id is None:
                    raise NotImplementedError('%s functional %s' % (ftype, key))
                else:
                    fn_facs.append((x_id, fac))
    def possible_x_for(key):
        return set((key,
                    'LDA_X_'+key, 'GGA_X_'+key, 'MGGA_X_'+key,
                    'HYB_GGA_X_'+key, 'HYB_MGGA_X_'+key))
    def possible_xc_for(key):
        return set((key, 'LDA_XC_'+key, 'GGA_XC_'+key, 'MGGA_XC_'+key,
                    'HYB_GGA_XC_'+key, 'HYB_MGGA_XC_'+key))
    def possible_k_for(key):
        return set((key,
                    'LDA_K_'+key, 'GGA_K_'+key,))
    def possible_x_k_for(key):
        return possible_x_for(key).union(possible_k_for(key))
    def possible_c_for(key):
        return set((key,
                    'LDA_C_'+key, 'GGA_C_'+key, 'MGGA_C_'+key))
    fpossible_dic = {'X': possible_x_for,
                     'C': possible_c_for,
                     'compound XC': possible_xc_for,
                     'K': possible_k_for,
                     'X or K': possible_x_k_for}

    description = format_xc_code(description)

    if '-' in description:  # To handle e.g. M06-L
        for key in _NAME_WITH_DASH:
            if key in description:
                description = description.replace(key, _NAME_WITH_DASH[key])

    if ',' in description:
        x_code, c_code = description.split(',')
        for token in x_code.replace('-', '+-').replace(';+', ';').split('+'):
            parse_token(token, 'X or K')
        for token in c_code.replace('-', '+-').replace(';+', ';').split('+'):
            parse_token(token, 'C')
    else:
        for token in description.replace('-', '+-').replace(';+', ';').split('+'):
            parse_token(token, 'compound XC', search_xc_alias=True)
    if hyb[2] == 0: # No omega is assigned. LR_HF is 0 for normal Coulomb operator
        hyb[1] = 0
    return hyb, remove_dup(fn_facs)

_NAME_WITH_DASH = {'SR-HF'    : 'SR_HF',
                   'LR-HF'    : 'LR_HF',
                   'OTPSS-D'  : 'OTPSS_D',
                   'B97-1'    : 'B97_1',
                   'B97-2'    : 'B97_2',
                   'B97-3'    : 'B97_3',
                   'B97-K'    : 'B97_K',
                   'B97-D'    : 'B97_D',
                   'HCTH-93'  : 'HCTH_93',
                   'HCTH-120' : 'HCTH_120',
                   'HCTH-147' : 'HCTH_147',
                   'HCTH-407' : 'HCTH_407',
                   'WB97X-D'  : 'WB97X_D',
                   'WB97X-V'  : 'WB97X_V',
                   'WB97M-V'  : 'WB97M_V',
                   'B97M-V'   : 'B97M_V',
                   'M05-2X'   : 'M05_2X',
                   'M06-L'    : 'M06_L',
                   'M06-HF'   : 'M06_HF',
                   'M06-2X'   : 'M06_2X',
                   'M08-HX'   : 'M08_HX',
                   'M08-SO'   : 'M08_SO',
                   'M11-L'    : 'M11_L',
                   'MN12-L'   : 'MN12_L',
                   'MN15-L'   : 'MN15_L',
                   'MN12-SX'  : 'MN12_SX',
                   'N12-SX'   : 'N12_SX',
                   'LRC-WPBE' : 'LRC_WPBE',
                   'LRC-WPBEH': 'LRC_WPBEH',
                   'LC-VV10'  : 'LC_VV10',
                   'CAM-B3LYP': 'CAM_B3LYP'}


def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=None):
    r'''Interface to call libxc library to evaluate XC functional, potential
    and functional derivatives.

    * The given functional xc_code must be a one-line string.
    * The functional xc_code is case-insensitive.
    * The functional xc_code string has two parts, separated by ",".  The
      first part describes the exchange functional, the second part sets the
      correlation functional.

      - If "," not appeared in string, the entire string is treated as the
        name of a compound functional (containing both the exchange and
        the correlation functional) which was declared in the functional
        aliases list. The full list of functional aliases can be obtained by
        calling the function pyscf.dft.xcfun.XC_ALIAS.keys() .

        If the string was not found in the aliased functional list, it is
        treated as X functional.

      - To input only X functional (without C functional), leave the second
        part blank. E.g. description='slater,' means a functional with LDA
        contribution only.

      - To neglect the contribution of X functional (just apply C functional),
        leave blank in the first part, e.g. description=',vwn' means a
        functional with VWN only.

      - If compound XC functional is specified, no matter whether it is in the
        X part (the string in front of comma) or the C part (the string behind
        comma), both X and C functionals of the compound XC functional will be
        used.

    * The functional name can be placed in arbitrary order.  Two names need to
      be separated by operators "+" or "-".  Blank spaces are ignored.
      NOTE the parser only reads operators "+" "-" "*".  / is not supported.

    * A functional name can have at most one factor.  If the factor is not
      given, it is set to 1.  Compound functional can be scaled as a unit. For
      example '0.5*b3lyp' is equivalent to
      'HF*0.1 + .04*LDA + .36*B88, .405*LYP + .095*VWN'

    * String "HF" stands for exact exchange (HF K matrix).  "HF" can be put in
      the correlation functional part (after comma). Putting "HF" in the
      correlation part is the same to putting "HF" in the exchange part.

    * String "RSH" means range-separated operator. Its format is
      RSH(omega, alpha, beta).  Another way to input RSH is to use keywords
      SR_HF and LR_HF: "SR_HF(0.1) * alpha_plus_beta" and "LR_HF(0.1) *
      alpha" where the number in parenthesis is the value of omega.

    * Be careful with the libxc convention of GGA functional, in which the LDA
      contribution is included.

    Args:
        xc_code : str
            A string to describe the linear combination of different XC functionals.
            The X and C functional are separated by comma like '.8*LDA+.2*B86,VWN'.
            If "HF" (exact exchange) is appeared in the string, the HF part will
            be skipped.  If an empty string "" is given, the returns exc, vxc,...
            will be vectors of zeros.
        rho : ndarray
            Shape of ((*,N)) for electron density (and derivatives) if spin = 0;
            Shape of ((*,N),(*,N)) for alpha/beta electron density (and derivatives) if spin > 0;
            where N is number of grids.
            rho (*,N) are ordered as (den,grad_x,grad_y,grad_z,laplacian,tau)
            where grad_x = d/dx den, laplacian = \nabla^2 den, tau = 1/2(\nabla f)^2
            In spin unrestricted case,
            rho is ((den_u,grad_xu,grad_yu,grad_zu,laplacian_u,tau_u)
                    (den_d,grad_xd,grad_yd,grad_zd,laplacian_d,tau_d))

    Kwargs:
        spin : int
            spin polarized if spin > 0
        relativity : int
            No effects.
        verbose : int or object of :class:`Logger`
            No effects.

    Returns:
        ex, vxc, fxc, kxc

        where

        * vxc = (vrho, vsigma, vlapl, vtau) for restricted case

        * vxc for unrestricted case
          | vrho[:,2]   = (u, d)
          | vsigma[:,3] = (uu, ud, dd)
          | vlapl[:,2]  = (u, d)
          | vtau[:,2]   = (u, d)

        * fxc for restricted case:
          (v2rho2, v2rhosigma, v2sigma2, v2lapl2, vtau2, v2rholapl, v2rhotau, v2lapltau, v2sigmalapl, v2sigmatau)

        * fxc for unrestricted case:
          | v2rho2[:,3]     = (u_u, u_d, d_d)
          | v2rhosigma[:,6] = (u_uu, u_ud, u_dd, d_uu, d_ud, d_dd)
          | v2sigma2[:,6]   = (uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd)
          | v2lapl2[:,3]
          | v2tau2[:,3]     = (u_u, u_d, d_d)
          | v2rholapl[:,4]
          | v2rhotau[:,4]   = (u_u, u_d, d_u, d_d)
          | v2lapltau[:,4]
          | v2sigmalapl[:,6]
          | v2sigmatau[:,6] = (uu_u, uu_d, ud_u, ud_d, dd_u, dd_d)

        * kxc for restricted case:
          (v3rho3, v3rho2sigma, v3rhosigma2, v3sigma3,
           v3rho2lapl, v3rho2tau,
           v3rhosigmalapl, v3rhosigmatau,
           v3rholapl2, v3rholapltau, v3rhotau2,
           v3sigma2lapl, v3sigma2tau,
           v3sigmalapl2, v3sigmalapltau, v3sigmatau2,
           v3lapl3, v3lapl2tau, v3lapltau2, v3tau3)

        * kxc for unrestricted case:
          | v3rho3[:,4]         = (u_u_u, u_u_d, u_d_d, d_d_d)
          | v3rho2sigma[:,9]    = (u_u_uu, u_u_ud, u_u_dd, u_d_uu, u_d_ud, u_d_dd, d_d_uu, d_d_ud, d_d_dd)
          | v3rhosigma2[:,12]   = (u_uu_uu, u_uu_ud, u_uu_dd, u_ud_ud, u_ud_dd, u_dd_dd, d_uu_uu, d_uu_ud, d_uu_dd, d_ud_ud, d_ud_dd, d_dd_dd)
          | v3sigma3[:,10]      = (uu_uu_uu, uu_uu_ud, uu_uu_dd, uu_ud_ud, uu_ud_dd, uu_dd_dd, ud_ud_ud, ud_ud_dd, ud_dd_dd, dd_dd_dd)
          | v3rho2lapl[:,6]
          | v3rho2tau[:,6]      = (u_u_u, u_u_d, u_d_u, u_d_d, d_d_u, d_d_d)
          | v3rhosigmalapl[:,12]
          | v3rhosigmatau[:,12] = (u_uu_u, u_uu_d, u_ud_u, u_ud_d, u_dd_u, u_dd_d,
                                   d_uu_u, d_uu_d, d_ud_u, d_ud_d, d_dd_u, d_dd_d)
          | v3rholapl2[:,6]
          | v3rholapltau[:,8]
          | v3rhotau2[:,6]      = (u_u_u, u_u_d, u_d_d, d_u_u, d_u_d, d_d_d)
          | v3sigma2lapl[:,12]
          | v3sigma2tau[:,12]   = (uu_uu_u, uu_uu_d, uu_ud_u, uu_ud_d, uu_dd_u, uu_dd_d,
                                   ud_ud_u, ud_ud_d, ud_dd_u, ud_dd_d, dd_dd_u, dd_dd_d)
          | v3sigmalapl2[:,9]
          | v3sigmalapltau[:,12]
          | v3sigmatau2[:,9]    = (uu_u_u, uu_u_d, uu_d_d, ud_u_u, ud_u_d, ud_d_d, dd_u_u, dd_u_d, dd_d_d)
          | v3lapl3[:,4]
          | v3lapl2tau[:,6]
          | v3lapltau2[:,6]
          | v3tau3[:,4]         = (u_u_u, u_u_d, u_d_d, d_d_d)

        see also libxc_itrf.c
    '''  # noqa: E501
    hyb, fn_facs = parse_xc(xc_code)
    if omega is not None:
        hyb[2] = float(omega)
    return _eval_xc(hyb, fn_facs, rho, spin, relativity, deriv, verbose)


def _eval_xc(hyb, fn_facs, rho, spin=0, relativity=0, deriv=1, verbose=None):
    assert(deriv <= 3)
    if spin == 0:
        nspin = 1
        rho_u = rho_d = numpy.asarray(rho, order='C')
    else:
        nspin = 2
        rho_u = numpy.asarray(rho[0], order='C')
        rho_d = numpy.asarray(rho[1], order='C')
    assert(rho_u.dtype == numpy.double)
    assert(rho_d.dtype == numpy.double)

    if rho_u.ndim == 1:
        rho_u = rho_u.reshape(1,-1)
        rho_d = rho_d.reshape(1,-1)
    ngrids = rho_u.shape[1]

    fn_ids = [x[0] for x in fn_facs]
    facs   = [x[1] for x in fn_facs]
    if hyb[2] != 0:
        # Current implementation does not support different omegas for
        # different RSH functionals if there are multiple RSHs
        omega = [hyb[2]] * len(facs)
    else:
        omega = [0] * len(facs)
    fn_ids_set = set(fn_ids)
    if fn_ids_set.intersection(PROBLEMATIC_XC):
        problem_xc = [PROBLEMATIC_XC[k]
                      for k in fn_ids_set.intersection(PROBLEMATIC_XC)]
        warnings.warn('Libxc functionals %s may have discrepancy to xcfun '
                      'library.\n' % problem_xc)

    if any([needs_laplacian(fid) for fid in fn_ids]):
        raise NotImplementedError('laplacian in meta-GGA method')

    n = len(fn_ids)
    if (n == 0 or  # xc_code = '' or xc_code = 'HF', an empty functional
        all((is_lda(x) for x in fn_ids))):
        if spin == 0:
            nvar = 1
            xctype = 'R-LDA'
        else:
            nvar = 2
            xctype = 'U-LDA'
    elif any((is_meta_gga(x) for x in fn_ids)):
        if spin == 0:
            nvar = 4
            xctype = 'R-MGGA'
        else:
            nvar = 9
            xctype = 'U-MGGA'
    else:  # GGA
        if spin == 0:
            nvar = 2
            xctype = 'R-GGA'
        else:
            nvar = 5
            xctype = 'U-GGA'

    # Check that the density rho has the appropriate shape
    # should it be >= or ==, in test test_xcfun.py, test_vs_libxc_rks
    # the density contain 6 rows independently of the functional
    if xctype[2:] == 'LDA':
        for rho_ud in [rho_u, rho_d]:
            assert rho_ud.shape[0] >= 1

    elif xctype[2:] == 'GGA':
        for rho_ud in [rho_u, rho_d]:
            assert rho_ud.shape[0] >= 4

    elif xctype[2:] == 'MGGA':
        for rho_ud in [rho_u, rho_d]:
            assert rho_ud.shape[0] >= 6
    else:
        raise ValueError("Unknow nvar {}".format(nvar))

    outlen = (math.factorial(nvar+deriv) //
              (math.factorial(nvar) * math.factorial(deriv)))
    outbuf = numpy.zeros((outlen,ngrids))

    density_threshold = 0
    _itrf.LIBXC_eval_xc(ctypes.c_int(n),
                        (ctypes.c_int*n)(*fn_ids),
                        (ctypes.c_double*n)(*facs),
                        (ctypes.c_double*n)(*omega),
                        ctypes.c_int(nspin),
                        ctypes.c_int(deriv), ctypes.c_int(rho_u.shape[1]),
                        rho_u.ctypes.data_as(ctypes.c_void_p),
                        rho_d.ctypes.data_as(ctypes.c_void_p),
                        outbuf.ctypes.data_as(ctypes.c_void_p),
                        ctypes.c_double(density_threshold))

    exc = outbuf[0]
    vxc = fxc = kxc = None
    if xctype == 'R-LDA':
        if deriv > 0:
            vxc = [outbuf[1]]
        if deriv > 1:
            fxc = [outbuf[2]]
        if deriv > 2:
            kxc = [outbuf[3]]
    elif xctype == 'R-GGA':
        if deriv > 0:
            vxc = [outbuf[1], outbuf[2]]
        if deriv > 1:
            fxc = [outbuf[3], outbuf[4], outbuf[5]]
        if deriv > 2:
            kxc = [outbuf[6], outbuf[7], outbuf[8], outbuf[9]]
    elif xctype == 'U-LDA':
        if deriv > 0:
            vxc = [outbuf[1:3].T]
        if deriv > 1:
            fxc = [outbuf[3:6].T]
        if deriv > 2:
            kxc = [outbuf[6:10].T]
    elif xctype == 'U-GGA':
        if deriv > 0:
            vxc = [outbuf[1:3].T, outbuf[3:6].T]
        if deriv > 1:
            fxc = [outbuf[6:9].T, outbuf[9:15].T, outbuf[15:21].T]
        if deriv > 2:
            kxc = [outbuf[21:25].T, outbuf[25:34].T, outbuf[34:46].T, outbuf[46:56].T]
    elif xctype == 'R-MGGA':
        if deriv > 0:
            vxc = [outbuf[1], outbuf[2], None, outbuf[4]]
        if deriv > 1:
            fxc = [
                # v2rho2, v2rhosigma, v2sigma2,
                outbuf[5], outbuf[6], outbuf[7],
                # v2lapl2, v2tau2,
                None, outbuf[9],
                # v2rholapl, v2rhotau,
                None, outbuf[11],
                # v2lapltau, v2sigmalapl, v2sigmatau,
                None, None, outbuf[14]]
        if deriv > 2:
            # v3lapltau2 might not be strictly 0
            # outbuf[18] = 0
            kxc = [
                # v3rho3, v3rho2sigma, v3rhosigma2, v3sigma3,
                outbuf[15], outbuf[16], outbuf[17], outbuf[18],
                # v3rho2lapl, v3rho2tau,
                None, outbuf[20],
                # v3rhosigmalapl, v3rhosigmatau,
                None, outbuf[22],
                # v3rholapl2, v3rholapltau, v3rhotau2,
                None, None, outbuf[25],
                # v3sigma2lapl, v3sigma2tau,
                None, outbuf[27],
                # v3sigmalapl2, v3sigmalapltau, v3sigmatau2,
                None, None, outbuf[30],
                # v3lapl3, v3lapl2tau, v3lapltau2, v3tau3)
                None, None, None, outbuf[34]]
    elif xctype == 'U-MGGA':
        if deriv > 0:
            vxc = [outbuf[1:3].T, outbuf[3:6].T, None, outbuf[8:10].T]
        if deriv > 1:
            # v2lapltau might not be strictly 0
            # outbuf[39:43] = 0
            fxc = [
                # v2rho2, v2rhosigma, v2sigma2,
                outbuf[10:13].T, outbuf[13:19].T, outbuf[19:25].T,
                # v2lapl2, v2tau2,
                None, outbuf[28:31].T,
                # v2rholapl, v2rhotau,
                None, outbuf[35:39].T,
                # v2lapltau, v2sigmalapl, v2sigmatau,
                None, None, outbuf[49:55].T]
        if deriv > 2:
            # v3lapltau2 might not be strictly 0
            # outbuf[204:216] = 0
            kxc = [
                # v3rho3, v3rho2sigma, v3rhosigma2, v3sigma3,
                outbuf[55:59].T, outbuf[59:68].T, outbuf[68:80].T, outbuf[80:90].T,
                # v3rho2lapl, v3rho2tau,
                None, outbuf[96:102].T,
                # v3rhosigmalapl, v3rhosigmatau,
                None, outbuf[114:126].T,
                # v3rholapl2, v3rholapltau, v3rhotau2,
                None, None, outbuf[140:146].T,
                # v3sigma2lapl, v3sigma2tau,
                None, outbuf[158:170].T,
                # v3sigmalapl2, v3sigmalapltau, v3sigmatau2,
                None, None, outbuf[191:200].T,
                # v3lapl3, v3lapl2tau, v3lapltau2, v3tau3)
                None, None, None, outbuf[216:220].T]
    return exc, vxc, fxc, kxc


def define_xc_(ni, description, xctype='LDA', hyb=0, rsh=(0,0,0)):
    '''Define XC functional.  See also :func:`eval_xc` for the rules of input description.

    Args:
        ni : an instance of :class:`NumInt`

        description : str
            A string to describe the linear combination of different XC functionals.
            The X and C functional are separated by comma like '.8*LDA+.2*B86,VWN'.
            If "HF" was appeared in the string, it stands for the exact exchange.

    Kwargs:
        xctype : str
            'LDA' or 'GGA' or 'MGGA'
        hyb : float
            hybrid functional coefficient
        rsh : a list of three floats
            coefficients (omega, alpha, beta) for range-separated hybrid functional.
            omega is the exponent factor in attenuated Coulomb operator e^{-omega r_{12}}/r_{12}
            alpha is the coefficient for long-range part, hybrid coefficient
            can be obtained by alpha + beta

    Examples:

    >>> mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='ccpvdz')
    >>> mf = dft.RKS(mol)
    >>> define_xc_(mf._numint, '.2*HF + .08*LDA + .72*B88, .81*LYP + .19*VWN')
    >>> mf.kernel()
    -76.3783361189611
    >>> define_xc_(mf._numint, 'LDA*.08 + .72*B88 + .2*HF, .81*LYP + .19*VWN')
    >>> mf.kernel()
    -76.3783361189611
    >>> def eval_xc(xc_code, rho, *args, **kwargs):
    ...     exc = 0.01 * rho**2
    ...     vrho = 0.01 * 2 * rho
    ...     vxc = (vrho, None, None, None)
    ...     fxc = None  # 2nd order functional derivative
    ...     kxc = None  # 3rd order functional derivative
    ...     return exc, vxc, fxc, kxc
    >>> define_xc_(mf._numint, eval_xc, xctype='LDA')
    >>> mf.kernel()
    48.8525211046668
    '''
    if isinstance(description, str):
        ni.eval_xc = lambda xc_code, rho, *args, **kwargs: \
                eval_xc(description, rho, *args, **kwargs)
        ni.hybrid_coeff = lambda *args, **kwargs: hybrid_coeff(description)
        ni.rsh_coeff = lambda *args: rsh_coeff(description)
        ni._xc_type = lambda *args: xc_type(description)

    elif callable(description):
        ni.eval_xc = description
        ni.hybrid_coeff = lambda *args, **kwargs: hyb
        ni.rsh_coeff = lambda *args, **kwargs: rsh
        ni._xc_type = lambda *args: xctype
    else:
        raise ValueError('Unknown description %s' % description)
    return ni

def define_xc(ni, description, xctype='LDA', hyb=0, rsh=(0,0,0)):
    return define_xc_(copy.copy(ni), description, xctype, hyb, rsh)
define_xc.__doc__ = define_xc_.__doc__


# pyscf/dft/numint.py

#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Numerical integration functions for RKS and UKS with real AO basis
'''

import warnings
import ctypes
import numpy
from pyscf import lib
from pyscf.dft.sap import sap_effective_charge
try:
    from pyscf.dft import libxc
except (ImportError, OSError):
    try:
        from pyscf.dft import xcfun
        libxc = xcfun
    except (ImportError, OSError):
        warnings.warn('XC functional libraries (libxc or XCfun) are not available.')
        raise

from pyscf.dft.gen_grid import make_mask, BLKSIZE
from pyscf import __config__

libdft = lib.load_library('libdft')
OCCDROP = getattr(__config__, 'dft_numint_OCCDROP', 1e-12)
# The system size above which to consider the sparsity of the density matrix.
# If the number of AOs in the system is less than this value, all tensors are
# treated as dense quantities and contracted by dgemm directly.
SWITCH_SIZE = getattr(__config__, 'dft_numint_SWITCH_SIZE', 800)

def eval_ao(mol, coords, deriv=0, shls_slice=None,
            non0tab=None, out=None, verbose=None):
    '''Evaluate AO function value on the given grids.

    Args:
        mol : an instance of :class:`Mole`

        coords : 2D array, shape (N,3)
            The coordinates of the grids.

    Kwargs:
        deriv : int
            AO derivative order.  It affects the shape of the return array.
            If deriv=0, the returned AO values are stored in a (N,nao) array.
            Otherwise the AO values are stored in an array of shape (M,N,nao).
            Here N is the number of grids, nao is the number of AO functions,
            M is the size associated to the derivative deriv.
        relativity : bool
            No effects.
        shls_slice : 2-element list
            (shl_start, shl_end).
            If given, only part of AOs (shl_start <= shell_id < shl_end) are
            evaluated.  By default, all shells defined in mol will be evaluated.
        non0tab : 2D bool array
            mask array to indicate whether the AO values are zero.  The mask
            array can be obtained by calling :func:`make_mask`
        out : ndarray
            If provided, results are written into this array.
        verbose : int or object of :class:`Logger`
            No effects.

    Returns:
        2D array of shape (N,nao) for AO values if deriv = 0.
        Or 3D array of shape (:,N,nao) for AO values and AO derivatives if deriv > 0.
        In the 3D array, the first (N,nao) elements are the AO values,
        followed by (3,N,nao) for x,y,z compoents;
        Then 2nd derivatives (6,N,nao) for xx, xy, xz, yy, yz, zz;
        Then 3rd derivatives (10,N,nao) for xxx, xxy, xxz, xyy, xyz, xzz, yyy, yyz, yzz, zzz;
        ...

    Examples:

    >>> mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='ccpvdz')
    >>> coords = numpy.random.random((100,3))  # 100 random points
    >>> ao_value = eval_ao(mol, coords)
    >>> print(ao_value.shape)
    (100, 24)
    >>> ao_value = eval_ao(mol, coords, deriv=1, shls_slice=(1,4))
    >>> print(ao_value.shape)
    (4, 100, 7)
    >>> ao_value = eval_ao(mol, coords, deriv=2, shls_slice=(1,4))
    >>> print(ao_value.shape)
    (10, 100, 7)
    '''
    comp = (deriv+1)*(deriv+2)*(deriv+3)//6
    if mol.cart:
        feval = 'GTOval_cart_deriv%d' % deriv
    else:
        feval = 'GTOval_sph_deriv%d' % deriv
    return mol.eval_gto(feval, coords, comp, shls_slice, non0tab, out=out)

#TODO: \nabla^2 rho and tau = 1/2 (\nabla f)^2
def eval_rho(mol, ao, dm, non0tab=None, xctype='LDA', hermi=0, verbose=None):
    r'''Calculate the electron density for LDA functional, and the density
    derivatives for GGA functional.

    Args:
        mol : an instance of :class:`Mole`

        ao : 2D array of shape (N,nao) for LDA, 3D array of shape (4,N,nao) for GGA
            or (5,N,nao) for meta-GGA.  N is the number of grids, nao is the
            number of AO functions.  If xctype is GGA, ao[0] is AO value
            and ao[1:3] are the AO gradients.  If xctype is meta-GGA, ao[4:10]
            are second derivatives of ao values.
        dm : 2D array
            Density matrix

    Kwargs:
        non0tab : 2D bool array
            mask array to indicate whether the AO values are zero.  The mask
            array can be obtained by calling :func:`make_mask`
        xctype : str
            LDA/GGA/mGGA.  It affects the shape of the return density.
        hermi : bool
            dm is hermitian or not
        verbose : int or object of :class:`Logger`
            No effects.

    Returns:
        1D array of size N to store electron density if xctype = LDA;  2D array
        of (4,N) to store density and "density derivatives" for x,y,z components
        if xctype = GGA;  (6,N) array for meta-GGA, where last two rows are
        \nabla^2 rho and tau = 1/2(\nabla f)^2

    Examples:

    >>> mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='ccpvdz')
    >>> coords = numpy.random.random((100,3))  # 100 random points
    >>> ao_value = eval_ao(mol, coords, deriv=0)
    >>> dm = numpy.random.random((mol.nao_nr(),mol.nao_nr()))
    >>> dm = dm + dm.T
    >>> rho, dx_rho, dy_rho, dz_rho = eval_rho(mol, ao, dm, xctype='LDA')
    '''
    xctype = xctype.upper()
    if xctype == 'LDA' or xctype == 'HF':
        ngrids, nao = ao.shape
    else:
        ngrids, nao = ao[0].shape

    if non0tab is None:
        non0tab = numpy.ones(((ngrids+BLKSIZE-1)//BLKSIZE,mol.nbas),
                             dtype=numpy.uint8)
    if not hermi:
        # (D + D.T)/2 because eval_rho computes 2*(|\nabla i> D_ij <j|) instead of
        # |\nabla i> D_ij <j| + |i> D_ij <\nabla j| for efficiency
        dm = (dm + dm.conj().T) * .5

    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_nr()
    if xctype == 'LDA' or xctype == 'HF':
        c0 = _dot_ao_dm(mol, ao, dm, non0tab, shls_slice, ao_loc)
        #:rho = numpy.einsum('pi,pi->p', ao, c0)
        rho = _contract_rho(ao, c0)
    elif xctype in ('GGA', 'NLC'):
        rho = numpy.empty((4,ngrids))
        c0 = _dot_ao_dm(mol, ao[0], dm, non0tab, shls_slice, ao_loc)
        #:rho[0] = numpy.einsum('pi,pi->p', c0, ao[0])
        rho[0] = _contract_rho(ao[0], c0)
        for i in range(1, 4):
            #:rho[i] = numpy.einsum('pi,pi->p', c0, ao[i])
            rho[i] = _contract_rho(ao[i], c0)
            rho[i] *= 2 # *2 for +c.c. in the next two lines
            #c1 = _dot_ao_dm(mol, ao[i], dm, non0tab, shls_slice, ao_loc)
            #rho[i] += numpy.einsum('pi,pi->p', c1, ao[0])
    else: # meta-GGA
        # rho[4] = \nabla^2 rho, rho[5] = 1/2 |nabla f|^2
        rho = numpy.empty((6,ngrids))
        c0 = _dot_ao_dm(mol, ao[0], dm, non0tab, shls_slice, ao_loc)
        #:rho[0] = numpy.einsum('pi,pi->p', ao[0], c0)
        rho[0] = _contract_rho(ao[0], c0)
        rho[5] = 0
        for i in range(1, 4):
            #:rho[i] = numpy.einsum('pi,pi->p', c0, ao[i]) * 2 # *2 for +c.c.
            rho[i] = _contract_rho(ao[i], c0) * 2
            c1 = _dot_ao_dm(mol, ao[i], dm.T, non0tab, shls_slice, ao_loc)
            #:rho[5] += numpy.einsum('pi,pi->p', c1, ao[i])
            rho[5] += _contract_rho(ao[i], c1)
        XX, YY, ZZ = 4, 7, 9
        ao2 = ao[XX] + ao[YY] + ao[ZZ]
        # \nabla^2 rho
        #:rho[4] = numpy.einsum('pi,pi->p', c0, ao2)
        rho[4] = _contract_rho(ao2, c0)
        rho[4] += rho[5]
        rho[4] *= 2
        # tau = 1/2 (\nabla f)^2
        rho[5] *= .5
    return rho

def eval_rho2(mol, ao, mo_coeff, mo_occ, non0tab=None, xctype='LDA',
              verbose=None):
    r'''Calculate the electron density for LDA functional, and the density
    derivatives for GGA functional.  This function has the same functionality
    as :func:`eval_rho` except that the density are evaluated based on orbital
    coefficients and orbital occupancy.  It is more efficient than
    :func:`eval_rho` in most scenario.

    Args:
        mol : an instance of :class:`Mole`

        ao : 2D array of shape (N,nao) for LDA, 3D array of shape (4,N,nao) for GGA
            or (5,N,nao) for meta-GGA.  N is the number of grids, nao is the
            number of AO functions.  If xctype is GGA, ao[0] is AO value
            and ao[1:3] are the AO gradients.  If xctype is meta-GGA, ao[4:10]
            are second derivatives of ao values.
        dm : 2D array
            Density matrix

    Kwargs:
        non0tab : 2D bool array
            mask array to indicate whether the AO values are zero.  The mask
            array can be obtained by calling :func:`make_mask`
        xctype : str
            LDA/GGA/mGGA.  It affects the shape of the return density.
        verbose : int or object of :class:`Logger`
            No effects.

    Returns:
        1D array of size N to store electron density if xctype = LDA;  2D array
        of (4,N) to store density and "density derivatives" for x,y,z components
        if xctype = GGA;  (6,N) array for meta-GGA, where last two rows are
        \nabla^2 rho and tau = 1/2(\nabla f)^2
    '''
    xctype = xctype.upper()
    if xctype == 'LDA' or xctype == 'HF':
        ngrids, nao = ao.shape
    else:
        ngrids, nao = ao[0].shape

    if non0tab is None:
        non0tab = numpy.ones(((ngrids+BLKSIZE-1)//BLKSIZE,mol.nbas),
                             dtype=numpy.uint8)
    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_nr()
    pos = mo_occ > OCCDROP
    if numpy.any(pos):
        cpos = numpy.einsum('ij,j->ij', mo_coeff[:,pos], numpy.sqrt(mo_occ[pos]))
        if xctype == 'LDA' or xctype == 'HF':
            c0 = _dot_ao_dm(mol, ao, cpos, non0tab, shls_slice, ao_loc)
            #:rho = numpy.einsum('pi,pi->p', c0, c0)
            rho = _contract_rho(c0, c0)
        elif xctype in ('GGA', 'NLC'):
            rho = numpy.empty((4,ngrids))
            c0 = _dot_ao_dm(mol, ao[0], cpos, non0tab, shls_slice, ao_loc)
            #:rho[0] = numpy.einsum('pi,pi->p', c0, c0)
            rho[0] = _contract_rho(c0, c0)
            for i in range(1, 4):
                c1 = _dot_ao_dm(mol, ao[i], cpos, non0tab, shls_slice, ao_loc)
                #:rho[i] = numpy.einsum('pi,pi->p', c0, c1) * 2 # *2 for +c.c.
                rho[i] = _contract_rho(c0, c1) * 2
        else: # meta-GGA
            # rho[4] = \nabla^2 rho, rho[5] = 1/2 |nabla f|^2
            rho = numpy.empty((6,ngrids))
            c0 = _dot_ao_dm(mol, ao[0], cpos, non0tab, shls_slice, ao_loc)
            #:rho[0] = numpy.einsum('pi,pi->p', c0, c0)
            rho[0] = _contract_rho(c0, c0)
            rho[5] = 0
            for i in range(1, 4):
                c1 = _dot_ao_dm(mol, ao[i], cpos, non0tab, shls_slice, ao_loc)
                #:rho[i] = numpy.einsum('pi,pi->p', c0, c1) * 2 # *2 for +c.c.
                #:rho[5] += numpy.einsum('pi,pi->p', c1, c1)
                rho[i] = _contract_rho(c0, c1) * 2
                rho[5] += _contract_rho(c1, c1)
            XX, YY, ZZ = 4, 7, 9
            ao2 = ao[XX] + ao[YY] + ao[ZZ]
            c1 = _dot_ao_dm(mol, ao2, cpos, non0tab, shls_slice, ao_loc)
            #:rho[4] = numpy.einsum('pi,pi->p', c0, c1)
            rho[4] = _contract_rho(c0, c1)
            rho[4] += rho[5]
            rho[4] *= 2

            rho[5] *= .5
    else:
        if xctype == 'LDA' or xctype == 'HF':
            rho = numpy.zeros(ngrids)
        elif xctype in ('GGA', 'NLC'):
            rho = numpy.zeros((4,ngrids))
        else:
            rho = numpy.zeros((6,ngrids))

    neg = mo_occ < -OCCDROP
    if numpy.any(neg):
        cneg = numpy.einsum('ij,j->ij', mo_coeff[:,neg], numpy.sqrt(-mo_occ[neg]))
        if xctype == 'LDA' or xctype == 'HF':
            c0 = _dot_ao_dm(mol, ao, cneg, non0tab, shls_slice, ao_loc)
            #:rho -= numpy.einsum('pi,pi->p', c0, c0)
            rho -= _contract_rho(c0, c0)
        elif xctype == 'GGA':
            c0 = _dot_ao_dm(mol, ao[0], cneg, non0tab, shls_slice, ao_loc)
            #:rho[0] -= numpy.einsum('pi,pi->p', c0, c0)
            rho[0] -= _contract_rho(c0, c0)
            for i in range(1, 4):
                c1 = _dot_ao_dm(mol, ao[i], cneg, non0tab, shls_slice, ao_loc)
                #:rho[i] -= numpy.einsum('pi,pi->p', c0, c1) * 2 # *2 for +c.c.
                rho[i] -= _contract_rho(c0, c1) * 2 # *2 for +c.c.
        else:
            c0 = _dot_ao_dm(mol, ao[0], cneg, non0tab, shls_slice, ao_loc)
            #:rho[0] -= numpy.einsum('pi,pi->p', c0, c0)
            rho[0] -= _contract_rho(c0, c0)
            rho5 = 0
            for i in range(1, 4):
                c1 = _dot_ao_dm(mol, ao[i], cneg, non0tab, shls_slice, ao_loc)
                #:rho[i] -= numpy.einsum('pi,pi->p', c0, c1) * 2 # *2 for +c.c.
                #:rho5 += numpy.einsum('pi,pi->p', c1, c1)
                rho[i] -= _contract_rho(c0, c1) * 2 # *2 for +c.c.
                rho5 += _contract_rho(c1, c1)
            XX, YY, ZZ = 4, 7, 9
            ao2 = ao[XX] + ao[YY] + ao[ZZ]
            c1 = _dot_ao_dm(mol, ao2, cneg, non0tab, shls_slice, ao_loc)
            #:rho[4] -= numpy.einsum('pi,pi->p', c0, c1) * 2
            rho[4] -= _contract_rho(c0, c1) * 2
            rho[4] -= rho5 * 2

            rho[5] -= rho5 * .5
    return rho

def _vv10nlc(rho, coords, vvrho, vvweight, vvcoords, nlc_pars):
    thresh=1e-8

    #output
    exc=numpy.zeros(rho[0,:].size)
    vxc=numpy.zeros([2,rho[0,:].size])

    #outer grid needs threshing
    threshind=rho[0,:]>=thresh
    coords=coords[threshind]
    R=rho[0,:][threshind]
    Gx=rho[1,:][threshind]
    Gy=rho[2,:][threshind]
    Gz=rho[3,:][threshind]
    G=Gx**2.+Gy**2.+Gz**2.

    #inner grid needs threshing
    innerthreshind=vvrho[0,:]>=thresh
    vvcoords=vvcoords[innerthreshind]
    vvweight=vvweight[innerthreshind]
    Rp=vvrho[0,:][innerthreshind]
    RpW=Rp*vvweight
    Gxp=vvrho[1,:][innerthreshind]
    Gyp=vvrho[2,:][innerthreshind]
    Gzp=vvrho[3,:][innerthreshind]
    Gp=Gxp**2.+Gyp**2.+Gzp**2.

    #constants and parameters
    Pi=numpy.pi
    Pi43=4.*Pi/3.
    Bvv, Cvv = nlc_pars
    Kvv=Bvv*1.5*Pi*((9.*Pi)**(-1./6.))
    Beta=((3./(Bvv*Bvv))**(0.75))/32.

    #inner grid
    W0p=Gp/(Rp*Rp)
    W0p=Cvv*W0p*W0p
    W0p=(W0p+Pi43*Rp)**0.5
    Kp=Kvv*(Rp**(1./6.))

    #outer grid
    W0tmp=G/(R**2)
    W0tmp=Cvv*W0tmp*W0tmp
    W0=(W0tmp+Pi43*R)**0.5
    dW0dR=(0.5*Pi43*R-2.*W0tmp)/W0
    dW0dG=W0tmp*R/(G*W0)
    K=Kvv*(R**(1./6.))
    dKdR=(1./6.)*K

    vvcoords = numpy.asarray(vvcoords, order='C')
    coords = numpy.asarray(coords, order='C')
    F = numpy.empty_like(R)
    U = numpy.empty_like(R)
    W = numpy.empty_like(R)
    #for i in range(R.size):
    #    DX=vvcoords[:,0]-coords[i,0]
    #    DY=vvcoords[:,1]-coords[i,1]
    #    DZ=vvcoords[:,2]-coords[i,2]
    #    R2=DX*DX+DY*DY+DZ*DZ
    #    gp=R2*W0p+Kp
    #    g=R2*W0[i]+K[i]
    #    gt=g+gp
    #    T=RpW/(g*gp*gt)
    #    F=numpy.sum(T)
    #    T*=(1./g+1./gt)
    #    U=numpy.sum(T)
    #    W=numpy.sum(T*R2)
    #    F*=-1.5
    libdft.VXC_vv10nlc(F.ctypes.data_as(ctypes.c_void_p),
                       U.ctypes.data_as(ctypes.c_void_p),
                       W.ctypes.data_as(ctypes.c_void_p),
                       vvcoords.ctypes.data_as(ctypes.c_void_p),
                       coords.ctypes.data_as(ctypes.c_void_p),
                       W0p.ctypes.data_as(ctypes.c_void_p),
                       W0.ctypes.data_as(ctypes.c_void_p),
                       K.ctypes.data_as(ctypes.c_void_p),
                       Kp.ctypes.data_as(ctypes.c_void_p),
                       RpW.ctypes.data_as(ctypes.c_void_p),
                       ctypes.c_int(vvcoords.shape[0]),
                       ctypes.c_int(coords.shape[0]))
    #exc is multiplied by Rho later
    exc[threshind] = Beta+0.5*F
    vxc[0,threshind] = Beta+F+1.5*(U*dKdR+W*dW0dR)
    vxc[1,threshind] = 1.5*W*dW0dG
    return exc,vxc

def eval_mat(mol, ao, weight, rho, vxc,
             non0tab=None, xctype='LDA', spin=0, verbose=None):
    r'''Calculate XC potential matrix.

    Args:
        mol : an instance of :class:`Mole`

        ao : ([4/10,] ngrids, nao) ndarray
            2D array of shape (N,nao) for LDA,
            3D array of shape (4,N,nao) for GGA
            or (10,N,nao) for meta-GGA.
            N is the number of grids, nao is the number of AO functions.
            If xctype is GGA, ao[0] is AO value and ao[1:3] are the real space
            gradients.  If xctype is meta-GGA, ao[4:10] are second derivatives
            of ao values.
        weight : 1D array
            Integral weights on grids.
        rho : ([4/6,] ngrids) ndarray
            Shape of ((*,N)) for electron density (and derivatives) if spin = 0;
            Shape of ((*,N),(*,N)) for alpha/beta electron density (and derivatives) if spin > 0;
            where N is number of grids.
            rho (*,N) are ordered as (den,grad_x,grad_y,grad_z,laplacian,tau)
            where grad_x = d/dx den, laplacian = \nabla^2 den, tau = 1/2(\nabla f)^2
            In spin unrestricted case,
            rho is ((den_u,grad_xu,grad_yu,grad_zu,laplacian_u,tau_u)
                    (den_d,grad_xd,grad_yd,grad_zd,laplacian_d,tau_d))
        vxc : ([4,] ngrids) ndarray
            XC potential value on each grid = (vrho, vsigma, vlapl, vtau)
            vsigma is GGA potential value on each grid.
            If the kwarg spin != 0, a list [vsigma_uu,vsigma_ud] is required.

    Kwargs:
        xctype : str
            LDA/GGA/mGGA.  It affects the shape of `ao` and `rho`
        non0tab : 2D bool array
            mask array to indicate whether the AO values are zero.  The mask
            array can be obtained by calling :func:`make_mask`
        spin : int
            If not 0, the returned matrix is the Vxc matrix of alpha-spin.  It
            is computed with the spin non-degenerated UKS formula.

    Returns:
        XC potential matrix in 2D array of shape (nao,nao) where nao is the
        number of AO functions.
    '''
    xctype = xctype.upper()
    if xctype == 'LDA' or xctype == 'HF':
        ngrids, nao = ao.shape
    else:
        ngrids, nao = ao[0].shape

    if non0tab is None:
        non0tab = numpy.ones(((ngrids+BLKSIZE-1)//BLKSIZE,mol.nbas),
                             dtype=numpy.uint8)
    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_nr()
    transpose_for_uks = False
    if xctype == 'LDA' or xctype == 'HF':
        if not isinstance(vxc, numpy.ndarray) or vxc.ndim == 2:
            vrho = vxc[0]
        else:
            vrho = vxc
        # *.5 because return mat + mat.T
        #:aow = numpy.einsum('pi,p->pi', ao, .5*weight*vrho)
        aow = _scale_ao(ao, .5*weight*vrho)
        mat = _dot_ao_ao(mol, ao, aow, non0tab, shls_slice, ao_loc)
    else:
        #wv = weight * vsigma * 2
        #aow  = numpy.einsum('pi,p->pi', ao[1], rho[1]*wv)
        #aow += numpy.einsum('pi,p->pi', ao[2], rho[2]*wv)
        #aow += numpy.einsum('pi,p->pi', ao[3], rho[3]*wv)
        #aow += numpy.einsum('pi,p->pi', ao[0], .5*weight*vrho)
        vrho, vsigma = vxc[:2]
        wv = numpy.empty((4,ngrids))
        if spin == 0:
            assert(vsigma is not None and rho.ndim==2)
            wv[0]  = weight * vrho * .5
            wv[1:4] = rho[1:4] * (weight * vsigma * 2)
        else:
            rho_a, rho_b = rho
            wv[0]  = weight * vrho * .5
            try:
                wv[1:4] = rho_a[1:4] * (weight * vsigma[0] * 2)  # sigma_uu
                wv[1:4]+= rho_b[1:4] * (weight * vsigma[1])      # sigma_ud
            except ValueError:
                warnings.warn('Note the output of libxc.eval_xc cannot be '
                              'directly used in eval_mat.\nvsigma from eval_xc '
                              'should be restructured as '
                              '(vsigma[:,0],vsigma[:,1])\n')
                transpose_for_uks = True
                vsigma = vsigma.T
                wv[1:4] = rho_a[1:4] * (weight * vsigma[0] * 2)  # sigma_uu
                wv[1:4]+= rho_b[1:4] * (weight * vsigma[1])      # sigma_ud
        #:aow = numpy.einsum('npi,np->pi', ao[:4], wv)
        aow = _scale_ao(ao[:4], wv)
        mat = _dot_ao_ao(mol, ao[0], aow, non0tab, shls_slice, ao_loc)

# JCP 138, 244108 (2013); DOI:10.1063/1.4811270
# JCP 112, 7002 (2000); DOI:10.1063/1.481298
    if xctype == 'MGGA':
        vlapl, vtau = vxc[2:]

        if vlapl is None:
            vlapl = 0
        else:
            if spin != 0:
                if transpose_for_uks:
                    vlapl = vlapl.T
                vlapl = vlapl[0]
            XX, YY, ZZ = 4, 7, 9
            ao2 = ao[XX] + ao[YY] + ao[ZZ]
            #:aow = numpy.einsum('pi,p->pi', ao2, .5 * weight * vlapl, out=aow)
            aow = _scale_ao(ao2, .5 * weight * vlapl, out=aow)
            mat += _dot_ao_ao(mol, ao[0], aow, non0tab, shls_slice, ao_loc)

        if spin != 0:
            if transpose_for_uks:
                vtau = vtau.T
            vtau = vtau[0]
        wv = weight * (.25*vtau + vlapl)
        mat += _tau_dot(mol, ao, ao, wv, non0tab, shls_slice, ao_loc)

    return mat + mat.T.conj()


def _dot_ao_ao(mol, ao1, ao2, non0tab, shls_slice, ao_loc, hermi=0):
    '''return numpy.dot(ao1.T, ao2)'''
    ngrids, nao = ao1.shape
    if nao < SWITCH_SIZE:
        return lib.dot(ao1.T.conj(), ao2)

    if not ao1.flags.f_contiguous:
        ao1 = lib.transpose(ao1)
    if not ao2.flags.f_contiguous:
        ao2 = lib.transpose(ao2)
    if ao1.dtype == ao2.dtype == numpy.double:
        fn = libdft.VXCdot_ao_ao
    else:
        fn = libdft.VXCzdot_ao_ao
        ao1 = numpy.asarray(ao1, numpy.complex128)
        ao2 = numpy.asarray(ao2, numpy.complex128)

    if non0tab is None or shls_slice is None or ao_loc is None:
        pnon0tab = pshls_slice = pao_loc = lib.c_null_ptr()
    else:
        pnon0tab    = non0tab.ctypes.data_as(ctypes.c_void_p)
        pshls_slice = (ctypes.c_int*2)(*shls_slice)
        pao_loc     = ao_loc.ctypes.data_as(ctypes.c_void_p)

    vv = numpy.empty((nao,nao), dtype=ao1.dtype)
    fn(vv.ctypes.data_as(ctypes.c_void_p),
       ao1.ctypes.data_as(ctypes.c_void_p),
       ao2.ctypes.data_as(ctypes.c_void_p),
       ctypes.c_int(nao), ctypes.c_int(ngrids),
       ctypes.c_int(mol.nbas), ctypes.c_int(hermi),
       pnon0tab, pshls_slice, pao_loc)
    return vv

def _dot_ao_dm(mol, ao, dm, non0tab, shls_slice, ao_loc, out=None):
    '''return numpy.dot(ao, dm)'''
    ngrids, nao = ao.shape
    if nao < SWITCH_SIZE:
        return lib.dot(numpy.asarray(dm, order='C').T, ao.T).T

    if not ao.flags.f_contiguous:
        ao = lib.transpose(ao)
    if ao.dtype == dm.dtype == numpy.double:
        fn = libdft.VXCdot_ao_dm
    else:
        fn = libdft.VXCzdot_ao_dm
        ao = numpy.asarray(ao, numpy.complex128)
        dm = numpy.asarray(dm, numpy.complex128)

    if non0tab is None or shls_slice is None or ao_loc is None:
        pnon0tab = pshls_slice = pao_loc = lib.c_null_ptr()
    else:
        pnon0tab    = non0tab.ctypes.data_as(ctypes.c_void_p)
        pshls_slice = (ctypes.c_int*2)(*shls_slice)
        pao_loc     = ao_loc.ctypes.data_as(ctypes.c_void_p)

    vm = numpy.ndarray((ngrids,dm.shape[1]), dtype=ao.dtype, order='F', buffer=out)
    dm = numpy.asarray(dm, order='C')
    fn(vm.ctypes.data_as(ctypes.c_void_p),
       ao.ctypes.data_as(ctypes.c_void_p),
       dm.ctypes.data_as(ctypes.c_void_p),
       ctypes.c_int(nao), ctypes.c_int(dm.shape[1]),
       ctypes.c_int(ngrids), ctypes.c_int(mol.nbas),
       pnon0tab, pshls_slice, pao_loc)
    return vm

def _scale_ao(ao, wv, out=None):
    #:aow = numpy.einsum('npi,np->pi', ao[:4], wv)
    if wv.ndim == 2:
        ao = ao.transpose(0,2,1)
    else:
        ngrids, nao = ao.shape
        ao = ao.T.reshape(1,nao,ngrids)
        wv = wv.reshape(1,ngrids)

    wv = numpy.asarray(wv, order='C')
    comp, nao, ngrids = ao.shape
    aow = numpy.ndarray((nao,ngrids), dtype=ao.dtype, buffer=out).T

    if not ao.flags.c_contiguous:
        aow = numpy.einsum('nip,np->pi', ao, wv)
    elif aow.dtype == numpy.double:
        libdft.VXC_dscale_ao(aow.ctypes.data_as(ctypes.c_void_p),
                             ao.ctypes.data_as(ctypes.c_void_p),
                             wv.ctypes.data_as(ctypes.c_void_p),
                             ctypes.c_int(comp), ctypes.c_int(nao),
                             ctypes.c_int(ngrids))
    elif aow.dtype == numpy.complex128:
        libdft.VXC_zscale_ao(aow.ctypes.data_as(ctypes.c_void_p),
                             ao.ctypes.data_as(ctypes.c_void_p),
                             wv.ctypes.data_as(ctypes.c_void_p),
                             ctypes.c_int(comp), ctypes.c_int(nao),
                             ctypes.c_int(ngrids))
    else:
        aow = numpy.einsum('nip,np->pi', ao, wv)
    return aow

def _contract_rho(bra, ket):
    #:rho  = numpy.einsum('pi,pi->p', bra.real, ket.real)
    #:rho += numpy.einsum('pi,pi->p', bra.imag, ket.imag)
    bra = bra.T
    ket = ket.T
    nao, ngrids = bra.shape
    rho = numpy.empty(ngrids)

    if not (bra.flags.c_contiguous and ket.flags.c_contiguous):
        rho  = numpy.einsum('ip,ip->p', bra.real, ket.real)
        rho += numpy.einsum('ip,ip->p', bra.imag, ket.imag)
    elif bra.dtype == numpy.double and ket.dtype == numpy.double:
        libdft.VXC_dcontract_rho(rho.ctypes.data_as(ctypes.c_void_p),
                                 bra.ctypes.data_as(ctypes.c_void_p),
                                 ket.ctypes.data_as(ctypes.c_void_p),
                                 ctypes.c_int(nao), ctypes.c_int(ngrids))
    elif bra.dtype == numpy.complex128 and ket.dtype == numpy.complex128:
        libdft.VXC_zcontract_rho(rho.ctypes.data_as(ctypes.c_void_p),
                                 bra.ctypes.data_as(ctypes.c_void_p),
                                 ket.ctypes.data_as(ctypes.c_void_p),
                                 ctypes.c_int(nao), ctypes.c_int(ngrids))
    else:
        rho  = numpy.einsum('ip,ip->p', bra.real, ket.real)
        rho += numpy.einsum('ip,ip->p', bra.imag, ket.imag)
    return rho

def _tau_dot(mol, bra, ket, wv, mask, shls_slice, ao_loc):
    # nabla_ao dot nabla_ao
    # numpy.einsum('p,xpi,xpj->ij', wv, bra[1:4].conj(), ket[1:4])
    aow = _scale_ao(ket[1], wv)
    mat = _dot_ao_ao(mol, bra[1], aow, mask, shls_slice, ao_loc)
    aow = _scale_ao(ket[2], wv, aow)
    mat += _dot_ao_ao(mol, bra[2], aow, mask, shls_slice, ao_loc)
    aow = _scale_ao(ket[3], wv, aow)
    mat += _dot_ao_ao(mol, bra[3], aow, mask, shls_slice, ao_loc)
    return mat

def nr_vxc(mol, grids, xc_code, dms, spin=0, relativity=0, hermi=0,
           max_memory=2000, verbose=None):
    '''
    Evaluate RKS/UKS XC functional and potential matrix on given meshgrids
    for a set of density matrices.  See :func:`nr_rks` and :func:`nr_uks`
    for more details.

    Args:
        mol : an instance of :class:`Mole`

        grids : an instance of :class:`Grids`
            grids.coords and grids.weights are needed for coordinates and weights of meshgrids.
        xc_code : str
            XC functional description.
            See :func:`parse_xc` of pyscf/dft/libxc.py for more details.
        dms : 2D array or a list of 2D arrays
            Density matrix or multiple density matrices

    Kwargs:
        hermi : int
            Input density matrices symmetric or not
        max_memory : int or float
            The maximum size of cache to use (in MB).

    Returns:
        nelec, excsum, vmat.
        nelec is the number of electrons generated by numerical integration.
        excsum is the XC functional value.  vmat is the XC potential matrix in
        2D array of shape (nao,nao) where nao is the number of AO functions.

    Examples:

    >>> from pyscf import gto, dft
    >>> mol = gto.M(atom='H 0 0 0; H 0 0 1.1')
    >>> grids = dft.gen_grid.Grids(mol)
    >>> grids.coords = numpy.random.random((100,3))  # 100 random points
    >>> grids.weights = numpy.random.random(100)
    >>> nao = mol.nao_nr()
    >>> dm = numpy.random.random((2,nao,nao))
    >>> nelec, exc, vxc = dft.numint.nr_vxc(mol, grids, 'lda,vwn', dm, spin=1)
    '''
    ni = NumInt()
    return ni.nr_vxc(mol, grids, xc_code, dms, spin, relativity,
                     hermi, max_memory, verbose)

def nr_sap_vxc(ni, mol, grids, max_memory=2000, verbose=None):
    '''Calculate superposition of atomic potentials matrix on given meshgrids.

    Args:
        ni : an instance of :class:`NumInt`

        mol : an instance of :class:`Mole`

        grids : an instance of :class:`Grids`
            grids.coords and grids.weights are needed for coordinates and weights of meshgrids.

    Kwargs:
        max_memory : int or float
            The maximum size of cache to use (in MB).

    Returns:
        vmat is the XC potential matrix in 2D array of shape (nao,nao)
        where nao is the number of AO functions.

    Examples:
    >>> import numpy
    >>> from pyscf import gto, dft
    >>> mol = gto.M(atom='H 0 0 0; H 0 0 1.1')
    >>> grids = dft.gen_grid.Grids(mol)
    >>> ni = dft.numint.NumInt()
    >>> vsap = ni.nr_sap(mol, grids)
    '''
    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_nr()
    nao = mol.nao
    vmat = numpy.zeros((nao,nao))
    aow = None
    vxcw = None
    ao_deriv = 0

    atom_coords = mol.atom_coords()
    atom_charges = mol.atom_charges()
    eps = numpy.finfo(float).eps

    for ao, mask, weight, coords in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
        aow = numpy.ndarray(ao.shape, order='F', buffer=aow)
        vxc = numpy.ndarray(coords.shape[0], buffer=vxcw)
        vxc.fill(0.0)
        # Form potential
        for igrid in range(vxc.size):
            for iatom in range(len(mol._atm)):
                # Distance from nucleus
                rnuc = numpy.linalg.norm(atom_coords[iatom,:] - coords[igrid,:])
                if rnuc > eps:
                    # Zeff(R)
                    Zeff = sap_effective_charge(atom_charges[iatom], rnuc)
                    vxc[igrid] -= Zeff/rnuc

        # *.5 because vmat + vmat.T
        #:aow = numpy.einsum('pi,p->pi', ao, .5*weight*vrho, out=aow)
        aow = _scale_ao(ao, .5*weight*vxc, out=aow)
        vmat += _dot_ao_ao(mol, ao, aow, mask, shls_slice, ao_loc)
        vxc = None

    vmat = vmat + vmat.conj().T
    return vmat

def nr_rks(ni, mol, grids, xc_code, dms, relativity=0, hermi=0,
           max_memory=2000, verbose=None):
    '''Calculate RKS XC functional and potential matrix on given meshgrids
    for a set of density matrices

    Args:
        ni : an instance of :class:`NumInt`

        mol : an instance of :class:`Mole`

        grids : an instance of :class:`Grids`
            grids.coords and grids.weights are needed for coordinates and weights of meshgrids.
        xc_code : str
            XC functional description.
            See :func:`parse_xc` of pyscf/dft/libxc.py for more details.
        dms : 2D array or a list of 2D arrays
            Density matrix or multiple density matrices

    Kwargs:
        hermi : int
            Input density matrices symmetric or not
        max_memory : int or float
            The maximum size of cache to use (in MB).

    Returns:
        nelec, excsum, vmat.
        nelec is the number of electrons generated by numerical integration.
        excsum is the XC functional value.  vmat is the XC potential matrix in
        2D array of shape (nao,nao) where nao is the number of AO functions.

    Examples:

    >>> from pyscf import gto, dft
    >>> mol = gto.M(atom='H 0 0 0; H 0 0 1.1')
    >>> grids = dft.gen_grid.Grids(mol)
    >>> grids.coords = numpy.random.random((100,3))  # 100 random points
    >>> grids.weights = numpy.random.random(100)
    >>> nao = mol.nao_nr()
    >>> dm = numpy.random.random((nao,nao))
    >>> ni = dft.numint.NumInt()
    >>> nelec, exc, vxc = ni.nr_rks(mol, grids, 'lda,vwn', dm)
    '''
    xctype = ni._xc_type(xc_code)
    make_rho, nset, nao = ni._gen_rho_evaluator(mol, dms, hermi)

    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_nr()

    nelec = numpy.zeros(nset)
    excsum = numpy.zeros(nset)
    if isinstance(dms, numpy.ndarray):
        vmat = numpy.zeros((nset,nao,nao), dtype=dms.dtype)
    else:
        vmat = numpy.zeros((nset,nao,nao), dtype=numpy.result_type(*dms))

    def block_loop(ao_deriv):
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            for i in range(nset):
                rho = make_rho(i, ao, mask, xctype)
                exc, vxc = ni.eval_xc(xc_code, rho, spin=0,
                                      relativity=relativity, deriv=1,
                                      verbose=verbose)[:2]
                if xctype == 'LDA':
                    den = rho * weight
                else:
                    den = rho[0] * weight
                nelec[i] += den.sum()
                excsum[i] += numpy.dot(den, exc)
                yield i, rho, ao, mask, weight, vxc

    aow = None
    if xctype == 'LDA':
        ao_deriv = 0
        for i, rho, ao, mask, weight, vxc in block_loop(ao_deriv):
            vrho = vxc[0]
            # *.5 because vmat + vmat.T
            #:aow = numpy.einsum('pi,p->pi', ao, .5*weight*vrho, out=aow)
            aow = _scale_ao(ao, .5*weight*vrho, out=aow)
            vmat[i] += _dot_ao_ao(mol, ao, aow, mask, shls_slice, ao_loc)

    elif xctype == 'GGA':
        ao_deriv = 1
        for i, rho, ao, mask, weight, vxc in block_loop(ao_deriv):
            wv = _rks_gga_wv0(rho, vxc, weight)
            #:aow = numpy.einsum('npi,np->pi', ao, wv, out=aow)
            aow = _scale_ao(ao, wv, out=aow)
            vmat[i] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)

    elif xctype == 'NLC':
        nlc_pars = ni.nlc_coeff(xc_code)
        ao_deriv = 1
        vvrho=numpy.empty([nset,4,0])
        vvweight=numpy.empty([nset,0])
        vvcoords=numpy.empty([nset,0,3])
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rhotmp = numpy.empty([0,4,weight.size])
            weighttmp = numpy.empty([0,weight.size])
            coordstmp = numpy.empty([0,weight.size,3])
            for idm in range(nset):
                rho = make_rho(idm, ao, mask, 'GGA')
                rho = numpy.expand_dims(rho,axis=0)
                rhotmp = numpy.concatenate((rhotmp,rho),axis=0)
                weighttmp = numpy.concatenate((weighttmp,numpy.expand_dims(weight,axis=0)),axis=0)
                coordstmp = numpy.concatenate((coordstmp,numpy.expand_dims(coords,axis=0)),axis=0)
                rho = None
            vvrho=numpy.concatenate((vvrho,rhotmp),axis=2)
            vvweight=numpy.concatenate((vvweight,weighttmp),axis=1)
            vvcoords=numpy.concatenate((vvcoords,coordstmp),axis=1)
            rhotmp = weighttmp = coordstmp = None
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            aow = numpy.ndarray(ao[0].shape, order='F', buffer=aow)
            for i in range(nset):
                rho = make_rho(i, ao, mask, 'GGA')
                exc, vxc = _vv10nlc(rho, coords, vvrho[i], vvweight[i], vvcoords[i], nlc_pars)
                den = rho[0] * weight
                nelec[i] += den.sum()
                excsum[i] += numpy.dot(den, exc)
                wv = _rks_gga_wv0(rho, vxc, weight)
                #:aow = numpy.einsum('npi,np->pi', ao, wv, out=aow)
                aow = _scale_ao(ao, wv, out=aow)
                vmat[i] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
        vvrho = vvweight = vvcoords = None

    elif xctype == 'MGGA':
        if (any(x in xc_code.upper() for x in ('CC06', 'CS', 'BR89', 'MK00'))):
            raise NotImplementedError('laplacian in meta-GGA method')
        ao_deriv = 2
        for i, rho, ao, mask, weight, vxc in block_loop(ao_deriv):
            wv = _rks_mgga_wv0(rho, vxc, weight)
            #:aow = numpy.einsum('npi,np->pi', ao[:4], wv, out=aow)
            aow = _scale_ao(ao[:4], wv[:4], out=aow)
            vmat[i] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
            vmat[i] += _tau_dot(mol, ao, ao, wv[5], mask, shls_slice, ao_loc)

    elif xctype == 'HF':
        pass
    else:
        raise NotImplementedError(f'numint.nr_uks for functional {xc_code}')

    for i in range(nset):
        vmat[i] = vmat[i] + vmat[i].conj().T
    if nset == 1:
        nelec = nelec[0]
        excsum = excsum[0]
        vmat = vmat[0]
    return nelec, excsum, vmat

def nr_uks(ni, mol, grids, xc_code, dms, relativity=0, hermi=0,
           max_memory=2000, verbose=None):
    '''Calculate UKS XC functional and potential matrix on given meshgrids
    for a set of density matrices

    Args:
        mol : an instance of :class:`Mole`

        grids : an instance of :class:`Grids`
            grids.coords and grids.weights are needed for coordinates and weights of meshgrids.
        xc_code : str
            XC functional description.
            See :func:`parse_xc` of pyscf/dft/libxc.py for more details.
        dms : a list of 2D arrays
            A list of density matrices, stored as (alpha,alpha,...,beta,beta,...)

    Kwargs:
        hermi : int
            Input density matrices symmetric or not
        max_memory : int or float
            The maximum size of cache to use (in MB).

    Returns:
        nelec, excsum, vmat.
        nelec is the number of (alpha,beta) electrons generated by numerical integration.
        excsum is the XC functional value.
        vmat is the XC potential matrix for (alpha,beta) spin.

    Examples:

    >>> from pyscf import gto, dft
    >>> mol = gto.M(atom='H 0 0 0; H 0 0 1.1')
    >>> grids = dft.gen_grid.Grids(mol)
    >>> grids.coords = numpy.random.random((100,3))  # 100 random points
    >>> grids.weights = numpy.random.random(100)
    >>> nao = mol.nao_nr()
    >>> dm = numpy.random.random((2,nao,nao))
    >>> ni = dft.numint.NumInt()
    >>> nelec, exc, vxc = ni.nr_uks(mol, grids, 'lda,vwn', dm)
    '''
    xctype = ni._xc_type(xc_code)
    if xctype == 'NLC':
        dms_sf = dms[0] + dms[1]
        nelec, excsum, vmat = nr_rks(ni, mol, grids, xc_code, dms_sf, relativity, hermi,
                                     max_memory, verbose)
        return [nelec,nelec], excsum, numpy.asarray([vmat,vmat])

    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_nr()

    dma, dmb = _format_uks_dm(dms)
    nao = dma.shape[-1]
    make_rhoa, nset = ni._gen_rho_evaluator(mol, dma, hermi)[:2]
    make_rhob       = ni._gen_rho_evaluator(mol, dmb, hermi)[0]

    nelec = numpy.zeros((2,nset))
    excsum = numpy.zeros(nset)
    vmat = numpy.zeros((2,nset,nao,nao), dtype=numpy.result_type(dma, dmb))

    def block_loop(ao_deriv):
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            for i in range(nset):
                rho_a = make_rhoa(i, ao, mask, xctype)
                rho_b = make_rhob(i, ao, mask, xctype)
                rho = (rho_a, rho_b)
                exc, vxc = ni.eval_xc(xc_code, rho, spin=1,
                                      relativity=relativity, deriv=1,
                                      verbose=verbose)[:2]
                if xctype == 'LDA':
                    den_a = rho_a * weight
                    den_b = rho_b * weight
                else:
                    den_a = rho_a[0] * weight
                    den_b = rho_b[0] * weight
                nelec[0,i] += den_a.sum()
                nelec[1,i] += den_b.sum()
                excsum[i] += numpy.dot(den_a, exc)
                excsum[i] += numpy.dot(den_b, exc)
                yield i, rho, ao, mask, weight, vxc

    aow = None
    if xctype == 'LDA':
        ao_deriv = 0
        for i, rho, ao, mask, weight, vxc in block_loop(ao_deriv):
            vrho = vxc[0]
            # *.5 due to +c.c. in the end
            #:aow = numpy.einsum('pi,p->pi', ao, .5*weight*vrho[:,0], out=aow)
            aow = _scale_ao(ao, .5*weight*vrho[:,0], out=aow)
            vmat[0,i] += _dot_ao_ao(mol, ao, aow, mask, shls_slice, ao_loc)
            #:aow = numpy.einsum('pi,p->pi', ao, .5*weight*vrho[:,1], out=aow)
            aow = _scale_ao(ao, .5*weight*vrho[:,1], out=aow)
            vmat[1,i] += _dot_ao_ao(mol, ao, aow, mask, shls_slice, ao_loc)

    elif xctype == 'GGA':
        ao_deriv = 1
        for i, rho, ao, mask, weight, vxc in block_loop(ao_deriv):
            wva, wvb = _uks_gga_wv0(rho, vxc, weight)
            #:aow = numpy.einsum('npi,np->pi', ao, wva, out=aow)
            aow = _scale_ao(ao, wva, out=aow)
            vmat[0,i] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
            #:aow = numpy.einsum('npi,np->pi', ao, wvb, out=aow)
            aow = _scale_ao(ao, wvb, out=aow)
            vmat[1,i] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)

    elif xctype == 'MGGA':
        if (any(x in xc_code.upper() for x in ('CC06', 'CS', 'BR89', 'MK00'))):
            raise NotImplementedError('laplacian in meta-GGA method')
        ao_deriv = 2
        for i, rho, ao, mask, weight, vxc in block_loop(ao_deriv):
            wva, wvb = _uks_mgga_wv0(rho, vxc, weight)
            #:aow = numpy.einsum('npi,np->pi', ao[:4], wva, out=aow)
            aow = _scale_ao(ao[:4], wva[:4], out=aow)
            vmat[0,i] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
            #:aow = numpy.einsum('npi,np->pi', ao[:4], wvb, out=aow)
            aow = _scale_ao(ao[:4], wvb[:4], out=aow)
            vmat[1,i] += _dot_ao_ao(mol, ao[0], aow, mask, shls_slice, ao_loc)
            vmat[0,i] += _tau_dot(mol, ao, ao, wva[5], mask, shls_slice, ao_loc)
            vmat[1,i] += _tau_dot(mol, ao, ao, wvb[5], mask, shls_slice, ao_loc)
    elif xctype == 'HF':
        pass
    else:
        raise NotImplementedError(f'numint.nr_uks for functional {xc_code}')

    for i in range(nset):
        vmat[0,i] = vmat[0,i] + vmat[0,i].conj().T
        vmat[1,i] = vmat[1,i] + vmat[1,i].conj().T
    if isinstance(dma, numpy.ndarray) and dma.ndim == 2:
        vmat = vmat[:,0]
        nelec = nelec.reshape(2)
        excsum = excsum[0]
    return nelec, excsum, vmat

def _format_uks_dm(dms):
    if isinstance(dms, numpy.ndarray) and dms.ndim == 2:  # RHF DM
        dma = dmb = dms * .5
    else:
        dma, dmb = dms
    if getattr(dms, 'mo_coeff', None) is not None:
        mo_coeff = dms.mo_coeff
        mo_occ = dms.mo_occ
        if mo_coeff[0].ndim < dma.ndim: # handle ROKS
            mo_occa = numpy.array(mo_occ> 0, dtype=numpy.double)
            mo_occb = numpy.array(mo_occ==2, dtype=numpy.double)
            dma = lib.tag_array(dma, mo_coeff=mo_coeff, mo_occ=mo_occa)
            dmb = lib.tag_array(dmb, mo_coeff=mo_coeff, mo_occ=mo_occb)
        else:
            dma = lib.tag_array(dma, mo_coeff=mo_coeff[0], mo_occ=mo_occ[0])
            dmb = lib.tag_array(dmb, mo_coeff=mo_coeff[1], mo_occ=mo_occ[1])
    return dma, dmb

nr_rks_vxc = nr_rks
nr_uks_vxc = nr_uks

def nr_rks_fxc(ni, mol, grids, xc_code, dm0, dms, relativity=0, hermi=0,
               rho0=None, vxc=None, fxc=None, max_memory=2000, verbose=None):
    '''Contract RKS XC (singlet hessian) kernel matrix with given density matrices

    Args:
        ni : an instance of :class:`NumInt`

        mol : an instance of :class:`Mole`

        grids : an instance of :class:`Grids`
            grids.coords and grids.weights are needed for coordinates and weights of meshgrids.
        xc_code : str
            XC functional description.
            See :func:`parse_xc` of pyscf/dft/libxc.py for more details.
        dms : 2D array a list of 2D arrays
            Density matrix or multiple density matrices

    Kwargs:
        hermi : int
            Input density matrices symmetric or not
        max_memory : int or float
            The maximum size of cache to use (in MB).
        rho0 : float array
            Zero-order density (and density derivative for GGA).  Giving kwargs rho0,
            vxc and fxc to improve better performance.
        vxc : float array
            First order XC derivatives
        fxc : float array
            Second order XC derivatives

    Returns:
        nelec, excsum, vmat.
        nelec is the number of electrons generated by numerical integration.
        excsum is the XC functional value.  vmat is the XC potential matrix in
        2D array of shape (nao,nao) where nao is the number of AO functions.

    Examples:

    '''
    xctype = ni._xc_type(xc_code)

    make_rho1, nset, nao = ni._gen_rho_evaluator(mol, dms, hermi)
    if rho0 is None and (xctype != 'LDA' or fxc is None):
        make_rho0 = ni._gen_rho_evaluator(mol, dm0, 1)[0]
    else:
        make_rho0 = None

    def block_loop(ao_deriv):
        p1 = 0
        _rho0 = None
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            p0, p1 = p1, p1 + weight.size
            if rho0 is not None:
                if xctype == 'LDA':
                    _rho0 = numpy.asarray(rho0[p0:p1], order='C')
                else:
                    _rho0 = numpy.asarray(rho0[:,p0:p1], order='C')
            elif make_rho0 is not None:
                _rho0 = make_rho0(0, ao, mask, xctype)
            if vxc is None or fxc is None:
                _vxc, _fxc = ni.eval_xc(xc_code, _rho0, spin=0,
                                        relativity=relativity, deriv=2,
                                        verbose=verbose)[1:3]
            else:
                _vxc = [None if x is None else x[p0:p1] for x in vxc]
                _fxc = [None if x is None else x[p0:p1] for x in fxc]

            for i in range(nset):
                rho1 = make_rho1(i, ao, mask, xctype)
                yield i, _rho0, rho1, ao, mask, weight, _vxc, _fxc

    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_nr()

    if isinstance(dms, numpy.ndarray):
        vmat = numpy.zeros((nset,nao,nao), dtype=dms.dtype)
    else:
        vmat = numpy.zeros((nset,nao,nao), dtype=numpy.result_type(*dms))

    if xctype == 'LDA':
        ao_deriv = 0
        aow = None
        for i, _rho0, rho1, ao, mask, weight, _vxc, _fxc in block_loop(ao_deriv):
            frr = _fxc[0]
            #:aow = numpy.einsum('pi,p->pi', ao, weight*frr*rho1, out=aow)
            aow = _scale_ao(ao, weight*frr*rho1, out=aow)
            vmat[i] += _dot_ao_ao(mol, aow, ao, mask, shls_slice, ao_loc)

    elif xctype == 'GGA':
        ao_deriv = 1
        aow = None
        for i, _rho0, rho1, ao, mask, weight, _vxc, _fxc in block_loop(ao_deriv):
            wv = _rks_gga_wv1(_rho0, rho1, _vxc, _fxc, weight)
            #:aow = numpy.einsum('npi,np->pi', ao, wv, out=aow)
            aow = _scale_ao(ao, wv, out=aow)
            vmat[i] += _dot_ao_ao(mol, aow, ao[0], mask, shls_slice, ao_loc)

        for i in range(nset):  # for (\nabla\mu) \nu + \mu (\nabla\nu)
            vmat[i] = vmat[i] + vmat[i].conj().T

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

    elif xctype == 'MGGA':
        ao_deriv = 2
        aow = None
        for i, _rho0, rho1, ao, mask, weight, _vxc, _fxc in block_loop(ao_deriv):
            wv = _rks_mgga_wv1(_rho0, rho1, _vxc, _fxc, weight)
            #:aow = numpy.einsum('npi,np->pi', ao, wv, out=aow)
            aow = _scale_ao(ao[:4], wv[:4], out=aow)
            vmat[i] += _dot_ao_ao(mol, aow, ao[0], mask, shls_slice, ao_loc)
            vmat[i] += _tau_dot(mol, ao, ao, wv[5], mask, shls_slice, ao_loc)

        for i in range(nset):  # for (\nabla\mu) \nu + \mu (\nabla\nu)
            vmat[i] = vmat[i] + vmat[i].conj().T

    if isinstance(dms, numpy.ndarray) and dms.ndim == 2:
        vmat = vmat[0]
    return vmat

def nr_rks_fxc_st(ni, mol, grids, xc_code, dm0, dms_alpha, relativity=0, singlet=True,
                  rho0=None, vxc=None, fxc=None, max_memory=2000, verbose=None):
    '''Associated to singlet or triplet Hessian
    Note the difference to nr_rks_fxc, dms_alpha is the response density
    matrices of alpha spin, alpha+/-beta DM is applied due to singlet/triplet
    coupling

    Ref. CPL, 256, 454
    '''
    xctype = ni._xc_type(xc_code)

    make_rho1, nset, nao = ni._gen_rho_evaluator(mol, dms_alpha, hermi=0)
    if rho0 is None and (xctype != 'LDA' or fxc is None):
        make_rho0a = ni._gen_rho_evaluator(mol, dm0*.5, hermi=1)[0]
    else:
        make_rho0a = None

    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_nr()

    if isinstance(dms_alpha, numpy.ndarray):
        vmat = numpy.zeros((nset,nao,nao), dtype=dms_alpha.dtype)
    else:
        vmat = numpy.zeros((nset,nao,nao), dtype=numpy.result_type(*dms_alpha))

    if xctype == 'LDA':
        ao_deriv = 0
        aow = None
        p1 = 0
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            p0, p1 = p1, p1 + weight.size
            if fxc is None:
                if rho0 is None:
                    rhoa = make_rho0a(0, ao, mask, xctype)
                else:
                    rhoa = numpy.asarray(rho0[0][p0:p1], order='C')
                _fxc = ni.eval_xc(xc_code, (rhoa,rhoa), spin=1, deriv=2)[2]
                u_u, u_d, d_d = _fxc[0].T
            else:
                if fxc[0].ndim == 1:
                    raise RuntimeError('cached (rho, vxc, fxc) need to be '
                                       'generated by cache_xc_kernel with flag spin=1')
                u_u, u_d, d_d = fxc[0][p0:p1].T
            if singlet:
                frho = u_u + u_d
            else:
                frho = u_u - u_d

            for i in range(nset):
                rho1 = make_rho1(i, ao, mask, 'LDA')
                #:aow = numpy.einsum('pi,p->pi', ao, weight*frho*rho1, out=aow)
                aow = _scale_ao(ao, weight*frho*rho1, out=aow)
                vmat[i] += _dot_ao_ao(mol, aow, ao, mask, shls_slice, ao_loc)
                rho1 = None

    elif xctype == 'GGA':
        ao_deriv = 1
        aow = None
        p1 = 0
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            p0, p1 = p1, p1 + weight.size
            if rho0 is None:
                rhoa = make_rho0a(0, ao, mask, xctype)
            else:
                rhoa = numpy.asarray(rho0[0][:,p0:p1], order='C')
            if vxc is None or fxc is None:
                _vxc, _fxc = ni.eval_xc(xc_code, (rhoa,rhoa), spin=1, deriv=2)[1:3]
            else:
                if rho0[0].ndim == 1:
                    raise RuntimeError('cached (rho, vxc, fxc) need to be '
                                       'generated by cache_xc_kernel with flag spin=1')
                _vxc = [None if x is None else x[p0:p1] for x in vxc]
                _fxc = [None if x is None else x[p0:p1] for x in fxc]
            vsigma = _vxc[1].T
            u_u, u_d, d_d = _fxc[0].T  # v2rho2
            u_uu, u_ud, u_dd, d_uu, d_ud, d_dd = _fxc[1].T  # v2rhosigma
            uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd = _fxc[2].T  # v2sigma2

            # Factorization differs to CPL, 256, 454, to use _rks_gga_wv1 function
            if singlet:
                vsigma = vsigma[0] + vsigma[1] * .5
                frr = u_u + u_d
                fgg = uu_uu + .5*ud_ud + 2*uu_ud + uu_dd
                frg = u_uu + u_dd + u_ud
            else:
                vsigma = vsigma[0] - vsigma[1] * .5
                frr = u_u - u_d
                fgg = uu_uu - uu_dd
                frg = u_uu - u_dd
            _vxc = (None, vsigma)
            _fxc = (frr, frg, fgg)

            for i in range(nset):
                # rho1[0 ] = |b><j| z_{bj}
                # rho1[1:] = \nabla(|b><j|) z_{bj}
                rho1 = make_rho1(i, ao, mask, xctype)
                wv = _rks_gga_wv1(rhoa, rho1, _vxc, _fxc, weight)
                #:aow = numpy.einsum('npi,np->pi', ao, wv, out=aow)
                aow = _scale_ao(ao, wv, out=aow)
                vmat[i] += _dot_ao_ao(mol, aow, ao[0], mask, shls_slice, ao_loc)
                rho1 = None

        for i in range(nset):  # for (\nabla\mu) \nu + \mu (\nabla\nu)
            vmat[i] = vmat[i] + vmat[i].T.conj()

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

    elif xctype == 'MGGA':
        ao_deriv = 2
        aow = None
        p1 = 0
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            p0, p1 = p1, p1 + weight.size
            if rho0 is None:
                rhoa = make_rho0a(0, ao, mask, xctype)
            else:
                rhoa = numpy.asarray(rho0[0][:,p0:p1], order='C')
            if vxc is None or fxc is None:
                _vxc, _fxc = ni.eval_xc(xc_code, (rhoa,rhoa), spin=1, deriv=2)[1:3]
            else:
                if rho0[0].ndim == 1:
                    raise RuntimeError('cached (rho, vxc, fxc) need to be '
                                       'generated by cache_xc_kernel with flag spin=1')
                _vxc = [None if x is None else x[p0:p1] for x in vxc]
                _fxc = [None if x is None else x[p0:p1] for x in fxc]
            vsigma = _vxc[1].T
            u_u, u_d, d_d = _fxc[0].T  # v2rho2
            u_uu, u_ud, u_dd, d_uu, d_ud, d_dd = _fxc[1].T  # v2rhosigma
            uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd = _fxc[2].T  # v2sigma2
            _tt = _fxc[4].T
            _rt = _fxc[6].T
            _gt = _fxc[9].T

            # Factorization differs to CPL, 256, 454, to use _rks_gga_wv1 function
            if singlet:
                vsigma = vsigma[0] + vsigma[1] * .5
                frr = u_u + u_d
                fgg = uu_uu + .5*ud_ud + 2*uu_ud + uu_dd
                frg = u_uu + u_dd + u_ud
                ftt = _tt[0] + _tt[1]  # u_u + u_d
                frt = _rt[0] + _rt[1]  # u_u + u_d
                fgt = _gt[0] + _gt[2] + _gt[4]  # uu_u + ud_u + dd_u
            else:
                vsigma = vsigma[0] - vsigma[1] * .5
                frr = u_u - u_d
                fgg = uu_uu - uu_dd
                frg = u_uu - u_dd
                ftt = _tt[0] - _tt[1]  # u_u - u_d
                frt = _rt[0] - _rt[1]  # u_u - u_d
                fgt = _gt[0] - _gt[4]  # uu_u - dd_u
            _vxc = (None, vsigma, None, None)
            _fxc = (frr, frg, fgg, None, ftt, None, frt, None, None, fgt)

            for i in range(nset):
                # rho1[0 ] = |b><j| z_{bj}
                # rho1[1:] = \nabla(|b><j|) z_{bj}
                rho1 = make_rho1(i, ao, mask, xctype)
                wv = _rks_mgga_wv1(rhoa, rho1, _vxc, _fxc, weight)
                #:aow = numpy.einsum('npi,np->pi', ao, wv, out=aow)
                aow = _scale_ao(ao[:4], wv[:4], out=aow)
                vmat[i] += _dot_ao_ao(mol, aow, ao[0], mask, shls_slice, ao_loc)
                vmat[i] += _tau_dot(mol, ao, ao, wv[5], mask, shls_slice, ao_loc)
                rho1 = None

        for i in range(nset):  # for (\nabla\mu) \nu + \mu (\nabla\nu)
            vmat[i] = vmat[i] + vmat[i].T.conj()

    if isinstance(dms_alpha, numpy.ndarray) and dms_alpha.ndim == 2:
        vmat = vmat[0]
    return vmat

def _rks_gga_wv0(rho, vxc, weight):
    vrho, vgamma = vxc[:2]
    ngrid = vrho.size
    wv = numpy.empty((4,ngrid))
    wv[0]  = vrho * .5  # v+v.T should be applied in the caller
    wv[1:] = 2 * vgamma * rho[1:4]
    wv[:] *= weight
    return wv

def _rks_gga_wv1(rho0, rho1, vxc, fxc, weight):
    vgamma = vxc[1]
    frho, frhogamma, fgg = fxc[:3]
    # sigma1 ~ \nabla(\rho_\alpha+\rho_\beta) dot \nabla(|b><j|) z_{bj}
    sigma1 = numpy.einsum('xi,xi->i', rho0[1:4], rho1[1:4])
    ngrid = sigma1.size
    wv = numpy.empty((4,ngrid), dtype=rho1.dtype)
    wv[0]  = frho * rho1[0]
    wv[0] += frhogamma * sigma1 * 2
    wv[1:] = (fgg * sigma1 * 4 + frhogamma * rho1[0] * 2) * rho0[1:4]
    wv[1:]+= vgamma * rho1[1:4] * 2
    wv *= weight
    wv[0] *= .5  # v+v.T should be applied in the caller
    return wv

def _rks_gga_wv2(rho0, rho1, fxc, kxc, weight):
    frr, frg, fgg = fxc[:3]
    frrr, frrg, frgg, fggg = kxc[:4]
    sigma1 = numpy.einsum('xi,xi->i', rho0[1:4], rho1[1:4])
    r1r1 = rho1[0]**2
    s1s1 = sigma1**2
    r1s1 = rho1[0] * sigma1
    sigma2 = numpy.einsum('xi,xi->i', rho1[1:4], rho1[1:4])
    ngrid = sigma1.size
    wv = numpy.empty((4,ngrid), dtype=rho1.dtype)
    wv[0]  = frrr * r1r1
    wv[0] += 4 * frrg * r1s1
    wv[0] += 4 * frgg * s1s1
    wv[0] += 2 * frg * sigma2
    wv[1:4]  = 2 * frrg * r1r1 * rho0[1:4]
    wv[1:4] += 8 * frgg * r1s1 * rho0[1:4]
    wv[1:4] += 4 * frg * rho1[0] * rho1[1:4]
    wv[1:4] += 4 * fgg * sigma2 * rho0[1:4]
    wv[1:4] += 8 * fgg * sigma1 * rho1[1:4]
    wv[1:4] += 8 * fggg * s1s1 * rho0[1:4]
    wv *= weight
    wv[0]*=.5  # v+v.T should be applied in the caller
    return wv

def _rks_mgga_wv0(rho, vxc, weight):
    vrho, vgamma, vlapl, vtau = vxc[:4]
    ngrid = vrho.size
    wv = numpy.zeros((6,ngrid))
    wv[0] = weight * vrho
    wv[1:4] = (weight * vgamma * 2) * rho[1:4]
    # *0.5 is for tau = 1/2 \nabla\phi\dot\nabla\phi
    wv[5] = weight * vtau * .5
    # *0.5 because v+v.T should be applied in the caller
    wv[0] *= .5
    wv[5] *= .5
    return wv

def _rks_mgga_wv1(rho0, rho1, vxc, fxc, weight):
    vsigma = vxc[1]
    frr, frg, fgg, fll, ftt, frl, frt, flt, fgl, fgt = fxc
    sigma1 = numpy.einsum('xi,xi->i', rho0[1:4], rho1[1:4])
    ngrids = sigma1.size
    wv = numpy.zeros((6,ngrids), dtype=rho1.dtype)
    wv[0]  = frr * rho1[0]
    wv[0] += frt * rho1[5]
    wv[0] += frg * sigma1 * 2
    wv[1:4] = (fgg * sigma1 * 4 + frg * rho1[0] * 2 + fgt * rho1[5] * 2) * rho0[1:4]
    wv[1:4]+= vsigma * rho1[1:4] * 2
    wv[5]  = ftt * rho1[5] * .5
    wv[5] += frt * rho1[0] * .5
    wv[5] += fgt * sigma1
    wv *= weight
    # v+v.T should be applied in the caller
    wv[0] *= .5
    wv[5] *= .5
    return wv

def _rks_mgga_wv2(rho0, rho1, fxc, kxc, weight):
    frr, frg, fgg, fll, ftt, frl, frt, flt, fgl, fgt = fxc
    frrr, frrg, frgg, fggg = kxc[:4]
    frrt = kxc[5]
    frgt = kxc[7]
    frtt = kxc[10]
    fggt = kxc[12]
    fgtt = kxc[15]
    fttt = kxc[19]
    sigma1 = numpy.einsum('xi,xi->i', rho0[1:4], rho1[1:4])
    r1r1 = rho1[0]**2
    t1t1 = rho1[5]**2
    r1t1 = rho1[0] * rho1[5]
    s1s1 = sigma1**2
    r1s1 = rho1[0] * sigma1
    s1t1 = sigma1 * rho1[5]
    sigma2 = numpy.einsum('xi,xi->i', rho1[1:4], rho1[1:4])

    ngrid = sigma1.size
    wv = numpy.zeros((6,ngrid), dtype=rho1.dtype)
    wv[0]  = frrr * r1r1
    wv[0] += 4 * frrg * r1s1
    wv[0] += 4 * frgg * s1s1
    wv[0] += 2 * frg * sigma2
    wv[0] += frtt * t1t1
    wv[0] += 2 * frrt * r1t1
    wv[0] += 4 * frgt * s1t1
    wv[1:4] += 2 * frrg * r1r1 * rho0[1:4]
    wv[1:4] += 8 * frgg * r1s1 * rho0[1:4]
    wv[1:4] += 4 * fgg * sigma2 * rho0[1:4]
    wv[1:4] += 8 * fggg * s1s1 * rho0[1:4]
    wv[1:4] += 2 * fgtt * t1t1 * rho0[1:4]
    wv[1:4] += 8 * fggt * s1t1 * rho0[1:4]
    wv[1:4] += 4 * frgt * r1t1 * rho0[1:4]
    wv[1:4] += 8 * fgg * sigma1 * rho1[1:4]
    wv[1:4] += 4 * frg * rho1[0] * rho1[1:4]
    wv[1:4] += 4 * fgt * rho1[5] * rho1[1:4]
    wv[5] += fttt * t1t1 * .5
    wv[5] += frtt * r1t1
    wv[5] += frrt * r1r1 * .5
    wv[5] += fgtt * s1t1 * 2
    wv[5] += fggt * s1s1 * 2
    wv[5] += frgt * r1s1 * 2
    wv[5] += fgt * sigma2
    wv *= weight
    # v+v.T should be applied in the caller
    wv[0] *= .5
    wv[5] *= .5
    return wv

def nr_uks_fxc(ni, mol, grids, xc_code, dm0, dms, relativity=0, hermi=0,
               rho0=None, vxc=None, fxc=None, max_memory=2000, verbose=None):
    '''Contract UKS XC kernel matrix with given density matrices

    Args:
        ni : an instance of :class:`NumInt`

        mol : an instance of :class:`Mole`

        grids : an instance of :class:`Grids`
            grids.coords and grids.weights are needed for coordinates and weights of meshgrids.
        xc_code : str
            XC functional description.
            See :func:`parse_xc` of pyscf/dft/libxc.py for more details.
        dms : 2D array a list of 2D arrays
            Density matrix or multiple density matrices

    Kwargs:
        hermi : int
            Input density matrices symmetric or not
        max_memory : int or float
            The maximum size of cache to use (in MB).
        rho0 : float array
            Zero-order density (and density derivative for GGA).  Giving kwargs rho0,
            vxc and fxc to improve better performance.
        vxc : float array
            First order XC derivatives
        fxc : float array
            Second order XC derivatives

    Returns:
        nelec, excsum, vmat.
        nelec is the number of electrons generated by numerical integration.
        excsum is the XC functional value.  vmat is the XC potential matrix in
        2D array of shape (nao,nao) where nao is the number of AO functions.

    Examples:

    '''
    xctype = ni._xc_type(xc_code)

    dma, dmb = _format_uks_dm(dms)
    nao = dms.shape[-1]
    make_rhoa, nset = ni._gen_rho_evaluator(mol, dma, hermi)[:2]
    make_rhob       = ni._gen_rho_evaluator(mol, dmb, hermi)[0]

    if rho0 is None and (xctype != 'LDA' or fxc is None):
        make_rho0 = ni._gen_rho_evaluator(mol, _format_uks_dm(dm0), 1)[0]
    else:
        make_rho0 = None

    def block_loop(ao_deriv):
        p1 = 0
        rho0a = rho0b = None
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            p0, p1 = p1, p1 + weight.size
            if rho0 is not None:
                if xctype == 'LDA':
                    rho0a = numpy.asarray(rho0[0][p0:p1], order='C')
                    rho0b = numpy.asarray(rho0[1][p0:p1], order='C')
                else:
                    rho0a = numpy.asarray(rho0[0][:,p0:p1], order='C')
                    rho0b = numpy.asarray(rho0[1][:,p0:p1], order='C')
            elif make_rho0 is not None:
                rho0a = make_rho0(0, ao, mask, xctype)
                rho0b = make_rho0(1, ao, mask, xctype)
            _rho0 = (rho0a, rho0b)
            if vxc is None or fxc is None:
                _vxc, _fxc = ni.eval_xc(xc_code, _rho0, spin=1,
                                        relativity=relativity, deriv=2,
                                        verbose=verbose)[1:3]
            else:
                _vxc = [None if x is None else x[p0:p1] for x in vxc]
                _fxc = [None if x is None else x[p0:p1] for x in fxc]

            for i in range(nset):
                rho1a = make_rhoa(i, ao, mask, xctype)
                rho1b = make_rhob(i, ao, mask, xctype)
                rho1 = (rho1a, rho1b)
                yield i, _rho0, rho1, ao, mask, weight, _vxc, _fxc

    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_nr()

    vmat = numpy.zeros((2,nset,nao,nao), dtype=numpy.result_type(dma, dmb))
    aow = None
    if xctype == 'LDA':
        ao_deriv = 0
        aow = None
        for i, _rho0, rho1, ao, mask, weight, _vxc, _fxc in block_loop(ao_deriv):
            u_u, u_d, d_d = _fxc[0].T
            rho1a, rho1b = rho1
            wva = (u_u * rho1a + u_d * rho1b) * weight
            wvb = (u_d * rho1a + d_d * rho1b) * weight
            #:aow = numpy.einsum('pi,p->pi', ao, wv, out=aow)
            aow = _scale_ao(ao, wva, out=aow)
            vmat[0,i] += _dot_ao_ao(mol, aow, ao, mask, shls_slice, ao_loc)
            #:aow = numpy.einsum('pi,p->pi', ao, wv, out=aow)
            aow = _scale_ao(ao, wvb, out=aow)
            vmat[1,i] += _dot_ao_ao(mol, aow, ao, mask, shls_slice, ao_loc)

    elif xctype == 'GGA':
        ao_deriv = 1
        aow = None
        for i, _rho0, rho1, ao, mask, weight, _vxc, _fxc in block_loop(ao_deriv):
            wva, wvb = _uks_gga_wv1(_rho0, rho1, _vxc, _fxc, weight)
            #:aow = numpy.einsum('npi,np->pi', ao, wva, out=aow)
            aow = _scale_ao(ao, wva, out=aow)
            vmat[0,i] += _dot_ao_ao(mol, aow, ao[0], mask, shls_slice, ao_loc)
            #:aow = numpy.einsum('npi,np->pi', ao, wvb, out=aow)
            aow = _scale_ao(ao, wvb, out=aow)
            vmat[1,i] += _dot_ao_ao(mol, aow, ao[0], mask, shls_slice, ao_loc)

        for i in range(nset):  # for (\nabla\mu) \nu + \mu (\nabla\nu)
            vmat[0,i] = vmat[0,i] + vmat[0,i].conj().T
            vmat[1,i] = vmat[1,i] + vmat[1,i].conj().T

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

    elif xctype == 'MGGA':
        ao_deriv = 2
        aow = None
        for i, _rho0, rho1, ao, mask, weight, _vxc, _fxc in block_loop(ao_deriv):
            wva, wvb = _uks_mgga_wv1(_rho0, rho1, _vxc, _fxc, weight)
            #:aow = numpy.einsum('npi,np->pi', ao, wva, out=aow)
            aow = _scale_ao(ao[:4], wva[:4], out=aow)
            vmat[0,i] += _dot_ao_ao(mol, aow, ao[0], mask, shls_slice, ao_loc)
            vmat[0,i] += _tau_dot(mol, ao, ao, wva[5], mask, shls_slice, ao_loc)
            #:aow = numpy.einsum('npi,np->pi', ao, wvb, out=aow)
            aow = _scale_ao(ao[:4], wvb[:4], out=aow)
            vmat[1,i] += _dot_ao_ao(mol, aow, ao[0], mask, shls_slice, ao_loc)
            vmat[1,i] += _tau_dot(mol, ao, ao, wvb[5], mask, shls_slice, ao_loc)

        for i in range(nset):  # for (\nabla\mu) \nu + \mu (\nabla\nu)
            vmat[0,i] = vmat[0,i] + vmat[0,i].conj().T
            vmat[1,i] = vmat[1,i] + vmat[1,i].conj().T

    if isinstance(dma, numpy.ndarray) and dma.ndim == 2:
        vmat = vmat[:,0]
    return vmat

def _uks_gga_wv0(rho, vxc, weight):
    rhoa, rhob = rho
    vrho, vsigma = vxc[:2]
    ngrids = vrho.shape[0]
    wva, wvb = numpy.empty((2, 4, ngrids))
    wva[0]  = vrho[:,0] * .5  # v+v.T should be applied in the caller
    wva[1:] = rhoa[1:4] * vsigma[:,0] * 2  # sigma_uu
    wva[1:]+= rhob[1:4] * vsigma[:,1]      # sigma_ud
    wva[:] *= weight
    wvb[0]  = vrho[:,1] * .5  # v+v.T should be applied in the caller
    wvb[1:] = rhob[1:4] * vsigma[:,2] * 2  # sigma_dd
    wvb[1:]+= rhoa[1:4] * vsigma[:,1]      # sigma_ud
    wvb[:] *= weight
    return wva, wvb

def _uks_gga_wv1(rho0, rho1, vxc, fxc, weight):
    uu, ud, dd = vxc[1].T
    u_u, u_d, d_d = fxc[0].T
    u_uu, u_ud, u_dd, d_uu, d_ud, d_dd = fxc[1].T
    uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd = fxc[2].T
    ngrid = uu.size

    rho0a, rho0b = rho0
    rho1a, rho1b = rho1
    a0a1 = numpy.einsum('xi,xi->i', rho0a[1:4], rho1a[1:4]) * 2
    a0b1 = numpy.einsum('xi,xi->i', rho0a[1:4], rho1b[1:4])
    b0a1 = numpy.einsum('xi,xi->i', rho0b[1:4], rho1a[1:4])
    b0b1 = numpy.einsum('xi,xi->i', rho0b[1:4], rho1b[1:4]) * 2
    ab_1 = a0b1 + b0a1

    wva, wvb = numpy.empty((2,4,ngrid), dtype=rho1a.dtype)
    # alpha = alpha-alpha * alpha
    wva[0]  = u_u * rho1a[0]
    wva[0] += u_uu * a0a1
    wva[0] += u_ud * ab_1
    wva[1:] = uu * rho1a[1:4] * 2
    wva[1:]+= u_uu * rho1a[0] * rho0a[1:4] * 2
    wva[1:]+= u_ud * rho1a[0] * rho0b[1:4]
    wva[1:]+= uu_uu * a0a1 * rho0a[1:4] * 2
    wva[1:]+= uu_ud * a0a1 * rho0b[1:4]
    wva[1:]+= uu_ud * ab_1 * rho0a[1:4] * 2
    wva[1:]+= ud_ud * ab_1 * rho0b[1:4]
    # alpha = alpha-beta  * beta
    wva[0] += u_d * rho1b[0]
    wva[0] += u_dd * b0b1
    wva[1:]+= ud * rho1b[1:4]
    wva[1:]+= d_uu * rho1b[0] * rho0a[1:4] * 2
    wva[1:]+= d_ud * rho1b[0] * rho0b[1:4]
    wva[1:]+= uu_dd * b0b1 * rho0a[1:4] * 2
    wva[1:]+= ud_dd * b0b1 * rho0b[1:4]
    wva *= weight
    wva[0] *= .5  # v+v.T should be applied in the caller

    # beta = beta-alpha * alpha
    wvb[0]  = u_d * rho1a[0]
    wvb[0] += d_ud * ab_1
    wvb[0] += d_uu * a0a1
    wvb[1:] = ud * rho1a[1:4]
    wvb[1:]+= u_dd * rho1a[0] * rho0b[1:4] * 2
    wvb[1:]+= u_ud * rho1a[0] * rho0a[1:4]
    wvb[1:]+= ud_dd * ab_1 * rho0b[1:4] * 2
    wvb[1:]+= ud_ud * ab_1 * rho0a[1:4]
    wvb[1:]+= uu_dd * a0a1 * rho0b[1:4] * 2
    wvb[1:]+= uu_ud * a0a1 * rho0a[1:4]
    # beta = beta-beta  * beta
    wvb[0] += d_d * rho1b[0]
    wvb[0] += d_dd * b0b1
    wvb[1:]+= dd * rho1b[1:4] * 2
    wvb[1:]+= d_dd * rho1b[0] * rho0b[1:4] * 2
    wvb[1:]+= d_ud * rho1b[0] * rho0a[1:4]
    wvb[1:]+= dd_dd * b0b1 * rho0b[1:4] * 2
    wvb[1:]+= ud_dd * b0b1 * rho0a[1:4]
    wvb *= weight
    wvb[0] *= .5  # v+v.T should be applied in the caller
    return wva, wvb

def _uks_gga_wv2(rho0, rho1, fxc, kxc, weight):
    u_u, u_d, d_d = fxc[0].T
    u_uu, u_ud, u_dd, d_uu, d_ud, d_dd = fxc[1].T
    uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd = fxc[2].T
    u_u_u, u_u_d, u_d_d, d_d_d = kxc[0].T
    u_u_uu, u_u_ud, u_u_dd, u_d_uu, u_d_ud, u_d_dd, d_d_uu, \
            d_d_ud, d_d_dd = kxc[1].T
    u_uu_uu, u_uu_ud, u_uu_dd, u_ud_ud, u_ud_dd, u_dd_dd, d_uu_uu, d_uu_ud, \
            d_uu_dd, d_ud_ud, d_ud_dd, d_dd_dd = kxc[2].T
    uu_uu_uu, uu_uu_ud, uu_uu_dd, uu_ud_ud, uu_ud_dd, uu_dd_dd, ud_ud_ud, \
            ud_ud_dd, ud_dd_dd, dd_dd_dd = kxc[3].T
    ngrid = u_u.size

    rho0a, rho0b = rho0
    rho1a, rho1b = rho1
    a0a1 = numpy.einsum('xi,xi->i', rho0a[1:4], rho1a[1:4]) * 2
    a0b1 = numpy.einsum('xi,xi->i', rho0a[1:4], rho1b[1:4])
    b0a1 = numpy.einsum('xi,xi->i', rho0b[1:4], rho1a[1:4])
    b0b1 = numpy.einsum('xi,xi->i', rho0b[1:4], rho1b[1:4]) * 2
    a1a1 = numpy.einsum('xi,xi->i', rho1a[1:4], rho1a[1:4]) * 2
    a1b1 = numpy.einsum('xi,xi->i', rho1a[1:4], rho1b[1:4]) * 2
    b1b1 = numpy.einsum('xi,xi->i', rho1b[1:4], rho1b[1:4]) * 2
    ab_1 = a0b1 + b0a1

    wva, wvb = numpy.zeros((2, 4, ngrid), dtype=rho1a.dtype)
    wva[0] += u_u_u * rho1a[0] * rho1a[0]
    wva[0] += u_u_d * rho1a[0] * rho1b[0] * 2
    wva[0] += u_d_d * rho1b[0] * rho1b[0]
    wva[0] += u_uu * a1a1
    wva[0] += u_ud * a1b1
    wva[0] += u_dd * b1b1
    wva[1:4] += u_uu * rho1a[0] * rho1a[1:4] * 4
    wva[1:4] += u_ud * rho1a[0] * rho1b[1:4] * 2
    wva[1:4] += d_uu * rho1b[0] * rho1a[1:4] * 4
    wva[1:4] += d_ud * rho1b[0] * rho1b[1:4] * 2
    wva[1:4] += uu_uu * a0a1 * rho1a[1:4] * 4
    wva[1:4] += uu_ud * a0a1 * rho1b[1:4] * 2
    wva[1:4] += uu_ud * ab_1 * rho1a[1:4] * 4
    wva[1:4] += uu_dd * b0b1 * rho1a[1:4] * 4
    wva[1:4] += ud_ud * ab_1 * rho1b[1:4] * 2
    wva[1:4] += ud_dd * b0b1 * rho1b[1:4] * 2
    wva[1:4] += uu_uu * a1a1 * rho0a[1:4] * 2
    wva[1:4] += uu_ud * a1b1 * rho0a[1:4] * 2
    wva[1:4] += uu_dd * b1b1 * rho0a[1:4] * 2
    wva[1:4] += uu_ud * a1a1 * rho0b[1:4]
    wva[1:4] += ud_ud * a1b1 * rho0b[1:4]
    wva[1:4] += ud_dd * b1b1 * rho0b[1:4]
    wva[0] += u_u_uu * rho1a[0] * a0a1 * 2
    wva[0] += u_d_uu * rho1b[0] * a0a1 * 2
    wva[0] += u_u_ud * rho1a[0] * ab_1 * 2
    wva[0] += u_d_ud * rho1b[0] * ab_1 * 2
    wva[0] += u_u_dd * rho1a[0] * b0b1 * 2
    wva[0] += u_d_dd * rho1b[0] * b0b1 * 2
    wva[1:4] += u_u_uu * rho1a[0] * rho1a[0] * rho0a[1:4] * 2
    wva[1:4] += u_d_uu * rho1a[0] * rho1b[0] * rho0a[1:4] * 4
    wva[1:4] += d_d_uu * rho1b[0] * rho1b[0] * rho0a[1:4] * 2
    wva[1:4] += u_u_ud * rho1a[0] * rho1a[0] * rho0a[1:4]
    wva[1:4] += u_d_ud * rho1a[0] * rho1b[0] * rho0a[1:4] * 2
    wva[1:4] += d_d_ud * rho1b[0] * rho1b[0] * rho0a[1:4]
    wva[1:4] += u_uu_uu * rho1a[0] * a0a1 * rho0a[1:4] * 4
    wva[1:4] += d_uu_uu * rho1b[0] * a0a1 * rho0a[1:4] * 4
    wva[1:4] += u_uu_ud * rho1a[0] * ab_1 * rho0a[1:4] * 4
    wva[1:4] += u_uu_ud * rho1a[0] * a0a1 * rho0b[1:4] * 2
    wva[1:4] += u_uu_dd * rho1a[0] * b0b1 * rho0a[1:4] * 4
    wva[1:4] += d_uu_dd * rho1b[0] * b0b1 * rho0a[1:4] * 4
    wva[1:4] += d_uu_ud * rho1b[0] * ab_1 * rho0a[1:4] * 4
    wva[1:4] += d_uu_ud * rho1b[0] * a0a1 * rho0b[1:4] * 2
    wva[1:4] += u_ud_ud * rho1a[0] * ab_1 * rho0b[1:4] * 2
    wva[1:4] += d_ud_ud * rho1b[0] * ab_1 * rho0b[1:4] * 2
    wva[1:4] += u_ud_dd * rho1a[0] * b0b1 * rho0b[1:4] * 2
    wva[1:4] += d_ud_dd * rho1b[0] * b0b1 * rho0b[1:4] * 2
    wva[0] += u_uu_uu * a0a1 * a0a1
    wva[0] += u_uu_ud * a0a1 * ab_1 * 2
    wva[0] += u_uu_dd * a0a1 * b0b1 * 2
    wva[0] += u_ud_ud * ab_1**2
    wva[0] += u_ud_dd * ab_1 * b0b1 * 2
    wva[0] += u_dd_dd * b0b1 * b0b1
    wva[1:4] += uu_uu_uu * a0a1 * a0a1 * rho0a[1:4] * 2
    wva[1:4] += uu_uu_ud * a0a1 * ab_1 * rho0a[1:4] * 4
    wva[1:4] += uu_uu_ud * a0a1 * a0a1 * rho0b[1:4]
    wva[1:4] += uu_uu_dd * a0a1 * b0b1 * rho0a[1:4] * 4
    wva[1:4] += uu_ud_ud * ab_1**2 * rho0a[1:4] * 2
    wva[1:4] += uu_ud_ud * a0a1 * ab_1 * rho0b[1:4] * 2
    wva[1:4] += uu_ud_dd * ab_1 * b0b1 * rho0a[1:4] * 4
    wva[1:4] += uu_ud_dd * a0a1 * b0b1 * rho0b[1:4] * 2
    wva[1:4] += uu_dd_dd * b0b1 * b0b1 * rho0a[1:4] * 2
    wva[1:4] += ud_ud_ud * ab_1**2 * rho0b[1:4]
    wva[1:4] += ud_ud_dd * ab_1 * b0b1 * rho0b[1:4] * 2
    wva[1:4] += ud_dd_dd * b0b1 * b0b1 * rho0b[1:4]
    wva *= weight
    wva[0]*=.5  # v+v.T should be applied in the caller

    wvb[0] += d_d_d * rho1b[0] * rho1b[0]
    wvb[0] += u_d_d * rho1b[0] * rho1a[0] * 2
    wvb[0] += u_u_d * rho1a[0] * rho1a[0]
    wvb[0] += d_dd * b1b1
    wvb[0] += d_ud * a1b1
    wvb[0] += d_uu * a1a1
    wvb[1:4] += u_ud * rho1a[0] * rho1a[1:4] * 2
    wvb[1:4] += u_dd * rho1a[0] * rho1b[1:4] * 4
    wvb[1:4] += d_ud * rho1b[0] * rho1a[1:4] * 2
    wvb[1:4] += d_dd * rho1b[0] * rho1b[1:4] * 4
    wvb[1:4] += dd_dd * b0b1 * rho1b[1:4] * 4
    wvb[1:4] += ud_dd * b0b1 * rho1a[1:4] * 2
    wvb[1:4] += ud_dd * ab_1 * rho1b[1:4] * 4
    wvb[1:4] += ud_ud * ab_1 * rho1a[1:4] * 2
    wvb[1:4] += uu_dd * a0a1 * rho1b[1:4] * 4
    wvb[1:4] += uu_ud * a0a1 * rho1a[1:4] * 2
    wvb[1:4] += dd_dd * b1b1 * rho0b[1:4] * 2
    wvb[1:4] += ud_dd * a1b1 * rho0b[1:4] * 2
    wvb[1:4] += uu_dd * a1a1 * rho0b[1:4] * 2
    wvb[1:4] += ud_dd * b1b1 * rho0a[1:4]
    wvb[1:4] += ud_ud * a1b1 * rho0a[1:4]
    wvb[1:4] += uu_ud * a1a1 * rho0a[1:4]
    wvb[0] += d_d_dd * rho1b[0] * b0b1 * 2
    wvb[0] += u_d_dd * rho1a[0] * b0b1 * 2
    wvb[0] += d_d_ud * rho1b[0] * ab_1 * 2
    wvb[0] += u_d_ud * rho1a[0] * ab_1 * 2
    wvb[0] += d_d_uu * rho1b[0] * a0a1 * 2
    wvb[0] += u_d_uu * rho1a[0] * a0a1 * 2
    wvb[1:4] += d_d_dd * rho1b[0] * rho1b[0] * rho0b[1:4] * 2
    wvb[1:4] += u_d_dd * rho1b[0] * rho1a[0] * rho0b[1:4] * 4
    wvb[1:4] += u_u_dd * rho1a[0] * rho1a[0] * rho0b[1:4] * 2
    wvb[1:4] += d_d_ud * rho1b[0] * rho1b[0] * rho0b[1:4]
    wvb[1:4] += u_d_ud * rho1b[0] * rho1a[0] * rho0b[1:4] * 2
    wvb[1:4] += u_u_ud * rho1a[0] * rho1a[0] * rho0b[1:4]
    wvb[1:4] += d_dd_dd * rho1b[0] * b0b1 * rho0b[1:4] * 4
    wvb[1:4] += u_dd_dd * rho1a[0] * b0b1 * rho0b[1:4] * 4
    wvb[1:4] += d_ud_dd * rho1b[0] * ab_1 * rho0b[1:4] * 4
    wvb[1:4] += u_ud_dd * rho1a[0] * ab_1 * rho0b[1:4] * 4
    wvb[1:4] += d_uu_dd * rho1b[0] * a0a1 * rho0b[1:4] * 4
    wvb[1:4] += u_uu_dd * rho1a[0] * a0a1 * rho0b[1:4] * 4
    wvb[1:4] += d_ud_dd * rho1b[0] * b0b1 * rho0a[1:4] * 2
    wvb[1:4] += u_ud_dd * rho1a[0] * b0b1 * rho0a[1:4] * 2
    wvb[1:4] += d_ud_ud * rho1b[0] * ab_1 * rho0a[1:4] * 2
    wvb[1:4] += u_ud_ud * rho1a[0] * ab_1 * rho0a[1:4] * 2
    wvb[1:4] += d_uu_ud * rho1b[0] * a0a1 * rho0a[1:4] * 2
    wvb[1:4] += u_uu_ud * rho1a[0] * a0a1 * rho0a[1:4] * 2
    wvb[0] += d_dd_dd * b0b1 * b0b1
    wvb[0] += d_ud_dd * ab_1 * b0b1 * 2
    wvb[0] += d_ud_ud * ab_1**2
    wvb[0] += d_uu_dd * b0b1 * a0a1 * 2
    wvb[0] += d_uu_ud * ab_1 * a0a1 * 2
    wvb[0] += d_uu_uu * a0a1 * a0a1
    wvb[1:4] += uu_uu_ud * a0a1 * a0a1 * rho0a[1:4]
    wvb[1:4] += uu_uu_dd * a0a1 * a0a1 * rho0b[1:4] * 2
    wvb[1:4] += uu_ud_ud * ab_1 * a0a1 * rho0a[1:4] * 2
    wvb[1:4] += uu_ud_dd * b0b1 * a0a1 * rho0a[1:4] * 2
    wvb[1:4] += uu_ud_dd * ab_1 * a0a1 * rho0b[1:4] * 4
    wvb[1:4] += uu_dd_dd * b0b1 * a0a1 * rho0b[1:4] * 4
    wvb[1:4] += ud_ud_ud * ab_1**2 * rho0a[1:4]
    wvb[1:4] += ud_ud_dd * ab_1 * b0b1 * rho0a[1:4] * 2
    wvb[1:4] += ud_ud_dd * ab_1**2 * rho0b[1:4] * 2
    wvb[1:4] += ud_dd_dd * b0b1 * b0b1 * rho0a[1:4]
    wvb[1:4] += ud_dd_dd * ab_1 * b0b1 * rho0b[1:4] * 4
    wvb[1:4] += dd_dd_dd * b0b1 * b0b1 * rho0b[1:4] * 2
    wvb *= weight
    wvb[0]*=.5

    return wva, wvb

def _uks_mgga_wv0(rho, vxc, weight):
    rhoa, rhob = rho
    vrho, vsigma, vlapl, vtau = vxc
    ngrid = vrho.shape[0]
    wva, wvb = numpy.zeros((2,6,ngrid))
    wva[0] = vrho[:,0] * .5  # v+v.T should be applied in the caller
    wva[1:4] = rhoa[1:4] * vsigma[:,0] * 2  # sigma_uu
    wva[1:4]+= rhob[1:4] * vsigma[:,1]      # sigma_ud
    wva[5] = vtau[:,0] * .25
    wva *= weight
    wvb[0] = vrho[:,1] * .5  # v+v.T should be applied in the caller
    wvb[1:4] = rhob[1:4] * vsigma[:,2] * 2  # sigma_dd
    wvb[1:4]+= rhoa[1:4] * vsigma[:,1]      # sigma_ud
    wvb[5] = vtau[:,1] * .25
    wvb *= weight
    return wva, wvb

def _uks_mgga_wv1(rho0, rho1, vxc, fxc, weight):
    uu, ud, dd = vxc[1].T
    u_u, u_d, d_d = fxc[0].T
    u_uu, u_ud, u_dd, d_uu, d_ud, d_dd = fxc[1].T
    uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd = fxc[2].T
    ftt = fxc[4].T
    frt = fxc[6].T
    fgt = fxc[9].T
    ngrids = uu.size

    rho0a, rho0b = rho0
    rho1a, rho1b = rho1
    a0a1 = numpy.einsum('xi,xi->i', rho0a[1:4], rho1a[1:4]) * 2
    a0b1 = numpy.einsum('xi,xi->i', rho0a[1:4], rho1b[1:4])
    b0a1 = numpy.einsum('xi,xi->i', rho0b[1:4], rho1a[1:4])
    b0b1 = numpy.einsum('xi,xi->i', rho0b[1:4], rho1b[1:4]) * 2
    ab_1 = a0b1 + b0a1

    wva, wvb = numpy.zeros((2, 6, ngrids), dtype=rho1a.dtype)

    # alpha = alpha-alpha * alpha
    wva[0] += u_u * rho1a[0]
    wva[0] += u_uu * a0a1
    wva[0] += u_ud * ab_1
    wva[0] += frt[0] * rho1a[5]
    wva[1:4]+= uu * rho1a[1:4] * 2
    wva[1:4]+= u_uu * rho1a[0] * rho0a[1:4] * 2
    wva[1:4]+= u_ud * rho1a[0] * rho0b[1:4]
    wva[1:4]+= uu_uu * a0a1 * rho0a[1:4] * 2
    wva[1:4]+= uu_ud * ab_1 * rho0a[1:4] * 2
    wva[1:4]+= uu_ud * a0a1 * rho0b[1:4]
    wva[1:4]+= ud_ud * ab_1 * rho0b[1:4]
    wva[1:4]+= fgt[0] * rho1a[5] * rho0a[1:4] * 2
    wva[1:4]+= fgt[2] * rho1a[5] * rho0b[1:4]
    wva[5] += ftt[0] * rho1a[5] * .5
    wva[5] += frt[0] * rho1a[0] * .5
    wva[5] += fgt[0] * a0a1 * .5
    wva[5] += fgt[2] * ab_1 * .5
    # alpha = alpha-beta  * beta
    wva[0] += u_d * rho1b[0]
    wva[0] += u_dd * b0b1
    wva[0] += frt[1] * rho1b[5]
    wva[1:4]+= ud * rho1b[1:4]
    wva[1:4]+= d_uu * rho1b[0] * rho0a[1:4] * 2
    wva[1:4]+= d_ud * rho1b[0] * rho0b[1:4]
    wva[1:4]+= uu_dd * b0b1 * rho0a[1:4] * 2
    wva[1:4]+= ud_dd * b0b1 * rho0b[1:4]
    # uu_d * rho1b[5] * rho0a[1:4]
    wva[1:4]+= fgt[1] * rho1b[5] * rho0a[1:4] * 2
    wva[1:4]+= fgt[3] * rho1b[5] * rho0b[1:4]
    wva[5] += ftt[1] * rho1b[5] * .5
    wva[5] += frt[2] * rho1b[0] * .5
    wva[5] += fgt[4] * b0b1 * .5
    wva *= weight
    # v+v.T should be applied in the caller
    wva[0] *= .5
    wva[5] *= .5

    # beta = beta-alpha * alpha
    wvb[0] += u_d * rho1a[0]
    wvb[0] += d_ud * ab_1
    wvb[0] += d_uu * a0a1
    wvb[0] += frt[2] * rho1a[5]
    wvb[1:4]+= ud * rho1a[1:4]
    wvb[1:4]+= u_dd * rho1a[0] * rho0b[1:4] * 2
    wvb[1:4]+= u_ud * rho1a[0] * rho0a[1:4]
    wvb[1:4]+= ud_dd * ab_1 * rho0b[1:4] * 2
    wvb[1:4]+= ud_ud * ab_1 * rho0a[1:4]
    wvb[1:4]+= uu_dd * a0a1 * rho0b[1:4] * 2
    wvb[1:4]+= uu_ud * a0a1 * rho0a[1:4]
    # dd_u * rho1a[5] * rho0b[1:4]
    wvb[1:4]+= fgt[4] * rho1a[5] * rho0b[1:4] * 2
    wvb[1:4]+= fgt[2] * rho1a[5] * rho0a[1:4]
    wvb[5] += ftt[1] * rho1a[5] * .5
    wvb[5] += frt[1] * rho1a[0] * .5
    wvb[5] += fgt[3] * ab_1 * .5
    wvb[5] += fgt[1] * a0a1 * .5
    # beta = beta-beta  * beta
    wvb[0] += d_d * rho1b[0]
    wvb[0] += d_dd * b0b1
    wvb[0] += frt[3] * rho1b[5]
    wvb[1:4]+= dd * rho1b[1:4] * 2
    wvb[1:4]+= d_dd * rho1b[0] * rho0b[1:4] * 2
    wvb[1:4]+= d_ud * rho1b[0] * rho0a[1:4]
    wvb[1:4]+= dd_dd * b0b1 * rho0b[1:4] * 2
    wvb[1:4]+= ud_dd * b0b1 * rho0a[1:4]
    wvb[1:4]+= fgt[5] * rho1b[5] * rho0b[1:4] * 2
    wvb[1:4]+= fgt[3] * rho1b[5] * rho0a[1:4]
    wvb[5] += ftt[2] * rho1b[5] * .5
    wvb[5] += frt[3] * rho1b[0] * .5
    wvb[5] += fgt[5] * b0b1 * .5
    wvb *= weight
    # v+v.T should be applied in the caller
    wvb[0] *= .5
    wvb[5] *= .5
    return wva, wvb

def _uks_mgga_wv2(rho0, rho1, fxc, kxc, weight):
    u_u, u_d, d_d = fxc[0].T
    u_uu, u_ud, u_dd, d_uu, d_ud, d_dd = fxc[1].T
    uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd = fxc[2].T
    u_u_u, u_u_d, u_d_d, d_d_d = kxc[0].T
    u_u_uu, u_u_ud, u_u_dd, u_d_uu, u_d_ud, u_d_dd, d_d_uu, \
            d_d_ud, d_d_dd = kxc[1].T
    u_uu_uu, u_uu_ud, u_uu_dd, u_ud_ud, u_ud_dd, u_dd_dd, d_uu_uu, d_uu_ud, \
            d_uu_dd, d_ud_ud, d_ud_dd, d_dd_dd = kxc[2].T
    uu_uu_uu, uu_uu_ud, uu_uu_dd, uu_ud_ud, uu_ud_dd, uu_dd_dd, ud_ud_ud, \
            ud_ud_dd, ud_dd_dd, dd_dd_dd = kxc[3].T
    fgt = fxc[9].T
    frrt = kxc[5].T
    frgt = kxc[7].T
    frtt = kxc[10].T
    fggt = kxc[12].T
    fgtt = kxc[15].T
    fttt = kxc[19].T
    ngrid = u_u.size

    rho0a, rho0b = rho0
    rho1a, rho1b = rho1
    a0a1 = numpy.einsum('xi,xi->i', rho0a[1:4], rho1a[1:4]) * 2
    a0b1 = numpy.einsum('xi,xi->i', rho0a[1:4], rho1b[1:4])
    b0a1 = numpy.einsum('xi,xi->i', rho0b[1:4], rho1a[1:4])
    b0b1 = numpy.einsum('xi,xi->i', rho0b[1:4], rho1b[1:4]) * 2
    a1a1 = numpy.einsum('xi,xi->i', rho1a[1:4], rho1a[1:4]) * 2
    a1b1 = numpy.einsum('xi,xi->i', rho1a[1:4], rho1b[1:4]) * 2
    b1b1 = numpy.einsum('xi,xi->i', rho1b[1:4], rho1b[1:4]) * 2
    ab_1 = a0b1 + b0a1
    rara = rho1a[0] * rho1a[0]
    rarb = rho1a[0] * rho1b[0]
    rbrb = rho1b[0] * rho1b[0]
    rata = rho1a[0] * rho1a[5]
    ratb = rho1a[0] * rho1b[5]
    rbta = rho1b[0] * rho1a[5]
    rbtb = rho1b[0] * rho1b[5]
    tata = rho1a[5] * rho1a[5]
    tatb = rho1a[5] * rho1b[5]
    tbtb = rho1b[5] * rho1b[5]

    wva, wvb = numpy.zeros((2, 6, ngrid))

    wva[0] += u_u_u * rara
    wva[0] += u_u_d * rarb * 2
    wva[0] += u_d_d * rbrb
    wva[0] += u_uu * a1a1
    wva[0] += u_ud * a1b1
    wva[0] += u_dd * b1b1
    wva[0] += u_u_uu * rho1a[0] * a0a1 * 2
    wva[0] += u_u_ud * rho1a[0] * ab_1 * 2
    wva[0] += u_u_dd * rho1a[0] * b0b1 * 2
    wva[0] += u_d_uu * rho1b[0] * a0a1 * 2
    wva[0] += u_d_ud * rho1b[0] * ab_1 * 2
    wva[0] += u_d_dd * rho1b[0] * b0b1 * 2
    wva[0] += u_uu_uu * a0a1 * a0a1
    wva[0] += u_uu_ud * a0a1 * ab_1 * 2
    wva[0] += u_uu_dd * a0a1 * b0b1 * 2
    wva[0] += u_ud_ud * ab_1**2
    wva[0] += u_ud_dd * b0b1 * ab_1 * 2
    wva[0] += u_dd_dd * b0b1 * b0b1
    wva[0] += frgt[0] * rho1a[5] * a0a1 * 2                     # u_uu_u
    wva[0] += frgt[1] * rho1b[5] * a0a1 * 2                     # u_uu_d
    wva[0] += frgt[2] * rho1a[5] * ab_1 * 2                     # u_ud_u
    wva[0] += frgt[3] * rho1b[5] * ab_1 * 2                     # u_ud_d
    wva[0] += frgt[4] * rho1a[5] * b0b1 * 2                     # u_dd_u
    wva[0] += frgt[5] * rho1b[5] * b0b1 * 2                     # u_dd_d
    wva[0] += frrt[0] * rata * 2                                # u_u_u
    wva[0] += frrt[1] * ratb * 2                                # u_u_d
    wva[0] += frrt[2] * rbta * 2                                # u_d_u
    wva[0] += frrt[3] * rbtb * 2                                # u_d_d
    wva[0] += frtt[0] * tata                                    # u_u_u
    wva[0] += frtt[1] * tatb * 2                                # u_u_d
    wva[0] += frtt[2] * tbtb                                    # u_d_d
    wva[1:4] += u_uu * rho1a[0] * rho1a[1:4] * 4
    wva[1:4] += u_ud * rho1a[0] * rho1b[1:4] * 2
    wva[1:4] += d_uu * rho1b[0] * rho1a[1:4] * 4
    wva[1:4] += d_ud * rho1b[0] * rho1b[1:4] * 2
    wva[1:4] += uu_uu * a1a1 * rho0a[1:4] * 2
    wva[1:4] += uu_uu * a0a1 * rho1a[1:4] * 4
    wva[1:4] += uu_ud * a0a1 * rho1b[1:4] * 2
    wva[1:4] += uu_ud * ab_1 * rho1a[1:4] * 4
    wva[1:4] += uu_ud * a1b1 * rho0a[1:4] * 2
    wva[1:4] += uu_ud * a1a1 * rho0b[1:4]
    wva[1:4] += uu_dd * b1b1 * rho0a[1:4] * 2
    wva[1:4] += uu_dd * b0b1 * rho1a[1:4] * 4
    wva[1:4] += ud_ud * ab_1 * rho1b[1:4] * 2
    wva[1:4] += ud_ud * a1b1 * rho0b[1:4]
    wva[1:4] += ud_dd * b0b1 * rho1b[1:4] * 2
    wva[1:4] += ud_dd * b1b1 * rho0b[1:4]
    wva[1:4] += u_u_uu * rara * rho0a[1:4] * 2
    wva[1:4] += u_u_ud * rara * rho0b[1:4]
    wva[1:4] += u_d_uu * rarb * rho0a[1:4] * 4
    wva[1:4] += u_d_ud * rarb * rho0b[1:4] * 2
    wva[1:4] += d_d_uu * rbrb * rho0a[1:4] * 2
    wva[1:4] += d_d_ud * rbrb * rho0b[1:4]
    wva[1:4] += u_uu_uu * rho1a[0] * a0a1 * rho0a[1:4] * 4
    wva[1:4] += u_uu_ud * rho1a[0] * ab_1 * rho0a[1:4] * 4
    wva[1:4] += u_uu_ud * rho1a[0] * a0a1 * rho0b[1:4] * 2
    wva[1:4] += u_uu_dd * rho1a[0] * b0b1 * rho0a[1:4] * 4
    wva[1:4] += u_ud_ud * rho1a[0] * ab_1 * rho0b[1:4] * 2
    wva[1:4] += u_ud_dd * rho1a[0] * b0b1 * rho0b[1:4] * 2
    wva[1:4] += d_uu_uu * rho1b[0] * a0a1 * rho0a[1:4] * 4
    wva[1:4] += d_uu_ud * rho1b[0] * ab_1 * rho0a[1:4] * 4
    wva[1:4] += d_uu_ud * rho1b[0] * a0a1 * rho0b[1:4] * 2
    wva[1:4] += d_uu_dd * rho1b[0] * b0b1 * rho0a[1:4] * 4
    wva[1:4] += d_ud_ud * rho1b[0] * ab_1 * rho0b[1:4] * 2
    wva[1:4] += d_ud_dd * rho1b[0] * b0b1 * rho0b[1:4] * 2
    wva[1:4] += uu_uu_uu * a0a1 * a0a1 * rho0a[1:4] * 2
    wva[1:4] += uu_uu_ud * a0a1 * ab_1 * rho0a[1:4] * 4
    wva[1:4] += uu_uu_ud * a0a1 * a0a1 * rho0b[1:4]
    wva[1:4] += uu_uu_dd * a0a1 * b0b1 * rho0a[1:4] * 4
    wva[1:4] += uu_ud_ud * ab_1**2 * rho0a[1:4] * 2
    wva[1:4] += uu_ud_ud * a0a1 * ab_1 * rho0b[1:4] * 2
    wva[1:4] += uu_ud_dd * ab_1 * b0b1 * rho0a[1:4] * 4
    wva[1:4] += uu_ud_dd * a0a1 * b0b1 * rho0b[1:4] * 2
    wva[1:4] += uu_dd_dd * b0b1 * b0b1 * rho0a[1:4] * 2
    wva[1:4] += ud_ud_ud * ab_1**2 * rho0b[1:4]
    wva[1:4] += ud_ud_dd * ab_1 * b0b1 * rho0b[1:4] * 2
    wva[1:4] += ud_dd_dd * b0b1 * b0b1 * rho0b[1:4]
    wva[1:4] += frgt[0] * rata * rho0a[1:4] * 4                 # u_uu_u
    wva[1:4] += frgt[1] * ratb * rho0a[1:4] * 4                 # u_uu_d
    wva[1:4] += frgt[2] * rata * rho0b[1:4] * 2                 # u_ud_u
    wva[1:4] += frgt[3] * ratb * rho0b[1:4] * 2                 # u_ud_d
    wva[1:4] += frgt[6] * rbta * rho0a[1:4] * 4                 # d_uu_u
    wva[1:4] += frgt[7] * rbtb * rho0a[1:4] * 4                 # d_uu_d
    wva[1:4] += frgt[8] * rbta * rho0b[1:4] * 2                 # d_ud_u
    wva[1:4] += frgt[9] * rbtb * rho0b[1:4] * 2                 # d_ud_d
    wva[1:4] += fgt[0] * rho1a[5] * rho1a[1:4] * 4              # uu_u
    wva[1:4] += fgt[1] * rho1b[5] * rho1a[1:4] * 4              # uu_d
    wva[1:4] += fgt[2] * rho1a[5] * rho1b[1:4] * 2              # ud_u
    wva[1:4] += fgt[3] * rho1b[5] * rho1b[1:4] * 2              # ud_d
    wva[1:4] += fggt[0] * rho1a[5] * a0a1 * rho0a[1:4] * 4      # uu_uu_u
    wva[1:4] += fggt[1] * rho1b[5] * a0a1 * rho0a[1:4] * 4      # uu_uu_d
    wva[1:4] += fggt[2] * rho1a[5] * a0a1 * rho0b[1:4] * 2      # uu_ud_u
    wva[1:4] += fggt[2] * rho1a[5] * ab_1 * rho0a[1:4] * 4      # uu_ud_u
    wva[1:4] += fggt[3] * rho1b[5] * a0a1 * rho0b[1:4] * 2      # uu_ud_d
    wva[1:4] += fggt[3] * rho1b[5] * ab_1 * rho0a[1:4] * 4      # uu_ud_d
    wva[1:4] += fggt[4] * rho1a[5] * b0b1 * rho0a[1:4] * 4      # uu_dd_u
    wva[1:4] += fggt[5] * rho1b[5] * b0b1 * rho0a[1:4] * 4      # uu_dd_d
    wva[1:4] += fggt[6] * rho1a[5] * ab_1 * rho0b[1:4] * 2      # ud_ud_u
    wva[1:4] += fggt[7] * rho1b[5] * ab_1 * rho0b[1:4] * 2      # ud_ud_d
    wva[1:4] += fggt[8] * rho1a[5] * b0b1 * rho0b[1:4] * 2      # ud_dd_u
    wva[1:4] += fggt[9] * rho1b[5] * b0b1 * rho0b[1:4] * 2      # ud_dd_d
    wva[1:4] += fgtt[0] * tata * rho0a[1:4] * 2                 # uu_u_u
    wva[1:4] += fgtt[1] * tatb * rho0a[1:4] * 4                 # uu_u_d
    wva[1:4] += fgtt[2] * tbtb * rho0a[1:4] * 2                 # uu_d_d
    wva[1:4] += fgtt[3] * tata * rho0b[1:4]                     # ud_u_u
    wva[1:4] += fgtt[4] * tatb * rho0b[1:4] * 2                 # ud_u_d
    wva[1:4] += fgtt[5] * tbtb * rho0b[1:4]                     # ud_d_d
    wva[5] += frgt[0 ] * rho1a[0] * a0a1                        # u_uu_u
    wva[5] += frgt[2 ] * rho1a[0] * ab_1                        # u_ud_u
    wva[5] += frgt[4 ] * rho1a[0] * b0b1                        # u_dd_u
    wva[5] += frgt[6 ] * rho1b[0] * a0a1                        # d_uu_u
    wva[5] += frgt[8 ] * rho1b[0] * ab_1                        # d_ud_u
    wva[5] += frgt[10] * rho1b[0] * b0b1                        # d_dd_u
    wva[5] += fggt[0 ] * a0a1 * a0a1 * .5                       # uu_uu_u
    wva[5] += fggt[2 ] * a0a1 * ab_1                            # uu_ud_u
    wva[5] += fggt[4 ] * a0a1 * b0b1                            # uu_dd_u
    wva[5] += fggt[6 ] * ab_1**2 * .5                           # ud_ud_u
    wva[5] += fggt[8 ] * ab_1 * b0b1                            # ud_dd_u
    wva[5] += fggt[10] * b0b1 * b0b1 * .5                       # dd_dd_u
    wva[5] += fgtt[0] * a0a1 * rho1a[5]                         # uu_u_u
    wva[5] += fgtt[1] * a0a1 * rho1b[5]                         # uu_u_d
    wva[5] += fgtt[3] * ab_1 * rho1a[5]                         # ud_u_u
    wva[5] += fgtt[4] * ab_1 * rho1b[5]                         # ud_u_d
    wva[5] += fgtt[6] * b0b1 * rho1a[5]                         # dd_u_u
    wva[5] += fgtt[7] * b0b1 * rho1b[5]                         # dd_u_d
    wva[5] += fgt[0] * a1a1 * .5                                # uu_u
    wva[5] += fgt[2] * a1b1 * .5                                # ud_u
    wva[5] += fgt[4] * b1b1 * .5                                # dd_u
    wva[5] += frrt[0] * rara * .5                               # u_u_u
    wva[5] += frrt[2] * rarb                                    # u_d_u
    wva[5] += frrt[4] * rbrb * .5                               # d_d_u
    wva[5] += frtt[0] * rata                                    # u_u_u
    wva[5] += frtt[1] * ratb                                    # u_u_d
    wva[5] += frtt[3] * rbta                                    # d_u_u
    wva[5] += frtt[4] * rbtb                                    # d_u_d
    wva[5] += fttt[0] * tata * .5                               # u_u_u
    wva[5] += fttt[1] * tatb                                    # u_u_d
    wva[5] += fttt[2] * tbtb * .5                               # u_d_d
    wva *= weight
    wva[0] *= .5
    wva[5] *= .5

    wvb[0] += u_u_d * rara
    wvb[0] += u_d_d * rarb * 2
    wvb[0] += d_d_d * rbrb
    wvb[0] += d_uu * a1a1
    wvb[0] += d_ud * a1b1
    wvb[0] += d_dd * b1b1
    wvb[0] += u_d_uu * rho1a[0] * a0a1 * 2
    wvb[0] += u_d_ud * rho1a[0] * ab_1 * 2
    wvb[0] += u_d_dd * rho1a[0] * b0b1 * 2
    wvb[0] += d_d_uu * rho1b[0] * a0a1 * 2
    wvb[0] += d_d_ud * rho1b[0] * ab_1 * 2
    wvb[0] += d_d_dd * rho1b[0] * b0b1 * 2
    wvb[0] += d_uu_uu * a0a1 * a0a1
    wvb[0] += d_uu_ud * a0a1 * ab_1 * 2
    wvb[0] += d_uu_dd * b0b1 * a0a1 * 2
    wvb[0] += d_ud_ud * ab_1**2
    wvb[0] += d_ud_dd * b0b1 * ab_1 * 2
    wvb[0] += d_dd_dd * b0b1 * b0b1
    wvb[0] += frgt[6 ] * rho1a[5] * a0a1 * 2                    # d_uu_u
    wvb[0] += frgt[7 ] * rho1b[5] * a0a1 * 2                    # d_uu_d
    wvb[0] += frgt[8 ] * rho1a[5] * ab_1 * 2                    # d_ud_u
    wvb[0] += frgt[9 ] * rho1b[5] * ab_1 * 2                    # d_ud_d
    wvb[0] += frgt[10] * rho1a[5] * b0b1 * 2                    # d_dd_u
    wvb[0] += frgt[11] * rho1b[5] * b0b1 * 2                    # d_dd_d
    wvb[0] += frrt[2] * rata * 2                                # u_d_u
    wvb[0] += frrt[3] * ratb * 2                                # u_d_d
    wvb[0] += frrt[4] * rbta * 2                                # d_d_u
    wvb[0] += frrt[5] * rbtb * 2                                # d_d_d
    wvb[0] += frtt[3] * tata                                    # d_u_u
    wvb[0] += frtt[4] * tatb * 2                                # d_u_d
    wvb[0] += frtt[5] * tbtb                                    # d_d_d
    wvb[1:4] += u_ud * rho1a[0] * rho1a[1:4] * 2
    wvb[1:4] += u_dd * rho1a[0] * rho1b[1:4] * 4
    wvb[1:4] += d_ud * rho1b[0] * rho1a[1:4] * 2
    wvb[1:4] += d_dd * rho1b[0] * rho1b[1:4] * 4
    wvb[1:4] += uu_ud * a1a1 * rho0a[1:4]
    wvb[1:4] += uu_ud * a0a1 * rho1a[1:4] * 2
    wvb[1:4] += uu_dd * a1a1 * rho0b[1:4] * 2
    wvb[1:4] += uu_dd * a0a1 * rho1b[1:4] * 4
    wvb[1:4] += ud_ud * a1b1 * rho0a[1:4]
    wvb[1:4] += ud_ud * ab_1 * rho1a[1:4] * 2
    wvb[1:4] += ud_dd * b1b1 * rho0a[1:4]
    wvb[1:4] += ud_dd * a1b1 * rho0b[1:4] * 2
    wvb[1:4] += ud_dd * b0b1 * rho1a[1:4] * 2
    wvb[1:4] += ud_dd * ab_1 * rho1b[1:4] * 4
    wvb[1:4] += dd_dd * b1b1 * rho0b[1:4] * 2
    wvb[1:4] += dd_dd * b0b1 * rho1b[1:4] * 4
    wvb[1:4] += u_u_ud * rara * rho0a[1:4]
    wvb[1:4] += u_u_dd * rara * rho0b[1:4] * 2
    wvb[1:4] += u_d_ud * rarb * rho0a[1:4] * 2
    wvb[1:4] += u_d_dd * rarb * rho0b[1:4] * 4
    wvb[1:4] += d_d_ud * rbrb * rho0a[1:4]
    wvb[1:4] += d_d_dd * rbrb * rho0b[1:4] * 2
    wvb[1:4] += u_uu_ud * rho1a[0] * a0a1 * rho0a[1:4] * 2
    wvb[1:4] += u_uu_dd * rho1a[0] * a0a1 * rho0b[1:4] * 4
    wvb[1:4] += u_ud_ud * rho1a[0] * ab_1 * rho0a[1:4] * 2
    wvb[1:4] += u_ud_dd * rho1a[0] * b0b1 * rho0a[1:4] * 2
    wvb[1:4] += u_ud_dd * rho1a[0] * ab_1 * rho0b[1:4] * 4
    wvb[1:4] += u_dd_dd * rho1a[0] * b0b1 * rho0b[1:4] * 4
    wvb[1:4] += d_uu_ud * rho1b[0] * a0a1 * rho0a[1:4] * 2
    wvb[1:4] += d_uu_dd * rho1b[0] * a0a1 * rho0b[1:4] * 4
    wvb[1:4] += d_ud_ud * rho1b[0] * ab_1 * rho0a[1:4] * 2
    wvb[1:4] += d_ud_dd * rho1b[0] * b0b1 * rho0a[1:4] * 2
    wvb[1:4] += d_ud_dd * rho1b[0] * ab_1 * rho0b[1:4] * 4
    wvb[1:4] += d_dd_dd * rho1b[0] * b0b1 * rho0b[1:4] * 4
    wvb[1:4] += uu_uu_ud * a0a1 * a0a1 * rho0a[1:4]
    wvb[1:4] += uu_uu_dd * a0a1 * a0a1 * rho0b[1:4] * 2
    wvb[1:4] += uu_ud_ud * ab_1 * a0a1 * rho0a[1:4] * 2
    wvb[1:4] += uu_ud_dd * b0b1 * a0a1 * rho0a[1:4] * 2
    wvb[1:4] += uu_ud_dd * ab_1 * a0a1 * rho0b[1:4] * 4
    wvb[1:4] += uu_dd_dd * b0b1 * a0a1 * rho0b[1:4] * 4
    wvb[1:4] += ud_ud_ud * ab_1**2 * rho0a[1:4]
    wvb[1:4] += ud_ud_dd * b0b1 * ab_1 * rho0a[1:4] * 2
    wvb[1:4] += ud_ud_dd * ab_1**2 * rho0b[1:4] * 2
    wvb[1:4] += ud_dd_dd * b0b1 * b0b1 * rho0a[1:4]
    wvb[1:4] += ud_dd_dd * b0b1 * ab_1 * rho0b[1:4] * 4
    wvb[1:4] += dd_dd_dd * b0b1 * b0b1 * rho0b[1:4] * 2
    wvb[1:4] += frgt[2 ] * rata * rho0a[1:4] * 2                # u_ud_u
    wvb[1:4] += frgt[3 ] * ratb * rho0a[1:4] * 2                # u_ud_d
    wvb[1:4] += frgt[4 ] * rata * rho0b[1:4] * 4                # u_dd_u
    wvb[1:4] += frgt[5 ] * ratb * rho0b[1:4] * 4                # u_dd_d
    wvb[1:4] += frgt[8 ] * rbta * rho0a[1:4] * 2                # d_ud_u
    wvb[1:4] += frgt[9 ] * rbtb * rho0a[1:4] * 2                # d_ud_d
    wvb[1:4] += frgt[10] * rbta * rho0b[1:4] * 4                # d_dd_u
    wvb[1:4] += frgt[11] * rbtb * rho0b[1:4] * 4                # d_dd_d
    wvb[1:4] += fgt[2] * rho1a[5] * rho1a[1:4] * 2              # ud_u
    wvb[1:4] += fgt[3] * rho1b[5] * rho1a[1:4] * 2              # ud_d
    wvb[1:4] += fgt[4] * rho1a[5] * rho1b[1:4] * 4              # dd_u
    wvb[1:4] += fgt[5] * rho1b[5] * rho1b[1:4] * 4              # dd_d
    wvb[1:4] += fggt[2 ] * rho1a[5] * a0a1 * rho0a[1:4] * 2     # uu_ud_u
    wvb[1:4] += fggt[3 ] * rho1b[5] * a0a1 * rho0a[1:4] * 2     # uu_ud_d
    wvb[1:4] += fggt[4 ] * rho1a[5] * a0a1 * rho0b[1:4] * 4     # uu_dd_u
    wvb[1:4] += fggt[5 ] * rho1b[5] * a0a1 * rho0b[1:4] * 4     # uu_dd_d
    wvb[1:4] += fggt[6 ] * rho1a[5] * ab_1 * rho0a[1:4] * 2     # ud_ud_u
    wvb[1:4] += fggt[7 ] * rho1b[5] * ab_1 * rho0a[1:4] * 2     # ud_ud_d
    wvb[1:4] += fggt[8 ] * rho1a[5] * ab_1 * rho0b[1:4] * 4     # ud_dd_u
    wvb[1:4] += fggt[8 ] * rho1a[5] * b0b1 * rho0a[1:4] * 2     # ud_dd_u
    wvb[1:4] += fggt[9 ] * rho1b[5] * ab_1 * rho0b[1:4] * 4     # ud_dd_d
    wvb[1:4] += fggt[9 ] * rho1b[5] * b0b1 * rho0a[1:4] * 2     # ud_dd_d
    wvb[1:4] += fggt[10] * rho1a[5] * b0b1 * rho0b[1:4] * 4     # dd_dd_u
    wvb[1:4] += fggt[11] * rho1b[5] * b0b1 * rho0b[1:4] * 4     # dd_dd_d
    wvb[1:4] += fgtt[3] * tata * rho0a[1:4]                     # ud_u_u
    wvb[1:4] += fgtt[4] * tatb * rho0a[1:4] * 2                 # ud_u_d
    wvb[1:4] += fgtt[5] * tbtb * rho0a[1:4]                     # ud_d_d
    wvb[1:4] += fgtt[6] * tata * rho0b[1:4] * 2                 # dd_u_u
    wvb[1:4] += fgtt[7] * tatb * rho0b[1:4] * 4                 # dd_u_d
    wvb[1:4] += fgtt[8] * tbtb * rho0b[1:4] * 2                 # dd_d_d
    wvb[5] += frgt[1 ] * rho1a[0] * a0a1                        # u_uu_d
    wvb[5] += frgt[3 ] * rho1a[0] * ab_1                        # u_ud_d
    wvb[5] += frgt[5 ] * rho1a[0] * b0b1                        # u_dd_d
    wvb[5] += frgt[7 ] * rho1b[0] * a0a1                        # d_uu_d
    wvb[5] += frgt[9 ] * rho1b[0] * ab_1                        # d_ud_d
    wvb[5] += frgt[11] * rho1b[0] * b0b1                        # d_dd_d
    wvb[5] += fggt[1 ] * a0a1 * a0a1 * .5                       # uu_uu_d
    wvb[5] += fggt[3 ] * ab_1 * a0a1 * .5                       # uu_ud_d
    wvb[5] += fggt[5 ] * b0b1 * a0a1                            # uu_dd_d
    wvb[5] += fggt[7 ] * ab_1**2 * .5                           # ud_ud_d
    wvb[5] += fggt[9 ] * b0b1 * ab_1                            # ud_dd_d
    wvb[5] += fggt[11] * b0b1 * b0b1 * .5                       # dd_dd_d
    wvb[5] += fgt[1] * a1a1 * .5                                # uu_d
    wvb[5] += fgt[3] * a1b1 * .5                                # ud_d
    wvb[5] += fgt[5] * b1b1 * .5                                # dd_d
    wvb[5] += fgtt[1] * a0a1 * rho1a[5]                         # uu_u_d
    wvb[5] += fgtt[2] * a0a1 * rho1b[5]                         # uu_d_d
    wvb[5] += fgtt[4] * ab_1 * rho1a[5]                         # ud_u_d
    wvb[5] += fgtt[5] * ab_1 * rho1b[5]                         # ud_d_d
    wvb[5] += fgtt[7] * b0b1 * rho1a[5]                         # dd_u_d
    wvb[5] += fgtt[8] * b0b1 * rho1b[5]                         # dd_d_d
    wvb[5] += frrt[1] * rara * .5                               # u_u_d
    wvb[5] += frrt[3] * rarb                                    # u_d_d
    wvb[5] += frrt[5] * rbrb * .5                               # d_d_d
    wvb[5] += frtt[1] * rata                                    # u_u_d
    wvb[5] += frtt[2] * ratb                                    # u_d_d
    wvb[5] += frtt[4] * rbta                                    # d_u_d
    wvb[5] += frtt[5] * rbtb                                    # d_d_d
    wvb[5] += fttt[1] * tata * .5                               # u_u_d
    wvb[5] += fttt[2] * tatb                                    # u_d_d
    wvb[5] += fttt[3] * tbtb * .5                               # d_d_d
    wvb *= weight
    wvb[0] *= .5
    wvb[5] *= .5

    return wva, wvb


def nr_fxc(mol, grids, xc_code, dm0, dms, spin=0, relativity=0, hermi=0,
           rho0=None, vxc=None, fxc=None, max_memory=2000, verbose=None):
    r'''Contract XC kernel matrix with given density matrices

    ... math::

            a_{pq} = f_{pq,rs} * x_{rs}

    '''
    ni = NumInt()
    return ni.nr_fxc(mol, grids, xc_code, dm0, dms, spin, relativity,
                     hermi, rho0, vxc, fxc, max_memory, verbose)


def cache_xc_kernel(ni, mol, grids, xc_code, mo_coeff, mo_occ, spin=0,
                    max_memory=2000):
    '''Compute the 0th order density, Vxc and fxc.  They can be used in TDDFT,
    DFT hessian module etc.
    '''
    xctype = ni._xc_type(xc_code)
    if xctype == 'MGGA':
        ao_deriv = 2
    elif xctype == 'GGA':
        ao_deriv = 1
    elif xctype == 'NLC':
        raise NotImplementedError('NLC')
    else:
        ao_deriv = 0

    if spin == 0:
        nao = mo_coeff.shape[0]
        rho = []
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory=max_memory):
            rho.append(ni.eval_rho2(mol, ao, mo_coeff, mo_occ, mask, xctype))
        rho = numpy.hstack(rho)
    else:
        nao = mo_coeff[0].shape[0]
        rhoa = []
        rhob = []
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            rhoa.append(ni.eval_rho2(mol, ao, mo_coeff[0], mo_occ[0], mask, xctype))
            rhob.append(ni.eval_rho2(mol, ao, mo_coeff[1], mo_occ[1], mask, xctype))
        rho = (numpy.hstack(rhoa), numpy.hstack(rhob))
    vxc, fxc = ni.eval_xc(xc_code, rho, spin=spin, relativity=0, deriv=2,
                          verbose=0)[1:3]
    return rho, vxc, fxc

def get_rho(ni, mol, dm, grids, max_memory=2000):
    '''Density in real space
    '''
    make_rho, nset, nao = ni._gen_rho_evaluator(mol, dm, 1)
    assert(nset == 1)
    rho = numpy.empty(grids.weights.size)
    p1 = 0
    for ao, mask, weight, coords \
            in ni.block_loop(mol, grids, nao, 0, max_memory):
        p0, p1 = p1, p1 + weight.size
        rho[p0:p1] = make_rho(0, ao, mask, 'LDA')
    return rho

class _NumIntMixin(lib.StreamObject):
    libxc = libxc

    @lib.with_doc(eval_ao.__doc__)
    def eval_ao(self, mol, coords, deriv=0, shls_slice=None,
                non0tab=None, out=None, verbose=None):
        return eval_ao(mol, coords, deriv, shls_slice, non0tab, out, verbose)

    get_rho = get_rho

    def block_loop(self, mol, grids, nao=None, deriv=0, max_memory=2000,
                   non0tab=None, blksize=None, buf=None):
        '''Define this macro to loop over grids by blocks.
        '''
        if grids.coords is None:
            grids.build(with_non0tab=True)
        if nao is None:
            nao = mol.nao
        ngrids = grids.coords.shape[0]
        comp = (deriv+1)*(deriv+2)*(deriv+3)//6
# NOTE to index grids.non0tab, the blksize needs to be the integer multiplier of BLKSIZE
        if blksize is None:
            blksize = int(max_memory*1e6/(comp*2*nao*8*BLKSIZE))*BLKSIZE
            blksize = max(BLKSIZE, min(blksize, ngrids, BLKSIZE*1200))
        if non0tab is None:
            non0tab = grids.non0tab
        if non0tab is None:
            non0tab = numpy.ones(((ngrids+BLKSIZE-1)//BLKSIZE,mol.nbas),
                                 dtype=numpy.uint8)
        if buf is None:
            buf = numpy.empty((comp,blksize,nao))
        for ip0 in range(0, ngrids, blksize):
            ip1 = min(ngrids, ip0+blksize)
            coords = grids.coords[ip0:ip1]
            weight = grids.weights[ip0:ip1]
            non0 = non0tab[ip0//BLKSIZE:]
            ao = self.eval_ao(mol, coords, deriv=deriv, non0tab=non0, out=buf)
            yield ao, non0, weight, coords

    def _gen_rho_evaluator(self, mol, dms, hermi=0):
        if getattr(dms, 'mo_coeff', None) is not None:
            #TODO: test whether dm.mo_coeff matching dm
            mo_coeff = dms.mo_coeff
            mo_occ = dms.mo_occ
            if isinstance(dms, numpy.ndarray) and dms.ndim == 2:
                mo_coeff = [mo_coeff]
                mo_occ = [mo_occ]
            nao = mo_coeff[0].shape[0]
            ndms = len(mo_occ)
            def make_rho(idm, ao, non0tab, xctype):
                return self.eval_rho2(mol, ao, mo_coeff[idm], mo_occ[idm],
                                      non0tab, xctype)
        else:
            if isinstance(dms, numpy.ndarray) and dms.ndim == 2:
                dms = [dms]
            if not hermi:
                # For eval_rho when xctype==GGA, which requires hermitian DMs
                dms = [(dm+dm.conj().T)*.5 for dm in dms]
            nao = dms[0].shape[0]
            ndms = len(dms)
            def make_rho(idm, ao, non0tab, xctype):
                return self.eval_rho(mol, ao, dms[idm], non0tab, xctype, hermi=1)
        return make_rho, ndms, nao

####################
# Overwrite following functions to use custom XC functional

    def hybrid_coeff(self, xc_code, spin=0):
        return self.libxc.hybrid_coeff(xc_code, spin)

    def nlc_coeff(self, xc_code):
        return self.libxc.nlc_coeff(xc_code)

    def rsh_coeff(self, xc_code):
        return self.libxc.rsh_coeff(xc_code)

    @lib.with_doc(libxc.eval_xc.__doc__)
    def eval_xc(self, xc_code, rho, spin=0, relativity=0, deriv=1, omega=None,
                verbose=None):
        if omega is None: omega = self.omega
        return self.libxc.eval_xc(xc_code, rho, spin, relativity, deriv,
                                  omega, verbose)

    def _xc_type(self, xc_code):
        return self.libxc.xc_type(xc_code)

    def rsh_and_hybrid_coeff(self, xc_code, spin=0):
        '''Range-separated parameter and HF exchange components: omega, alpha, beta

        Exc_RSH = c_SR * SR_HFX + c_LR * LR_HFX + (1-c_SR) * Ex_SR + (1-c_LR) * Ex_LR + Ec
                = alpha * HFX + beta * SR_HFX + (1-c_SR) * Ex_SR + (1-c_LR) * Ex_LR + Ec
                = alpha * LR_HFX + hyb * SR_HFX + (1-c_SR) * Ex_SR + (1-c_LR) * Ex_LR + Ec

        SR_HFX = < pi | e^{-omega r_{12}}/r_{12} | iq >
        LR_HFX = < pi | (1-e^{-omega r_{12}})/r_{12} | iq >
        alpha = c_LR
        beta = c_SR - c_LR
        '''
        omega, alpha, beta = self.rsh_coeff(xc_code)
        if self.omega is not None:
            omega = self.omega

        if abs(omega) > 1e-10:
            hyb = alpha + beta
        else:
            hyb = self.hybrid_coeff(xc_code, spin)
        return omega, alpha, hyb

class NumInt(_NumIntMixin):
    '''Numerical integration methods for non-relativistic RKS and UKS'''

    def __init__(self):
        self.omega = None  # RSH paramter

    @lib.with_doc(nr_vxc.__doc__)
    def nr_vxc(self, mol, grids, xc_code, dms, spin=0, relativity=0, hermi=0,
               max_memory=2000, verbose=None):
        if spin == 0:
            return self.nr_rks(mol, grids, xc_code, dms, relativity, hermi,
                               max_memory, verbose)
        else:
            return self.nr_uks(mol, grids, xc_code, dms, relativity, hermi,
                               max_memory, verbose)
    get_vxc = nr_vxc

    @lib.with_doc(nr_fxc.__doc__)
    def nr_fxc(self, mol, grids, xc_code, dm0, dms, spin=0, relativity=0, hermi=0,
               rho0=None, vxc=None, fxc=None, max_memory=2000, verbose=None):
        if spin == 0:
            return self.nr_rks_fxc(mol, grids, xc_code, dm0, dms, relativity,
                                   hermi, rho0, vxc, fxc, max_memory, verbose)
        else:
            return self.nr_uks_fxc(mol, grids, xc_code, dm0, dms, relativity,
                                   hermi, rho0, vxc, fxc, max_memory, verbose)
    get_fxc = nr_fxc

    @lib.with_doc(nr_sap_vxc.__doc__)
    def nr_sap(self, mol, grids, max_memory=2000, verbose=None):
        return self.nr_sap_vxc(mol, grids, max_memory, verbose)

    nr_rks = nr_rks
    nr_uks = nr_uks
    nr_sap_vxc = nr_sap_vxc
    nr_rks_fxc = nr_rks_fxc
    nr_uks_fxc = nr_uks_fxc
    cache_xc_kernel  = cache_xc_kernel

    @lib.with_doc(make_mask.__doc__)
    def make_mask(self, mol, coords, relativity=0, shls_slice=None,
                  verbose=None):
        return make_mask(mol, coords, relativity, shls_slice, verbose)

    @lib.with_doc(eval_rho.__doc__)
    def eval_rho(self, mol, ao, dm, non0tab=None, xctype='LDA', hermi=0, verbose=None):
        return eval_rho(mol, ao, dm, non0tab, xctype, hermi, verbose)

    @lib.with_doc(eval_rho2.__doc__)
    def eval_rho2(self, mol, ao, mo_coeff, mo_occ, non0tab=None, xctype='LDA',
                  verbose=None):
        return eval_rho2(mol, ao, mo_coeff, mo_occ, non0tab, xctype, verbose)

_NumInt = NumInt


if __name__ == '__main__':
    from pyscf import gto
    from pyscf import dft

    mol = gto.M(atom=[
        ["O" , (0. , 0.     , 0.)],
        [1   , (0. , -0.757 , 0.587)],
        [1   , (0. , 0.757  , 0.587)] ],
        basis='6311g**')
    mf = dft.RKS(mol)
    mf.grids.atom_grid = {"H": (30, 194), "O": (30, 194),}
    mf.grids.prune = None
    mf.grids.build()
    dm = mf.get_init_guess(key='minao')

    numpy.random.seed(1)
    dm1 = numpy.random.random((dm.shape))
    dm1 = lib.hermi_triu(dm1)
    res = mf._numint.nr_vxc(mol, mf.grids, mf.xc, dm1, spin=0)
    print(res[1] - -37.084047825971282)
    res = mf._numint.nr_vxc(mol, mf.grids, mf.xc, (dm1,dm1), spin=1)
    print(res[1] - -92.436362308687094)
    res = mf._numint.nr_vxc(mol, mf.grids, mf.xc, dm, spin=0)
    print(res[1] - -8.6313329288394947)
    res = mf._numint.nr_vxc(mol, mf.grids, mf.xc, (dm,dm), spin=1)
    print(res[1] - -21.520301399504582)


# pyscf/dft/numint2c.py

#!/usr/bin/env python
# Copyright 2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Numerical integration functions for (2-component) GKS with real AO basis
'''

import numpy as np
import scipy.linalg
from pyscf import lib
from pyscf.dft import numint


class NumInt2C(numint._NumIntMixin):
    '''Numerical integration methods for 2-component basis (used by GKS)'''

    def __init__(self):
        self.omega = None  # RSH paramter
        self.collinear = True

    def eval_rho(self, mol, ao, dm, non0tab=None, xctype='LDA', hermi=0, verbose=None):
        '''Calculate the electron density for LDA functional and the density
        derivatives for GGA functional in the framework of 2-component basis.
        '''
        nao = ao.shape[-1]
        assert dm.ndim == 2 and nao * 2 == dm.shape[0]
        dm_a = dm[:nao,:nao].real.copy()
        dm_b = dm[nao:,nao:].real.copy()
        rho = numint.eval_rho(mol, ao, dm_a, non0tab, xctype, hermi, verbose)
        rho += numint.eval_rho(mol, ao, dm_b, non0tab, xctype, hermi, verbose)
        if dm.dtype == np.complex128:
            dm_a = dm[:nao,:nao].imag
            dm_b = dm[nao:,nao:].imag
            if abs(dm_a).max() > 1e-8 or abs(dm_b).max() > 1e-8:
                rhoI = numint.eval_rho(mol, ao, dm_a, non0tab, xctype, hermi, verbose)
                rhoI += numint.eval_rho(mol, ao, dm_b, non0tab, xctype, hermi, verbose)
                rho = rho + rhoI * 1j
        return rho

    def eval_rho2(self, mol, ao, mo_coeff, mo_occ, non0tab=None, xctype='LDA',
                  verbose=None):
        '''Calculate the electron density for LDA functional and the density
        derivatives for GGA functional in the framework of 2-component basis.
        '''
        if mo_coeff.dtype == np.double:
            nao = ao.shape[-1]
            assert nao * 2 == mo_coeff.shape[0]
            mo_aR = mo_coeff[:nao]
            mo_bR = mo_coeff[nao:]
            hermi = 1
            rho  = numint.eval_rho2(mol, ao, mo_aR, mo_occ, non0tab, xctype, hermi, verbose)
            rho += numint.eval_rho2(mol, ao, mo_bR, mo_occ, non0tab, xctype, hermi, verbose)
        else:
            dm = np.dot(mo_coeff * mo_occ, mo_coeff.conj().T)
            hermi = 1
            rho = self.eval_rho(mol, dm, ao, dm, non0tab, xctype, hermi, verbose)
        return rho

    def cache_xc_kernel(self, mol, grids, xc_code, mo_coeff, mo_occ, spin=0,
                        max_memory=2000):
        '''Compute the 0th order density, Vxc and fxc.  They can be used in TDDFT,
        DFT hessian module etc.
        '''
        xctype = self._xc_type(xc_code)
        if xctype == 'MGGA':
            ao_deriv = 2
        elif xctype == 'GGA':
            ao_deriv = 1
        elif xctype == 'NLC':
            raise NotImplementedError('NLC')
        else:
            ao_deriv = 0

        nao = mo_coeff.shape[0] // 2
        dm = np.dot(mo_coeff * mo_occ, mo_coeff.conj().T)
        dm_a = dm[:nao,:nao].real.copy()
        dm_b = dm[nao:,nao:].real.copy()
        rhoa = []
        rhob = []
        for ao, mask, weight, coords \
                in self.block_loop(mol, grids, nao, ao_deriv, max_memory):
            # rhoa and rhob have to be real
            rhoa.append(numint.eval_rho(mol, ao, dm_a, mask, xctype))
            rhob.append(numint.eval_rho(mol, ao, dm_b, mask, xctype))
        rho = (np.hstack(rhoa), np.hstack(rhob))
        vxc, fxc = self.eval_xc(xc_code, rho, spin=1, relativity=0, deriv=2,
                                verbose=0)[1:3]
        return rho, vxc, fxc

    def get_rho(self, mol, dm, grids, max_memory=2000):
        '''Density in real space
        '''
        nao = dm.shape[-1] // 2
        dm_a = dm[:nao,:nao].real
        dm_b = dm[nao:,nao:].real
        ni = self.view(numint.NumInt)
        return numint.get_rho(ni, mol, dm_a+dm_b, grids, max_memory)

    def nr_vxc(self, mol, grids, xc_code, dms, spin=0, relativity=0, hermi=1,
               max_memory=2000, verbose=None):
        dms = np.asarray(dms)
        nao = dms.shape[-1] // 2
        # ground state density is real
        dm_a = dms[...,:nao,:nao].real.copy()
        dm_b = dms[...,nao:,nao:].real.copy()
        ni = self.view(numint.NumInt)
        n, exc, vxc = numint.nr_uks(ni, mol, grids, xc_code, (dm_a, dm_b),
                                    max_memory=max_memory)
        vmat = np.zeros_like(dms)
        vmat[...,:nao,:nao] = vxc[0]
        vmat[...,nao:,nao:] = vxc[1]
        return n, exc, vmat
    get_vxc = nr_gks_vxc = nr_vxc

    def nr_fxc(self, mol, grids, xc_code, dm0, dms, spin=0, relativity=0, hermi=0,
               rho0=None, vxc=None, fxc=None, max_memory=2000, verbose=None):
        dms = np.asarray(dms)
        nao = dms.shape[-1] // 2
        dm0a = dm0[:nao,:nao].real.copy()
        dm0b = dm0[nao:,nao:].real.copy()
        # dms_a and dms_b may be complex if they are TDDFT amplitudes
        dms_a = dms[...,:nao,:nao].copy()
        dms_b = dms[...,nao:,nao:].copy()
        ni = self.view(numint.NumInt)
        vmat = numint.nr_uks_fxc(
            ni, mol, grids, xc_code, (dm0a, dm0b), (dms_a, dms_b), relativity=0,
            hermi=0, rho0=None, vxc=None, fxc=None, max_memory=2000, verbose=None)
        fxcmat = np.zeros_like(dms)
        fxcmat[...,:nao,:nao] = vmat[0]
        fxcmat[...,nao:,nao:] = vmat[1]
        return fxcmat
    get_fxc = nr_gks_fxc = nr_fxc


# pyscf/dft/r_numint.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Numerical integration functions for (4-component) DKS with j-adapted AO basis
'''

import numpy
from pyscf import lib
from pyscf.dft import numint
from pyscf.dft.numint import _dot_ao_dm, _dot_ao_ao, _scale_ao, BLKSIZE


def eval_ao(mol, coords, deriv=0, with_s=True, shls_slice=None,
            non0tab=None, out=None, verbose=None):
    '''Evaluates the value of 2-component or 4-component j-adapted basis on grids.
    '''
    comp = (deriv+1)*(deriv+2)*(deriv+3)//6
    feval = 'GTOval_spinor_deriv%d' % deriv
    if not with_s:
        # aoLa, aoLb = aoL
        ao = aoL = mol.eval_gto(feval, coords, comp, shls_slice, non0tab, out=out)
    else:
        assert(deriv <= 1)  # only GTOval_ipsp_spinor
        ngrids = coords.shape[0]
        nao = mol.nao_2c()
        ao = numpy.ndarray((4,comp,nao,ngrids), dtype=numpy.complex128, buffer=out)
        aoL = mol.eval_gto(feval, coords, comp, shls_slice, non0tab, out=ao[:2])  # noqa
        ao = ao.transpose(0,1,3,2)
        aoS = ao[2:]
        aoSa, aoSb = aoS
        feval_gto = ['GTOval_sp_spinor', 'GTOval_ipsp_spinor']
        p1 = 0
        for n in range(deriv+1):
            comp = (n+1)*(n+2)//2
            p0, p1 = p1, p1 + comp
            aoSa[p0:p1], aoSb[p0:p1] = mol.eval_gto(
                feval_gto[n], coords, comp, shls_slice, non0tab)

        if deriv == 0:
            ao = ao[:,0]
    return ao

def _dot_spinor_dm(mol, ket, dm, non0tab, shls_slice, ao_loc):
    ket_a, ket_b = ket
    outa = _dot_ao_dm(mol, ket_a, dm, non0tab, shls_slice, ao_loc)
    outb = _dot_ao_dm(mol, ket_b, dm, non0tab, shls_slice, ao_loc)
    return outa, outb

def _contract_rho2x2(bra, ket):
    '''rho2x2 = sum_i |ket_i> <bra_i|
    '''
    ket_a, ket_b = ket
    bra_a, bra_b = bra
    rhoaa = numpy.einsum('pi,pi->p', ket_a.real, bra_a.real)
    rhoaa+= numpy.einsum('pi,pi->p', ket_a.imag, bra_a.imag)
    rhoab = numpy.einsum('pi,pi->p', ket_a, bra_b.conj())
    rhoba = numpy.einsum('pi,pi->p', ket_b, bra_a.conj())
    rhobb = numpy.einsum('pi,pi->p', ket_b.real, bra_b.real)
    rhobb+= numpy.einsum('pi,pi->p', ket_b.imag, bra_b.imag)
    return rhoaa, rhoab, rhoba, rhobb

def _rho2x2_to_rho_m(rho2x2):
    # rho = einsum('xgi,ij,xgj->g', ket, dm, bra.conj())
    # mx = einsum('xy,ygi,ij,xgj->g', sx, ket, dm, bra.conj())
    # my = einsum('xy,ygi,ij,xgj->g', sy, ket, dm, bra.conj())
    # mz = einsum('xy,ygi,ij,xgj->g', sz, ket, dm, bra.conj())
    raa, rab, rba, rbb = rho2x2
    ngrids = raa.size
    rho_m = numpy.empty((4, ngrids))
    rho, mx, my, mz = rho_m
    rho[:] = raa.real + rbb.real
    mx[:] = rab.real + rba.real
    my[:] = rba.imag - rab.imag
    mz[:] = raa.real - rbb.real
    return rho_m

def _get_rho_m(mol, ket, dm, bra, non0tab, shls_slice, ao_loc):
    out = _dot_spinor_dm(mol, ket, dm, non0tab, shls_slice, ao_loc)
    tmp = _contract_rho2x2(bra, out)
    return _rho2x2_to_rho_m(tmp)

def eval_rho(mol, ao, dm, non0tab=None, xctype='LDA', hermi=0, verbose=None):
    aoa, aob = ao
    ngrids, nao = aoa.shape[-2:]
    xctype = xctype.upper()

    if non0tab is None:
        non0tab = numpy.ones(((ngrids+BLKSIZE-1)//BLKSIZE,mol.nbas),
                             dtype=numpy.uint8)
    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_2c()

    if xctype == 'LDA':
        rho_m = _get_rho_m(mol, ao, dm, ao, non0tab, shls_slice, ao_loc)
    elif xctype == 'GGA':
        # first 4 ~ (rho, m), second 4 ~ (rho0, dx, dy, dz)
        rho_m = numpy.empty((4, 4, ngrids))
        c0 = _dot_spinor_dm(mol, ao[:,0], dm, non0tab, shls_slice, ao_loc)
        rho_m[:,0] = _rho2x2_to_rho_m(_contract_rho2x2(ao[:,0], c0))
        for i in range(1, 4):
            rho_m[:,i] = _rho2x2_to_rho_m(_contract_rho2x2(ao[:,i], c0))
            rho_m[:,i] *= 2  # *2 for +c.c. corresponding to |dx ao> dm < ao|
    else: # meta-GGA
        rho_m = numpy.zeros((4, 6, ngrids))
        c0 = _dot_spinor_dm(mol, ao[:,0], dm, non0tab, shls_slice, ao_loc)
        rho_m[:,0] = _rho2x2_to_rho_m(_contract_rho2x2(ao[:,0], c0))
        for i in range(1, 4):
            rho_m[:,i] = _rho2x2_to_rho_m(_contract_rho2x2(ao[:,i], c0))
            rho_m[:,i] *= 2  # *2 for +c.c. corresponding to |dx ao> dm < ao|
            c1 = _dot_spinor_dm(mol, ao[:,1], dm, non0tab, shls_slice, ao_loc)
            rho_m[:,5] += _contract_rho2x2(ao[i], c1)
        XX, YY, ZZ = 4, 7, 9
        ao2 = ao[XX] + ao[YY] + ao[ZZ]
        # \nabla^2 rho
        rho_m[:,4] = _rho2x2_to_rho_m(_contract_rho2x2(ao2, c0))
        rho_m[:,4] += rho_m[:,5]
        rho_m[:,4] *= 2
        # tau = 1/2 (\nabla f)^2
        rho_m[:,5] *= .5

    rho, m = rho_m[0], rho_m[1:]
    return rho, m

def _ncol_lda_vxc_mat(mol, ao, weight, rho, vxc, non0tab, shls_slice, ao_loc,
                      on_LL=True):
    '''Vxc matrix of non-collinear LDA'''
    aoa, aob = ao
    r, m = rho
    vrho = vxc[0]
    vr, vm = vrho.T
    s = lib.norm(m, axis=0)

    idx = s < 1e-20
    with numpy.errstate(divide='ignore',invalid='ignore'):
        ws = vm * weight / s
    ws[idx] = 0

    # * .5 because of v+v.conj().T in r_vxc
    wv = .5 * weight * vr
    if on_LL:
        ws *= .5
    else:  # for SS block
        # flip the sign for small components because
        # M = \beta\Sigma = [sigma,  0    ]
        #                   [0    , -sigma]
        ws *= -.5

    # einsum('g,g,xgi,xgj->ij', vr, weight, ao, ao)
    # + einsum('xy,g,g,xgi,ygj->ij', sx, vm*m[0]/s, weight, ao, ao)
    # + einsum('xy,g,g,xgi,ygj->ij', sy, vm*m[1]/s, weight, ao, ao)
    # + einsum('xy,g,g,xgi,ygj->ij', sz, vm*m[2]/s, weight, ao, ao)
    aow = None
    aow = _scale_ao(aoa, ws*m[0], out=aow)  # Mx
    tmp = _dot_ao_ao(mol, aob, aow, non0tab, shls_slice, ao_loc)
    mat = tmp + tmp.conj().T
    aow = _scale_ao(aoa, ws*m[1], out=aow)  # My
    tmp = _dot_ao_ao(mol, aob, aow, non0tab, shls_slice, ao_loc)
    mat+= (tmp - tmp.conj().T) * 1j
    aow = _scale_ao(aoa, wv, out=aow)
    aow+= _scale_ao(aoa, ws*m[2])  # Mz
    mat+= _dot_ao_ao(mol, aoa, aow, non0tab, shls_slice, ao_loc)
    aow = _scale_ao(aob, wv, out=aow)
    aow-= _scale_ao(aob, ws*m[2])  # Mz
    mat+= _dot_ao_ao(mol, aob, aow, non0tab, shls_slice, ao_loc)
    return mat

def _ncol_gga_vxc_mat(mol, ao, weight, rho, vxc, non0tab, shls_slice, ao_loc,
                      on_LL=True):
    '''Vxc matrix of non-collinear GGA'''
    raise NotImplementedError

def _ncol_mgga_vxc_mat(mol, ao, weight, rho, vxc, non0tab, shls_slice, ao_loc,
                       on_LL=True):
    '''Vxc matrix of non-collinear MGGA'''
    raise NotImplementedError

def _dks_gga_wv0(rho, vxc, weight):
    r, (mx, my, mz) = rho
    rhoa = (r + mz) * .5
    rhob = (r - mz) * .5
    return numint._uks_gga_wv0((rhoa, rhob), vxc, weight)

def _col_lda_vxc_mat(mol, ao, weight, rho, vxc, non0tab, shls_slice, ao_loc,
                     on_LL=True):
    '''Vxc matrix of collinear LDA'''
    vrho = vxc[0]
    aoa, aob = ao
    if on_LL:
        wva, wvb = .5 * weight * vrho.T
    else:  # for SS block
        # v_rho = (vxc_a + vxc_b) * .5
        # v_mz  = (vxc_a - vxc_b) * .5
        # For small components, M = \beta\Sigma leads to
        # (v_rho - sigma_z*v_mz) = [vxc_b, 0    ]
        #                          [0    , vxc_a]
        wvb, wva = .5 * weight * vrho.T
    mat  = _dot_ao_ao(mol, aoa, _scale_ao(aoa, wva), non0tab, shls_slice, ao_loc)
    mat += _dot_ao_ao(mol, aob, _scale_ao(aob, wvb), non0tab, shls_slice, ao_loc)
    return mat

def _col_gga_vxc_mat(mol, ao, weight, rho, vxc, non0tab, shls_slice, ao_loc,
                     on_LL=True):
    '''Vxc matrix of collinear GGA'''
    if on_LL:
        wva, wvb = _dks_gga_wv0(rho, vxc, weight)
    else:  # for SS block
        wvb, wva = _dks_gga_wv0(rho, vxc, weight)
    aoa, aob = ao
    mat  = _dot_ao_ao(mol, aoa[0], _scale_ao(aoa[:4], wva[:4]), non0tab, shls_slice, ao_loc)
    mat += _dot_ao_ao(mol, aob[0], _scale_ao(aob[:4], wvb[:4]), non0tab, shls_slice, ao_loc)
    return mat

def _col_mgga_vxc_mat(mol, ao, weight, rho, vxc, non0tab, shls_slice, ao_loc,
                      on_LL=True):
    '''Vxc matrix of collinear MGGA'''
    if on_LL:
        wva, wvb = _dks_gga_wv0(rho, vxc, weight)
    else:
        wvb, wva = _dks_gga_wv0(rho, vxc, weight)
    aoa, aob = ao
    mat  = _dot_ao_ao(mol, aoa[0], _scale_ao(aoa[:4], wva[:4]), non0tab, shls_slice, ao_loc)
    mat += _dot_ao_ao(mol, aob[0], _scale_ao(aob[:4], wvb[:4]), non0tab, shls_slice, ao_loc)

    vtau = vxc[3]
    if on_LL:
        wva, wvb = .25 * weight * vtau.T
    else:
        wvb, wva = .25 * weight * vtau.T
    mat += _dot_ao_ao(mol, aoa[1], _scale_ao(aoa[1], wva), non0tab, shls_slice, ao_loc)
    mat += _dot_ao_ao(mol, aoa[2], _scale_ao(aoa[2], wva), non0tab, shls_slice, ao_loc)
    mat += _dot_ao_ao(mol, aoa[3], _scale_ao(aoa[3], wva), non0tab, shls_slice, ao_loc)
    mat += _dot_ao_ao(mol, aob[1], _scale_ao(aob[1], wvb), non0tab, shls_slice, ao_loc)
    mat += _dot_ao_ao(mol, aob[2], _scale_ao(aob[2], wvb), non0tab, shls_slice, ao_loc)
    mat += _dot_ao_ao(mol, aob[3], _scale_ao(aob[3], wvb), non0tab, shls_slice, ao_loc)
    return mat

def r_vxc(ni, mol, grids, xc_code, dms, spin=0, relativity=1, hermi=1,
          max_memory=2000, verbose=None):
    '''Calculate 2-component or 4-component Vxc matrix in j-adapted basis
    '''
    xctype = ni._xc_type(xc_code)
    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_2c()
    n2c = ao_loc[-1]

    make_rho, nset, nao = ni._gen_rho_evaluator(mol, dms, hermi)
    with_s = (nao == n2c*2)  # 4C DM
    nelec = numpy.zeros(nset)
    excsum = numpy.zeros(nset)
    matLL = numpy.zeros((nset,n2c,n2c), dtype=numpy.complex128)
    matSS = numpy.zeros_like(matLL)

    if xctype in ('LDA', 'GGA', 'MGGA'):
        f_eval_mat = {
            ('LDA' , True ): (_col_lda_vxc_mat  , 0),
            ('GGA' , True ): (_col_gga_vxc_mat  , 1),
            ('MGGA', True ): (_col_mgga_vxc_mat , 2),
            ('LDA' , False): (_ncol_lda_vxc_mat , 0),
            ('GGA' , False): (_ncol_gga_vxc_mat , 1),
            ('MGGA', False): (_ncol_mgga_vxc_mat, 2)
        }
        fmat, ao_deriv = f_eval_mat[(xctype, ni.collinear)]

        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, with_s, max_memory):
            for idm in range(nset):
                rho = make_rho(idm, ao, mask, xctype)
                exc, vxc = ni.eval_xc(xc_code, rho, spin=1,
                                      relativity=relativity, deriv=1,
                                      verbose=verbose)[:2]
                den = rho[0] * weight
                nelec[idm] += den.sum()
                excsum[idm] += (den*exc).sum()

                matLL[idm] += fmat(mol, ao[:2], weight, rho, vxc,
                                   mask, shls_slice, ao_loc, True)
                if with_s:
                    matSS[idm] += fmat(mol, ao[2:], weight, rho, vxc,
                                       mask, shls_slice, ao_loc, False)

        matLL = matLL + matLL.conj().transpose(0,2,1)
        if with_s:
            matSS = matSS + matSS.conj().transpose(0,2,1)
    elif xctype == 'HF':
        pass
    else:
        raise NotImplementedError(f'r_vxc for functional {xc_code}')

    if with_s:
        matSS *= (.5 / lib.param.LIGHT_SPEED)**2
        vmat = numpy.zeros((nset,nao,nao), dtype=numpy.complex128)
        vmat[:,:n2c,:n2c] = matLL
        vmat[:,n2c:,n2c:] = matSS
    else:
        vmat = matLL

    if nset == 1:
        nelec = nelec[0]
        excsum = excsum[0]
    return nelec, excsum, vmat.reshape(dms.shape)

def _dks_lda_wv1(rho0, rho1, vxc, fxc, weight):
    u_u, u_d, d_d = fxc[0].T
    rho1, (m1x, m1y, m1z) = rho1
    rho1a = (rho1 + m1z) * .5
    rho1b = (rho1 - m1z) * .5
    # * .5 because v+v.T is applied in the caller
    wva = (u_u * rho1a + u_d * rho1b) * weight * .5
    wvb = (u_d * rho1a + d_d * rho1b) * weight * .5
    return wva, wvb

def _dks_gga_wv1(rho0, rho1, vxc, fxc, weight):
    r0, (m0x, m0y, m0z) = rho0
    r1, (m1x, m1y, m1z) = rho1
    rho0 = ((r0 + m0z) * .5, (r0 - m0z) * .5)
    rho1 = ((r1 + m1z) * .5, (r1 - m1z) * .5)
    return numint._uks_gga_wv1(rho0, rho1, vxc, fxc, weight)

def _dks_mgga_wv1(rho0, rho1, vxc, fxc, weight):
    r0, (m0x, m0y, m0z) = rho0
    r1, (m1x, m1y, m1z) = rho1
    rho0 = ((r0 + m0z) * .5, (r0 - m0z) * .5)
    rho1 = ((r1 + m1z) * .5, (r1 - m1z) * .5)
    return numint._uks_mgga_wv1(rho0, rho1, vxc, fxc, weight)

def _col_lda_fxc_mat(mol, ao, weight, rho0, rho1, vxc, fxc,
                     non0tab, shls_slice, ao_loc, on_LL=True):
    '''Kernel matrix of collinear LDA'''
    if on_LL:
        wva, wvb = _dks_lda_wv1(rho0, rho1, vxc, fxc, weight)
    else:
        wvb, wva = _dks_lda_wv1(rho0, rho1, vxc, fxc, weight)
    aoa, aob = ao
    mat  = _dot_ao_ao(mol, _scale_ao(aoa, wva), aoa, non0tab, shls_slice, ao_loc)
    mat += _dot_ao_ao(mol, _scale_ao(aob, wvb), aob, non0tab, shls_slice, ao_loc)
    return mat

def _col_gga_fxc_mat(mol, ao, weight, rho0, rho1, vxc, fxc,
                     non0tab, shls_slice, ao_loc, on_LL=True):
    '''Kernel matrix of collinear GGA'''
    if on_LL:
        wva, wvb = _dks_gga_wv1(rho0, rho1, vxc, fxc, weight)
    else:
        wvb, wva = _dks_gga_wv1(rho0, rho1, vxc, fxc, weight)
    aoa, aob = ao
    aow = None
    aow = _scale_ao(aoa[:4], wva[:4], aow)
    mat = _dot_ao_ao(mol, aow, aoa[0], non0tab, shls_slice, ao_loc)
    aow = _scale_ao(aob[:4], wvb[:4], aow)
    mat += _dot_ao_ao(mol, aow, aob[0], non0tab, shls_slice, ao_loc)
    return mat

def _col_mgga_fxc_mat(mol, ao, weight, rho0, rho1, vxc, fxc,
                      non0tab, shls_slice, ao_loc, on_LL=True):
    '''Kernel matrix of collinear MGGA'''
    if on_LL:
        wva, wvb = _dks_mgga_wv1(rho0, rho1, vxc, fxc, weight)
    else:
        wvb, wva = _dks_mgga_wv1(rho0, rho1, vxc, fxc, weight)
    aoa, aob = ao
    aow = None
    aow = _scale_ao(aoa[:4], wva[:4], aow)
    mat = _dot_ao_ao(mol, aow, aoa[0], non0tab, shls_slice, ao_loc)
    aow = _scale_ao(aob[:4], wvb[:4], aow)
    mat += _dot_ao_ao(mol, aow, aob[0], non0tab, shls_slice, ao_loc)

    mat += _dot_ao_ao(mol, ao[1], _scale_ao(ao[1], wva[5], aow), non0tab, shls_slice, ao_loc)
    mat += _dot_ao_ao(mol, ao[2], _scale_ao(ao[2], wva[5], aow), non0tab, shls_slice, ao_loc)
    mat += _dot_ao_ao(mol, ao[3], _scale_ao(ao[3], wva[5], aow), non0tab, shls_slice, ao_loc)
    mat += _dot_ao_ao(mol, ao[1], _scale_ao(ao[1], wvb[5], aow), non0tab, shls_slice, ao_loc)
    mat += _dot_ao_ao(mol, ao[2], _scale_ao(ao[2], wvb[5], aow), non0tab, shls_slice, ao_loc)
    mat += _dot_ao_ao(mol, ao[3], _scale_ao(ao[3], wvb[5], aow), non0tab, shls_slice, ao_loc)
    return mat

def r_fxc(ni, mol, grids, xc_code, dm0, dms, relativity=1, hermi=0,
          rho0=None, vxc=None, fxc=None, max_memory=2000, verbose=None):
    '''Calculate 2-component or 4-component Vxc matrix in j-adapted basis
    '''
    xctype = ni._xc_type(xc_code)
    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_2c()
    n2c = ao_loc[-1]
    if not ni.collinear:
        raise NotImplementedError('non-collinear fxc')

    make_rho1, nset, nao = ni._gen_rho_evaluator(mol, dms, hermi)
    with_s = (nao == n2c*2)  # 4C DM

    if rho0 is None and (xctype != 'LDA' or fxc is None):
        make_rho0 = ni._gen_rho_evaluator(mol, dm0, 1)[0]
    else:
        make_rho0 = None

    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_2c()

    matLL = numpy.zeros((nset,n2c,n2c), dtype=dms.dtype)
    matSS = numpy.zeros_like(matLL)
    if xctype in ('LDA', 'GGA', 'MGGA'):
        f_eval_mat = {
            'LDA' : (_col_lda_fxc_mat  , 0),
            'GGA' : (_col_gga_fxc_mat  , 1),
            'MGGA': (_col_mgga_fxc_mat , 2),
        }
        fmat, ao_deriv = f_eval_mat[xctype]

        _rho0 = None
        p1 = 0
        for ao, mask, weight, coords \
                in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
            p0, p1 = p1, p1 + weight.size
            if rho0 is not None:
                if xctype == 'LDA':
                    _rho0 = [numpy.asarray(rho0[0][p0:p1], order='C'),
                             numpy.asarray(rho0[1][:,p0:p1], order='C')]
                else:
                    _rho0 = [numpy.asarray(rho0[0][:,p0:p1], order='C'),
                             numpy.asarray(rho0[1][:,:,p0:p1], order='C')]
            elif make_rho0 is not None:
                _rho0 = make_rho0(0, ao, mask, xctype)

            if vxc is None or fxc is None:
                _vxc, _fxc = ni.eval_xc(xc_code, _rho0, spin=1,
                                        relativity=relativity, deriv=2,
                                        verbose=verbose)[1:3]
            else:
                _vxc = [None if x is None else x[p0:p1] for x in vxc]
                _fxc = [None if x is None else x[p0:p1] for x in fxc]

            for i in range(nset):
                rho1 = make_rho1(i, ao, mask, xctype)
                matLL[i] += fmat(mol, ao[:2], weight, _rho0, rho1, _vxc, _fxc,
                                 mask, shls_slice, ao_loc)
                if with_s:
                    matSS[i] += fmat(mol, ao[2:], weight, _rho0, rho1, _vxc, _fxc,
                                     mask, shls_slice, ao_loc, False)

        # for (\nabla\mu) \nu + \mu (\nabla\nu)
        matLL = matLL + matLL.conj().transpose(0,2,1)
        if with_s:
            matSS = matSS + matSS.conj().transpose(0,2,1)

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

    if with_s:
        matSS *= (.5 / lib.param.LIGHT_SPEED)**2
        vmat = numpy.zeros((nset,nao,nao), dtype=numpy.complex128)
        vmat[:,:n2c,:n2c] = matLL
        vmat[:,n2c:,n2c:] = matSS
    else:
        vmat = matLL

    if isinstance(dms, numpy.ndarray) and dms.ndim == 2:
        vmat = vmat[0]
    return vmat

def cache_xc_kernel(ni, mol, grids, xc_code, mo_coeff, mo_occ, spin=1,
                    max_memory=2000):
    '''Compute the 0th order density, Vxc and fxc.  They can be used in TDDFT,
    DFT hessian module etc.
    '''
    xctype = ni._xc_type(xc_code)
    if xctype == 'MGGA':
        ao_deriv = 2
    elif xctype == 'GGA':
        ao_deriv = 1
    elif xctype == 'NLC':
        raise NotImplementedError('NLC')
    else:
        ao_deriv = 0

    dm = numpy.dot(mo_coeff * mo_occ, mo_coeff.conj().T)
    hermi = 1
    make_rho, nset, nao = ni._gen_rho_evaluator(mol, dm, hermi)
    rho = []
    m = []
    for ao, mask, weight, coords \
            in ni.block_loop(mol, grids, nao, ao_deriv, max_memory):
        rho_m = make_rho(0, ao, mask, xctype)
        rho.append(rho_m[0])
        m.append(rho_m[1])
    rho = (numpy.hstack(rho), numpy.hstack(m))
    vxc, fxc = ni.eval_xc(xc_code, rho, spin=spin, relativity=1, deriv=2,
                          verbose=0)[1:3]
    return rho, vxc, fxc

def get_rho(ni, mol, dm, grids, max_memory=2000):
    make_rho, nset, nao = ni._gen_rho_evaluator(mol, dm, hermi=1)
    n2c = mol.nao_2c()
    with_s = (nao == n2c*2)  # 4C DM
    rho = numpy.empty(grids.weights.size)
    p1 = 0
    for ao, mask, weight, coords \
            in ni.block_loop(mol, grids, nao, 0, with_s, max_memory):
        p0, p1 = p1, p1 + weight.size
        rho[p0:p1] = make_rho(0, ao, mask, 'LDA')[0]
    return rho


class RNumInt(numint._NumIntMixin):
    '''NumInt for j-adapted (spinor) basis'''
    def __init__(self):
        self.omega = None  # RSH paramter
        self.collinear = False

    get_rho = get_rho
    cache_xc_kernel = cache_xc_kernel
    get_vxc = r_vxc = r_vxc
    get_fxc = r_fxc = r_fxc

    def eval_ao(self, mol, coords, deriv=0, with_s=True, shls_slice=None,
                non0tab=None, out=None, verbose=None):
        return eval_ao(mol, coords, deriv, with_s, shls_slice, non0tab, out, verbose)

    def eval_rho2(self, mol, ao, mo_coeff, mo_occ, non0tab=None, xctype='LDA',
                  verbose=None):
        raise NotImplementedError

    @lib.with_doc(eval_rho.__doc__)
    def eval_rho(self, mol, ao, dm, non0tab=None, xctype='LDA', verbose=None):
        return eval_rho(mol, ao, dm, non0tab, xctype, verbose)

    def block_loop(self, mol, grids, nao, deriv=0, with_s=False, max_memory=2000,
                   non0tab=None, blksize=None, buf=None):
        '''Define this macro to loop over grids by blocks.
        '''
        if grids.coords is None:
            grids.build(with_non0tab=True)
        ngrids = grids.weights.size
        comp = (deriv+1)*(deriv+2)*(deriv+3)//6
# NOTE to index ni.non0tab, the blksize needs to be the integer multiplier of BLKSIZE
        if blksize is None:
            blksize = min(int(max_memory*1e6/((comp*4+4)*nao*16*BLKSIZE))*BLKSIZE, ngrids)
            blksize = max(blksize, BLKSIZE)
        if non0tab is None:
            non0tab = grids.non0tab
        if non0tab is None:
            non0tab = numpy.ones(((ngrids+BLKSIZE-1)//BLKSIZE,mol.nbas),
                                 dtype=numpy.uint8)

        if buf is None:
            buf = numpy.empty((4,comp,blksize,nao), dtype=numpy.complex128)
        for ip0 in range(0, ngrids, blksize):
            ip1 = min(ngrids, ip0+blksize)
            coords = grids.coords[ip0:ip1]
            weight = grids.weights[ip0:ip1]
            non0 = non0tab[ip0//BLKSIZE:]
            ao = self.eval_ao(mol, coords, deriv=deriv, with_s=with_s,
                              non0tab=non0, out=buf)
            yield ao, non0, weight, coords

    def _gen_rho_evaluator(self, mol, dms, hermi=1):
        dms = numpy.asarray(dms)
        nao = dms.shape[-1]
        if isinstance(dms, numpy.ndarray) and dms.ndim == 2:
            dms = dms.reshape(1,nao,nao)
        ndms = len(dms)
        n2c = mol.nao_2c()
        with_s = (nao == n2c*2)  # 4C DM
        if with_s:
            c1 = .5 / lib.param.LIGHT_SPEED
            dmLL = dms[:,:n2c,:n2c].copy('C')
            dmSS = dms[:,n2c:,n2c:] * c1**2

            def make_rho(idm, ao, non0tab, xctype):
                rho , m  = self.eval_rho(mol, ao[:2], dmLL[idm], non0tab, xctype)
                rhoS, mS = self.eval_rho(mol, ao[2:], dmSS[idm], non0tab, xctype)
                rho += rhoS
                # M = |\beta\Sigma|
                m[0] -= mS[0]
                m[1] -= mS[1]
                m[2] -= mS[2]
                return rho, m
        else:
            def make_rho(idm, ao, non0tab, xctype):
                return self.eval_rho(mol, ao, dms[idm], non0tab, xctype)
        return make_rho, ndms, nao

    def eval_xc(self, xc_code, rho, spin=1, relativity=0, deriv=1, omega=None,
                verbose=None):
        if omega is None: omega = self.omega
        if self.collinear:
            r, (mx, my, mz) = rho
            rhou = (r + mz) * .5
            rhod = (r - mz) * .5
            rho = (rhou, rhod)
            xc = self.libxc.eval_xc(xc_code, rho, 1, relativity, deriv,
                                    omega, verbose)
        else:
            # JTCC, 2, 257
            r, m = rho
            s = lib.norm(m, axis=0)
            rhou = (r + s) * .5
            rhod = (r - s) * .5
            rho = (rhou, rhod)
            xc = self.libxc.eval_xc(xc_code, rho, 1, relativity, deriv,
                                    omega, verbose)
            exc, vxc = xc[:2]
            # update vxc[0] inplace
            vrho = vxc[0]
            vr, vm = (vrho[:,0]+vrho[:,1])*.5, (vrho[:,0]-vrho[:,1])*.5
            vrho[:,0] = vr
            vrho[:,1] = vm
            # TODO:
            # vsigma_rr = (vxc[1][0] + vxc[1][1] + fxc[1][2]) / 4
            # vsigma_rs = (vxc[1][0] - vxc[1][2]) / 2
            # vsigma_ss = (vxc[1][0] - vxc[1][1] + fxc[1][2]) / 4
            # vtau_r = (vxc[3][0] + vxc[3][1]) / 2
            # vtau_s = (vxc[3][0] - vxc[3][1]) / 2
        return xc

_RNumInt = RNumInt


# pyscf/dft/radi.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''radii grids'''

import numpy
from pyscf.data import radii
from pyscf.data.elements import charge as elements_proton

BRAGG_RADII = radii.BRAGG
COVALENT_RADII = radii.COVALENT

# P.M.W. Gill, B.G. Johnson, J.A. Pople, Chem. Phys. Letters 209 (1993) 506-512
SG1RADII = numpy.array((
    0,
    1.0000,                                                 0.5882,
    3.0769, 2.0513, 1.5385, 1.2308, 1.0256, 0.8791, 0.7692, 0.6838,
    4.0909, 3.1579, 2.5714, 2.1687, 1.8750, 1.6514, 1.4754, 1.3333))


# Murray, N.C. Handy, G.J. Laming,  Mol. Phys. 78, 997(1993)
def murray(n, *args, **kwargs):
    raise RuntimeError('Not implemented')

# Gauss-Chebyshev of the first kind,  and the transformed interval [0,\infty)
def becke(n, charge, *args, **kwargs):
    '''Becke, JCP 88, 2547 (1988); DOI:10.1063/1.454033'''
    if charge == 1:
        rm = BRAGG_RADII[charge]
    else:
        rm = BRAGG_RADII[charge] * .5
    t, w = numpy.polynomial.chebyshev.chebgauss(n)
    r = (1+t)/(1-t) * rm
    w *= 2/(1-t)**2 * rm
    return r, w

# scale rad and rad_weight if necessary
# gauss-legendre
def delley(n, *args, **kwargs):
    '''B. Delley radial grids. Ref. JCP 104, 9848 (1996); DOI:10.1063/1.471749. log2 algorithm'''
    r = numpy.empty(n)
    dr = numpy.empty(n)
    r_outer = 12.
    step = 1. / (n+1)
    rfac = r_outer / numpy.log(1 - (n*step)**2)
    for i in range(1, n+1):
        xi = rfac * numpy.log(1-(i*step)**2)
        r[i-1] = xi
        dri = rfac * (-2.0*i*(step)**2) / ((1-(i*step)**2)) # d xi / dr
        dr[i-1] = dri
    return r, dr
gauss_legendre = delley

def mura_knowles(n, charge=None, *args, **kwargs):
    '''Mura-Knowles [JCP 104, 9848 (1996); DOI:10.1063/1.471749] log3 quadrature radial grids'''
    r = numpy.empty(n)
    dr = numpy.empty(n)
# 7 for Li, Be, Na, Mg, K, Ca, otherwise 5
    if charge in (3, 4, 11, 12, 19, 20):
        far = 7
    else:
        far = 5.2
    for i in range(n):
        x = (i+.5) / n
        r[i] = -far * numpy.log(1-x**3)
        dr[i] = far * 3*x*x/((1-x**3)*n)
    return r, dr

# Gauss-Chebyshev of the second kind,  and the transformed interval [0,\infty)
# Ref  Matthias Krack and Andreas M. Koster,  J. Chem. Phys. 108 (1998), 3226
def gauss_chebyshev(n, *args, **kwargs):
    '''Gauss-Chebyshev [JCP 108, 3226 (1998); DOI:10.1063/1.475719) radial grids'''
    ln2 = 1 / numpy.log(2)
    fac = 16./3 / (n+1)
    x1 = numpy.arange(1,n+1) * numpy.pi / (n+1)
    xi = ((n-1-numpy.arange(n)*2) / (n+1.) +
          (1+2./3*numpy.sin(x1)**2) * numpy.sin(2*x1) / numpy.pi)
    xi = (xi - xi[::-1])/2
    r = 1 - numpy.log(1+xi) * ln2
    dr = fac * numpy.sin(x1)**4 * ln2/(1+xi)
    return r, dr


def treutler_ahlrichs(n, *args, **kwargs):
    '''
    Treutler-Ahlrichs [JCP 102, 346 (1995); DOI:10.1063/1.469408] (M4) radial grids
    '''
    r = numpy.empty(n)
    dr = numpy.empty(n)
    step = numpy.pi / (n+1)
    ln2 = 1 / numpy.log(2)
    for i in range(n):
        x = numpy.cos((i+1)*step)
        r [i] = -ln2*(1+x)**.6 * numpy.log((1-x)/2)
        dr[i] = step * numpy.sin((i+1)*step) \
                * ln2*(1+x)**.6 *(-.6/(1+x)*numpy.log((1-x)/2)+1/(1-x))
    return r[::-1], dr[::-1]
treutler = treutler_ahlrichs




def becke_atomic_radii_adjust(mol, atomic_radii):
    '''Becke atomic radii adjust function'''
# Becke atomic size adjustment.  J. Chem. Phys. 88, 2547
# i > j
# fac(i,j) = \frac{1}{4} ( \frac{ra(j)}{ra(i)} - \frac{ra(i)}{ra(j)}
# fac(j,i) = -fac(i,j)
    charges = [elements_proton(x) for x in mol.elements]
    rad = atomic_radii[charges] + 1e-200
    rr = rad.reshape(-1,1) * (1./rad)
    a = .25 * (rr.T - rr)
    a[a<-.5] = -.5
    a[a>0.5] = 0.5
    #:return lambda i,j,g: g + a[i,j]*(1-g**2)
    def fadjust(i, j, g):
        g1 = g**2
        g1 -= 1.
        g1 *= -a[i,j]
        g1 += g
        return g1
    return fadjust

def treutler_atomic_radii_adjust(mol, atomic_radii):
    '''Treutler atomic radii adjust function: [JCP 102, 346 (1995); DOI:10.1063/1.469408]'''
# JCP 102, 346 (1995)
# i > j
# fac(i,j) = \frac{1}{4} ( \frac{ra(j)}{ra(i)} - \frac{ra(i)}{ra(j)}
# fac(j,i) = -fac(i,j)
    charges = [elements_proton(x) for x in mol.elements]
    rad = numpy.sqrt(atomic_radii[charges]) + 1e-200
    rr = rad.reshape(-1,1) * (1./rad)
    a = .25 * (rr.T - rr)
    a[a<-.5] = -.5
    a[a>0.5] = 0.5
    #:return lambda i,j,g: g + a[i,j]*(1-g**2)
    def fadjust(i, j, g):
        g1 = g**2
        g1 -= 1.
        g1 *= -a[i,j]
        g1 += g
        return g1
    return fadjust



# pyscf/dft/rks.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Non-relativistic restricted Kohn-Sham
'''


import textwrap
import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf import scf
from pyscf.scf import hf
from pyscf.scf import _vhf
from pyscf.scf import jk
from pyscf.dft import gen_grid
from pyscf.dft import numint
from pyscf import __config__


def get_veff(ks, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
    '''Coulomb + XC functional

    .. note::
        This function will modify the input ks object.

    Args:
        ks : an instance of :class:`RKS`
            XC functional are controlled by ks.xc attribute.  Attribute
            ks.grids might be initialized.
        dm : ndarray or list of ndarrays
            A density matrix or a list of density matrices

    Kwargs:
        dm_last : ndarray or a list of ndarrays or 0
            The density matrix baseline.  If not 0, this function computes the
            increment of HF potential w.r.t. the reference HF potential matrix.
        vhf_last : ndarray or a list of ndarrays or 0
            The reference Vxc potential matrix.
        hermi : int
            Whether J, K matrix is hermitian

            | 0 : no hermitian or symmetric
            | 1 : hermitian
            | 2 : anti-hermitian

    Returns:
        matrix Veff = J + Vxc.  Veff can be a list matrices, if the input
        dm is a list of density matrices.
    '''
    if mol is None: mol = ks.mol
    if dm is None: dm = ks.make_rdm1()
    ks.initialize_grids(mol, dm)

    t0 = (logger.process_clock(), logger.perf_counter())

    ground_state = (isinstance(dm, numpy.ndarray) and dm.ndim == 2)

    ni = ks._numint
    if hermi == 2:  # because rho = 0
        n, exc, vxc = 0, 0, 0
    else:
        max_memory = ks.max_memory - lib.current_memory()[0]
        n, exc, vxc = ni.nr_rks(mol, ks.grids, ks.xc, dm, max_memory=max_memory)
        if ks.nlc:
            assert 'VV10' in ks.nlc.upper()
            _, enlc, vnlc = ni.nr_rks(mol, ks.nlcgrids, ks.xc+'__'+ks.nlc, dm,
                                      max_memory=max_memory)
            exc += enlc
            vxc += vnlc
        logger.debug(ks, 'nelec by numeric integration = %s', n)
        t0 = logger.timer(ks, 'vxc', *t0)

    #enabling range-separated hybrids
    omega, alpha, hyb = ni.rsh_and_hybrid_coeff(ks.xc, spin=mol.spin)

    if abs(hyb) < 1e-10 and abs(alpha) < 1e-10:
        vk = None
        if (ks._eri is None and ks.direct_scf and
            getattr(vhf_last, 'vj', None) is not None):
            ddm = numpy.asarray(dm) - numpy.asarray(dm_last)
            vj = ks.get_j(mol, ddm, hermi)
            vj += vhf_last.vj
        else:
            vj = ks.get_j(mol, dm, hermi)
        vxc += vj
    else:
        if (ks._eri is None and ks.direct_scf and
            getattr(vhf_last, 'vk', None) is not None):
            ddm = numpy.asarray(dm) - numpy.asarray(dm_last)
            vj, vk = ks.get_jk(mol, ddm, hermi)
            vk *= hyb
            if abs(omega) > 1e-10:  # For range separated Coulomb operator
                vklr = ks.get_k(mol, ddm, hermi, omega=omega)
                vklr *= (alpha - hyb)
                vk += vklr
            vj += vhf_last.vj
            vk += vhf_last.vk
        else:
            vj, vk = ks.get_jk(mol, dm, hermi)
            vk *= hyb
            if abs(omega) > 1e-10:
                vklr = ks.get_k(mol, dm, hermi, omega=omega)
                vklr *= (alpha - hyb)
                vk += vklr
        vxc += vj - vk * .5

        if ground_state:
            exc -= numpy.einsum('ij,ji', dm, vk).real * .5 * .5

    if ground_state:
        ecoul = numpy.einsum('ij,ji', dm, vj).real * .5
    else:
        ecoul = None

    vxc = lib.tag_array(vxc, ecoul=ecoul, exc=exc, vj=vj, vk=vk)
    return vxc

def get_vsap(ks, mol=None):
    '''Superposition of atomic potentials

    S. Lehtola, Assessment of initial guesses for self-consistent
    field calculations. Superposition of Atomic Potentials: simple yet
    efficient, J. Chem. Theory Comput. 15, 1593 (2019). DOI:
    10.1021/acs.jctc.8b01089. arXiv:1810.11659.

    This function evaluates the effective charge of a neutral atom,
    given by exchange-only LDA on top of spherically symmetric
    unrestricted Hartree-Fock calculations as described in

    S. Lehtola, L. Visscher, E. Engel, Efficient implementation of the
    superposition of atomic potentials initial guess for electronic
    structure calculations in Gaussian basis sets, J. Chem. Phys., in
    press (2020).

    The potentials have been calculated for the ground-states of
    spherically symmetric atoms at the non-relativistic level of theory
    as described in

    S. Lehtola, "Fully numerical calculations on atoms with fractional
    occupations and range-separated exchange functionals", Phys. Rev. A
    101, 012516 (2020). DOI: 10.1103/PhysRevA.101.012516

    using accurate finite-element calculations as described in

    S. Lehtola, "Fully numerical Hartree-Fock and density functional
    calculations. I. Atoms", Int. J. Quantum Chem. e25945 (2019).
    DOI: 10.1002/qua.25945

    .. note::
        This function will modify the input ks object.

    Args:
        ks : an instance of :class:`RKS`
            XC functional are controlled by ks.xc attribute.  Attribute
            ks.grids might be initialized.

    Returns:
        matrix Vsap = Vnuc + J + Vxc.
    '''
    if mol is None: mol = ks.mol
    t0 = (logger.process_clock(), logger.perf_counter())

    if ks.grids.coords is None:
        ks.grids.build(with_non0tab=True)
        t0 = logger.timer(ks, 'setting up grids', *t0)

    ni = ks._numint
    max_memory = ks.max_memory - lib.current_memory()[0]
    vsap = ni.nr_sap(mol, ks.grids, max_memory=max_memory)
    return vsap

# The vhfopt of standard Coulomb operator can be used here as an approximate
# opt since long-range part Coulomb is always smaller than standard Coulomb.
# It's safe to prescreen LR integrals with the integral estimation from
# standard Coulomb.
def _get_k_lr(mol, dm, omega=0, hermi=0, vhfopt=None):
    import sys
    sys.stderr.write('This function is deprecated. '
                     'It is replaced by mol.get_k(mol, dm, omege=omega)')
    dm = numpy.asarray(dm)
# Note, ks object caches the ERIs for small systems. The cached eris are
# computed with regular Coulomb operator. ks.get_jk or ks.get_k do not evalute
# the K matrix with the range separated Coulomb operator.  Here jk.get_jk
# function computes the K matrix with the modified Coulomb operator.
    nao = dm.shape[-1]
    dms = dm.reshape(-1,nao,nao)
    with mol.with_range_coulomb(omega):
        # Compute the long range part of ERIs temporarily with omega. Restore
        # the original omega when the block ends
        if vhfopt is None:
            contents = lambda: None # just a place_holder
        else:
            contents = vhfopt._this.contents
        with lib.temporary_env(contents,
                               fprescreen=_vhf._fpointer('CVHFnrs8_vk_prescreen')):
            intor = mol._add_suffix('int2e')
            vklr = jk.get_jk(mol, dms, ['ijkl,jk->il']*len(dms), intor=intor,
                             vhfopt=vhfopt)
    return numpy.asarray(vklr).reshape(dm.shape)


def energy_elec(ks, dm=None, h1e=None, vhf=None):
    r'''Electronic part of RKS energy.

    Note this function has side effects which cause mf.scf_summary updated.

    Args:
        ks : an instance of DFT class

        dm : 2D ndarray
            one-partical density matrix
        h1e : 2D ndarray
            Core hamiltonian

    Returns:
        RKS electronic energy and the 2-electron contribution
    '''
    if dm is None: dm = ks.make_rdm1()
    if h1e is None: h1e = ks.get_hcore()
    if vhf is None or getattr(vhf, 'ecoul', None) is None:
        vhf = ks.get_veff(ks.mol, dm)
    e1 = numpy.einsum('ij,ji->', h1e, dm).real
    ecoul = vhf.ecoul.real
    exc = vhf.exc.real
    e2 = ecoul + exc
    ks.scf_summary['e1'] = e1
    ks.scf_summary['coul'] = ecoul
    ks.scf_summary['exc'] = exc
    logger.debug(ks, 'E1 = %s  Ecoul = %s  Exc = %s', e1, ecoul, exc)
    return e1+e2, e2


NELEC_ERROR_TOL = getattr(__config__, 'dft_rks_prune_error_tol', 0.02)
def prune_small_rho_grids_(ks, mol, dm, grids):
    rho = ks._numint.get_rho(mol, dm, grids, ks.max_memory)
    n = numpy.dot(rho, grids.weights)
    if abs(n-mol.nelectron) < NELEC_ERROR_TOL*n:
        rho *= grids.weights
        idx = abs(rho) > ks.small_rho_cutoff / grids.weights.size
        logger.debug(ks, 'Drop grids %d',
                     grids.weights.size - numpy.count_nonzero(idx))
        grids.coords  = numpy.asarray(grids.coords [idx], order='C')
        grids.weights = numpy.asarray(grids.weights[idx], order='C')
        grids.non0tab = grids.make_mask(mol, grids.coords)
    return grids

def define_xc_(ks, description, xctype='LDA', hyb=0, rsh=(0,0,0)):
    libxc = ks._numint.libxc
    ks._numint = libxc.define_xc_(ks._numint, description, xctype, hyb, rsh)
    return ks


def _dft_common_init_(mf, xc='LDA,VWN'):
    mf.xc = xc
    mf.nlc = ''
    mf.grids = gen_grid.Grids(mf.mol)
    mf.grids.level = getattr(__config__, 'dft_rks_RKS_grids_level',
                             mf.grids.level)
    mf.nlcgrids = gen_grid.Grids(mf.mol)
    mf.nlcgrids.level = getattr(__config__, 'dft_rks_RKS_nlcgrids_level',
                                mf.nlcgrids.level)
    # Use rho to filter grids
    mf.small_rho_cutoff = getattr(__config__, 'dft_rks_RKS_small_rho_cutoff', 1e-7)
##################################################
# don't modify the following attributes, they are not input options
    mf._numint = numint.NumInt()
    mf._keys = mf._keys.union(['xc', 'nlc', 'omega', 'grids', 'nlcgrids',
                               'small_rho_cutoff'])

class KohnShamDFT(object):
    '''
    Attributes for Kohn-Sham DFT:
        xc : str
            'X_name,C_name' for the XC functional.  Default is 'lda,vwn'
        nlc : str
            'NLC_name' for the NLC functional.  Default is '' (i.e., None)
        omega : float
            Omega of the range-separated Coulomb operator e^{-omega r_{12}^2} / r_{12}
        grids : Grids object
            grids.level (0 - 9)  big number for large mesh grids. Default is 3

            radii_adjust
                | radi.treutler_atomic_radii_adjust (default)
                | radi.becke_atomic_radii_adjust
                | None : to switch off atomic radii adjustment

            grids.atomic_radii
                | radi.BRAGG_RADII  (default)
                | radi.COVALENT_RADII
                | None : to switch off atomic radii adjustment

            grids.radi_method  scheme for radial grids
                | radi.treutler  (default)
                | radi.delley
                | radi.mura_knowles
                | radi.gauss_chebyshev

            grids.becke_scheme  weight partition function
                | gen_grid.original_becke  (default)
                | gen_grid.stratmann

            grids.prune  scheme to reduce number of grids
                | gen_grid.nwchem_prune  (default)
                | gen_grid.sg1_prune
                | gen_grid.treutler_prune
                | None : to switch off grids pruning

            grids.symmetry  True/False  to symmetrize mesh grids (TODO)

            grids.atom_grid  Set (radial, angular) grids for particular atoms.
            Eg, grids.atom_grid = {'H': (20,110)} will generate 20 radial
            grids and 110 angular grids for H atom.

        small_rho_cutoff : float
            Drop grids if their contribution to total electrons smaller than
            this cutoff value.  Default is 1e-7.

    Examples:

    >>> mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='ccpvdz', verbose=0)
    >>> mf = dft.RKS(mol)
    >>> mf.xc = 'b3lyp'
    >>> mf.kernel()
    -76.415443079840458
    '''

    __init__ = _dft_common_init_

    @property
    def omega(self):
        return self._numint.omega
    @omega.setter
    def omega(self, v):
        self._numint.omega = float(v)

    def dump_flags(self, verbose=None):
        log = logger.new_logger(self, verbose)
        log.info('XC library %s version %s\n    %s',
                 self._numint.libxc.__name__,
                 self._numint.libxc.__version__,
                 self._numint.libxc.__reference__)

        if log.verbose >= logger.INFO:
            log.info('XC functionals = %s', self.xc)
            if hasattr(self._numint.libxc, 'xc_reference'):
                log.info(textwrap.indent('\n'.join(self._numint.libxc.xc_reference(self.xc)), '    '))

        if self.nlc!='':
            log.info('NLC functional = %s', self.nlc)

        self.grids.dump_flags(verbose)
        if self.nlc!='':
            log.info('** Following is NLC Grids **')
            self.nlcgrids.dump_flags(verbose)

        log.info('small_rho_cutoff = %g', self.small_rho_cutoff)
        return self

    define_xc_ = define_xc_

    def to_rhf(self):
        '''Convert the input mean-field object to a RHF/ROHF object.

        Note this conversion only changes the class of the mean-field object.
        The total energy and wave-function are the same as them in the input
        mean-field object.
        '''
        mf = scf.RHF(self.mol)
        mf.__dict__.update(self.to_rks().__dict__)
        mf.converged = False
        return mf

    def to_uhf(self):
        '''Convert the input mean-field object to a UHF object.

        Note this conversion only changes the class of the mean-field object.
        The total energy and wave-function are the same as them in the input
        mean-field object.
        '''
        mf = scf.UHF(self.mol)
        mf.__dict__.update(self.to_uks().__dict__)
        mf.converged = False
        return mf

    def to_ghf(self):
        '''Convert the input mean-field object to a GHF object.

        Note this conversion only changes the class of the mean-field object.
        The total energy and wave-function are the same as them in the input
        mean-field object.
        '''
        mf = scf.GHF(self.mol)
        mf.__dict__.update(self.to_gks().__dict__)
        mf.converged = False
        return mf

    def to_rks(self, xc=None):
        '''Convert the input mean-field object to a RKS/ROKS object.

        Note this conversion only changes the class of the mean-field object.
        The total energy and wave-function are the same as them in the input
        mean-field object.
        '''
        mf = scf.addons.convert_to_rhf(self)
        if xc is not None:
            mf.xc = xc
        if xc != self.xc or not isinstance(self, RKS):
            mf.converged = False
        return mf

    def to_uks(self, xc=None):
        '''Convert the input mean-field object to a UKS object.

        Note this conversion only changes the class of the mean-field object.
        The total energy and wave-function are the same as them in the input
        mean-field object.
        '''
        mf = scf.addons.convert_to_uhf(self)
        if xc is not None:
            mf.xc = xc
        if xc != self.xc:
            mf.converged = False
        return mf

    def to_gks(self, xc=None):
        '''Convert the input mean-field object to a GKS object.

        Note this conversion only changes the class of the mean-field object.
        The total energy and wave-function are the same as them in the input
        mean-field object.
        '''
        mf = scf.addons.convert_to_ghf(self)
        if xc is not None:
            mf.xc = xc
        if xc != self.xc:
            mf.converged = False
        return mf

    def reset(self, mol=None):
        hf.SCF.reset(self, mol)
        self.grids.reset(mol)
        self.nlcgrids.reset(mol)
        return self

    def initialize_grids(self, mol=None, dm=None):
        '''Initialize self.grids the first time call get_veff'''
        if mol is None: mol = self.mol

        if self.grids.coords is None:
            t0 = (logger.process_clock(), logger.perf_counter())
            self.grids.build(with_non0tab=True)
            if (self.small_rho_cutoff > 1e-20 and
                # dm.ndim == 2 indicates ground state
                isinstance(dm, numpy.ndarray) and dm.ndim == 2):
                # Filter grids the first time setup grids
                self.grids = prune_small_rho_grids_(self, self.mol, dm,
                                                    self.grids)
            t0 = logger.timer(self, 'setting up grids', *t0)

        if self.nlc != '':
            if self.nlcgrids.coords is None:
                t0 = (logger.process_clock(), logger.perf_counter())
                self.nlcgrids.build(with_non0tab=True)
                if (self.small_rho_cutoff > 1e-20 and
                    # dm.ndim == 2 indicates ground state
                    isinstance(dm, numpy.ndarray) and dm.ndim == 2):
                    # Filter grids the first time setup grids
                    self.nlcgrids = prune_small_rho_grids_(self, self.mol, dm,
                                                           self.nlcgrids)
                t0 = logger.timer(self, 'setting up nlc grids', *t0)
        return self

# Update the KohnShamDFT label in scf.hf module
hf.KohnShamDFT = KohnShamDFT

def init_guess_by_vsap(mf, mol=None):
    '''Form SAP guess'''
    if mol is None: mol = mf.mol

    vsap = mf.get_vsap()
    t = mol.intor_symmetric('int1e_kin')
    s = mf.get_ovlp(mol)
    hsap = t + vsap

    # Form guess orbitals
    mo_energy, mo_coeff = mf.eig(hsap, s)
    logger.debug(mf, 'VSAP mo energies\n{}'.format(mo_energy))

    # and guess density
    mo_occ = mf.get_occ(mo_energy, mo_coeff)
    return mf.make_rdm1(mo_coeff, mo_occ)


class RKS(KohnShamDFT, hf.RHF):
    __doc__ = '''Restricted Kohn-Sham\n''' + hf.SCF.__doc__ + KohnShamDFT.__doc__

    def __init__(self, mol, xc='LDA,VWN'):
        hf.RHF.__init__(self, mol)
        KohnShamDFT.__init__(self, xc)

    def dump_flags(self, verbose=None):
        hf.RHF.dump_flags(self, verbose)
        return KohnShamDFT.dump_flags(self, verbose)

    get_veff = get_veff
    get_vsap = get_vsap
    energy_elec = energy_elec

    init_guess_by_vsap = init_guess_by_vsap

    def nuc_grad_method(self):
        from pyscf.grad import rks as rks_grad
        return rks_grad.Gradients(self)


# pyscf/dft/rks_symm.py

#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Non-relativistic Restricted Kohn-Sham
'''

from pyscf.scf import hf_symm
from pyscf.dft import rks
from pyscf.dft import uks


class SymAdaptedRKS(hf_symm.SymAdaptedRHF, rks.KohnShamDFT):
    ''' Restricted Kohn-Sham '''
    def __init__(self, mol, xc='LDA,VWN'):
        hf_symm.RHF.__init__(self, mol)
        rks.KohnShamDFT.__init__(self, xc)

    def dump_flags(self, verbose=None):
        hf_symm.RHF.dump_flags(self, verbose)
        rks.KohnShamDFT.dump_flags(self, verbose)
        return self

    get_veff = rks.get_veff
    get_vsap = rks.get_vsap
    energy_elec = rks.energy_elec

    init_guess_by_vsap = rks.init_guess_by_vsap

    reset = rks.KohnShamDFT.reset

    def nuc_grad_method(self):
        from pyscf.grad import rks
        return rks.Gradients(self)

RKS = SymAdaptedRKS


class SymAdaptedROKS(hf_symm.SymAdaptedROHF, rks.KohnShamDFT):
    ''' Restricted Kohn-Sham '''
    def __init__(self, mol, xc='LDA,VWN'):
        hf_symm.ROHF.__init__(self, mol)
        rks.KohnShamDFT.__init__(self, xc)

    def dump_flags(self, verbose=None):
        hf_symm.ROHF.dump_flags(self, verbose)
        rks.KohnShamDFT.dump_flags(self, verbose)
        return self

    get_veff = uks.get_veff
    get_vsap = rks.get_vsap
    energy_elec = uks.energy_elec

    init_guess_by_vsap = rks.init_guess_by_vsap

    reset = rks.KohnShamDFT.reset

    def nuc_grad_method(self):
        from pyscf.grad import roks
        return roks.Gradients(self)

ROKS = SymAdaptedROKS


if __name__ == '__main__':
    from pyscf import gto
    mol = gto.Mole()
    mol.verbose = 2
    mol.output = '/dev/null'

    mol.atom.extend([['He', (0.,0.,0.)], ])
    mol.basis = { 'He': 'cc-pvdz'}
    mol.symmetry = 1
    mol.build()

    m = RKS(mol)
    m.xc = 'b3lyp'
    print(m.scf())  # -2.89992555753



# pyscf/dft/roks.py

#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Non-relativistic restricted open-shell Kohn-Sham
'''

import numpy
from pyscf import lib
from pyscf.scf import rohf
from pyscf.dft.uks import energy_elec
from pyscf.dft import rks
from pyscf.dft import uks


@lib.with_doc(uks.get_veff.__doc__)
def get_veff(ks, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
    if getattr(dm, 'mo_coeff', None) is not None:
        mo_coeff = dm.mo_coeff
        mo_occ_a = (dm.mo_occ > 0).astype(numpy.double)
        mo_occ_b = (dm.mo_occ ==2).astype(numpy.double)
        dm = lib.tag_array(dm, mo_coeff=(mo_coeff,mo_coeff),
                           mo_occ=(mo_occ_a,mo_occ_b))
    return uks.get_veff(ks, mol, dm, dm_last, vhf_last, hermi)


class ROKS(rks.KohnShamDFT, rohf.ROHF):
    '''Restricted open-shell Kohn-Sham
    See pyscf/dft/rks.py RKS class for the usage of the attributes'''
    def __init__(self, mol, xc='LDA,VWN'):
        rohf.ROHF.__init__(self, mol)
        rks.KohnShamDFT.__init__(self, xc)

    def dump_flags(self, verbose=None):
        rohf.ROHF.dump_flags(self, verbose)
        rks.KohnShamDFT.dump_flags(self, verbose)
        return self

    get_veff = get_veff
    get_vsap = rks.get_vsap
    energy_elec = energy_elec

    init_guess_by_vsap = rks.init_guess_by_vsap

    def nuc_grad_method(self):
        from pyscf.grad import roks
        return roks.Gradients(self)


if __name__ == '__main__':
    from pyscf import gto
    from pyscf.dft import xcfun
    mol = gto.Mole()
    mol.verbose = 7
    mol.output = '/dev/null'#'out_rks'

    mol.atom.extend([['He', (0.,0.,0.)], ])
    mol.basis = { 'He': 'cc-pvdz'}
    #mol.grids = { 'He': (10, 14),}
    mol.build()

    m = ROKS(mol).run()
    m.xc = 'b88,lyp'
    print(m.scf())  # -2.8978518405

    m = ROKS(mol)
    m._numint.libxc = xcfun
    m.xc = 'b88,lyp'
    print(m.scf())  # -2.8978518405



# pyscf/dft/sap.py

# Copyright (c) 2020, Susi Lehtola
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
# * Redistributions in binary form must reproduce the above copyright
# notice, this list of conditions and the following disclaimer in the
# documentation and/or other materials provided with the distribution.
# * Neither the name of the <organization> nor the
# names of its contributors may be used to endorse or promote products
# derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
# USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
# OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.
#
# Routines for the implementation of the superposition of atomic
# potentials guess for electronic structure calculations, see
#
# S. Lehtola, "Assessment of Initial Guesses for Self-Consistent Field
# Calculations. Superposition of Atomic Potentials: Simple yet
# Efficient", J. Chem. Theory Comput. 15, 1593 (2019).
# DOI: 10.1021/acs.jctc.8b01089
#
# This function evaluates the effective charge of a neutral atom,
# given by exchange-only LDA on top of spherically symmetric
# unrestricted Hartree-Fock calculations as described in
#
# S. Lehtola, L. Visscher, E. Engel, Efficient implementation of the
# superposition of atomic potentials initial guess for electronic
# structure calculations in Gaussian basis sets, J. Chem. Phys., in
# press (2020).
#
# The potentials have been calculated for the ground-states of
# spherically symmetric atoms at the non-relativistic level of theory
# as described in
#
# S. Lehtola, "Fully numerical calculations on atoms with fractional
# occupations and range-separated exchange functionals", Phys. Rev. A
# 101, 012516 (2020). DOI: 10.1103/PhysRevA.101.012516
#
# using accurate finite-element calculations as described in
#
# S. Lehtola, "Fully numerical Hartree-Fock and density functional
# calculations. I. Atoms", Int. J. Quantum Chem. e25945 (2019).
# DOI: 10.1002/qua.25945

import numpy
from pyscf.dft.sap_data import sap_Zeff

def sap_effective_charge(Z, r):
    '''
    Calculates the effective charge for the superposition of atomic potentials.

    S. Lehtola, "Assessment of Initial Guesses for Self-Consistent Field
    Calculations. Superposition of Atomic Potentials: Simple yet
    Efficient", J. Chem. Theory Comput. 15, 1593 (2019).
    DOI: 10.1021/acs.jctc.8b01089

    This function evaluates the effective charge of a neutral atom,
    given by exchange-only LDA on top of spherically symmetric
    unrestricted Hartree-Fock calculations as described in

    S. Lehtola, L. Visscher, E. Engel, Efficient implementation of the
    superposition of atomic potentials initial guess for electronic
    structure calculations in Gaussian basis sets, J. Chem. Phys., in
    press (2020).

    The potentials have been calculated for the ground-states of
    spherically symmetric atoms at the non-relativistic level of theory
    as described in

    S. Lehtola, "Fully numerical calculations on atoms with fractional
    occupations and range-separated exchange functionals", Phys. Rev. A
    101, 012516 (2020). DOI: 10.1103/PhysRevA.101.012516

    using accurate finite-element calculations as described in

    S. Lehtola, "Fully numerical Hartree-Fock and density functional
    calculations. I. Atoms", Int. J. Quantum Chem. e25945 (2019).
    DOI: 10.1002/qua.25945

    Input:
       Z: atomic charge
       r: distance from nucleus
    Output:
       Z(r): screened charge
    '''

    if Z < 1:
        return 0.0
    if Z >= sap_Zeff.shape[1]:
        raise ValueError('Atoms beyond Og are not supported')
    if r < 0.0:
        raise ValueError('Distance cannot be negative')
    if r >= sap_Zeff.shape[0]:
        return 0.0

    # Linear interpolation
    return numpy.interp(r, sap_Zeff[0,:], sap_Zeff[Z,:])


# pyscf/dft/uks.py

#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Non-relativistic Unrestricted Kohn-Sham
'''


import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf.scf import uhf
from pyscf.dft import rks

def get_veff(ks, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
    '''Coulomb + XC functional for UKS.  See pyscf/dft/rks.py
    :func:`get_veff` fore more details.
    '''
    if mol is None: mol = ks.mol
    if dm is None: dm = ks.make_rdm1()
    if not isinstance(dm, numpy.ndarray):
        dm = numpy.asarray(dm)
    if dm.ndim == 2:  # RHF DM
        dm = numpy.asarray((dm*.5,dm*.5))
    ks.initialize_grids(mol, dm)

    t0 = (logger.process_clock(), logger.perf_counter())

    ground_state = (dm.ndim == 3 and dm.shape[0] == 2)

    ni = ks._numint
    if hermi == 2:  # because rho = 0
        n, exc, vxc = (0,0), 0, 0
    else:
        max_memory = ks.max_memory - lib.current_memory()[0]
        n, exc, vxc = ni.nr_uks(mol, ks.grids, ks.xc, dm, max_memory=max_memory)
        if ks.nlc:
            assert 'VV10' in ks.nlc.upper()
            _, enlc, vnlc = ni.nr_rks(mol, ks.nlcgrids, ks.xc+'__'+ks.nlc, dm[0]+dm[1],
                                      max_memory=max_memory)
            exc += enlc
            vxc += vnlc
        logger.debug(ks, 'nelec by numeric integration = %s', n)
        t0 = logger.timer(ks, 'vxc', *t0)

    #enabling range-separated hybrids
    omega, alpha, hyb = ni.rsh_and_hybrid_coeff(ks.xc, spin=mol.spin)

    if abs(hyb) < 1e-10 and abs(alpha) < 1e-10:
        vk = None
        if (ks._eri is None and ks.direct_scf and
            getattr(vhf_last, 'vj', None) is not None):
            ddm = numpy.asarray(dm) - numpy.asarray(dm_last)
            vj = ks.get_j(mol, ddm[0]+ddm[1], hermi)
            vj += vhf_last.vj
        else:
            vj = ks.get_j(mol, dm[0]+dm[1], hermi)
        vxc += vj
    else:
        if (ks._eri is None and ks.direct_scf and
            getattr(vhf_last, 'vk', None) is not None):
            ddm = numpy.asarray(dm) - numpy.asarray(dm_last)
            vj, vk = ks.get_jk(mol, ddm, hermi)
            vk *= hyb
            if abs(omega) > 1e-10:
                vklr = ks.get_k(mol, ddm, hermi, omega)
                vklr *= (alpha - hyb)
                vk += vklr
            vj = vj[0] + vj[1] + vhf_last.vj
            vk += vhf_last.vk
        else:
            vj, vk = ks.get_jk(mol, dm, hermi)
            vj = vj[0] + vj[1]
            vk *= hyb
            if abs(omega) > 1e-10:
                vklr = ks.get_k(mol, dm, hermi, omega)
                vklr *= (alpha - hyb)
                vk += vklr
        vxc += vj - vk

        if ground_state:
            exc -=(numpy.einsum('ij,ji', dm[0], vk[0]).real +
                   numpy.einsum('ij,ji', dm[1], vk[1]).real) * .5
    if ground_state:
        ecoul = numpy.einsum('ij,ji', dm[0]+dm[1], vj).real * .5
    else:
        ecoul = None

    vxc = lib.tag_array(vxc, ecoul=ecoul, exc=exc, vj=vj, vk=vk)
    return vxc

def get_vsap(ks, mol=None):
    '''Superposition of atomic potentials

    S. Lehtola, Assessment of initial guesses for self-consistent
    field calculations. Superposition of Atomic Potentials: simple yet
    efficient, J. Chem. Theory Comput. 15, 1593 (2019). DOI:
    10.1021/acs.jctc.8b01089. arXiv:1810.11659.

    This function evaluates the effective charge of a neutral atom,
    given by exchange-only LDA on top of spherically symmetric
    unrestricted Hartree-Fock calculations as described in

    S. Lehtola, L. Visscher, E. Engel, Efficient implementation of the
    superposition of atomic potentials initial guess for electronic
    structure calculations in Gaussian basis sets, J. Chem. Phys., in
    press (2020).

    The potentials have been calculated for the ground-states of
    spherically symmetric atoms at the non-relativistic level of theory
    as described in

    S. Lehtola, "Fully numerical calculations on atoms with fractional
    occupations and range-separated exchange functionals", Phys. Rev. A
    101, 012516 (2020). DOI: 10.1103/PhysRevA.101.012516

    using accurate finite-element calculations as described in

    S. Lehtola, "Fully numerical Hartree-Fock and density functional
    calculations. I. Atoms", Int. J. Quantum Chem. e25945 (2019).
    DOI: 10.1002/qua.25945

    .. note::
        This function will modify the input ks object.

    Args:
        ks : an instance of :class:`RKS`
            XC functional are controlled by ks.xc attribute.  Attribute
            ks.grids might be initialized.

    Returns:
        matrix Vsap = Vnuc + J + Vxc.
    '''
    Vsap = rks.get_vsap(ks, mol)
    return numpy.asarray([Vsap, Vsap])

def energy_elec(ks, dm=None, h1e=None, vhf=None):
    if dm is None: dm = ks.make_rdm1()
    if h1e is None: h1e = ks.get_hcore()
    if vhf is None or getattr(vhf, 'ecoul', None) is None:
        vhf = ks.get_veff(ks.mol, dm)
    if not (isinstance(dm, numpy.ndarray) and dm.ndim == 2):
        dm = dm[0] + dm[1]
    return rks.energy_elec(ks, dm, h1e, vhf)


class UKS(rks.KohnShamDFT, uhf.UHF):
    '''Unrestricted Kohn-Sham
    See pyscf/dft/rks.py RKS class for document of the attributes'''
    def __init__(self, mol, xc='LDA,VWN'):
        uhf.UHF.__init__(self, mol)
        rks.KohnShamDFT.__init__(self, xc)

    def dump_flags(self, verbose=None):
        uhf.UHF.dump_flags(self, verbose)
        rks.KohnShamDFT.dump_flags(self, verbose)
        return self

    def initialize_grids(self, mol=None, dm=None):
        ground_state = (isinstance(dm, numpy.ndarray)
                        and dm.ndim == 3 and dm.shape[0] == 2)
        if ground_state:
            super().initialize_grids(mol, dm[0]+dm[1])
        else:
            super().initialize_grids(mol)
        return self

    get_veff = get_veff
    get_vsap = get_vsap
    energy_elec = energy_elec

    init_guess_by_vsap = rks.init_guess_by_vsap

    def nuc_grad_method(self):
        from pyscf.grad import uks
        return uks.Gradients(self)


# pyscf/dft/uks_symm.py

#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Non-relativistic Unrestricted Kohn-Sham
'''

from pyscf.lib import logger
from pyscf.scf import uhf_symm
from pyscf.dft import uks
from pyscf.dft import rks


class SymAdaptedUKS(uhf_symm.UHF, rks.KohnShamDFT):
    ''' Restricted Kohn-Sham '''
    def __init__(self, mol, xc='LDA,VWN'):
        uhf_symm.UHF.__init__(self, mol)
        rks.KohnShamDFT.__init__(self, xc)

    def dump_flags(self, verbose=None):
        uhf_symm.UHF.dump_flags(self, verbose)
        rks.KohnShamDFT.dump_flags(self, verbose)
        return self

    get_veff = uks.get_veff
    get_vsap = uks.get_vsap
    energy_elec = uks.energy_elec

    init_guess_by_vsap = rks.init_guess_by_vsap

    reset = rks.KohnShamDFT.reset

    def nuc_grad_method(self):
        from pyscf.grad import uks
        return uks.Gradients(self)

UKS = SymAdaptedUKS


if __name__ == '__main__':
    from pyscf import gto
    mol = gto.Mole()
    mol.verbose = 7
    mol.output = '/dev/null'

    mol.atom.extend([['He', (0.,0.,0.)], ])
    mol.basis = { 'He': 'cc-pvdz'}
    mol.symmetry = 1
    mol.build()

    m = UKS(mol)
    m.xc = 'b3lyp'
    print(m.scf())  # -2.89992555753



# pyscf/dft/xcfun.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
XC functional, the interface to xcfun (https://github.com/dftlibs/xcfun)
U. Ekstrom et al, J. Chem. Theory Comput., 6, 1971
'''

import copy
import ctypes
import math
import numpy
from pyscf import lib
from pyscf.dft.xc.utils import remove_dup, format_xc_code

_itrf = lib.load_library('libxcfun_itrf')

_itrf.xcfun_splash.restype = ctypes.c_char_p
_itrf.xcfun_version.restype = ctypes.c_char_p

__version__ = _itrf.xcfun_version().decode("UTF-8")
__reference__ = _itrf.xcfun_splash().decode("UTF-8")

XC = XC_CODES = {
'SLATERX'       :  0,  #Slater LDA exchange
'PW86X'         :  1,  #PW86 exchange
'VWN3C'         :  2,  #VWN3 LDA Correlation functional
'VWN5C'         :  3,  #VWN5 LDA Correlation functional
'PBEC'          :  4,  #PBE correlation functional
'PBEX'          :  5,  #PBE Exchange Functional
'BECKEX'        :  6,  #Becke 88 exchange
'BECKECORRX'    :  7,  #Becke 88 exchange correction
'BECKESRX'      :  8,  #Short range Becke 88 exchange
'BECKECAMX'     :  9,  #CAM Becke 88 exchange
'BRX'           : 10,  #Becke-Roussells exchange with jp dependence
'BRC'           : 11,  #Becke-Roussells correlation with jp dependence
'BRXC'          : 12,  #Becke-Roussells correlation with jp dependence
'LDAERFX'       : 13,  #Short-range spin-dependent LDA exchange functional
'LDAERFC'       : 14,  #Short-range spin-dependent LDA correlation functional
'LDAERFC_JT'    : 15,  #Short-range spin-unpolarized LDA correlation functional
'LYPC'          : 16,  #LYP correlation
'OPTX'          : 17,  #OPTX Handy & Cohen exchange
'OPTXCORR'      : 18,  #OPTX Handy & Cohen exchange -- correction part only
'REVPBEX'       : 19,  #Revised PBE Exchange Functional
'RPBEX'         : 20,  #RPBE Exchange Functional
'SPBEC'         : 21,  #sPBE correlation functional
'VWN_PBEC'      : 22,  #PBE correlation functional using VWN LDA correlation.
'KTX'           : 23,  #KT exchange GGA correction
'TFK'           : 24,  #Thomas-Fermi Kinetic Energy Functional
'TW'            : 25,  #von Weizsacker Kinetic Energy Functional
'PW91X'         : 26,  #Perdew-Wang 1991 GGA Exchange Functional
'PW91K'         : 27,  #PW91 GGA Kinetic Energy Functional
'PW92C'         : 28,  #PW92 LDA correlation
'M05X'          : 29,  #M05 exchange
'M05X2X'        : 30,  #M05-2X exchange
'M06X'          : 31,  #M06 exchange
'M06X2X'        : 32,  #M06-2X exchange
'M06LX'         : 33,  #M06-L exchange
'M06HFX'        : 34,  #M06-HF exchange
'M05X2C'        : 35,  #M05-2X Correlation
'M05C'          : 36,  #M05 Correlation
'M06C'          : 37,  #M06 Correlation
'M06HFC'        : 38,  #M06-HF Correlation
'M06LC'         : 39,  #M06-L Correlation
'M06X2C'        : 40,  #M06-2X Correlation
'TPSSC'         : 41,  #TPSS original correlation functional
'TPSSX'         : 42,  #TPSS original exchange functional
'REVTPSSC'      : 43,  #Revised TPSS correlation functional
'REVTPSSX'      : 44,  #Reviewed TPSS exchange functional
'SCANC'         : 45,  #SCAN correlation functional
'SCANX'         : 46,  #SCAN exchange functional
'RSCANC'        : 47,  #rSCAN correlation functional
'RSCANX'        : 48,  #rSCAN exchange functional
'RPPSCANC'      : 49,  #r++SCAN correlation functional
'RPPSCANX'      : 50,  #r++SCAN exchange functional
'R2SCANC'       : 51,  #r2SCAN correlation functional
'R2SCANX'       : 52,  #r2SCAN exchange functional
'R4SCANC'       : 53,  #r4SCAN correlation functional
'R4SCANX'       : 54,  #r4SCAN exchange functional
'PZ81C'         : 55,  #PZ81 LDA correlation
'P86C'          : 56,  #P86C GGA correlation
'P86CORRC'      : 57,  #P86C GGA correlation
'BTK'           : 58,  #Borgoo-Tozer TS
'VWK'           : 59,  #von Weizsaecker kinetic energy
'B97X'          : 60,  #B97 exchange
'B97C'          : 61,  #B97 correlation
'B97_1X'        : 62,  #B97-1 exchange
'B97_1C'        : 63,  #B97-1 correlation
'B97_2X'        : 64,  #B97-2 exchange
'B97_2C'        : 65,  #B97-2 correlation
'CSC'           : 66,  #Colle-Salvetti correlation functional
'APBEC'         : 67,  #APBE correlation functional.
'APBEX'         : 68,  #APBE Exchange Functional
'ZVPBESOLC'     : 69,  #zvPBEsol correlation Functional
#'BLOCX'         : 70,  #BLOC exchange functional
'PBEINTC'       : 71,  #PBEint correlation Functional
'PBEINTX'       : 72,  #PBEint Exchange Functional
'PBELOCC'       : 73,  #PBEloc correlation functional.
'PBESOLX'       : 74,  #PBEsol Exchange Functional
'TPSSLOCC'      : 75,  #TPSSloc correlation functional
'ZVPBEINTC'     : 76,  #zvPBEint correlation Functional
'PW91C'         : 77,  #PW91 Correlation
#
# alias
#
'SLATER'        : 0,  # SLATERX
'LDA'           : 0,  # SLATERX
'VWN'           : 3,  # VWN5C
'VWN5'          : 3,  # VWN5C
'VWN3'          : 2,  # VWN3C
'SVWN'          : 'SLATERX + VWN5',
'B88'           : 6,  # BECKEX
'LYP'           : 16,
'P86'           : 56,
'M052XX'        : 30,  # M05-2X exchange
'M062XX'        : 32,  # M06-2X exchange
'M052XC'        : 35,  # M05-2X Correlation
'M062XC'        : 40,  # M06-2X Correlation
'BLYP'          : 'B88 + LYP',
'BP86'          : 'B88 + P86',  # Becke-Perdew 1986
'BPW91'         : 'B88 + PW91C',
'BPW92'         : 'B88 + PW92C',
'OLYP'          : '2.4832*SLATER - 1.43169*OPTX + LYP',  # CPL, 341, 319
'KT1X'           : 'SLATERX - 0.006*KTX',  # Keal-Tozer 1, JCP, 119, 3015
'KT2XC'         : '1.07173*SLATER - .006*KTX + 0.576727*VWN5',  # Keal-Tozer 2, JCP, 119, 3015
'KT3XC'         : 'SLATERX*1.092 + KTX*-0.004 + OPTXCORR*-0.925452 + LYPC*0.864409',  # Keal-Tozer 3, JCP, 121, 5654
# == '2.021452*SLATER - .004*KTX - .925452*OPTX + .864409*LYP',
'PBE0'          : '.25*HF + .75*PBEX + PBEC',  # Perdew-Burke-Ernzerhof, JCP, 110, 6158
'PBE1PBE'       : 'PBE0',
'PBEH'          : 'PBE0',
'B3P86'         : '.2*HF + .08*SLATER + .72*B88 + .81*P86C + .19*VWN5C',
'B3P86G'        : '.2*HF + .08*SLATER + .72*B88 + .81*P86C + .19*VWN3C',
'B3PW91'        : '.2*HF + .08*SLATER + .72*B88 + .81*PW91C + .19*VWN5C',
'B3PW91G'       : '.2*HF + .08*SLATER + .72*B88 + .81*PW91C + .19*VWN3C',
# Note, use VWN5 for B3LYP. It is different to the libxc default B3LYP
'B3LYP'         : 'B3LYP5',
'B3LYP5'        : '.2*HF + .08*SLATER + .72*B88 + .81*LYP + .19*VWN5C',
'B3LYPG'        : '.2*HF + .08*SLATER + .72*B88 + .81*LYP + .19*VWN3C', # B3LYP-VWN3 used by Gaussian and libxc
#'O3LYP'         : '.1161*HF + .9262*SLATER + .8133*OPTXCORR + .81*LYP + .19*VWN5C',  # Mol. Phys. 99 607
#'O3LYPG'        : '.1161*HF + .9262*SLATER + .8133*OPTXCORR + .81*LYP + .19*VWN3C',
# libxc implementation as below, see also discussion in https://gitlab.com/libxc/libxc/issues/47
#'O3LYP'         : '.1161*HF + .9262*SLATER + 1.164393477*OPTXCORR + .81*LYP + .19*VWN5C', #1.164393477 = .8133*1.43169
#'O3LYPG'        : '.1161*HF + .9262*SLATER + 1.164393477*OPTXCORR + .81*LYP + .19*VWN3C',
'O3LYP'         : '.1161*HF + 0.071006917*SLATER + .8133*OPTX, .81*LYP + .19*VWN5',  # libxc implementation
'X3LYP'         : '.218*HF + .073*SLATER + 0.542385*B88 + .166615*PW91X + .871*LYP + .129*VWN5C',  # Xu, PNAS, 101, 2673
'X3LYPG'        : '.218*HF + .073*SLATER + 0.542385*B88 + .166615*PW91X + .871*LYP + .129*VWN3C',
# Range-separated-hybrid functional: (alpha+beta)*SR_HF(0.33) + alpha*LR_HF(0.33)
# Note default mu of xcfun is 0.4. It can cause discrepancy for CAMB3LYP
'CAMB3LYP'      : '0.19*SR_HF(0.33) + 0.65*LR_HF(0.33) + 0.46*BECKESRX + 0.35*B88 + VWN5C*0.19 + LYPC*0.81',
'CAM_B3LYP'     : 'CAMB3LYP',
'LDAERF'        : 'LDAERFX + LDAERFC',  # Short-range exchange and correlation LDA functional
'B97XC'         : 'B97X + B97C + HF*0.1943',
'B97_1XC'       : 'B97_1X + B97_1C + HF*0.21',
'B97_2XC'       : 'B97_2X + B97_2C + HF*0.21',
'TPSSH'         : '0.1*HF + 0.9*TPSSX + TPSSC',
'TF'            : 'TFK',
}

# Some XC functionals have conventional name, like M06-L means M06-L for X
# functional and M06-L for C functional, PBE mean PBE-X plus PBE-C. If the
# conventional name was placed in the XC_CODES, it may lead to recursive
# reference when parsing the xc description.  These names (as exceptions of
# XC_CODES) are listed in XC_ALIAS below and they should be treated as a
# shortcut for XC functional.
XC_ALIAS = {
    # Conventional name : name in XC_CODES
    'BLYP'              : 'B88,LYP',
    'BP86'              : 'B88,P86',
    'PW91'              : 'PW91,PW91',
    'PBE'               : 'PBE,PBE',
    'REVPBE'            : 'REVPBE,PBE',
    'PBESOL'            : 'PBESOL,PBESOL',
#    'PKZB'              : 'PKZB,PKZB',
    'TPSS'              : 'TPSS,TPSS',
    'REVTPSS'           : 'REVTPSS,REVTPSS',
    'SCAN'              : 'SCAN,SCAN',
#    'SOGGA'             : 'SOGGA,PBE',
    #'BLOC'              : 'BLOC,TPSSLOC',
    'OLYP'              : 'OPTX,LYP',
    'RPBE'              : 'RPBE,PBE',
    'BPBE'              : 'B88,PBE',
#    'MPW91'             : 'MPW91,PW91',
    'HFLYP'             : 'HF,LYP',
#    'HFPW92'            : 'HF,PWMOD',
#    'SPW92'             : 'SLATER,PWMOD',
    'SVWN'              : 'SLATER,VWN',
#    'MS0'               : 'MS0,REGTPSS',
#    'MS1'               : 'MS1,REGTPSS',
#    'MS2'               : 'MS2,REGTPSS',
#    'MS2H'              : 'MS2H,REGTPSS',
#    'MVS'               : 'MVS,REGTPSS',
#    'MVSH'              : 'MVSH,REGTPSS',
#    'SOGGA11'           : 'SOGGA11,SOGGA11',
#    'SOGGA11-X'         : 'SOGGA11X,SOGGA11X',
    'KT1'               : 'KT1X,VWN',
#    'DLDF'              : 'DLDF,DLDF',
#    'GAM'               : 'GAM,GAM',
    'M06-L'             : 'M06L,M06L',
#    'M11-L'             : 'M11L,M11L',
#    'MN12-L'            : 'MN12L,MN12L',
#    'MN15-L'            : 'MN15L,MN15L',
#    'N12'               : 'N12,N12',
#    'N12-SX'            : 'N12SX,N12SX',
#    'MN12-SX'           : 'MN12SX,MN12SX',
#    'MN15'              : 'MN15,MN15',
#    'MBEEF'             : 'MBEEF,PBESOL',
#    'SCAN0'             : 'SCAN0,SCAN',
#    'PBEOP'             : 'PBE,OPPBE',
#    'BOP'               : 'B88,OPB88',
    'M05'               : '.28*HF + .72*M05X + M05C',
    'M06'               : '.27*HF +     M06X + M06C',
    #'M05_2X'            : '.56*HF + .44*M05X2X + M06C2X',
    #'M06_2X'            : '.54*HF +     M06X2X + M06C2X',
}
XC_ALIAS.update([(key.replace('-',''), XC_ALIAS[key])
                 for key in XC_ALIAS if '-' in key])

'''
LDA_IDS = set([0, 2, 3, 13, 14, 15, 24, 28, 55])
GGA_IDS = set([1, 4, 5, 6, 7, 8, 9, 16, 17, 18, 19, 20, 21, 22, 23, 25, 26,
               27, 44, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 67, 68, 69, 71,
               72, 73, 74, 76, 77])
MGGA_IDS =set([10, 11, 12, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41,
               42, 43, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 66, 70, 75])
'''
HYB_XC = set(('PBE0'    , 'PBE1PBE' , 'B3PW91'  , 'B3P86'   , 'B3LYP'   ,
              'B3PW91G' , 'B3P86G'  , 'B3LYPG'  , 'O3LYP'   , 'CAMB3LYP',
              'B97XC'   , 'B97_1XC' , 'B97_2XC' , 'M05XC'   , 'TPSSH'   ,
              'HFLYP'))
RSH_XC = set(('CAMB3LYP',))
MAX_DERIV_ORDER = 3

VV10_XC = {
    'B97M_V'    : [6.0, 0.01],
    'WB97M_V'   : [6.0, 0.01],
    'WB97X_V'   : [6.0, 0.01],
    'VV10'      : [5.9, 0.0093],
    'LC_VV10'   : [6.3, 0.0089],
    'REVSCAN_VV10': [9.8, 0.0093],
    'SCAN_RVV10'  : [15.7, 0.0093],
    'SCAN_VV10'   : [14.0, 0.0093],
    'SCANL_RVV10' : [15.7, 0.0093],
    'SCANL_VV10'  : [14.0, 0.0093],
}
VV10_XC.update([(key.replace('_', ''), val) for key, val in VV10_XC.items()])

def xc_type(xc_code):
    if xc_code is None:
        return None
    elif isinstance(xc_code, str):
        if is_nlc(xc_code):
            return 'NLC'
        hyb, fn_facs = parse_xc(xc_code)
    else:
        fn_facs = [(xc_code, 1)]  # mimic fn_facs
    if not fn_facs:
        return 'HF'
    elif all(_itrf.XCFUN_xc_type(ctypes.c_int(xid)) == 0 for xid, val in fn_facs):
        return 'LDA'
    elif any(_itrf.XCFUN_xc_type(ctypes.c_int(xid)) == 2 for xid, val in fn_facs):
        return 'MGGA'
    else:
        # all((xid in GGA_IDS or xid in LDA_IDS for xid, val in fn_fns)):
        # include hybrid_xc
        return 'GGA'

def is_lda(xc_code):
    return xc_type(xc_code) == 'LDA'

def is_hybrid_xc(xc_code):
    if isinstance(xc_code, str):
        xc_code = xc_code.replace(' ','').upper()
        return ('HF' in xc_code or xc_code in HYB_XC or
                hybrid_coeff(xc_code) != 0)
    elif isinstance(xc_code, int):
        return False
    else:
        return any((is_hybrid_xc(x) for x in xc_code))

def is_meta_gga(xc_code):
    return xc_type(xc_code) == 'MGGA'

def is_gga(xc_code):
    return xc_type(xc_code) == 'GGA'

def is_nlc(xc_code):
    return '__VV10' in xc_code.upper()

def nlc_coeff(xc_code):
    '''Get NLC coefficients
    '''
    xc_code = xc_code.upper()

    nlc_part = None
    if '__VV10' in xc_code:
        xc_code, nlc_part = xc_code.split('__', 1)

    if xc_code in VV10_XC:
        return VV10_XC[xc_code]
    elif nlc_part is not None:
        # Use VV10 NLC parameters by default for the general case
        return VV10_XC[nlc_part]
    else:
        raise NotImplementedError(
            '%s does not have NLC part. Available functionals are %s' %
            (xc_code, ', '.join(VV10_XC.keys())))

def rsh_coeff(xc_code):
    '''Get Range-separated-hybrid coefficients
    '''
    hyb, fn_facs = parse_xc(xc_code)
    hyb, alpha, omega = hyb
    beta = hyb - alpha
    return omega, alpha, beta

def max_deriv_order(xc_code):
    hyb, fn_facs = parse_xc(xc_code)
    return MAX_DERIV_ORDER

def test_deriv_order(xc_code, deriv, raise_error=False):
    support = deriv <= max_deriv_order(xc_code)
    if not support and raise_error:
        raise NotImplementedError('xcfun library does not support derivative '
                                  'order %d for  %s' % (deriv, xc_code))
    return support

def hybrid_coeff(xc_code, spin=0):
    if is_nlc(xc_code):
        return 0
    hyb, fn_facs = parse_xc(xc_code)
    return hyb[0]

def parse_xc_name(xc_name):
    fn_facs = parse_xc(xc_name)[1]
    return fn_facs[0][0], fn_facs[1][0]

def parse_xc(description):
    r'''Rules to input functional description:

    * The given functional description must be a one-line string.
    * The functional description is case-insensitive.
    * The functional description string has two parts, separated by ",".  The
      first part describes the exchange functional, the second is the correlation
      functional.


      - If "," not appeared in string, the entire string is treated as the
        name of a compound functional (containing both the exchange and
        the correlation functional) which was declared in the functional
        aliases list. The full list of functional aliases can be obtained by
        calling the function pyscf.dft.xcfun.XC_ALIAS.keys() .

      - To input only X functional (without C functional), leave the second
        part blank. E.g. description='slater,' means a functional with LDA
        contribution only.

      - To neglect the contribution of X functional (just apply C functional),
        leave blank in the first part, e.g. description=',vwn' means a
        functional with VWN only.

      - If compound XC functional is specified, no matter whether it is in the
        X part (the string in front of comma) or the C part (the string behind
        comma), both X and C functionals of the compound XC functional will be
        used.

    * The functional name can be placed in arbitrary order.  Two names need to
      be separated by operators "+" or "-".  Blank spaces are ignored.
      NOTE the parser only reads operators "+" "-" "*".  / is not supported.

    * A functional name can have at most one factor.  If the factor is not
      given, it is set to 1.  Compound functional can be scaled as a unit. For
      example '0.5*b3lyp' is equivalent to
      'HF*0.1 + .04*LDA + .36*B88, .405*LYP + .095*VWN'

    * String "HF" stands for exact exchange (HF K matrix).  "HF" can be put in
      the correlation functional part (after comma). Putting "HF" in the
      correlation part is the same to putting "HF" in the exchange part.

    * String "RSH" means range-separated operator. Its format is
      RSH(omega, alpha, beta).  Another way to input RSH is to use keywords
      SR_HF and LR_HF: "SR_HF(0.1) * alpha_plus_beta" and "LR_HF(0.1) *
      alpha" where the number in parenthesis is the value of omega.

    * Be careful with the convention of GGA functional, in which the LDA
      contribution has been included.
    '''
    hyb = [0, 0, 0]  # hybrid, alpha, omega
    if description is None:
        return hyb, []
    elif isinstance(description, int):
        return hyb, [(description, 1.)]
    elif not isinstance(description, str): #isinstance(description, (tuple,list)):
        return parse_xc('%s,%s' % tuple(description))

    def assign_omega(omega, hyb_or_sr, lr=0):
        if hyb[2] == omega or omega == 0:
            hyb[0] += hyb_or_sr
            hyb[1] += lr
        elif hyb[2] == 0:
            hyb[0] += hyb_or_sr
            hyb[1] += lr
            hyb[2] = omega
        else:
            raise ValueError('Different values of omega found for RSH functionals')
    fn_facs = []
    def parse_token(token, suffix, search_xc_alias=False):
        if token:
            if token[0] == '-':
                sign = -1
                token = token[1:]
            else:
                sign = 1
            if '*' in token:
                fac, key = token.split('*')
                if fac[0].isalpha():
                    fac, key = key, fac
                fac = sign * float(fac)
            else:
                fac, key = sign, token

            if key[:3] == 'RSH':
                # RSH(alpha; beta; omega): Range-separated-hybrid functional
                # See also utils.format_xc_code
                alpha, beta, omega = [float(x) for x in key[4:-1].split(';')]
                assign_omega(omega, fac*(alpha+beta), fac*alpha)
            elif key == 'HF':
                hyb[0] += fac
                hyb[1] += fac  # also add to LR_HF
            elif 'SR_HF' in key or 'SRHF' in key:
                if '(' in key:
                    omega = float(key.split('(')[1].split(')')[0])
                    assign_omega(omega, fac, 0)
                else:  # Assuming this omega the same to the existing omega
                    hyb[0] += fac
            elif 'LR_HF' in key:
                if '(' in key:
                    omega = float(key.split('(')[1].split(')')[0])
                    assign_omega(omega, 0, fac)
                else:
                    hyb[1] += fac  # == alpha
            elif key.isdigit():
                fn_facs.append((int(key), fac))
            else:
                if search_xc_alias and key in XC_ALIAS:
                    x_id = XC_ALIAS[key]
                elif key in XC_CODES:
                    x_id = XC_CODES[key]
                elif key+suffix in XC_CODES:
                    x_id = XC_CODES[key+suffix]
                else:
                    raise KeyError('Unknown %s functional  %s' % (suffix, key))
                if isinstance(x_id, str):
                    hyb1, fn_facs1 = parse_xc(x_id)
# Recursively scale the composed functional, to support e.g. '0.5*b3lyp'
                    if hyb1[0] != 0 or hyb1[1] != 0:
                        assign_omega(hyb1[2], hyb1[0]*fac, hyb1[1]*fac)
                    fn_facs.extend([(xid, c*fac) for xid, c in fn_facs1])
                elif x_id is None:
                    raise NotImplementedError('Unknown %s functional  %s' % (suffix, key))
                else:
                    fn_facs.append((x_id, fac))

    description = format_xc_code(description)

    if '-' in description:  # To handle e.g. M06-L
        for key in _NAME_WITH_DASH:
            if key in description:
                description = description.replace(key, _NAME_WITH_DASH[key])

    if ',' in description:
        x_code, c_code = description.split(',')
        for token in x_code.replace('-', '+-').replace(';+', ';').split('+'):
            parse_token(token, 'X')
        for token in c_code.replace('-', '+-').replace(';+', ';').split('+'):
            parse_token(token, 'C')
    else:
        for token in description.replace('-', '+-').replace(';+', ';').split('+'):
            parse_token(token, 'XC', search_xc_alias=True)
    if hyb[2] == 0: # No omega is assigned. LR_HF is 0 for normal Coulomb operator
        hyb[1] = 0
    return hyb, remove_dup(fn_facs)

_NAME_WITH_DASH = {'SR-HF'  : 'SR_HF',
                   'LR-HF'  : 'LR_HF',
                   'M06-L'  : 'M06L',
                   'M05-2X' : 'M052X',
                   'M06-HF' : 'M06HF',
                   'M06-2X' : 'M062X',}


def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, omega=None, verbose=None):
    r'''Interface to call xcfun library to evaluate XC functional, potential
    and functional derivatives.

    See also :func:`pyscf.dft.libxc.eval_xc`
    '''
    hyb, fn_facs = parse_xc(xc_code)
    if omega is not None:
        hyb[2] = float(omega)
    return _eval_xc(hyb, fn_facs, rho, spin, relativity, deriv, verbose)

XC_D0 = 0
XC_D1 = 1
XC_D2 = 2
XC_D3 = 3
XC_D4 = 4

XC_D00 = 0
XC_D10 = 1
XC_D01 = 2
XC_D20 = 3
XC_D11 = 4
XC_D02 = 5
XC_D30 = 6
XC_D21 = 7
XC_D12 = 8
XC_D03 = 9
XC_D40 = 10
XC_D31 = 11
XC_D22 = 12
XC_D13 = 13
XC_D04 = 14

XC_D000 = 0
XC_D100 = 1
XC_D010 = 2
XC_D001 = 3
XC_D200 = 4
XC_D110 = 5
XC_D101 = 6
XC_D020 = 7
XC_D011 = 8
XC_D002 = 9
XC_D300 = 10
XC_D210 = 11
XC_D201 = 12
XC_D120 = 13
XC_D111 = 14
XC_D102 = 15
XC_D030 = 16
XC_D021 = 17
XC_D012 = 18
XC_D003 = 19
XC_D400 = 20
XC_D310 = 21
XC_D301 = 22
XC_D220 = 23
XC_D211 = 24
XC_D202 = 25
XC_D130 = 26
XC_D121 = 27
XC_D112 = 28
XC_D103 = 29
XC_D040 = 30
XC_D031 = 31
XC_D022 = 32
XC_D013 = 33
XC_D004 = 34

XC_D00000 = 0
XC_D10000 = 1
XC_D01000 = 2
XC_D00100 = 3
XC_D00010 = 4
XC_D00001 = 5
XC_D20000 = 6
XC_D11000 = 7
XC_D10100 = 8
XC_D10010 = 9
XC_D10001 = 10
XC_D02000 = 11
XC_D01100 = 12
XC_D01010 = 13
XC_D01001 = 14
XC_D00200 = 15
XC_D00110 = 16
XC_D00101 = 17
XC_D00020 = 18
XC_D00011 = 19
XC_D00002 = 20
XC_D30000 = 21
XC_D21000 = 22
XC_D20100 = 23
XC_D20010 = 24
XC_D20001 = 25
XC_D12000 = 26
XC_D11100 = 27
XC_D11010 = 28
XC_D11001 = 29
XC_D10200 = 30
XC_D10110 = 31
XC_D10101 = 32
XC_D10020 = 33
XC_D10011 = 34
XC_D10002 = 35
XC_D03000 = 36
XC_D02100 = 37
XC_D02010 = 38
XC_D02001 = 39
XC_D01200 = 40
XC_D01110 = 41
XC_D01101 = 42
XC_D01020 = 43
XC_D01011 = 44
XC_D01002 = 45
XC_D00300 = 46
XC_D00210 = 47
XC_D00201 = 48
XC_D00120 = 49
XC_D00111 = 50
XC_D00102 = 51
XC_D00030 = 52
XC_D00021 = 53
XC_D00012 = 54
XC_D00003 = 55
XC_D40000 = 56
XC_D31000 = 57
XC_D30100 = 58
XC_D30010 = 59
XC_D30001 = 60
XC_D22000 = 61
XC_D21100 = 62
XC_D21010 = 63
XC_D21001 = 64
XC_D20200 = 65
XC_D20110 = 66
XC_D20101 = 67
XC_D20020 = 68
XC_D20011 = 69
XC_D20002 = 70
XC_D13000 = 71
XC_D12100 = 72
XC_D12010 = 73
XC_D12001 = 74
XC_D11200 = 75
XC_D11110 = 76
XC_D11101 = 77
XC_D11020 = 78
XC_D11011 = 79
XC_D11002 = 80
XC_D10300 = 81
XC_D10210 = 82
XC_D10201 = 83
XC_D10120 = 84
XC_D10111 = 85
XC_D10102 = 86
XC_D10030 = 87
XC_D10021 = 88
XC_D10012 = 89
XC_D10003 = 90
XC_D04000 = 91
XC_D03100 = 92
XC_D03010 = 93
XC_D03001 = 94
XC_D02200 = 95
XC_D02110 = 96
XC_D02101 = 97
XC_D02020 = 98
XC_D02011 = 99
XC_D02002 = 100
XC_D01300 = 101
XC_D01210 = 102
XC_D01201 = 103
XC_D01120 = 104
XC_D01111 = 105
XC_D01102 = 106
XC_D01030 = 107
XC_D01021 = 108
XC_D01012 = 109
XC_D01003 = 110
XC_D00400 = 111
XC_D00310 = 112
XC_D00301 = 113
XC_D00220 = 114
XC_D00211 = 115
XC_D00202 = 116
XC_D00130 = 117
XC_D00121 = 118
XC_D00112 = 119
XC_D00103 = 120
XC_D00040 = 121
XC_D00031 = 122
XC_D00022 = 123
XC_D00013 = 124
XC_D00004 = 125

XC_D0000000 = 0
XC_D1000000 = 1
XC_D0100000 = 2
XC_D0010000 = 3
XC_D0001000 = 4
XC_D0000100 = 5
XC_D0000010 = 6
XC_D0000001 = 7
XC_D2000000 = 8
XC_D1100000 = 9
XC_D1010000 = 10
XC_D1001000 = 11
XC_D1000100 = 12
XC_D1000010 = 13
XC_D1000001 = 14
XC_D0200000 = 15
XC_D0110000 = 16
XC_D0101000 = 17
XC_D0100100 = 18
XC_D0100010 = 19
XC_D0100001 = 20
XC_D0020000 = 21
XC_D0011000 = 22
XC_D0010100 = 23
XC_D0010010 = 24
XC_D0010001 = 25
XC_D0002000 = 26
XC_D0001100 = 27
XC_D0001010 = 28
XC_D0001001 = 29
XC_D0000200 = 30
XC_D0000110 = 31
XC_D0000101 = 32
XC_D0000020 = 33
XC_D0000011 = 34
XC_D0000002 = 35
XC_D3000000 = 36
XC_D2100000 = 37
XC_D2010000 = 38
XC_D2001000 = 39
XC_D2000100 = 40
XC_D2000010 = 41
XC_D2000001 = 42
XC_D1200000 = 43
XC_D1110000 = 44
XC_D1101000 = 45
XC_D1100100 = 46
XC_D1100010 = 47
XC_D1100001 = 48
XC_D1020000 = 49
XC_D1011000 = 50
XC_D1010100 = 51
XC_D1010010 = 52
XC_D1010001 = 53
XC_D1002000 = 54
XC_D1001100 = 55
XC_D1001010 = 56
XC_D1001001 = 57
XC_D1000200 = 58
XC_D1000110 = 59
XC_D1000101 = 60
XC_D1000020 = 61
XC_D1000011 = 62
XC_D1000002 = 63
XC_D0300000 = 64
XC_D0210000 = 65
XC_D0201000 = 66
XC_D0200100 = 67
XC_D0200010 = 68
XC_D0200001 = 69
XC_D0120000 = 70
XC_D0111000 = 71
XC_D0110100 = 72
XC_D0110010 = 73
XC_D0110001 = 74
XC_D0102000 = 75
XC_D0101100 = 76
XC_D0101010 = 77
XC_D0101001 = 78
XC_D0100200 = 79
XC_D0100110 = 80
XC_D0100101 = 81
XC_D0100020 = 82
XC_D0100011 = 83
XC_D0100002 = 84
XC_D0030000 = 85
XC_D0021000 = 86
XC_D0020100 = 87
XC_D0020010 = 88
XC_D0020001 = 89
XC_D0012000 = 90
XC_D0011100 = 91
XC_D0011010 = 92
XC_D0011001 = 93
XC_D0010200 = 94
XC_D0010110 = 95
XC_D0010101 = 96
XC_D0010020 = 97
XC_D0010011 = 98
XC_D0010002 = 99
XC_D0003000 = 100
XC_D0002100 = 101
XC_D0002010 = 102
XC_D0002001 = 103
XC_D0001200 = 104
XC_D0001110 = 105
XC_D0001101 = 106
XC_D0001020 = 107
XC_D0001011 = 108
XC_D0001002 = 109
XC_D0000300 = 110
XC_D0000210 = 111
XC_D0000201 = 112
XC_D0000120 = 113
XC_D0000111 = 114
XC_D0000102 = 115
XC_D0000030 = 116
XC_D0000021 = 117
XC_D0000012 = 118
XC_D0000003 = 119

def _eval_xc(hyb, fn_facs, rho, spin=0, relativity=0, deriv=1, verbose=None):
    assert(deriv < 4)
    if spin == 0:
        rho_u = rho_d = numpy.asarray(rho, order='C')
    else:
        rho_u = numpy.asarray(rho[0], order='C')
        rho_d = numpy.asarray(rho[1], order='C')
    assert(rho_u.dtype == numpy.double)
    assert(rho_d.dtype == numpy.double)

    if rho_u.ndim == 1:
        rho_u = rho_u.reshape(1,-1)
        rho_d = rho_d.reshape(1,-1)
    ngrids = rho_u.shape[1]

    fn_ids = [x[0] for x in fn_facs]
    facs   = [x[1] for x in fn_facs]
    if hyb[2] != 0:
        # Current implementation does not support different omegas for
        # different functionals
        omega = [hyb[2]] * len(facs)
    else:
        omega = [0] * len(facs)

    n = len(fn_ids)
    if (n == 0 or  # xc_code = '' or xc_code = 'HF', an empty functional
        all((is_lda(x) for x in fn_ids))):  # LDA
        if spin == 0:
            nvar = 1
            xctype = 'R-LDA'
        else:
            nvar = 2
            xctype = 'U-LDA'
    elif any((is_meta_gga(x) for x in fn_ids)):
        if spin == 0:
            nvar = 3
            xctype = 'R-MGGA'
        else:
            nvar = 7
            xctype = 'U-MGGA'
    else:  # GGA
        if spin == 0:
            nvar = 2
            xctype = 'R-GGA'
        else:
            nvar = 5
            xctype = 'U-GGA'
    outlen = (math.factorial(nvar+deriv) //
              (math.factorial(nvar) * math.factorial(deriv)))
    outbuf = numpy.zeros((ngrids,outlen))

    if n > 0:
        _itrf.XCFUN_eval_xc(ctypes.c_int(n),
                            (ctypes.c_int*n)(*fn_ids),
                            (ctypes.c_double*n)(*facs),
                            (ctypes.c_double*n)(*omega),
                            ctypes.c_int(spin),
                            ctypes.c_int(deriv), ctypes.c_int(rho_u.shape[1]),
                            rho_u.ctypes.data_as(ctypes.c_void_p),
                            rho_d.ctypes.data_as(ctypes.c_void_p),
                            outbuf.ctypes.data_as(ctypes.c_void_p))

    outbuf = lib.transpose(outbuf)
    exc = outbuf[0]
    vxc = fxc = kxc = None
    if xctype == 'R-LDA':
        if deriv > 0:
            vxc = [outbuf[1]]
        if deriv > 1:
            fxc = [outbuf[2]]
        if deriv > 2:
            kxc = [outbuf[3]]
    elif xctype == 'R-GGA':
        if deriv > 0:
            vxc = [outbuf[1], outbuf[2]]
        if deriv > 1:
            fxc = [outbuf[3], outbuf[4], outbuf[5]]
        if deriv > 2:
            kxc = [outbuf[6], outbuf[7], outbuf[8], outbuf[9]]
    elif xctype == 'U-LDA':
        if deriv > 0:
            vxc = [outbuf[1:3].T]
        if deriv > 1:
            fxc = [outbuf[3:6].T]
        if deriv > 2:
            kxc = [outbuf[6:10].T]
    elif xctype == 'U-GGA':
        if deriv > 0:
            vxc = [outbuf[1:3].T, outbuf[3:6].T]
        if deriv > 1:
            fxc = [outbuf[[XC_D20000,XC_D11000,XC_D02000]].T,
                   outbuf[[XC_D10100,XC_D10010,XC_D10001,
                           XC_D01100,XC_D01010,XC_D01001]].T,
                   outbuf[[XC_D00200,XC_D00110,XC_D00101,XC_D00020,XC_D00011,XC_D00002]].T]
        if deriv > 2:
            kxc = [outbuf[[XC_D30000,XC_D21000,XC_D12000,XC_D03000]].T,
                   outbuf[[XC_D20100,XC_D20010,XC_D20001,
                           XC_D11100,XC_D11010,XC_D11001,
                           XC_D02100,XC_D02010,XC_D02001]].T,
                   outbuf[[XC_D10200,XC_D10110,XC_D10101,XC_D10020,XC_D10011,XC_D10002,
                           XC_D01200,XC_D01110,XC_D01101,XC_D01020,XC_D01011,XC_D01002]].T,
                   outbuf[[XC_D00300,XC_D00210,XC_D00201,XC_D00120,XC_D00111,
                           XC_D00102,XC_D00030,XC_D00021,XC_D00012,XC_D00003]].T]
# MGGA/MLGGA: Note the MLGGA interface are not implemented. MGGA only needs 3
# input arguments.  To make the interface compatible with libxc, treat MGGA as
# MLGGA
    elif xctype == 'R-MGGA':
        if deriv > 0:
            vxc = [outbuf[1], outbuf[2], None, outbuf[3]]
        if deriv > 1:
            fxc = [
                # v2rho2, v2rhosigma, v2sigma2,
                outbuf[XC_D200], outbuf[XC_D110], outbuf[XC_D020],
                # v2lapl2, v2tau2,
                None, outbuf[XC_D002],
                # v2rholapl, v2rhotau,
                None, outbuf[XC_D101],
                # v2lapltau, v2sigmalapl, v2sigmatau,
                None, None, outbuf[XC_D011]]
        if deriv > 2:
            kxc = [
                # v3rho3, v3rho2sigma, v3rhosigma2, v3sigma3,
                outbuf[XC_D300], outbuf[XC_D210], outbuf[XC_D120], outbuf[XC_D030],
                # v3rho2lapl, v3rho2tau,
                None, outbuf[XC_D201],
                # v3rhosigmalapl, v3rhosigmatau,
                None, outbuf[XC_D111],
                # v3rholapl2, v3rholapltau, v3rhotau2,
                None, None, outbuf[XC_D102],
                # v3sigma2lapl, v3sigma2tau,
                None, outbuf[XC_D021],
                # v3sigmalapl2, v3sigmalapltau, v3sigmatau2,
                None, None, outbuf[XC_D012],
                # v3lapl3, v3lapl2tau, v3lapltau2, v3tau3)
                None, None, None, outbuf[XC_D003]]
    elif xctype == 'U-MGGA':
        if deriv > 0:
            vxc = (outbuf[1:3].T, outbuf[3:6].T, None, outbuf[6:8].T)
        if deriv > 1:
            fxc = [
                # v2rho2, v2rhosigma, v2sigma2,
                outbuf[[XC_D2000000,XC_D1100000,XC_D0200000]].T,
                outbuf[[XC_D1010000,XC_D1001000,XC_D1000100,
                        XC_D0110000,XC_D0101000,XC_D0100100]].T,
                outbuf[[XC_D0020000,XC_D0011000,XC_D0010100,
                        XC_D0002000,XC_D0001100,XC_D0000200]].T,
                # v2lapl2, v2tau2,
                None,
                outbuf[[XC_D0000020,XC_D0000011,XC_D0000002]].T,
                # v2rholapl, v2rhotau,
                None,
                outbuf[[XC_D1000010,XC_D1000001,XC_D0100010,XC_D0100001]].T,
                # v2lapltau, v2sigmalapl, v2sigmatau,
                None, None,
                outbuf[[XC_D0010010,XC_D0010001,XC_D0001010,XC_D0001001,
                        XC_D0000110,XC_D0000101]].T]
        if deriv > 2:
            kxc = [
                # v3rho3, v3rho2sigma, v3rhosigma2, v3sigma3,
                outbuf[[XC_D3000000,XC_D2100000,XC_D1200000,XC_D0300000]].T,
                outbuf[[XC_D2010000,XC_D2001000,XC_D2000100,
                        XC_D1110000,XC_D1101000,XC_D1100100,
                        XC_D0210000,XC_D0201000,XC_D0200100]].T,
                outbuf[[XC_D1020000,XC_D1011000,XC_D1010100,XC_D1002000,XC_D1001100,XC_D1000200,
                        XC_D0120000,XC_D0111000,XC_D0110100,XC_D0102000,XC_D0101100,XC_D0100200]].T,
                outbuf[[XC_D0030000,XC_D0021000,XC_D0020100,XC_D0012000,XC_D0011100,
                        XC_D0010200,XC_D0003000,XC_D0002100,XC_D0001200,XC_D0000300]].T,
                # v3rho2lapl, v3rho2tau,
                None,
                outbuf[[XC_D2000010,XC_D2000001,XC_D1100010,XC_D1100001,XC_D0200010,XC_D0200001]].T,
                # v3rhosigmalapl, v3rhosigmatau,
                None,
                outbuf[[XC_D1010010,XC_D1010001,XC_D1001010,XC_D1001001,XC_D1000110,XC_D1000101,
                        XC_D0110010,XC_D0110001,XC_D0101010,XC_D0101001,XC_D0100110,XC_D0100101]].T,
                # v3rholapl2, v3rholapltau, v3rhotau2,
                None, None,
                outbuf[[XC_D1000020,XC_D1000011,XC_D1000002,XC_D0100020,XC_D0100011,XC_D0100002]].T,
                # v3sigma2lapl, v3sigma2tau,
                None,
                outbuf[[XC_D0020010,XC_D0020001,XC_D0011010,XC_D0011001,XC_D0010110,XC_D0010101,
                        XC_D0002010,XC_D0002001,XC_D0001110,XC_D0001101,XC_D0000210,XC_D0000201]].T,
                # v3sigmalapl2, v3sigmalapltau, v3sigmatau2,
                None, None,
                outbuf[[XC_D0010020,XC_D0010011,XC_D0010002,
                        XC_D0001020,XC_D0001011,XC_D0001002,
                        XC_D0000120,XC_D0000111,XC_D0000102]].T,
                # v3lapl3, v3lapl2tau, v3lapltau2, v3tau3)
                None, None, None,
                outbuf[[XC_D0000030,XC_D0000021,XC_D0000012,XC_D0000003]].T]
    return exc, vxc, fxc, kxc


def define_xc_(ni, description, xctype='LDA', hyb=0, rsh=(0,0,0)):
    '''Define XC functional.  See also :func:`eval_xc` for the rules of input description.

    Args:
        ni : an instance of :class:`NumInt`

        description : str
            A string to describe the linear combination of different XC functionals.
            The X and C functional are separated by comma like '.8*LDA+.2*B86,VWN'.
            If "HF" was appeared in the string, it stands for the exact exchange.

    Kwargs:
        xctype : str
            'LDA' or 'GGA' or 'MGGA'
        hyb : float
            hybrid functional coefficient
        rsh : a list of three floats
            coefficients (omega, alpha, beta) for range-separated hybrid functional.
            omega is the exponent factor in attenuated Coulomb operator e^{-omega r_{12}}/r_{12}
            alpha is the coefficient for long-range part, hybrid coefficient
            can be obtained by alpha + beta

    Examples:

    >>> mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='ccpvdz')
    >>> mf = dft.RKS(mol)
    >>> define_xc_(mf._numint, '.2*HF + .08*LDA + .72*B88, .81*LYP + .19*VWN')
    >>> mf.kernel()
    -76.3783361189611
    >>> define_xc_(mf._numint, 'LDA*.08 + .72*B88 + .2*HF, .81*LYP + .19*VWN')
    >>> mf.kernel()
    -76.3783361189611
    >>> def eval_xc(xc_code, rho, *args, **kwargs):
    ...     exc = 0.01 * rho**2
    ...     vrho = 0.01 * 2 * rho
    ...     vxc = (vrho, None, None, None)
    ...     fxc = None  # 2nd order functional derivative
    ...     kxc = None  # 3rd order functional derivative
    ...     return exc, vxc, fxc, kxc
    >>> define_xc_(mf._numint, eval_xc, xctype='LDA')
    >>> mf.kernel()
    48.8525211046668
    '''
    if isinstance(description, str):
        ni.eval_xc = lambda xc_code, rho, *args, **kwargs: \
                eval_xc(description, rho, *args, **kwargs)
        ni.hybrid_coeff = lambda *args, **kwargs: hybrid_coeff(description)
        ni.rsh_coeff = lambda *args: rsh_coeff(description)
        ni._xc_type = lambda *args: xc_type(description)

    elif callable(description):
        ni.eval_xc = description
        ni.hybrid_coeff = lambda *args, **kwargs: hyb
        ni.rsh_coeff = lambda *args, **kwargs: rsh
        ni._xc_type = lambda *args: xctype
    else:
        raise ValueError('Unknown description %s' % description)
    return ni

def define_xc(ni, description, xctype='LDA', hyb=0, rsh=(0,0,0)):
    return define_xc_(copy.copy(ni), description, xctype, hyb, rsh)
define_xc.__doc__ = define_xc_.__doc__


# pyscf/dft/xc_deriv.py
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 2022 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

'''
Transform XC functional derivatives between different representations
'''

import ctypes
import numpy as np
from pyscf import lib

libdft = lib.load_library('libdft')


def transform_vxc(rho, vxc, xctype, spin=0):
    r'''
    Transform libxc functional derivatives to the derivative tensor of
    parameters in rho:
        [[density_a, (nabla_x)_a, (nabla_y)_a, (nabla_z)_a, tau_a],
         [density_b, (nabla_x)_b, (nabla_y)_b, (nabla_z)_b, tau_b]].
    The output tensor has the shape:
        * spin polarized
            LDA : [2,1,N]
            GGA : [2,4,N]
            MGGA: [2,5,N]
        * spin unpolarized
            LDA : [1,N]
            GGA : [4,N]
            MGGA: [5,N]
    '''
    rho = np.asarray(rho, order='C')
    if xctype == 'GGA':
        order = 1
        nvar = 4
        fr = vxc[0].T
        fg = vxc[1].T
    elif xctype == 'MGGA':
        order = 2
        nvar = 5
        fr = vxc[0].T
        fg = vxc[1].T
        ft = vxc[3].T
    else:  # LDA
        order = 0
        nvar = 1
        fr = vxc[0].T

    ngrids = rho.shape[-1]
    if spin == 1:
        if order == 0:
            vp = fr.reshape(2, nvar, ngrids)
        else:
            vp = np.empty((2,nvar, ngrids))
            vp[:,0] = fr
            #:vp[:,1:4] = np.einsum('abg,bxg->axg', _stack_fg(fg), rho[:,1:4])
            vp[:,1:4] = _stack_fg(fg, rho=rho)
        if order > 1:
            vp[:,4] = ft
    else:
        if order == 0:
            vp = fr.reshape(nvar, ngrids)
        else:
            vp = np.empty((nvar, ngrids))
            vp[0] = fr
            vp[1:4] = 2 * fg * rho[1:4]
        if order > 1:
            vp[4] = ft
    return vp

def transform_fxc(rho, vxc, fxc, xctype, spin=0):
    r'''
    Transform libxc functional derivatives to the derivative tensor of
    parameters in rho:
        [[density_a, (nabla_x)_a, (nabla_y)_a, (nabla_z)_a, tau_a],
         [density_b, (nabla_x)_b, (nabla_y)_b, (nabla_z)_b, tau_b]].
    The output tensor has the shape:
        * spin polarized
            LDA : [2,1,2,1,N]
            GGA : [2,4,2,4,N]
            MGGA: [2,5,2,5,N]
        * spin unpolarized
            LDA : [1,1,N]
            GGA : [4,4,N]
            MGGA: [5,5,N]
    '''
    rho = np.asarray(rho, order='C')
    if xctype == 'GGA':
        order = 1
        nvar = 4
        fg = vxc[1].T
        frr = fxc[0].T
        frg = fxc[1].T
        fgg = fxc[2].T
    elif xctype == 'MGGA':
        order = 2
        nvar = 5
        fg = vxc[1].T
        frr, frg, fgg, ftt, frt, fgt = [fxc[i].T for i in [0, 1, 2, 4, 6, 9]]
    else:  # LDA
        order = 0
        nvar = 1
        frr = fxc[0].T

    ngrids = rho.shape[-1]
    if spin == 1:
        if order == 0:
            vp = _stack_frr(frr).reshape(2,nvar, 2,nvar, ngrids).transpose(1,3,0,2,4)
        else:
            vp = np.empty((2,nvar, 2,nvar, ngrids)).transpose(1,3,0,2,4)
            vp[0,0] = _stack_frr(frr)
            i3 = np.arange(3)
            #:qgg = _stack_fgg(fgg)
            #:qgg = np.einsum('abcdg,axg->xbcdg', qgg, rho[:,1:4])
            #:qgg = np.einsum('xbcdg,cyg->xybdg', qgg, rho[:,1:4])
            qgg = _stack_fgg(fgg, rho=rho).transpose(1,3,0,2,4)
            qgg[i3,i3] += _stack_fg(fg)
            vp[1:4,1:4] = qgg

            frg = frg.reshape(2,3,ngrids)
            #:qrg = _stack_fg(frg, axis=1)
            #:qrg = np.einsum('rabg,axg->xrbg', qrg, rho[:,1:4])
            qrg = _stack_fg(frg, axis=1, rho=rho).transpose(2,0,1,3)
            vp[0,1:4] = qrg
            vp[1:4,0] = qrg.transpose(0,2,1,3)

        if order > 1:
            fgt = fgt.reshape(3,2,ngrids)
            #:qgt = _stack_fg(fgt, axis=0)
            #:qgt = np.einsum('abrg,axg->xbrg', qgt, rho[:,1:4])
            qgt = _stack_fg(fgt, axis=0, rho=rho).transpose(1,0,2,3)
            vp[1:4,4] = qgt
            vp[4,1:4] = qgt.transpose(0,2,1,3)

            qrt = frt.reshape(2,2,ngrids)
            vp[0,4] = qrt
            vp[4,0] = qrt.transpose(1,0,2)

            vp[4,4] = _stack_frr(ftt)

        vp = vp.transpose(2,0,3,1,4)

    else:
        if order == 0:
            vp = frr.reshape(nvar, nvar, ngrids)
        else:
            vp = np.empty((nvar, nvar, ngrids))
            vp[0,0] = frr
            i3 = np.arange(3)
            qgg = 4 * fgg * rho[1:4] * rho[1:4,None]
            qgg[i3,i3] += fg * 2
            vp[1:4,1:4] = qgg
            vp[0,1:4] = vp[1:4,0] = 2 * frg * rho[1:4]
        if order > 1:
            vp[4,1:4] = vp[1:4,4] = 2 * fgt * rho[1:4]
            vp[0,4] = frt
            vp[4,0] = frt
            vp[4,4] = ftt
    return vp

def transform_kxc(rho, fxc, kxc, xctype, spin=0):
    r'''
    Transform libxc functional derivatives to the derivative tensor of
    parameters in rho:
        [[density_a, (nabla_x)_a, (nabla_y)_a, (nabla_z)_a, tau_a],
         [density_b, (nabla_x)_b, (nabla_y)_b, (nabla_z)_b, tau_b]].
    The output tensor has the shape:
        * spin polarized
            LDA : [2,1,2,1,2,1,N]
            GGA : [2,4,2,4,2,4,N]
            MGGA: [2,5,2,5,2,5,N]
        * spin unpolarized
            LDA : [1,1,1,N]
            GGA : [4,4,4,N]
            MGGA: [5,5,5,N]
    '''
    rho = np.asarray(rho, order='C')
    if xctype == 'GGA':
        order = 1
        nvar = 4
        frg = fxc[1].T
        fgg = fxc[2].T
        frrr, frrg, frgg, fggg = [x.T for x in kxc[:4]]
    elif xctype == 'MGGA':
        order = 2
        nvar = 5
        frg = fxc[1].T
        fgg = fxc[2].T
        fgt = fxc[9].T
        frrr, frrg, frgg, fggg, frrt, frgt, frtt, fggt, fgtt, fttt = \
                [kxc[i].T for i in [0, 1, 2, 3, 5, 7, 10, 12, 15, 19]]
    else:  # LDA
        order = 0
        nvar = 1
        frrr = kxc[0].T

    ngrids = rho.shape[-1]
    if spin == 1:
        if order == 0:
            vp = _stack_frrr(frrr).reshape(2,nvar, 2,nvar, 2,nvar, ngrids).transpose(1,3,5,0,2,4,6)
        else:
            vp = np.empty((2,nvar, 2,nvar, 2,nvar, ngrids)).transpose(1,3,5,0,2,4,6)
            vp[0,0,0] = _stack_frrr(frrr)
            i3 = np.arange(3)
            #:qggg = _stack_fggg(fggg)
            #:qggg = np.einsum('abcdefg,axg->xbcdefg', qggg, rho[:,1:4])
            #:qggg = np.einsum('xbcdefg,cyg->xybdefg', qggg, rho[:,1:4])
            #:qggg = np.einsum('xybdefg,ezg->xyzbdfg', qggg, rho[:,1:4])
            qggg = _stack_fggg(fggg, rho=rho).transpose(1,3,5,0,2,4,6)
            qgg = _stack_fgg(fgg)
            qgg = np.einsum('abcdg,axg->xbcdg', qgg, rho[:,1:4])
            for i in range(3):
                qggg[:,i,i] += qgg
                qggg[i,:,i] += qgg.transpose(0,2,1,3,4)
                qggg[i,i,:] += qgg.transpose(0,2,3,1,4)
            vp[1:4,1:4,1:4] = qggg

            frgg = frgg.reshape(2,6,ngrids)
            #:qrgg = _stack_fgg(frgg, axis=1)
            #:qrgg = np.einsum('rabcdg,axg->xrbcdg', qrgg, rho[:,1:4])
            #:qrgg = np.einsum('xrbcdg,cyg->xyrbdg', qrgg, rho[:,1:4])
            qrgg = _stack_fgg(frgg, axis=1, rho=rho).transpose(2,4,0,1,3,5)
            qrg = _stack_fg(frg.reshape(2,3,ngrids), axis=1)
            qrgg[i3,i3] += qrg
            vp[0,1:4,1:4] = qrgg
            vp[1:4,0,1:4] = qrgg.transpose(0,1,3,2,4,5)
            vp[1:4,1:4,0] = qrgg.transpose(0,1,3,4,2,5)

            frrg = frrg.reshape(3,3,ngrids)
            qrrg = _stack_frr(frrg, axis=0)
            #:qrrg = _stack_fg(qrrg, axis=2)
            #:qrrg = np.einsum('rsabg,axg->rsxbg', qrrg, rho[:,1:4])
            qrrg = _stack_fg(qrrg, axis=2, rho=rho).transpose(3,0,1,2,4)
            vp[0,0,1:4] = qrrg
            vp[0,1:4,0] = qrrg.transpose(0,1,3,2,4)
            vp[1:4,0,0] = qrrg.transpose(0,3,1,2,4)

        if order > 1:
            fggt = fggt.reshape(6,2,ngrids)
            #:qggt = _stack_fgg(fggt, axis=0)
            #:qggt = np.einsum('abcdrg,axg->xbcdrg', qggt, rho[:,1:4])
            #:qggt = np.einsum('xbcdrg,cyg->xybdrg', qggt, rho[:,1:4])
            qggt = _stack_fgg(fggt, axis=0, rho=rho).transpose(1,3,0,2,4,5)
            qgt = _stack_fg(fgt.reshape(3,2,ngrids), axis=0)
            i3 = np.arange(3)
            qggt[i3,i3] += qgt
            vp[1:4,1:4,4] = qggt
            vp[1:4,4,1:4] = qggt.transpose(0,1,2,4,3,5)
            vp[4,1:4,1:4] = qggt.transpose(0,1,4,2,3,5)

            qgtt = _stack_frr(fgtt.reshape(3,3,ngrids), axis=1)
            #:qgtt = _stack_fg(qgtt, axis=0)
            #:qgtt = np.einsum('abrsg,axg->xbrsg', qgtt, rho[:,1:4])
            qgtt = _stack_fg(qgtt, axis=0, rho=rho).transpose(1,0,2,3,4)
            vp[1:4,4,4] = qgtt
            vp[4,1:4,4] = qgtt.transpose(0,2,1,3,4)
            vp[4,4,1:4] = qgtt.transpose(0,2,3,1,4)

            frgt = frgt.reshape(2,3,2,ngrids)
            #:qrgt = _stack_fg(frgt, axis=1)
            #:qrgt = np.einsum('rabsg,axg->xrbsg', qrgt, rho[:,1:4])
            qrgt = _stack_fg(frgt, axis=1, rho=rho).transpose(2,0,1,3,4)
            vp[0,1:4,4] = qrgt
            vp[0,4,1:4] = qrgt.transpose(0,1,3,2,4)
            vp[1:4,0,4] = qrgt.transpose(0,2,1,3,4)
            vp[4,0,1:4] = qrgt.transpose(0,3,1,2,4)
            vp[1:4,4,0] = qrgt.transpose(0,2,3,1,4)
            vp[4,1:4,0] = qrgt.transpose(0,3,2,1,4)

            qrrt = _stack_frr(frrt.reshape(3,2,ngrids), axis=0)
            vp[0,0,4] = qrrt
            vp[0,4,0] = qrrt.transpose(0,2,1,3)
            vp[4,0,0] = qrrt.transpose(2,0,1,3)

            qrtt = _stack_frr(frtt.reshape(2,3,ngrids), axis=1)
            vp[0,4,4] = qrtt
            vp[4,0,4] = qrtt.transpose(1,0,2,3)
            vp[4,4,0] = qrtt.transpose(1,2,0,3)

            vp[4,4,4] = _stack_frrr(fttt, axis=0)

        vp = vp.transpose(3,0,4,1,5,2,6)

    else:
        if order == 0:
            vp = frrr.reshape(nvar, nvar, nvar, ngrids)
        else:
            vp = np.empty((nvar, nvar, nvar, ngrids))
            vp[0,0,0] = frrr
            i3 = np.arange(3)
            qggg = 8 * fggg * rho[1:4] * rho[1:4,None] * rho[1:4,None,None]
            qgg = 4 * fgg * rho[1:4]
            for i in range(3):
                qggg[i,i,:] += qgg
                qggg[i,:,i] += qgg
                qggg[:,i,i] += qgg
            vp[1:4,1:4,1:4] = qggg

            qrgg = 4 * frgg * rho[1:4] * rho[1:4,None]
            qrgg[i3,i3] += frg * 2
            vp[0,1:4,1:4] = qrgg
            vp[1:4,0,1:4] = qrgg
            vp[1:4,1:4,0] = qrgg

            qrrg = 2 * frrg * rho[1:4]
            vp[0,0,1:4] = qrrg
            vp[0,1:4,0] = qrrg
            vp[1:4,0,0] = qrrg

        if order > 1:
            qggt = 4 * fggt * rho[1:4] * rho[1:4,None]
            qggt[i3,i3] += fgt * 2
            vp[1:4,1:4,4] = qggt
            vp[1:4,4,1:4] = qggt
            vp[4,1:4,1:4] = qggt

            qgtt = 2 * fgtt * rho[1:4]
            vp[1:4,4,4] = qgtt
            vp[4,1:4,4] = qgtt
            vp[4,4,1:4] = qgtt

            qrgt = 2 * frgt * rho[1:4]
            vp[0,1:4,4] = qrgt
            vp[0,4,1:4] = qrgt
            vp[1:4,0,4] = qrgt
            vp[4,0,1:4] = qrgt
            vp[1:4,4,0] = qrgt
            vp[4,1:4,0] = qrgt

            vp[0,0,4] = frrt
            vp[0,4,0] = frrt
            vp[4,0,0] = frrt
            vp[0,4,4] = frtt
            vp[4,0,4] = frtt
            vp[4,4,0] = frtt
            vp[4,4,4] = fttt
    return vp

def transform_lxc(rho, fxc, kxc, lxc, xctype, spin=0):
    r'''
    Transform libxc vxc functional output to the derivative tensor of parameters
    in rho: [density, nabla_x, nabla_y, nabla_z, tau]. The output tensor has the
    shape:
        * spin polarized
            LDA : [2,1,2,1,2,1,2,1,N]
            GGA : [2,4,2,4,2,4,2,4,N]
            MGGA: [2,5,2,5,2,5,2,5,N]
        * spin unpolarized
            LDA : [1,1,1,1,N]
            GGA : [4,4,4,4,N]
            MGGA: [5,5,5,5,N]
    '''
    raise NotImplementedError

def _stack_fg(fg, axis=0, rho=None, out=None):
    '''fg [uu, ud, dd] -> [[uu*2, ud], [du, dd*2]]'''
    if rho is None:
        qg = _stack_frr(fg, axis)
        if axis == 0:
            qg[0,0] *= 2
            qg[1,1] *= 2
        elif axis == 1:
            qg[:,0,0] *= 2
            qg[:,1,1] *= 2
        elif axis == 2:
            qg[:,:,0,0] *= 2
            qg[:,:,1,1] *= 2
        else:
            raise NotImplementedError
        return qg
    else:
        rho = np.asarray(rho, order='C')
        fg = np.asarray(fg, order='C')
        if fg.shape[axis] != 3:
            fg = fg.reshape(fg.shape[:axis] + (3, -1) + fg.shape[axis+1:])
        nvar, ngrids = rho.shape[1:]
        ncounts = int(np.prod(fg.shape[:axis]))
        mcounts = int(np.prod(fg.shape[axis+1:-1]))
        qg = np.empty(fg.shape[:axis] + (2, 3) + fg.shape[axis+1:])
        rho = libdft.VXCfg_to_direct_deriv(
            qg.ctypes.data_as(ctypes.c_void_p),
            fg.ctypes.data_as(ctypes.c_void_p),
            rho.ctypes.data_as(ctypes.c_void_p),
            ctypes.c_int(ncounts), ctypes.c_int(nvar), ctypes.c_int(mcounts),
            ctypes.c_int(ngrids))
        return qg


def _stack_frr(frr, axis=0):
    '''frr [u_u, u_d, d_d] -> [[u_u, u_d], [d_u, d_d]]'''
    if frr.shape[axis] != 3:
        frr = frr.reshape(frr.shape[:axis] + (3, -1) + frr.shape[axis+1:])
    slices = [slice(None)] * frr.ndim
    slices[axis] = [[0,1],[1,2]]
    return frr[tuple(slices)]

def _stack_fgg(fgg, axis=0, rho=None):
    '''
    fgg [uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd] ->
    [[uu_uu, ud_ud, ud_dd],
     [ud_uu, ud_ud, ud_dd],
     [dd_uu, dd_ud, dd_dd]] -> tensor with shape [2,2, 2,2, ...]
    '''
    if fgg.shape[axis] != 6:
        fgg = fgg.reshape(fgg.shape[:axis] + (6, -1) + fgg.shape[axis+1:])
    slices = [slice(None)] * fgg.ndim
    slices[axis] = [[0, 1, 2], [1, 3, 4], [2, 4, 5]]
    fgg = fgg[tuple(slices)]
    return _stack_fg(_stack_fg(fgg, axis=axis+1, rho=rho), axis=axis, rho=rho)

def _stack_frrr(frrr, axis=0):
    '''
    frrr [u_u_u, u_u_d, u_d_d, d_d_d]
    -> tensor with shape [2, 2, 2, ...]
    '''
    if frrr.shape[axis] != 4:
        frrr = frrr.reshape(frrr.shape[:axis] + (4, -1) + frrr.shape[axis+1:])
    slices = [slice(None)] * frrr.ndim
    slices[axis] = [[[0, 1], [1, 2]],
                    [[1, 2], [2, 3]]]
    return frrr[tuple(slices)]

def _stack_fggg(fggg, axis=0, rho=None):
    '''
    fggg [uu_uu_uu, uu_uu_ud, uu_uu_dd, uu_ud_ud, uu_ud_dd, uu_dd_dd, ud_ud_ud, ud_ud_dd, ud_dd_dd, dd_dd_dd]
    -> tensor with shape [2,2, 2,2, 2,2, ...]
    '''
    if fggg.shape[axis] != 10:
        fggg = fggg.reshape(fggg.shape[:axis] + (10, 2) + fggg.shape[axis+1:])
    slices = [slice(None)] * fggg.ndim
    slices[axis] = [[[0, 1, 2], [1, 3, 4], [2, 4, 5]],
                    [[1, 3, 4], [3, 6, 7], [4, 7, 8]],
                    [[2, 4, 5], [4, 7, 8], [5, 8, 9]]]
    fggg = fggg[tuple(slices)]
    return _stack_fg(_stack_fg(_stack_fg(
        fggg, axis=axis+2, rho=rho), axis=axis+1, rho=rho), axis=axis, rho=rho)

def compress(vp, spin=0):
    if spin != 0:  # spin polarized
        shape = vp.shape
        comp_shape = [shape[i*2]*shape[i*2+1] for i in range(0, vp.ndim-1, 2)]
        vp = vp.reshape(comp_shape + [shape[-1]])

    order = vp.ndim - 1
    if order < 2:
        pass
    elif order == 2:  # 2nd derivatives
        vp = vp[np.tril_indices(vp.shape[0])]
    elif order == 3:  # 2nd derivatives
        nd = vp.shape[0]
        vp = np.vstack([vp[i][np.tril_indices(i+1)] for i in range(nd)])
    else:
        raise NotImplementedError('High order derivatives')
    return vp

def decompress(vp, spin=0):
    order = vp.ndim - 1
    ngrids = vp.shape[-1]
    if order < 2:
        out = vp
    elif order == 2:  # 2nd derivatives
        nd = vp.shape[0]
        out = np.empty((nd, nd, ngrids))
        idx, idy = np.tril_indices(nd)
        out[idx,idy] = vp
        out[idy,idx] = vp
    elif order == 3:  # 2nd derivatives
        nd = vp.shape[0]
        out = np.empty((nd, nd, nd, ngrids))
        p1 = 0
        for i in range(nd):
            idx, idy = np.tril_indices(i+1)
            p0, p1 = p1, p1+idx.size
            out[i,idx,idy] = vp[p0:p1]
            out[i,idy,idx] = vp[p0:p1]
            out[idx,i,idy] = vp[p0:p1]
            out[idy,i,idx] = vp[p0:p1]
            out[idx,idy,i] = vp[p0:p1]
            out[idy,idx,i] = vp[p0:p1]
    else:
        raise NotImplementedError('High order derivatives')

    if spin != 0:  # spin polarized
        shape = out.shape
        nvar = shape[0] // 2
        decomp_shape = [(2, nvar)] * order + [ngrids]
        # reshape to (2,n,2,n,...,ngrids)
        out = out.reshape(np.hstack(decomp_shape))
    return out

def ud2ts(v_ud):
    '''XC derivatives spin-up/spin-down representations to
    total-density/spin-density representations'''
    v_ud = np.asarray(v_ud, order='C')
    nvar, ngrids = v_ud.shape[-2:]
    #:c = np.array([[.5,  .5],    # vrho = (va + vb) / 2
    #:              [.5, -.5]])   # vs   = (va - vb) / 2
    if v_ud.ndim == 3:  # vxc
        #:v_ts = np.einsum('ra,axg->rxg', c, v_ud)
        drv = libdft.VXCud2ts_deriv1
    elif v_ud.ndim == 5:  # fxc
        #:v_ts = np.einsum('ra,axbyg->rxbyg', c, v_ud)
        #:v_ts = np.einsum('sb,rxbyg->rxsyg', c, v_ts)
        drv = libdft.VXCud2ts_deriv2
    elif v_ud.ndim == 7:  # kxc
        #:v_ts = np.einsum('ra,axbyczg->rxbyczg', c, v_ud)
        #:v_ts = np.einsum('sb,rxbyczg->rxsyczg', c, v_ts)
        #:v_ts = np.einsum('tc,rxsyczg->rxsytzg', c, v_ts)
        drv = libdft.VXCud2ts_deriv3
    else:
        raise NotImplementedError(f'Shape {v_ud.shape} not supported')
    v_ts = np.empty_like(v_ud)
    drv(v_ts.ctypes.data_as(ctypes.c_void_p),
        v_ud.ctypes.data_as(ctypes.c_void_p),
        ctypes.c_int(nvar), ctypes.c_int(ngrids))
    return v_ts

def ts2ud(v_ts):
    '''XC derivatives total-density/spin-density representations to
    spin-up/spin-down representations'''
    c = np.array([[1,  1],    # va = vrho + vs
                  [1, -1]])   # vb = vrho - vs
    if v_ts.ndim == 3:  # vxc
        v_ud = np.einsum('ra,axg->rxg', c, v_ts)
    elif v_ts.ndim == 5:  # fxc
        v_ud = np.einsum('ra,axbyg->rxbyg', c, v_ts)
        v_ud = np.einsum('sb,rxbyg->rxsyg', c, v_ud)
    elif v_ts.ndim == 7:  # kxc
        v_ud = np.einsum('ra,axbyczg->rxbyczg', c, v_ts)
        v_ud = np.einsum('sb,rxbyczg->rxsyczg', c, v_ud)
        v_ud = np.einsum('tc,rxsyczg->rxsytzg', c, v_ud)
    else:
        raise NotImplementedError(f'Shape {v_ts.shape} not supported')
    return v_ud
