      veff = vj[1] * 2 + f1vo[0] * 2
        else:
            veff = f1vo[0] * 2
        veff0mop = reduce(numpy.dot, (mo_coeff.T, veff, mo_coeff))
        wvo -= numpy.einsum('ki,ai->ak', veff0mop[:nocc,:nocc], xpy) * 2
        wvo += numpy.einsum('ac,ai->ci', veff0mop[nocc:,nocc:], xpy) * 2
        veff0mom = numpy.zeros((nmo,nmo))

    with lib.temporary_env(mf.with_solvent, equilibrium_solvation=True):
        # set singlet=None, generate function for CPHF type response kernel
        vresp = mf.gen_response(singlet=None, hermi=1)
        def fvind(x):
            dm = reduce(numpy.dot, (orbv, x.reshape(nvir,nocc)*2, orbo.T))
            v1ao = vresp(dm+dm.T)
            return reduce(numpy.dot, (orbv.T, v1ao, orbo)).ravel()
        z1 = cphf.solve(fvind, mo_energy, mo_occ, wvo,
                        max_cycle=td_grad.cphf_max_cycle,
                        tol=td_grad.cphf_conv_tol)[0]
        z1 = z1.reshape(nvir,nocc)
        time1 = log.timer('Z-vector using CPHF solver', *time0)

        z1ao  = reduce(numpy.dot, (orbv, z1, orbo.T))
        veff = vresp(z1ao+z1ao.T)

    im0 = numpy.zeros((nmo,nmo))
    im0[:nocc,:nocc] = reduce(numpy.dot, (orbo.T, veff0doo+veff, orbo))
    im0[:nocc,:nocc]+= numpy.einsum('ak,ai->ki', veff0mop[nocc:,:nocc], xpy)
    im0[:nocc,:nocc]+= numpy.einsum('ak,ai->ki', veff0mom[nocc:,:nocc], xmy)
    im0[nocc:,nocc:] = numpy.einsum('ci,ai->ac', veff0mop[nocc:,:nocc], xpy)
    im0[nocc:,nocc:]+= numpy.einsum('ci,ai->ac', veff0mom[nocc:,:nocc], xmy)
    im0[nocc:,:nocc] = numpy.einsum('ki,ai->ak', veff0mop[:nocc,:nocc], xpy)*2
    im0[nocc:,:nocc]+= numpy.einsum('ki,ai->ak', veff0mom[:nocc,:nocc], xmy)*2

    zeta = lib.direct_sum('i+j->ij', mo_energy, mo_energy) * .5
    zeta[nocc:,:nocc] = mo_energy[:nocc]
    zeta[:nocc,nocc:] = mo_energy[nocc:]
    dm1 = numpy.zeros((nmo,nmo))
    dm1[:nocc,:nocc] = doo
    dm1[nocc:,nocc:] = dvv
    dm1[nocc:,:nocc] = z1
    dm1[:nocc,:nocc] += numpy.eye(nocc)*2 # for ground state
    im0 = reduce(numpy.dot, (mo_coeff, im0+zeta*dm1, mo_coeff.T))

    # Initialize hcore_deriv with the underlying SCF object because some
    # extensions (e.g. QM/MM, solvent) modifies the SCF object only.
    mf_grad = td_grad.base._scf.nuc_grad_method()
    hcore_deriv = mf_grad.hcore_generator(mol)
    s1 = mf_grad.get_ovlp(mol)

    dmz1doo = z1ao + dmzoo
    oo0 = reduce(numpy.dot, (orbo, orbo.T))
    if abs(hyb) > 1e-10:
        dm = (oo0, dmz1doo+dmz1doo.T, dmxpy+dmxpy.T, dmxmy-dmxmy.T)
        vj, vk = td_grad.get_jk(mol, dm)
        vk *= hyb
        if abs(omega) > 1e-10:
            with mol.with_range_coulomb(omega):
                vk += td_grad.get_k(mol, dm) * (alpha-hyb)
        vj = vj.reshape(-1,3,nao,nao)
        vk = vk.reshape(-1,3,nao,nao)
        if singlet:
            veff1 = vj * 2 - vk
        else:
            veff1 = numpy.vstack((vj[:2]*2-vk[:2], -vk[2:]))
    else:
        vj = td_grad.get_j(mol, (oo0, dmz1doo+dmz1doo.T, dmxpy+dmxpy.T))
        vj = vj.reshape(-1,3,nao,nao)
        veff1 = numpy.zeros((4,3,nao,nao))
        if singlet:
            veff1[:3] = vj * 2
        else:
            veff1[:2] = vj[:2] * 2

    fxcz1 = tdrks_grad._contract_xc_kernel(td_grad, mf.xc, z1ao, None,
                                           False, False, True, max_memory)[0]

    veff1[0] += vxc1[1:]
    veff1[1] +=(f1oo[1:] + fxcz1[1:] + k1ao[1:]*2)*2 # *2 for dmz1doo+dmz1oo.T
    veff1[2] += f1vo[1:] * 2
    time1 = log.timer('2e AO integral derivatives', *time1)

    if atmlst is None:
        atmlst = range(mol.natm)
    offsetdic = mol.offset_nr_by_atom()
    de = numpy.zeros((len(atmlst),3))
    for k, ia in enumerate(atmlst):
        shl0, shl1, p0, p1 = offsetdic[ia]

        # Ground state gradients
        h1ao = hcore_deriv(ia)
        h1ao[:,p0:p1]   += veff1[0,:,p0:p1]
        h1ao[:,:,p0:p1] += veff1[0,:,p0:p1].transpose(0,2,1)
        # oo0*2 for doubly occupied orbitals
        e1  = numpy.einsum('xpq,pq->x', h1ao, oo0) * 2

        e1 += numpy.einsum('xpq,pq->x', h1ao, dmz1doo)
        e1 -= numpy.einsum('xpq,pq->x', s1[:,p0:p1], im0[p0:p1])
        e1 -= numpy.einsum('xqp,pq->x', s1[:,p0:p1], im0[:,p0:p1])

        e1 += numpy.einsum('xij,ij->x', veff1[1,:,p0:p1], oo0[p0:p1])
        e1 += numpy.einsum('xij,ij->x', veff1[2,:,p0:p1], dmxpy[p0:p1,:]) * 2
        e1 += numpy.einsum('xij,ij->x', veff1[3,:,p0:p1], dmxmy[p0:p1,:]) * 2
        e1 += numpy.einsum('xji,ij->x', veff1[2,:,p0:p1], dmxpy[:,p0:p1]) * 2
        e1 -= numpy.einsum('xji,ij->x', veff1[3,:,p0:p1], dmxmy[:,p0:p1]) * 2

        de[k] = e1

    de += _grad_solvent(with_solvent, oo0*2, dmz1doo, dmxpy*2, singlet)

    log.timer('TDDFT nuclear gradients', *time0)
    return de

def tduhf_grad_elec(td_grad, x_y, atmlst=None, max_memory=2000, verbose=logger.INFO):
    '''
    See also function pyscf.grad.tduhf.grad_elec
    '''
    log = logger.new_logger(td_grad, verbose)
    time0 = logger.process_clock(), logger.perf_counter()

    mol = td_grad.mol
    mf = td_grad.base._scf
    mo_coeff = mf.mo_coeff
    mo_energy = mf.mo_energy
    mo_occ = mf.mo_occ

    with_solvent = getattr(td_grad.base, 'with_solvent', mf.with_solvent)

    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]
    nao = mo_coeff[0].shape[0]
    nmoa = nocca + nvira
    nmob = noccb + nvirb

    (xa, xb), (ya, yb) = x_y
    xpya = (xa+ya).reshape(nocca,nvira).T
    xpyb = (xb+yb).reshape(noccb,nvirb).T
    xmya = (xa-ya).reshape(nocca,nvira).T
    xmyb = (xb-yb).reshape(noccb,nvirb).T

    dvva = numpy.einsum('ai,bi->ab', xpya, xpya) + numpy.einsum('ai,bi->ab', xmya, xmya)
    dvvb = numpy.einsum('ai,bi->ab', xpyb, xpyb) + numpy.einsum('ai,bi->ab', xmyb, xmyb)
    dooa =-numpy.einsum('ai,aj->ij', xpya, xpya) - numpy.einsum('ai,aj->ij', xmya, xmya)
    doob =-numpy.einsum('ai,aj->ij', xpyb, xpyb) - numpy.einsum('ai,aj->ij', xmyb, xmyb)
    dmxpya = reduce(numpy.dot, (orbva, xpya, orboa.T))
    dmxpyb = reduce(numpy.dot, (orbvb, xpyb, orbob.T))
    dmxmya = reduce(numpy.dot, (orbva, xmya, orboa.T))
    dmxmyb = reduce(numpy.dot, (orbvb, xmyb, orbob.T))
    dmzooa = reduce(numpy.dot, (orboa, dooa, orboa.T))
    dmzoob = reduce(numpy.dot, (orbob, doob, orbob.T))
    dmzooa+= reduce(numpy.dot, (orbva, dvva, orbva.T))
    dmzoob+= reduce(numpy.dot, (orbvb, dvvb, orbvb.T))

    vj, vk = mf.get_jk(mol, (dmzooa, dmxpya+dmxpya.T, dmxmya-dmxmya.T,
                             dmzoob, dmxpyb+dmxpyb.T, dmxmyb-dmxmyb.T), hermi=0)
    vj = vj.reshape(2,3,nao,nao)
    vk = vk.reshape(2,3,nao,nao)
    if with_solvent.equilibrium_solvation:
        dmxpy = dmxpya + dmxpyb
        vj[0,:2] += mf.with_solvent._B_dot_x((dmzooa+dmzoob, dmxpy+dmxpy.T))
    else:
        vj[0,0] += mf.with_solvent._B_dot_x(dmzooa+dmzoob)

    veff0doo = vj[0,0]+vj[1,0] - vk[:,0]
    wvoa = reduce(numpy.dot, (orbva.T, veff0doo[0], orboa)) * 2
    wvob = reduce(numpy.dot, (orbvb.T, veff0doo[1], orbob)) * 2
    veff = vj[0,1]+vj[1,1] - vk[:,1]
    veff0mopa = reduce(numpy.dot, (mo_coeff[0].T, veff[0], mo_coeff[0]))
    veff0mopb = reduce(numpy.dot, (mo_coeff[1].T, veff[1], mo_coeff[1]))
    wvoa -= numpy.einsum('ki,ai->ak', veff0mopa[:nocca,:nocca], xpya) * 2
    wvob -= numpy.einsum('ki,ai->ak', veff0mopb[:noccb,:noccb], xpyb) * 2
    wvoa += numpy.einsum('ac,ai->ci', veff0mopa[nocca:,nocca:], xpya) * 2
    wvob += numpy.einsum('ac,ai->ci', veff0mopb[noccb:,noccb:], xpyb) * 2
    veff = -vk[:,2]
    veff0moma = reduce(numpy.dot, (mo_coeff[0].T, veff[0], mo_coeff[0]))
    veff0momb = reduce(numpy.dot, (mo_coeff[1].T, veff[1], mo_coeff[1]))
    wvoa -= numpy.einsum('ki,ai->ak', veff0moma[:nocca,:nocca], xmya) * 2
    wvob -= numpy.einsum('ki,ai->ak', veff0momb[:noccb,:noccb], xmyb) * 2
    wvoa += numpy.einsum('ac,ai->ci', veff0moma[nocca:,nocca:], xmya) * 2
    wvob += numpy.einsum('ac,ai->ci', veff0momb[noccb:,noccb:], xmyb) * 2

    with lib.temporary_env(mf.with_solvent, equilibrium_solvation=True):
        vresp = mf.gen_response(hermi=1)
        def fvind(x):
            dm1 = numpy.empty((2,nao,nao))
            xa = x[0,:nvira*nocca].reshape(nvira,nocca)
            xb = x[0,nvira*nocca:].reshape(nvirb,noccb)
            dma = reduce(numpy.dot, (orbva, xa, orboa.T))
            dmb = reduce(numpy.dot, (orbvb, xb, orbob.T))
            dm1[0] = dma + dma.T
            dm1[1] = dmb + dmb.T
            v1 = vresp(dm1)
            v1a = reduce(numpy.dot, (orbva.T, v1[0], orboa))
            v1b = reduce(numpy.dot, (orbvb.T, v1[1], orbob))
            return numpy.hstack((v1a.ravel(), v1b.ravel()))
        z1a, z1b = ucphf.solve(fvind, mo_energy, mo_occ, (wvoa,wvob),
                               max_cycle=td_grad.cphf_max_cycle,
                               tol=td_grad.cphf_conv_tol)[0]
        time1 = log.timer('Z-vector using UCPHF solver', *time0)

        z1ao = numpy.empty((2,nao,nao))
        z1ao[0] = reduce(numpy.dot, (orbva, z1a, orboa.T))
        z1ao[1] = reduce(numpy.dot, (orbvb, z1b, orbob.T))
        veff = vresp((z1ao+z1ao.transpose(0,2,1)) * .5)

    im0a = numpy.zeros((nmoa,nmoa))
    im0b = numpy.zeros((nmob,nmob))
    im0a[:nocca,:nocca] = reduce(numpy.dot, (orboa.T, veff0doo[0]+veff[0], orboa)) * .5
    im0b[:noccb,:noccb] = reduce(numpy.dot, (orbob.T, veff0doo[1]+veff[1], orbob)) * .5
    im0a[:nocca,:nocca]+= numpy.einsum('ak,ai->ki', veff0mopa[nocca:,:nocca], xpya) * .5
    im0b[:noccb,:noccb]+= numpy.einsum('ak,ai->ki', veff0mopb[noccb:,:noccb], xpyb) * .5
    im0a[:nocca,:nocca]+= numpy.einsum('ak,ai->ki', veff0moma[nocca:,:nocca], xmya) * .5
    im0b[:noccb,:noccb]+= numpy.einsum('ak,ai->ki', veff0momb[noccb:,:noccb], xmyb) * .5
    im0a[nocca:,nocca:] = numpy.einsum('ci,ai->ac', veff0mopa[nocca:,:nocca], xpya) * .5
    im0b[noccb:,noccb:] = numpy.einsum('ci,ai->ac', veff0mopb[noccb:,:noccb], xpyb) * .5
    im0a[nocca:,nocca:]+= numpy.einsum('ci,ai->ac', veff0moma[nocca:,:nocca], xmya) * .5
    im0b[noccb:,noccb:]+= numpy.einsum('ci,ai->ac', veff0momb[noccb:,:noccb], xmyb) * .5
    im0a[nocca:,:nocca] = numpy.einsum('ki,ai->ak', veff0mopa[:nocca,:nocca], xpya)
    im0b[noccb:,:noccb] = numpy.einsum('ki,ai->ak', veff0mopb[:noccb,:noccb], xpyb)
    im0a[nocca:,:nocca]+= numpy.einsum('ki,ai->ak', veff0moma[:nocca,:nocca], xmya)
    im0b[noccb:,:noccb]+= numpy.einsum('ki,ai->ak', veff0momb[:noccb,:noccb], xmyb)

    zeta_a = (mo_energy[0][:,None] + mo_energy[0]) * .5
    zeta_b = (mo_energy[1][:,None] + mo_energy[1]) * .5
    zeta_a[nocca:,:nocca] = mo_energy[0][:nocca]
    zeta_b[noccb:,:noccb] = mo_energy[1][:noccb]
    zeta_a[:nocca,nocca:] = mo_energy[0][nocca:]
    zeta_b[:noccb,noccb:] = mo_energy[1][noccb:]
    dm1a = numpy.zeros((nmoa,nmoa))
    dm1b = numpy.zeros((nmob,nmob))
    dm1a[:nocca,:nocca] = dooa * .5
    dm1b[:noccb,:noccb] = doob * .5
    dm1a[nocca:,nocca:] = dvva * .5
    dm1b[noccb:,noccb:] = dvvb * .5
    dm1a[nocca:,:nocca] = z1a * .5
    dm1b[noccb:,:noccb] = z1b * .5
    dm1a[:nocca,:nocca] += numpy.eye(nocca) # for ground state
    dm1b[:noccb,:noccb] += numpy.eye(noccb)
    im0a = reduce(numpy.dot, (mo_coeff[0], im0a+zeta_a*dm1a, mo_coeff[0].T))
    im0b = reduce(numpy.dot, (mo_coeff[1], im0b+zeta_b*dm1b, mo_coeff[1].T))
    im0 = im0a + im0b

    # Initialize hcore_deriv with the underlying SCF object because some
    # extensions (e.g. QM/MM, solvent) modifies the SCF object only.
    mf_grad = td_grad.base._scf.nuc_grad_method()
    hcore_deriv = mf_grad.hcore_generator(mol)
    s1 = mf_grad.get_ovlp(mol)

    dmz1dooa = z1ao[0] + dmzooa
    dmz1doob = z1ao[1] + dmzoob
    oo0a = reduce(numpy.dot, (orboa, orboa.T))
    oo0b = reduce(numpy.dot, (orbob, orbob.T))
    as_dm1 = oo0a + oo0b + (dmz1dooa + dmz1doob) * .5
    vj, vk = td_grad.get_jk(mol, (oo0a, dmz1dooa+dmz1dooa.T, dmxpya+dmxpya.T, dmxmya-dmxmya.T,
                                  oo0b, dmz1doob+dmz1doob.T, dmxpyb+dmxpyb.T, dmxmyb-dmxmyb.T))
    vj = vj.reshape(2,4,3,nao,nao)
    vk = vk.reshape(2,4,3,nao,nao)
    vhf1a, vhf1b = vj[0] + vj[1] - vk
    time1 = log.timer('2e AO integral derivatives', *time1)

    if atmlst is None:
        atmlst = range(mol.natm)
    offsetdic = mol.offset_nr_by_atom()
    de = numpy.zeros((len(atmlst),3))
    for k, ia in enumerate(atmlst):
        shl0, shl1, p0, p1 = offsetdic[ia]

        # Ground state gradients
        h1ao = hcore_deriv(ia)
        de[k] = numpy.einsum('xpq,pq->x', h1ao, as_dm1)

        de[k] += numpy.einsum('xpq,pq->x', vhf1a[0,:,p0:p1], oo0a[p0:p1])
        de[k] += numpy.einsum('xpq,pq->x', vhf1b[0,:,p0:p1], oo0b[p0:p1])
        de[k] += numpy.einsum('xpq,qp->x', vhf1a[0,:,p0:p1], oo0a[:,p0:p1])
        de[k] += numpy.einsum('xpq,qp->x', vhf1b[0,:,p0:p1], oo0b[:,p0:p1])

        de[k] += numpy.einsum('xpq,pq->x', vhf1a[0,:,p0:p1], dmz1dooa[p0:p1]) * .5
        de[k] += numpy.einsum('xpq,pq->x', vhf1b[0,:,p0:p1], dmz1doob[p0:p1]) * .5
        de[k] += numpy.einsum('xpq,qp->x', vhf1a[0,:,p0:p1], dmz1dooa[:,p0:p1]) * .5
        de[k] += numpy.einsum('xpq,qp->x', vhf1b[0,:,p0:p1], dmz1doob[:,p0:p1]) * .5

        de[k] -= numpy.einsum('xpq,pq->x', s1[:,p0:p1], im0[p0:p1])
        de[k] -= numpy.einsum('xqp,pq->x', s1[:,p0:p1], im0[:,p0:p1])

        de[k] += numpy.einsum('xij,ij->x', vhf1a[1,:,p0:p1], oo0a[p0:p1]) * .5
        de[k] += numpy.einsum('xij,ij->x', vhf1b[1,:,p0:p1], oo0b[p0:p1]) * .5
        de[k] += numpy.einsum('xij,ij->x', vhf1a[2,:,p0:p1], dmxpya[p0:p1,:])
        de[k] += numpy.einsum('xij,ij->x', vhf1b[2,:,p0:p1], dmxpyb[p0:p1,:])
        de[k] += numpy.einsum('xij,ij->x', vhf1a[3,:,p0:p1], dmxmya[p0:p1,:])
        de[k] += numpy.einsum('xij,ij->x', vhf1b[3,:,p0:p1], dmxmyb[p0:p1,:])
        de[k] += numpy.einsum('xji,ij->x', vhf1a[2,:,p0:p1], dmxpya[:,p0:p1])
        de[k] += numpy.einsum('xji,ij->x', vhf1b[2,:,p0:p1], dmxpyb[:,p0:p1])
        de[k] -= numpy.einsum('xji,ij->x', vhf1a[3,:,p0:p1], dmxmya[:,p0:p1])
        de[k] -= numpy.einsum('xji,ij->x', vhf1b[3,:,p0:p1], dmxmyb[:,p0:p1])

    dm0 = oo0a + oo0b
    dmz1doo = (dmz1dooa + dmz1doob) * .5
    dmxpy = dmxpya + dmxpyb
    de += _grad_solvent(with_solvent, dm0, dmz1doo, dmxpy)

    log.timer('TDUHF nuclear gradients', *time0)
    return de

def tduks_grad_elec(td_grad, x_y, atmlst=None, max_memory=2000, verbose=logger.INFO):
    '''
    See also function pyscf.grad.tduks.grad_elec
    '''
    log = logger.new_logger(td_grad, verbose)
    time0 = logger.process_clock(), logger.perf_counter()

    mol = td_grad.mol
    mf = td_grad.base._scf
    mo_coeff = mf.mo_coeff
    mo_energy = mf.mo_energy
    mo_occ = mf.mo_occ

    with_solvent = getattr(td_grad.base, 'with_solvent', mf.with_solvent)

    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]
    nao = mo_coeff[0].shape[0]
    nmoa = nocca + nvira
    nmob = noccb + nvirb

    (xa, xb), (ya, yb) = x_y
    xpya = (xa+ya).reshape(nocca,nvira).T
    xpyb = (xb+yb).reshape(noccb,nvirb).T
    xmya = (xa-ya).reshape(nocca,nvira).T
    xmyb = (xb-yb).reshape(noccb,nvirb).T

    dvva = numpy.einsum('ai,bi->ab', xpya, xpya) + numpy.einsum('ai,bi->ab', xmya, xmya)
    dvvb = numpy.einsum('ai,bi->ab', xpyb, xpyb) + numpy.einsum('ai,bi->ab', xmyb, xmyb)
    dooa =-numpy.einsum('ai,aj->ij', xpya, xpya) - numpy.einsum('ai,aj->ij', xmya, xmya)
    doob =-numpy.einsum('ai,aj->ij', xpyb, xpyb) - numpy.einsum('ai,aj->ij', xmyb, xmyb)
    dmxpya = reduce(numpy.dot, (orbva, xpya, orboa.T))
    dmxpyb = reduce(numpy.dot, (orbvb, xpyb, orbob.T))
    dmxmya = reduce(numpy.dot, (orbva, xmya, orboa.T))
    dmxmyb = reduce(numpy.dot, (orbvb, xmyb, orbob.T))
    dmzooa = reduce(numpy.dot, (orboa, dooa, orboa.T))
    dmzoob = reduce(numpy.dot, (orbob, doob, orbob.T))
    dmzooa+= reduce(numpy.dot, (orbva, dvva, orbva.T))
    dmzoob+= reduce(numpy.dot, (orbvb, dvvb, orbvb.T))

    ni = mf._numint
    ni.libxc.test_deriv_order(mf.xc, 3, raise_error=True)
    omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, mol.spin)
    # dm0 = mf.make_rdm1(mo_coeff, mo_occ), but it is not used when computing
    # fxc since rho0 is passed to fxc function.
    dm0 = None
    rho0, vxc, fxc = ni.cache_xc_kernel(mf.mol, mf.grids, mf.xc,
                                        mo_coeff, mo_occ, spin=1)
    f1vo, f1oo, vxc1, k1ao = \
            tduks_grad._contract_xc_kernel(td_grad, mf.xc, (dmxpya,dmxpyb),
                                           (dmzooa,dmzoob), True, True, max_memory)

    if abs(hyb) > 1e-10:
        dm = (dmzooa, dmxpya+dmxpya.T, dmxmya-dmxmya.T,
              dmzoob, dmxpyb+dmxpyb.T, dmxmyb-dmxmyb.T)
        vj, vk = mf.get_jk(mol, dm, hermi=0)
        vk *= hyb
        if abs(omega) > 1e-10:
            vk += mf.get_k(mol, dm, hermi=0, omega=omega) * (alpha-hyb)
        vj = vj.reshape(2,3,nao,nao)
        vk = vk.reshape(2,3,nao,nao)

        if with_solvent.equilibrium_solvation:
            dmxpy = dmxpya + dmxpyb
            vj[0,:2] += mf.with_solvent._B_dot_x((dmzooa+dmzoob, dmxpy+dmxpy.T))
        else:
            vj[0,0] += mf.with_solvent._B_dot_x(dmzooa+dmzoob)

        veff0doo = vj[0,0]+vj[1,0] - vk[:,0] + f1oo[:,0] + k1ao[:,0] * 2
        wvoa = reduce(numpy.dot, (orbva.T, veff0doo[0], orboa)) * 2
        wvob = reduce(numpy.dot, (orbvb.T, veff0doo[1], orbob)) * 2
        veff = vj[0,1]+vj[1,1] - vk[:,1] + f1vo[:,0] * 2
        veff0mopa = reduce(numpy.dot, (mo_coeff[0].T, veff[0], mo_coeff[0]))
        veff0mopb = reduce(numpy.dot, (mo_coeff[1].T, veff[1], mo_coeff[1]))
        wvoa -= numpy.einsum('ki,ai->ak', veff0mopa[:nocca,:nocca], xpya) * 2
        wvob -= numpy.einsum('ki,ai->ak', veff0mopb[:noccb,:noccb], xpyb) * 2
        wvoa += numpy.einsum('ac,ai->ci', veff0mopa[nocca:,nocca:], xpya) * 2
        wvob += numpy.einsum('ac,ai->ci', veff0mopb[noccb:,noccb:], xpyb) * 2
        veff = -vk[:,2]
        veff0moma = reduce(numpy.dot, (mo_coeff[0].T, veff[0], mo_coeff[0]))
        veff0momb = reduce(numpy.dot, (mo_coeff[1].T, veff[1], mo_coeff[1]))
        wvoa -= numpy.einsum('ki,ai->ak', veff0moma[:nocca,:nocca], xmya) * 2
        wvob -= numpy.einsum('ki,ai->ak', veff0momb[:noccb,:noccb], xmyb) * 2
        wvoa += numpy.einsum('ac,ai->ci', veff0moma[nocca:,nocca:], xmya) * 2
        wvob += numpy.einsum('ac,ai->ci', veff0momb[noccb:,noccb:], xmyb) * 2
    else:
        dm = (dmzooa, dmxpya+dmxpya.T,
              dmzoob, dmxpyb+dmxpyb.T)
        vj = mf.get_j(mol, dm, hermi=1).reshape(2,2,nao,nao)

        if with_solvent.equilibrium_solvation:
            dmxpy = dmxpya + dmxpyb
            vj[0,:2] += mf.with_solvent._B_dot_x((dmzooa+dmzoob, dmxpy+dmxpy.T))
        else:
            vj[0,0] += mf.with_solvent._B_dot_x(dmzooa+dmzoob)

        veff0doo = vj[0,0]+vj[1,0] + f1oo[:,0] + k1ao[:,0] * 2
        wvoa = reduce(numpy.dot, (orbva.T, veff0doo[0], orboa)) * 2
        wvob = reduce(numpy.dot, (orbvb.T, veff0doo[1], orbob)) * 2
        veff = vj[0,1]+vj[1,1] + f1vo[:,0] * 2
        veff0mopa = reduce(numpy.dot, (mo_coeff[0].T, veff[0], mo_coeff[0]))
        veff0mopb = reduce(numpy.dot, (mo_coeff[1].T, veff[1], mo_coeff[1]))
        wvoa -= numpy.einsum('ki,ai->ak', veff0mopa[:nocca,:nocca], xpya) * 2
        wvob -= numpy.einsum('ki,ai->ak', veff0mopb[:noccb,:noccb], xpyb) * 2
        wvoa += numpy.einsum('ac,ai->ci', veff0mopa[nocca:,nocca:], xpya) * 2
        wvob += numpy.einsum('ac,ai->ci', veff0mopb[noccb:,noccb:], xpyb) * 2
        veff0moma = numpy.zeros((nmoa,nmoa))
        veff0momb = numpy.zeros((nmob,nmob))

    with lib.temporary_env(mf.with_solvent, equilibrium_solvation=True):
        vresp = mf.gen_response(hermi=1)
        def fvind(x):
            dm1 = numpy.empty((2,nao,nao))
            xa = x[0,:nvira*nocca].reshape(nvira,nocca)
            xb = x[0,nvira*nocca:].reshape(nvirb,noccb)
            dma = reduce(numpy.dot, (orbva, xa, orboa.T))
            dmb = reduce(numpy.dot, (orbvb, xb, orbob.T))
            dm1[0] = dma + dma.T
            dm1[1] = dmb + dmb.T
            v1 = vresp(dm1)
            v1a = reduce(numpy.dot, (orbva.T, v1[0], orboa))
            v1b = reduce(numpy.dot, (orbvb.T, v1[1], orbob))
            return numpy.hstack((v1a.ravel(), v1b.ravel()))
        z1a, z1b = ucphf.solve(fvind, mo_energy, mo_occ, (wvoa,wvob),
                               max_cycle=td_grad.cphf_max_cycle,
                               tol=td_grad.cphf_conv_tol)[0]
        time1 = log.timer('Z-vector using UCPHF solver', *time0)

        z1ao = numpy.empty((2,nao,nao))
        z1ao[0] = reduce(numpy.dot, (orbva, z1a, orboa.T))
        z1ao[1] = reduce(numpy.dot, (orbvb, z1b, orbob.T))
        veff = vresp((z1ao+z1ao.transpose(0,2,1)) * .5)

    im0a = numpy.zeros((nmoa,nmoa))
    im0b = numpy.zeros((nmob,nmob))
    im0a[:nocca,:nocca] = reduce(numpy.dot, (orboa.T, veff0doo[0]+veff[0], orboa)) * .5
    im0b[:noccb,:noccb] = reduce(numpy.dot, (orbob.T, veff0doo[1]+veff[1], orbob)) * .5
    im0a[:nocca,:nocca]+= numpy.einsum('ak,ai->ki', veff0mopa[nocca:,:nocca], xpya) * .5
    im0b[:noccb,:noccb]+= numpy.einsum('ak,ai->ki', veff0mopb[noccb:,:noccb], xpyb) * .5
    im0a[:nocca,:nocca]+= numpy.einsum('ak,ai->ki', veff0moma[nocca:,:nocca], xmya) * .5
    im0b[:noccb,:noccb]+= numpy.einsum('ak,ai->ki', veff0momb[noccb:,:noccb], xmyb) * .5
    im0a[nocca:,nocca:] = numpy.einsum('ci,ai->ac', veff0mopa[nocca:,:nocca], xpya) * .5
    im0b[noccb:,noccb:] = numpy.einsum('ci,ai->ac', veff0mopb[noccb:,:noccb], xpyb) * .5
    im0a[nocca:,nocca:]+= numpy.einsum('ci,ai->ac', veff0moma[nocca:,:nocca], xmya) * .5
    im0b[noccb:,noccb:]+= numpy.einsum('ci,ai->ac', veff0momb[noccb:,:noccb], xmyb) * .5
    im0a[nocca:,:nocca] = numpy.einsum('ki,ai->ak', veff0mopa[:nocca,:nocca], xpya)
    im0b[noccb:,:noccb] = numpy.einsum('ki,ai->ak', veff0mopb[:noccb,:noccb], xpyb)
    im0a[nocca:,:nocca]+= numpy.einsum('ki,ai->ak', veff0moma[:nocca,:nocca], xmya)
    im0b[noccb:,:noccb]+= numpy.einsum('ki,ai->ak', veff0momb[:noccb,:noccb], xmyb)

    zeta_a = (mo_energy[0][:,None] + mo_energy[0]) * .5
    zeta_b = (mo_energy[1][:,None] + mo_energy[1]) * .5
    zeta_a[nocca:,:nocca] = mo_energy[0][:nocca]
    zeta_b[noccb:,:noccb] = mo_energy[1][:noccb]
    zeta_a[:nocca,nocca:] = mo_energy[0][nocca:]
    zeta_b[:noccb,noccb:] = mo_energy[1][noccb:]
    dm1a = numpy.zeros((nmoa,nmoa))
    dm1b = numpy.zeros((nmob,nmob))
    dm1a[:nocca,:nocca] = dooa * .5
    dm1b[:noccb,:noccb] = doob * .5
    dm1a[nocca:,nocca:] = dvva * .5
    dm1b[noccb:,noccb:] = dvvb * .5
    dm1a[nocca:,:nocca] = z1a * .5
    dm1b[noccb:,:noccb] = z1b * .5
    dm1a[:nocca,:nocca] += numpy.eye(nocca) # for ground state
    dm1b[:noccb,:noccb] += numpy.eye(noccb)
    im0a = reduce(numpy.dot, (mo_coeff[0], im0a+zeta_a*dm1a, mo_coeff[0].T))
    im0b = reduce(numpy.dot, (mo_coeff[1], im0b+zeta_b*dm1b, mo_coeff[1].T))
    im0 = im0a + im0b

    # Initialize hcore_deriv with the underlying SCF object because some
    # extensions (e.g. QM/MM, solvent) modifies the SCF object only.
    mf_grad = td_grad.base._scf.nuc_grad_method()
    hcore_deriv = mf_grad.hcore_generator(mol)
    s1 = mf_grad.get_ovlp(mol)

    dmz1dooa = z1ao[0] + dmzooa
    dmz1doob = z1ao[1] + dmzoob
    oo0a = reduce(numpy.dot, (orboa, orboa.T))
    oo0b = reduce(numpy.dot, (orbob, orbob.T))
    as_dm1 = oo0a + oo0b + (dmz1dooa + dmz1doob) * .5

    if abs(hyb) > 1e-10:
        dm = (oo0a, dmz1dooa+dmz1dooa.T, dmxpya+dmxpya.T, dmxmya-dmxmya.T,
              oo0b, dmz1doob+dmz1doob.T, dmxpyb+dmxpyb.T, dmxmyb-dmxmyb.T)
        vj, vk = td_grad.get_jk(mol, dm)
        vj = vj.reshape(2,4,3,nao,nao)
        vk = vk.reshape(2,4,3,nao,nao) * hyb
        if abs(omega) > 1e-10:
            with mol.with_range_coulomb(omega):
                vk += td_grad.get_k(mol, dm).reshape(2,4,3,nao,nao) * (alpha-hyb)
        veff1 = vj[0] + vj[1] - vk
    else:
        dm = (oo0a, dmz1dooa+dmz1dooa.T, dmxpya+dmxpya.T,
              oo0b, dmz1doob+dmz1doob.T, dmxpyb+dmxpyb.T)
        vj = td_grad.get_j(mol, dm).reshape(2,3,3,nao,nao)
        veff1 = numpy.zeros((2,4,3,nao,nao))
        veff1[:,:3] = vj[0] + vj[1]

    fxcz1 = tduks_grad._contract_xc_kernel(td_grad, mf.xc, z1ao, None,
                                           False, False, max_memory)[0]

    veff1[:,0] += vxc1[:,1:]
    veff1[:,1] +=(f1oo[:,1:] + fxcz1[:,1:] + k1ao[:,1:]*2)*2 # *2 for dmz1doo+dmz1oo.T
    veff1[:,2] += f1vo[:,1:] * 2
    veff1a, veff1b = veff1
    time1 = log.timer('2e AO integral derivatives', *time1)

    if atmlst is None:
        atmlst = range(mol.natm)
    offsetdic = mol.offset_nr_by_atom()
    de = numpy.zeros((len(atmlst),3))
    for k, ia in enumerate(atmlst):
        shl0, shl1, p0, p1 = offsetdic[ia]

        # Ground state gradients
        h1ao = hcore_deriv(ia)
        de[k] = numpy.einsum('xpq,pq->x', h1ao, as_dm1)

        de[k] += numpy.einsum('xpq,pq->x', veff1a[0,:,p0:p1], oo0a[p0:p1])
        de[k] += numpy.einsum('xpq,pq->x', veff1b[0,:,p0:p1], oo0b[p0:p1])
        de[k] += numpy.einsum('xpq,qp->x', veff1a[0,:,p0:p1], oo0a[:,p0:p1])
        de[k] += numpy.einsum('xpq,qp->x', veff1b[0,:,p0:p1], oo0b[:,p0:p1])

        de[k] += numpy.einsum('xpq,pq->x', veff1a[0,:,p0:p1], dmz1dooa[p0:p1]) * .5
        de[k] += numpy.einsum('xpq,pq->x', veff1b[0,:,p0:p1], dmz1doob[p0:p1]) * .5
        de[k] += numpy.einsum('xpq,qp->x', veff1a[0,:,p0:p1], dmz1dooa[:,p0:p1]) * .5
        de[k] += numpy.einsum('xpq,qp->x', veff1b[0,:,p0:p1], dmz1doob[:,p0:p1]) * .5

        de[k] -= numpy.einsum('xpq,pq->x', s1[:,p0:p1], im0[p0:p1])
        de[k] -= numpy.einsum('xqp,pq->x', s1[:,p0:p1], im0[:,p0:p1])

        de[k] += numpy.einsum('xij,ij->x', veff1a[1,:,p0:p1], oo0a[p0:p1]) * .5
        de[k] += numpy.einsum('xij,ij->x', veff1b[1,:,p0:p1], oo0b[p0:p1]) * .5
        de[k] += numpy.einsum('xij,ij->x', veff1a[2,:,p0:p1], dmxpya[p0:p1,:])
        de[k] += numpy.einsum('xij,ij->x', veff1b[2,:,p0:p1], dmxpyb[p0:p1,:])
        de[k] += numpy.einsum('xij,ij->x', veff1a[3,:,p0:p1], dmxmya[p0:p1,:])
        de[k] += numpy.einsum('xij,ij->x', veff1b[3,:,p0:p1], dmxmyb[p0:p1,:])
        de[k] += numpy.einsum('xji,ij->x', veff1a[2,:,p0:p1], dmxpya[:,p0:p1])
        de[k] += numpy.einsum('xji,ij->x', veff1b[2,:,p0:p1], dmxpyb[:,p0:p1])
        de[k] -= numpy.einsum('xji,ij->x', veff1a[3,:,p0:p1], dmxmya[:,p0:p1])
        de[k] -= numpy.einsum('xji,ij->x', veff1b[3,:,p0:p1], dmxmyb[:,p0:p1])

    dm0 = oo0a + oo0b
    dmz1doo = (dmz1dooa + dmz1doob) * .5
    dmxpy = dmxpya + dmxpyb
    de += _grad_solvent(with_solvent, dm0, dmz1doo, dmxpy)

    log.timer('TDUHF nuclear gradients', *time0)
    return de


def _grad_solvent(pcmobj, dm0, dmz1doo, dmxpy, singlet=True):
    '''Energy derivatives associated to derivatives of B tensor'''
    dielectric = pcmobj.eps
    if dielectric > 0:
        f_epsilon = (dielectric-1.)/dielectric
    else:
        f_epsilon = 1

    r_vdw      = pcmobj._intermediates['r_vdw'     ]
    ylm_1sph   = pcmobj._intermediates['ylm_1sph'  ]
    ui         = pcmobj._intermediates['ui'        ]
    Lmat       = pcmobj._intermediates['Lmat'      ]
    cached_pol = pcmobj._intermediates['cached_pol']

    # First order nuclei-solvent-electron contribution
    tmp = _grad_ne(pcmobj, dmz1doo,
                   r_vdw, ui, ylm_1sph, cached_pol, Lmat)
    de = .5 * f_epsilon * tmp

    # First order electron-solvent-electron contribution
    tmp = _grad_ee(pcmobj, (dm0, dmxpy), (dmz1doo, dmxpy),
                   r_vdw, ui, ylm_1sph, cached_pol, Lmat)
    de += .5 * f_epsilon * tmp[0]  # (dm0 * dmz1doo)

    if singlet and pcmobj.equilibrium_solvation:
        de += .5 * f_epsilon * tmp[1]  # (dmxpy * dmxpy)

    return de

def _grad_nn(pcmobj, r_vdw, ui, ylm_1sph, cached_pol, L):
    '''nuclei-solvent-nuclei term'''
    mol = pcmobj.mol
    natm = mol.natm
    coords_1sph, weights_1sph = ddcosmo.make_grids_one_sphere(pcmobj.lebedev_order)
    ngrid_1sph = coords_1sph.shape[0]
    lmax = pcmobj.lmax
    nlm = (lmax+1)**2
    atom_coords = mol.atom_coords()
    atom_charges = mol.atom_charges()

    fi0 = ddcosmo.make_fi(pcmobj, r_vdw)
    fi1 = ddcosmo_grad.make_fi1(pcmobj, pcmobj.get_atomic_radii())
    fi1[:,:,ui==0] = 0
    ui1 = -fi1

    de = numpy.zeros((natm,3))

    cav_coords = (atom_coords.reshape(natm,1,3)
                  + numpy.einsum('r,gx->rgx', r_vdw, coords_1sph))

    v_phi = numpy.zeros((natm, ngrid_1sph))
    for ia in range(natm):
        # Note (-) sign is not applied to atom_charges, because (-) is explicitly
        # included in rhs and L matrix
        d_rs = atom_coords.reshape(-1,1,3) - cav_coords[ia]
        v_phi[ia] = numpy.einsum('z,zp->p', atom_charges, 1./lib.norm(d_rs,axis=2))
    phi0 = -numpy.einsum('n,xn,jn,jn->jx', weights_1sph, ylm_1sph, ui, v_phi)

    psi0 = numpy.zeros((natm, nlm))
    for ia in range(natm):
        psi0[ia,0] += numpy.sqrt(4*numpy.pi)/r_vdw[ia] * mol.atom_charge(ia)

    v_phi0 = numpy.empty((natm,ngrid_1sph))
    for ia in range(natm):
        cav_coords = atom_coords[ia] + r_vdw[ia] * coords_1sph
        d_rs = atom_coords.reshape(-1,1,3) - cav_coords
        v_phi0[ia] = numpy.einsum('z,zp->p', atom_charges, 1./lib.norm(d_rs,axis=2))
    phi1 = -numpy.einsum('n,ln,azjn,jn->azjl', weights_1sph, ylm_1sph, ui1, v_phi0)

    for ia in range(natm):
        cav_coords = atom_coords[ia] + r_vdw[ia] * coords_1sph
        for ja in range(natm):
            rs = atom_coords[ja] - cav_coords
            d_rs = lib.norm(rs, axis=1)
            v_phi = atom_charges[ja] * numpy.einsum('px,p->px', rs, 1./d_rs**3)
            tmp = numpy.einsum('n,ln,n,nx->xl', weights_1sph, ylm_1sph, ui[ia], v_phi)
            phi1[ja,:,ia] += tmp  # response of the other atoms
            phi1[ia,:,ia] -= tmp  # response of cavity grids

    L1 = ddcosmo_grad.make_L1(pcmobj, r_vdw, ylm_1sph, fi0)
    Xvec0 = numpy.linalg.solve(L.reshape(natm*nlm,-1), phi0.ravel())
    Xvec0 = Xvec0.reshape(natm,nlm)
    phi1 -= numpy.einsum('aziljm,jm->azil', L1, Xvec0)

    LS0 = numpy.linalg.solve(L.reshape(natm*nlm,-1).T, psi0.ravel())
    LS0 = LS0.reshape(natm,nlm)
    de += numpy.einsum('il,azil->az', LS0, phi1)

    return de

def _grad_ne(pcmobj, dm, r_vdw, ui, ylm_1sph, cached_pol, L):
    '''nuclear charge-electron density cross term'''
    mol = pcmobj.mol
    natm = mol.natm
    coords_1sph, weights_1sph = ddcosmo.make_grids_one_sphere(pcmobj.lebedev_order)
    ngrid_1sph = coords_1sph.shape[0]
    lmax = pcmobj.lmax
    nlm = (lmax+1)**2
    nao = mol.nao
    atom_coords = mol.atom_coords()
    atom_charges = mol.atom_charges()
    grids = pcmobj.grids
    aoslices = mol.aoslice_by_atom()

    #extern_point_idx = ui > 0
    fi0 = ddcosmo.make_fi(pcmobj, r_vdw)
    fi1 = ddcosmo_grad.make_fi1(pcmobj, pcmobj.get_atomic_radii())
    fi1[:,:,ui==0] = 0
    ui1 = -fi1

    dms = numpy.asarray(dm)
    is_single_dm = dms.ndim == 2
    dms = dms.reshape(-1,nao,nao)
    n_dm = dms.shape[0]

    de = numpy.zeros((n_dm,natm,3))

    cav_coords = (atom_coords.reshape(natm,1,3)
                  + numpy.einsum('r,gx->rgx', r_vdw, coords_1sph))

    v_phi = numpy.zeros((natm, ngrid_1sph))
    for ia in range(natm):
        # Note (-) sign is not applied to atom_charges, because (-) is explicitly
        # included in rhs and L matrix
        d_rs = atom_coords.reshape(-1,1,3) - cav_coords[ia]
        v_phi[ia] = numpy.einsum('z,zp->p', atom_charges, 1./lib.norm(d_rs,axis=2))
    phi0 = -numpy.einsum('n,xn,jn,jn->jx', weights_1sph, ylm_1sph, ui, v_phi)

    Xvec0 = numpy.linalg.solve(L.reshape(natm*nlm,-1), phi0.ravel())
    Xvec0 = Xvec0.reshape(natm,nlm)

    v_phi0 = numpy.empty((natm,ngrid_1sph))
    for ia in range(natm):
        cav_coords = atom_coords[ia] + r_vdw[ia] * coords_1sph
        d_rs = atom_coords.reshape(-1,1,3) - cav_coords
        v_phi0[ia] = numpy.einsum('z,zp->p', atom_charges, 1./lib.norm(d_rs,axis=2))
    phi1 = -numpy.einsum('n,ln,azjn,jn->azjl', weights_1sph, ylm_1sph, ui1, v_phi0)

    for ia in range(natm):
        cav_coords = atom_coords[ia] + r_vdw[ia] * coords_1sph
        for ja in range(natm):
            rs = atom_coords[ja] - cav_coords
            d_rs = lib.norm(rs, axis=1)
            v_phi = atom_charges[ja] * numpy.einsum('px,p->px', rs, 1./d_rs**3)
            tmp = numpy.einsum('n,ln,n,nx->xl', weights_1sph, ylm_1sph, ui[ia], v_phi)
            phi1[ja,:,ia] += tmp  # response of the other atoms
            phi1[ia,:,ia] -= tmp  # response of cavity grids

    L1 = ddcosmo_grad.make_L1(pcmobj, r_vdw, ylm_1sph, fi0)

    phi1 -= numpy.einsum('aziljm,jm->azil', L1, Xvec0)
    Xvec1 = numpy.linalg.solve(L.reshape(natm*nlm,-1), phi1.reshape(-1,natm*nlm).T)
    Xvec1 = Xvec1.T.reshape(natm,3,natm,nlm)

    for ia, (coords, weight, weight1) in enumerate(rks_grad.grids_response_cc(grids)):
        ao = mol.eval_gto('GTOval_sph_deriv1', coords)

        fak_pol, leak_idx = cached_pol[mol.atom_symbol(ia)]
        fac_pol = ddcosmo._vstack_factor_fak_pol(fak_pol, lmax)
        scaled_weights = numpy.einsum('azm,mn->azn', Xvec1[:,:,ia], fac_pol)
        scaled_weights *= weight
        aow = numpy.einsum('gi,azg->azgi', ao[0], scaled_weights)
        de -= numpy.einsum('nij,gi,azgj->naz', dms, ao[0], aow)

        aow0 = numpy.einsum('gi,g->gi', ao[0], weight)
        aow1 = numpy.einsum('gi,zxg->zxgi', ao[0], weight1)

        den0 = numpy.einsum('nij,gi,zxgj->nzxg', dms, ao[0], aow1)
        de -= numpy.einsum('m,mg,nzxg->nzx', Xvec0[ia], fac_pol, den0)

        eta_nj = numpy.einsum('m,mg->g', Xvec0[ia], fac_pol)
        dm_ao  = lib.einsum('nij,gj->ngi', dms, aow0)
        dm_ao += lib.einsum('nji,gj->ngi', dms, aow0)
        for ja in range(natm):
            shl0, shl1, p0, p1 = aoslices[ja]
            den1 = numpy.einsum('ngi,xgi->nxg', dm_ao[:,:,p0:p1], ao[1:,:,p0:p1])
            detmp = numpy.einsum('g,nxg->nx', eta_nj, den1)
            de[:,ja] += detmp
            de[:,ia] -= detmp

    psi0 = numpy.zeros((natm, nlm))
    for ia in range(natm):
        psi0[ia,0] += numpy.sqrt(4*numpy.pi)/r_vdw[ia] * mol.atom_charge(ia)

    LS0 = numpy.linalg.solve(L.reshape(natm*nlm,-1).T, psi0.ravel())
    LS0 = LS0.reshape(natm,nlm)

    LS1 = numpy.einsum('il,aziljm->azjm', LS0, L1)
    LS1 = numpy.linalg.solve(L.reshape(natm*nlm,-1).T, LS1.reshape(-1,natm*nlm).T)
    LS1 = LS1.T.reshape(natm,3,natm,nlm)

    int3c2e = mol._add_suffix('int3c2e')
    int3c2e_ip1 = mol._add_suffix('int3c2e_ip1')
    cintopt = gto.moleintor.make_cintopt(mol._atm, mol._bas, mol._env, int3c2e)
    cintopt_ip1 = gto.moleintor.make_cintopt(mol._atm, mol._bas, mol._env, int3c2e_ip1)
    for ia in range(natm):
        cav_coords = atom_coords[ia] + r_vdw[ia] * coords_1sph
        #fakemol = gto.fakemol_for_charges(cav_coords[ui[ia]>0])
        fakemol = gto.fakemol_for_charges(cav_coords)
        wtmp = numpy.einsum('l,n,ln->ln', LS0[ia], weights_1sph, ylm_1sph)
        v_nj = df.incore.aux_e2(mol, fakemol, intor=int3c2e, aosym='s1', cintopt=cintopt)
        jaux = numpy.einsum('ijg,nij->ng', v_nj, dms)
        de -= numpy.einsum('azl,g,lg,g,ng->naz', LS1[:,:,ia], weights_1sph, ylm_1sph, ui[ia], jaux)
        de += numpy.einsum('lg,azg,ng->naz', wtmp, ui1[:,:,ia], jaux)

        v_e1_nj = df.incore.aux_e2(mol, fakemol, intor=int3c2e_ip1, comp=3,
                                   aosym='s1', cintopt=cintopt_ip1)
        for ja in range(natm):
            shl0, shl1, p0, p1 = aoslices[ja]
            jaux1  = numpy.einsum('xijg,nij->nxg', v_e1_nj[:,p0:p1], dms[:,p0:p1])
            jaux1 += numpy.einsum('xijg,nji->nxg', v_e1_nj[:,p0:p1], dms[:,:,p0:p1])
            detmp = numpy.einsum('lg,g,nxg->nx', wtmp, ui[ia], jaux1)
            de[:,ja] -= detmp
            de[:,ia] += detmp

    if is_single_dm:
        de = de[0]
    return de

def _grad_ee(pcmobj, dm1, dm2, r_vdw, ui, ylm_1sph, cached_pol, L):
    '''electron density-electorn density term'''
    mol = pcmobj.mol
    mol = pcmobj.mol
    natm = mol.natm
    nao  = mol.nao
    lmax = pcmobj.lmax
    nlm = (lmax+1)**2

    atom_coords = mol.atom_coords()
    aoslices = mol.aoslice_by_atom()
    grids = pcmobj.grids
    coords_1sph, weights_1sph = ddcosmo.make_grids_one_sphere(pcmobj.lebedev_order)

    #extern_point_idx = ui > 0
    fi0 = ddcosmo.make_fi(pcmobj, r_vdw)
    fi1 = ddcosmo_grad.make_fi1(pcmobj, pcmobj.get_atomic_radii())
    fi1[:,:,ui==0] = 0
    ui1 = -fi1

    dm1s = numpy.asarray(dm1)
    dm2s = numpy.asarray(dm2)
    is_single_dm = dm1s.ndim == 2
    dm1s = dm1s.reshape(-1,nao,nao)
    dm2s = dm2s.reshape(-1,nao,nao)
    n_dm = dm1s.shape[0]
    assert dm2s.shape[0] == n_dm

    de = numpy.zeros((n_dm,natm,3))

    ni = numint.NumInt()
    make_rho, nset, nao = ni._gen_rho_evaluator(mol, dm1s, hermi=0)
    den = numpy.empty((n_dm,grids.weights.size))
    p1 = 0
    for ao, mask, weight, coords in ni.block_loop(mol, grids, nao, 0):
        p0, p1 = p1, p1 + weight.size
        for i in range(n_dm):
            den[i,p0:p1] = make_rho(i, ao, mask, 'LDA') * weight

    psi0_dm1 = numpy.zeros((n_dm, natm, nlm))
    i1 = 0
    for ia in range(natm):
        fak_pol, leak_idx = cached_pol[mol.atom_symbol(ia)]
        fac_pol = ddcosmo._vstack_factor_fak_pol(fak_pol, lmax)
        i0, i1 = i1, i1 + fac_pol.shape[1]
        psi0_dm1[:,ia] = -numpy.einsum('mg,ng->nm', fac_pol, den[:,i0:i1])
    LS0 = numpy.linalg.solve(L.reshape(natm*nlm,-1).T, psi0_dm1.reshape(n_dm,-1).T)
    LS0 = LS0.T.reshape(n_dm,natm,nlm)

    phi0_dm1 = numpy.zeros((n_dm,natm,nlm))
    phi0_dm2 = numpy.zeros((n_dm,natm,nlm))
    phi1_dm1 = numpy.zeros((n_dm,natm,3,natm,nlm))
    int3c2e = mol._add_suffix('int3c2e')
    int3c2e_ip1 = mol._add_suffix('int3c2e_ip1')
    cintopt = gto.moleintor.make_cintopt(mol._atm, mol._bas, mol._env, int3c2e)
    cintopt_ip1 = gto.moleintor.make_cintopt(mol._atm, mol._bas, mol._env, int3c2e_ip1)
    for ia in range(natm):
        cav_coords = atom_coords[ia] + r_vdw[ia] * coords_1sph
        #fakemol = gto.fakemol_for_charges(cav_coords[ui[ia]>0])
        fakemol = gto.fakemol_for_charges(cav_coords)
        v_nj = df.incore.aux_e2(mol, fakemol, intor=int3c2e, aosym='s1', cintopt=cintopt)

        v_phi = numpy.einsum('ijg,nij->ng', v_nj, dm1s)
        phi0_dm1[:,ia] = numpy.einsum('g,lg,g,ng->nl', weights_1sph, ylm_1sph, ui[ia], v_phi)
        phi1_dm1[:,:,:,ia] += numpy.einsum('g,lg,azg,ng->nazl', weights_1sph, ylm_1sph, ui1[:,:,ia], v_phi)

        jaux = numpy.einsum('ijg,nij->ng', v_nj, dm2s)
        de += numpy.einsum('nl,g,lg,azg,ng->naz', LS0[:,ia], weights_1sph, ylm_1sph, ui1[:,:,ia], jaux)

        v_e1_nj = df.incore.aux_e2(mol, fakemol, intor=int3c2e_ip1, comp=3,
                                   aosym='s1', cintopt=cintopt_ip1)
        wtmp = numpy.einsum('g,lg,g->lg', weights_1sph, ylm_1sph, ui[ia])
        phi0_dm2[:,ia] = numpy.einsum('lg,ng->nl', wtmp, jaux)
        for ja in range(natm):
            shl0, shl1, p0, p1 = aoslices[ja]
            jaux1  = numpy.einsum('xijg,nij->nxg', v_e1_nj[:,p0:p1], dm2s[:,p0:p1])
            jaux1 += numpy.einsum('xijg,nji->nxg', v_e1_nj[:,p0:p1], dm2s[:,:,p0:p1])
            detmp = numpy.einsum('nl,lg,nxg->nx', LS0[:,ia], wtmp, jaux1)
            de[:,ja] -= detmp
            de[:,ia] += detmp

            tmp  = numpy.einsum('xijg,nij->nxg', v_e1_nj[:,p0:p1], dm1s[:,p0:p1])
            tmp += numpy.einsum('xijg,nji->nxg', v_e1_nj[:,p0:p1], dm1s[:,:,p0:p1])
            phitmp = numpy.einsum('lg,nxg->nxl', wtmp, tmp)
            phi1_dm1[:,ja,:,ia] -= phitmp
            phi1_dm1[:,ia,:,ia] += phitmp

    Xvec0 = numpy.linalg.solve(L.reshape(natm*nlm,-1), phi0_dm1.reshape(n_dm,-1).T)
    Xvec0 = Xvec0.T.reshape(n_dm,natm,nlm)

    L1 = ddcosmo_grad.make_L1(pcmobj, r_vdw, ylm_1sph, fi0)

    phi1_dm1 -= numpy.einsum('aziljm,njm->nazil', L1, Xvec0)
    Xvec1 = numpy.linalg.solve(L.reshape(natm*nlm,-1), phi1_dm1.reshape(-1,natm*nlm).T)
    Xvec1 = Xvec1.T.reshape(n_dm,natm,3,natm,nlm)

    psi1_dm1 = numpy.zeros((n_dm,natm,3,natm,nlm))
    for ia, (coords, weight, weight1) in enumerate(rks_grad.grids_response_cc(grids)):
        ao = mol.eval_gto('GTOval_sph_deriv1', coords)

        fak_pol, leak_idx = cached_pol[mol.atom_symbol(ia)]
        fac_pol = ddcosmo._vstack_factor_fak_pol(fak_pol, lmax)
        scaled_weights = numpy.einsum('nazm,mg->nazg', Xvec1[:,:,:,ia], fac_pol)
        scaled_weights *= weight
        aow = numpy.einsum('gi,nazg->nazgi', ao[0], scaled_weights)
        de -= numpy.einsum('nij,gi,nazgj->naz', dm2s, ao[0], aow)

        aow0 = numpy.einsum('gi,g->gi', ao[0], weight)
        aow1 = numpy.einsum('gi,zxg->zxgi', ao[0], weight1)

        den0 = numpy.einsum('nij,gi,zxgj->nzxg', dm1s, ao[0], aow1)
        psi1_dm1[:,:,:,ia] -= numpy.einsum('mg,nzxg->nzxm', fac_pol, den0)

        dm_ao  = lib.einsum('nij,gj->ngi', dm1s, aow0)
        dm_ao += lib.einsum('nji,gj->ngi', dm1s, aow0)
        for ja in range(natm):
            shl0, shl1, p0, p1 = aoslices[ja]
            den1 = numpy.einsum('ngi,xgi->nxg', dm_ao[:,:,p0:p1], ao[1:,:,p0:p1])
            psitmp = numpy.einsum('mg,nxg->nxm', fac_pol, den1)
            psi1_dm1[:,ja,:,ia] += psitmp
            psi1_dm1[:,ia,:,ia] -= psitmp

        eta_nj = numpy.einsum('nm,mg->ng', Xvec0[:,ia], fac_pol)
        den0 = numpy.einsum('nij,gi,zxgj->nzxg', dm2s, ao[0], aow1)
        de -= numpy.einsum('ng,nzxg->nzx', eta_nj, den0)

        dm_ao  = lib.einsum('nij,gj->ngi', dm2s, aow0)
        dm_ao += lib.einsum('nji,gj->ngi', dm2s, aow0)
        for ja in range(natm):
            shl0, shl1, p0, p1 = aoslices[ja]
            den1 = lib.einsum('ngi,xgi->nxg', dm_ao[:,:,p0:p1], ao[1:,:,p0:p1])
            detmp = numpy.einsum('ng,nxg->nx', eta_nj, den1)
            de[:,ja] += detmp
            de[:,ia] -= detmp

    psi1_dm1 -= numpy.einsum('nil,aziljm->nazjm', LS0, L1)
    LS1 = numpy.linalg.solve(L.reshape(natm*nlm,-1).T, psi1_dm1.reshape(-1,natm*nlm).T)
    LS1 = LS1.T.reshape(n_dm,natm,3,natm,nlm)
    de += numpy.einsum('nazjx,njx->naz', LS1, phi0_dm2)

    if is_single_dm:
        de = de[0]
    return de


if __name__ == '__main__':
    mol0 = gto.M(atom='H  0.  0.  1.804; F  0.  0.  0.', verbose=0, unit='B')
    mol1 = gto.M(atom='H  0.  0.  1.803; F  0.  0.  0.', verbose=0, unit='B')
    mol2 = gto.M(atom='H  0.  0.  1.805; F  0.  0.  0.', verbose=0, unit='B')

    # TDA with equilibrium_solvation
    mf = mol0.RHF().ddCOSMO().run()
    td = mf.TDA().ddCOSMO().run(equilibrium_solvation=True)
    g1 = td.nuc_grad_method().kernel() # 0  0  -0.5116214042

    mf = mol1.RHF().ddCOSMO().run()
    td1 = mf.TDA().ddCOSMO().run(equilibrium_solvation=True)
    mf = mol2.RHF().ddCOSMO().run()
    td2 = mf.TDA().ddCOSMO().run(equilibrium_solvation=True)
    print((td2.e_tot[0]-td1.e_tot[0])/0.002, g1[0,2])
    print((td2.e_tot[0]-td1.e_tot[0])/0.002 - g1[0,2])

    # TDA without equilibrium_solvation
    mf = mol0.RHF().ddCOSMO().run()
    td = mf.TDA().ddCOSMO().run()
    g1 = td.nuc_grad_method().kernel()

    mf = mol1.RHF().ddCOSMO().run()
    td1 = mf.TDA().ddCOSMO().run()
    mf = mol2.RHF().ddCOSMO().run()
    td2 = mf.TDA().ddCOSMO().run()
    print((td2.e_tot[0]-td1.e_tot[0])/0.002, g1[0,2])
    print((td2.e_tot[0]-td1.e_tot[0])/0.002 - g1[0,2])

    # TDA lda with equilibrium_solvation
    mf = mol0.RKS().ddCOSMO().run(xc='svwn')
    td = mf.TDA().ddCOSMO().run(equilibrium_solvation=True)
    g1 = td.nuc_grad_method().kernel()

    mf = mol1.RKS().ddCOSMO().run(xc='svwn')
    td1 = mf.TDA().ddCOSMO().run(equilibrium_solvation=True)
    mf = mol2.RKS().ddCOSMO().run(xc='svwn')
    td2 = mf.TDA().ddCOSMO().run(equilibrium_solvation=True)
    print((td2.e_tot[0]-td1.e_tot[0])/0.002, g1[0,2])
    print((td2.e_tot[0]-td1.e_tot[0])/0.002 - g1[0,2])


# pyscf/solvent/ddcosmo.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>
#

'''
domain decomposition COSMO

See also the code on github

https://github.com/filippolipparini/ddPCM

and the papers

[1] Domain decomposition for implicit solvation models.
E. Cances, Y. Maday, B. Stamm
J. Chem. Phys., 139, 054111 (2013)
http://dx.doi.org/10.1063/1.4816767

[2] Fast Domain Decomposition Algorithm for Continuum Solvation Models: Energy and First Derivatives.
F. Lipparini, B. Stamm, E. Cances, Y. Maday, B. Mennucci
J. Chem. Theory Comput., 9, 3637-3648 (2013)
http://dx.doi.org/10.1021/ct400280b

[3] Quantum, classical, and hybrid QM/MM calculations in solution: General implementation of the ddCOSMO linear scaling strategy.
F. Lipparini, G. Scalmani, L. Lagardere, B. Stamm, E. Cances, Y. Maday, J.-P.Piquemal, M. J. Frisch, B. Mennucci
J. Chem. Phys., 141, 184108 (2014)
http://dx.doi.org/10.1063/1.4901304

-- Dielectric constants (from https://gaussian.com/scrf/) --
More dataset can be found in Minnesota Solvent Descriptor Database
(https://comp.chem.umn.edu/solvation)
Water                                  78.3553
Acetonitrile                           35.688
Methanol                               32.613
Ethanol                                24.852
IsoQuinoline                           11.00
Quinoline                              9.16
Chloroform                             4.7113
DiethylEther                           4.2400
Dichloromethane                        8.93
DiChloroEthane                         10.125
CarbonTetraChloride                    2.2280
Benzene                                2.2706
Toluene                                2.3741
ChloroBenzene                          5.6968
NitroMethane                           36.562
Heptane                                1.9113
CycloHexane                            2.0165
Aniline                                6.8882
Acetone                                20.493
TetraHydroFuran                        7.4257
DiMethylSulfoxide                      46.826
Argon                                  1.430
Krypton                                1.519
Xenon                                  1.706
n-Octanol                              9.8629
1,1,1-TriChloroEthane                  7.0826
1,1,2-TriChloroEthane                  7.1937
1,2,4-TriMethylBenzene                 2.3653
1,2-DiBromoEthane                      4.9313
1,2-EthaneDiol                         40.245
1,4-Dioxane                            2.2099
1-Bromo-2-MethylPropane                7.7792
1-BromoOctane                          5.0244
1-BromoPentane                         6.269
1-BromoPropane                         8.0496
1-Butanol                              17.332
1-ChloroHexane                         5.9491
1-ChloroPentane                        6.5022
1-ChloroPropane                        8.3548
1-Decanol                              7.5305
1-FluoroOctane                         3.89
1-Heptanol                             11.321
1-Hexanol                              12.51
1-Hexene                               2.0717
1-Hexyne                               2.615
1-IodoButane                           6.173
1-IodoHexaDecane                       3.5338
1-IodoPentane                          5.6973
1-IodoPropane                          6.9626
1-NitroPropane                         23.73
1-Nonanol                              8.5991
1-Pentanol                             15.13
1-Pentene                              1.9905
1-Propanol                             20.524
2,2,2-TriFluoroEthanol                 26.726
2,2,4-TriMethylPentane                 1.9358
2,4-DiMethylPentane                    1.8939
2,4-DiMethylPyridine                   9.4176
2,6-DiMethylPyridine                   7.1735
2-BromoPropane                         9.3610
2-Butanol                              15.944
2-ChloroButane                         8.3930
2-Heptanone                            11.658
2-Hexanone                             14.136
2-MethoxyEthanol                       17.2
2-Methyl-1-Propanol                    16.777
2-Methyl-2-Propanol                    12.47
2-MethylPentane                        1.89
2-MethylPyridine                       9.9533
2-NitroPropane                         25.654
2-Octanone                             9.4678
2-Pentanone                            15.200
2-Propanol                             19.264
2-Propen-1-ol                          19.011
3-MethylPyridine                       11.645
3-Pentanone                            16.78
4-Heptanone                            12.257
4-Methyl-2-Pentanone                   12.887
4-MethylPyridine                       11.957
5-Nonanone                             10.6
AceticAcid                             6.2528
AcetoPhenone                           17.44
a-ChloroToluene                        6.7175
Anisole                                4.2247
Benzaldehyde                           18.220
BenzoNitrile                           25.592
BenzylAlcohol                          12.457
BromoBenzene                           5.3954
BromoEthane                            9.01
Bromoform                              4.2488
Butanal                                13.45
ButanoicAcid                           2.9931
Butanone                               18.246
ButanoNitrile                          24.291
ButylAmine                             4.6178
ButylEthanoate                         4.9941
CarbonDiSulfide                        2.6105
Cis-1,2-DiMethylCycloHexane            2.06
Cis-Decalin                            2.2139
CycloHexanone                          15.619
CycloPentane                           1.9608
CycloPentanol                          16.989
CycloPentanone                         13.58
Decalin-mixture                        2.196
DiBromomEthane                         7.2273
DiButylEther                           3.0473
DiEthylAmine                           3.5766
DiEthylSulfide                         5.723
DiIodoMethane                          5.32
DiIsoPropylEther                       3.38
DiMethylDiSulfide                      9.6
DiPhenylEther                          3.73
DiPropylAmine                          2.9112
e-1,2-DiChloroEthene                   2.14
e-2-Pentene                            2.051
EthaneThiol                            6.667
EthylBenzene                           2.4339
EthylEthanoate                         5.9867
EthylMethanoate                        8.3310
EthylPhenylEther                       4.1797
FluoroBenzene                          5.42
Formamide                              108.94
FormicAcid                             51.1
HexanoicAcid                           2.6
IodoBenzene                            4.5470
IodoEthane                             7.6177
IodoMethane                            6.8650
IsoPropylBenzene                       2.3712
m-Cresol                               12.44
Mesitylene                             2.2650
MethylBenzoate                         6.7367
MethylButanoate                        5.5607
MethylCycloHexane                      2.024
MethylEthanoate                        6.8615
MethylMethanoate                       8.8377
MethylPropanoate                       6.0777
m-Xylene                               2.3478
n-ButylBenzene                         2.36
n-Decane                               1.9846
n-Dodecane                             2.0060
n-Hexadecane                           2.0402
n-Hexane                               1.8819
NitroBenzene                           34.809
NitroEthane                            28.29
n-MethylAniline                        5.9600
n-MethylFormamide-mixture              181.56
n,n-DiMethylAcetamide                  37.781
n,n-DiMethylFormamide                  37.219
n-Nonane                               1.9605
n-Octane                               1.9406
n-Pentadecane                          2.0333
n-Pentane                              1.8371
n-Undecane                             1.9910
o-ChloroToluene                        4.6331
o-Cresol                               6.76
o-DiChloroBenzene                      9.9949
o-NitroToluene                         25.669
o-Xylene                               2.5454
Pentanal                               10.0
PentanoicAcid                          2.6924
PentylAmine                            4.2010
PentylEthanoate                        4.7297
PerFluoroBenzene                       2.029
p-IsoPropylToluene                     2.2322
Propanal                               18.5
PropanoicAcid                          3.44
PropanoNitrile                         29.324
PropylAmine                            4.9912
PropylEthanoate                        5.5205
p-Xylene                               2.2705
Pyridine                               12.978
sec-ButylBenzene                       2.3446
tert-ButylBenzene                      2.3447
TetraChloroEthene                      2.268
TetraHydroThiophene-s,s-dioxide        43.962
Tetralin                               2.771
Thiophene                              2.7270
Thiophenol                             4.2728
trans-Decalin                          2.1781
TriButylPhosphate                      8.1781
TriChloroEthene                        3.422
TriEthylAmine                          2.3832
Xylene-mixture                         2.3879
z-1,2-DiChloroEthene                   9.2
'''  # noqa: E501

import ctypes
import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf import gto
from pyscf import df
from pyscf.dft import gen_grid, numint
from pyscf.data import radii
from pyscf.symm import sph

from pyscf.solvent import _attach_solvent

@lib.with_doc(_attach_solvent._for_scf.__doc__)
def ddcosmo_for_scf(mf, solvent_obj=None, dm=None):
    if solvent_obj is None:
        solvent_obj = DDCOSMO(mf.mol)
    return _attach_solvent._for_scf(mf, solvent_obj, dm)

@lib.with_doc(_attach_solvent._for_casscf.__doc__)
def ddcosmo_for_casscf(mc, solvent_obj=None, dm=None):
    if solvent_obj is None:
        if isinstance(getattr(mc._scf, 'with_solvent', None), DDCOSMO):
            solvent_obj = mc._scf.with_solvent
        else:
            solvent_obj = DDCOSMO(mc.mol)
    return _attach_solvent._for_casscf(mc, solvent_obj, dm)

@lib.with_doc(_attach_solvent._for_casci.__doc__)
def ddcosmo_for_casci(mc, solvent_obj=None, dm=None):
    if solvent_obj is None:
        if isinstance(getattr(mc._scf, 'with_solvent', None), DDCOSMO):
            solvent_obj = mc._scf.with_solvent
        else:
            solvent_obj = DDCOSMO(mc.mol)
    return _attach_solvent._for_casci(mc, solvent_obj, dm)

@lib.with_doc(_attach_solvent._for_post_scf.__doc__)
def ddcosmo_for_post_scf(method, solvent_obj=None, dm=None):
    if solvent_obj is None:
        if isinstance(getattr(method._scf, 'with_solvent', None), DDCOSMO):
            solvent_obj = method._scf.with_solvent
        else:
            solvent_obj = DDCOSMO(method.mol)
    return _attach_solvent._for_post_scf(method, solvent_obj, dm)

@lib.with_doc(_attach_solvent._for_tdscf.__doc__)
def ddcosmo_for_tdscf(method, solvent_obj=None, dm=None):
    scf_solvent = getattr(method._scf, 'with_solvent', None)
    assert scf_solvent is None or isinstance(scf_solvent, DDCOSMO)

    if solvent_obj is None:
        solvent_obj = DDCOSMO(method.mol)
    return _attach_solvent._for_tdscf(method, solvent_obj, dm)


# Inject ddCOSMO into other methods
from pyscf import scf
from pyscf import mcscf
from pyscf import mp, ci, cc
from pyscf import tdscf
scf.hf.SCF.ddCOSMO    = scf.hf.SCF.DDCOSMO    = ddcosmo_for_scf
mp.mp2.MP2.ddCOSMO    = mp.mp2.MP2.DDCOSMO    = ddcosmo_for_post_scf
ci.cisd.CISD.ddCOSMO  = ci.cisd.CISD.DDCOSMO  = ddcosmo_for_post_scf
cc.ccsd.CCSD.ddCOSMO  = cc.ccsd.CCSD.DDCOSMO  = ddcosmo_for_post_scf
tdscf.rhf.TDMixin.ddCOSMO = tdscf.rhf.TDMixin.DDCOSMO = ddcosmo_for_tdscf
mcscf.casci.CASCI.ddCOSMO = mcscf.casci.CASCI.DDCOSMO = ddcosmo_for_casci
mcscf.mc1step.CASSCF.ddCOSMO = mcscf.mc1step.CASSCF.DDCOSMO = ddcosmo_for_casscf


# Keep gen_ddcosmo_solver for backward compatibility
def gen_ddcosmo_solver(pcmobj, verbose=None):
    '''Generate ddcosmo function to compute energy and potential matrix
    '''
    return pcmobj._get_vind

def energy(pcmobj, dm):
    r'''
    ddCOSMO energy
    Es = 1/2 f(eps) \int rho(r) W(r) dr
    '''
    epcm = pcmobj._get_vind(dm)[0]
    return epcm

def get_atomic_radii(pcmobj):
    mol = pcmobj.mol
    vdw_radii = pcmobj.radii_table
    atom_radii = pcmobj.atom_radii

    atom_symb = [mol.atom_symbol(i) for i in range(mol.natm)]
    r_vdw = [vdw_radii[gto.charge(x)] for x in atom_symb]
    if atom_radii is not None:
        for i in range(mol.natm):
            if atom_symb[i] in atom_radii:
                r_vdw[i] = atom_radii[atom_symb[i]]
    return numpy.asarray(r_vdw)


def regularize_xt(t, eta):
    xt = numpy.zeros_like(t)
    inner = t <= 1-eta
    on_shell = (1-eta < t) & (t < 1)
    xt[inner] = 1
    ti = t[on_shell]
# JCTC, 9, 3637
    xt[on_shell] = 1./eta**5 * (1-ti)**3 * (6*ti**2 + (15*eta-12)*ti
                                            + 10*eta**2 - 15*eta + 6)
# JCP, 139, 054111
#    xt[on_shell] = 1./eta**4 * (1-ti)**2 * (ti-1+2*eta)**2
    return xt

def make_grids_one_sphere(lebedev_order):
    ngrid_1sph = gen_grid.LEBEDEV_ORDER[lebedev_order]
    leb_grid = numpy.empty((ngrid_1sph,4))
    gen_grid.libdft.MakeAngularGrid(leb_grid.ctypes.data_as(ctypes.c_void_p),
                                    ctypes.c_int(ngrid_1sph))
    coords_1sph = leb_grid[:,:3]
    # Note the Lebedev angular grids are normalized to 1 in pyscf
    weights_1sph = 4*numpy.pi * leb_grid[:,3]
    return coords_1sph, weights_1sph

def make_L(pcmobj, r_vdw, ylm_1sph, fi):
    # See JCTC, 9, 3637, Eq (18)
    mol = pcmobj.mol
    natm = mol.natm
    lmax = pcmobj.lmax
    eta = pcmobj.eta
    nlm = (lmax+1)**2

    coords_1sph, weights_1sph = make_grids_one_sphere(pcmobj.lebedev_order)
    ngrid_1sph = weights_1sph.size
    atom_coords = mol.atom_coords()
    ylm_1sph = ylm_1sph.reshape(nlm,ngrid_1sph)

# JCP, 141, 184108 Eq (9), (12) is incorrect
# L_diag = <lm|(1/|s-s'|)|l'm'>
# Using Laplace expansion for electrostatic potential 1/r
# L_diag = 4pi/(2l+1)/|s| <lm|l'm'>
    L_diag = numpy.zeros((natm,nlm))
    p1 = 0
    for l in range(lmax+1):
        p0, p1 = p1, p1 + (l*2+1)
        L_diag[:,p0:p1] = 4*numpy.pi/(l*2+1)
    L_diag *= 1./r_vdw.reshape(-1,1)
    Lmat = numpy.diag(L_diag.ravel()).reshape(natm,nlm,natm,nlm)

    for ja in range(natm):
        # scale the weight, precontract d_nj and w_n
        # see JCTC 9, 3637, Eq (16) - (18)
        # Note all values are scaled by 1/r_vdw to make the formulas
        # consistent to Psi in JCP, 141, 184108
        part_weights = weights_1sph.copy()
        part_weights[fi[ja]>1] /= fi[ja,fi[ja]>1]
        for ka in atoms_with_vdw_overlap(ja, atom_coords, r_vdw):
            vjk = r_vdw[ja] * coords_1sph + atom_coords[ja] - atom_coords[ka]
            tjk = lib.norm(vjk, axis=1) / r_vdw[ka]
            wjk = pcmobj.regularize_xt(tjk, eta, r_vdw[ka])
            wjk *= part_weights
            pol = sph.multipoles(vjk, lmax)
            p1 = 0
            for l in range(lmax+1):
                fac = 4*numpy.pi/(l*2+1) / r_vdw[ka]**(l+1)
                p0, p1 = p1, p1 + (l*2+1)
                a = numpy.einsum('xn,n,mn->xm', ylm_1sph, wjk, pol[l])
                Lmat[ja,:,ka,p0:p1] += -fac * a
    return Lmat

def make_fi(pcmobj, r_vdw):
    coords_1sph, weights_1sph = make_grids_one_sphere(pcmobj.lebedev_order)
    mol = pcmobj.mol
    eta = pcmobj.eta
    natm = mol.natm
    atom_coords = mol.atom_coords()
    ngrid_1sph = coords_1sph.shape[0]
    fi = numpy.zeros((natm,ngrid_1sph))
    for ia in range(natm):
        for ja in atoms_with_vdw_overlap(ia, atom_coords, r_vdw):
            v = r_vdw[ia]*coords_1sph + atom_coords[ia] - atom_coords[ja]
            rv = lib.norm(v, axis=1)
            t = rv / r_vdw[ja]
            xt = pcmobj.regularize_xt(t, eta, r_vdw[ja])
            fi[ia] += xt
    fi[fi < 1e-20] = 0
    return fi

def make_phi(pcmobj, dm, r_vdw, ui, ylm_1sph, with_nuc=True):
    '''
    Induced potential of ddCOSMO model

    Kwargs:
        with_nuc (bool): Mute the contribution of nuclear charges when
            computing the second order derivatives of energy
    '''
    mol = pcmobj.mol
    natm = mol.natm
    coords_1sph, weights_1sph = make_grids_one_sphere(pcmobj.lebedev_order)
    ngrid_1sph = coords_1sph.shape[0]

    dms = numpy.asarray(dm)
    is_single_dm = dms.ndim == 2

    nao = dms.shape[-1]
    dms = dms.reshape(-1,nao,nao)
    n_dm = dms.shape[0]
    diagidx = numpy.arange(nao)
    diagidx = diagidx*(diagidx+1)//2 + diagidx
    tril_dm = lib.pack_tril(dms+dms.transpose(0,2,1))
    tril_dm[:,diagidx] *= .5

    atom_coords = mol.atom_coords()
    atom_charges = mol.atom_charges()

    extern_point_idx = ui > 0
    cav_coords = (atom_coords.reshape(natm,1,3)
                  + numpy.einsum('r,gx->rgx', r_vdw, coords_1sph))

    v_phi = numpy.zeros((n_dm, natm, ngrid_1sph))

    if with_nuc:
        for ia in range(natm):
            # Note (-) sign is not applied to atom_charges, because (-) is explicitly
            # included in rhs and L matrix
            d_rs = atom_coords.reshape(-1,1,3) - cav_coords[ia]
            v_phi[:,ia] = numpy.einsum('z,zp->p', atom_charges, 1./lib.norm(d_rs,axis=2))

    max_memory = pcmobj.max_memory - lib.current_memory()[0]
    blksize = int(max(max_memory*.9e6/8/nao**2, 400))

    cav_coords = cav_coords[extern_point_idx]
    v_phi_e = numpy.empty((n_dm, cav_coords.shape[0]))
    int3c2e = mol._add_suffix('int3c2e')
    cintopt = gto.moleintor.make_cintopt(mol._atm, mol._bas, mol._env, int3c2e)
    for i0, i1 in lib.prange(0, cav_coords.shape[0], blksize):
        fakemol = gto.fakemol_for_charges(cav_coords[i0:i1])
        v_nj = df.incore.aux_e2(mol, fakemol, intor=int3c2e, aosym='s2ij', cintopt=cintopt)
        v_phi_e[:,i0:i1] = numpy.einsum('nx,xk->nk', tril_dm, v_nj)
    v_phi[:,extern_point_idx] -= v_phi_e

    phi = -numpy.einsum('n,xn,jn,ijn->ijx', weights_1sph, ylm_1sph, ui, v_phi)
    if is_single_dm:
        phi = phi[0]
    return phi

def make_psi_vmat(pcmobj, dm, r_vdw, ui, ylm_1sph, cached_pol, Xvec, L,
                  with_nuc=True):
    '''
    The first order derivative of E_ddCOSMO wrt density matrix

    Kwargs:
        with_nuc (bool): Mute the contribution of nuclear charges when
            computing the second order derivatives of energy.
    '''
    mol = pcmobj.mol
    natm = mol.natm
    lmax = pcmobj.lmax
    nlm = (lmax+1)**2

    dms = numpy.asarray(dm)
    is_single_dm = dms.ndim == 2
    grids = pcmobj.grids

    ni = numint.NumInt()
    max_memory = pcmobj.max_memory - lib.current_memory()[0]
    make_rho, n_dm, nao = ni._gen_rho_evaluator(mol, dms)
    dms = dms.reshape(n_dm,nao,nao)
    Xvec = Xvec.reshape(n_dm, natm, nlm)

    i1 = 0
    scaled_weights = numpy.empty((n_dm, grids.weights.size))
    for ia in range(natm):
        fak_pol, leak_idx = cached_pol[mol.atom_symbol(ia)]
        fac_pol = _vstack_factor_fak_pol(fak_pol, lmax)
        i0, i1 = i1, i1 + fac_pol.shape[1]
        scaled_weights[:,i0:i1] = numpy.einsum('mn,im->in', fac_pol, Xvec[:,ia])
    scaled_weights *= grids.weights

    shls_slice = (0, mol.nbas)
    ao_loc = mol.ao_loc_nr()
    den = numpy.empty((n_dm, grids.weights.size))
    vmat = numpy.zeros((n_dm, nao, nao))
    p1 = 0
    aow = None
    for ao, mask, weight, coords \
            in ni.block_loop(mol, grids, nao, 0, max_memory):
        p0, p1 = p1, p1 + weight.size
        for i in range(n_dm):
            den[i,p0:p1] = make_rho(i, ao, mask, 'LDA')
            aow = numint._scale_ao(ao, scaled_weights[i,p0:p1], out=aow)
            vmat[i] -= numint._dot_ao_ao(mol, ao, aow, mask, shls_slice, ao_loc)
    den *= grids.weights
    ao = aow = scaled_weights = None

    nelec_leak = 0
    psi = numpy.zeros((n_dm, natm, nlm))
    i1 = 0
    for ia in range(natm):
        fak_pol, leak_idx = cached_pol[mol.atom_symbol(ia)]
        fac_pol = _vstack_factor_fak_pol(fak_pol, lmax)
        i0, i1 = i1, i1 + fac_pol.shape[1]
        nelec_leak += den[:,i0:i1][:,leak_idx].sum(axis=1)
        psi[:,ia] = -numpy.einsum('in,mn->im', den[:,i0:i1], fac_pol)
    logger.debug(pcmobj, 'electron leaks %s', nelec_leak)

# Contribution of nuclear charges to the total density
# The factor numpy.sqrt(4*numpy.pi) is due to the product of 4*pi * Y_0^0
    if with_nuc:
        for ia in range(natm):
            psi[:,ia,0] += numpy.sqrt(4*numpy.pi)/r_vdw[ia] * mol.atom_charge(ia)

    # <Psi, L^{-1}g> -> Psi = SL the adjoint equation to LX = g
    L_S = numpy.linalg.solve(L.reshape(natm*nlm,-1).T, psi.reshape(n_dm,-1).T)
    L_S = L_S.reshape(natm,nlm,n_dm).transpose(2,0,1)
    coords_1sph, weights_1sph = make_grids_one_sphere(pcmobj.lebedev_order)
    # JCP, 141, 184108, Eq (39)
    xi_jn = numpy.einsum('n,jn,xn,ijx->ijn', weights_1sph, ui, ylm_1sph, L_S)
    extern_point_idx = ui > 0
    cav_coords = (mol.atom_coords().reshape(natm,1,3)
                  + numpy.einsum('r,gx->rgx', r_vdw, coords_1sph))
    cav_coords = cav_coords[extern_point_idx]
    xi_jn = xi_jn[:,extern_point_idx]

    max_memory = pcmobj.max_memory - lib.current_memory()[0]
    blksize = int(max(max_memory*.9e6/8/nao**2, 400))

    cintopt = gto.moleintor.make_cintopt(mol._atm, mol._bas, mol._env, 'int3c2e')
    vmat_tril = 0
    for i0, i1 in lib.prange(0, cav_coords.shape[0], blksize):
        fakemol = gto.fakemol_for_charges(cav_coords[i0:i1])
        v_nj = df.incore.aux_e2(mol, fakemol, intor='int3c2e', aosym='s2ij',
                                cintopt=cintopt)
        vmat_tril += numpy.einsum('xn,in->ix', v_nj, xi_jn[:,i0:i1])
    vmat += lib.unpack_tril(vmat_tril)

    if is_single_dm:
        psi = psi[0]
        L_S = L_S[0]
        vmat = vmat[0]
    return psi, vmat, L_S

def cache_fake_multipoles(grids, r_vdw, lmax):
    # For each type of atoms, cache the product of last two terms in
    # JCP, 141, 184108, Eq (31):
    # x_{<}^{l} / x_{>}^{l+1} Y_l^m
    mol = grids.mol
    atom_grids_tab = grids.gen_atomic_grids(mol)
    r_vdw_type = {}
    for ia in range(mol.natm):
        symb = mol.atom_symbol(ia)
        if symb not in r_vdw_type:
            r_vdw_type[symb] = r_vdw[ia]

    cached_pol = {}
    for symb in atom_grids_tab:
        x_nj, w = atom_grids_tab[symb]
        r = lib.norm(x_nj, axis=1)
        # Different equations are used in JCTC, 9, 3637. r*Ys (the fake_pole)
        # is computed as r^l/r_vdw. "leak_idx" is not needed.
        # Here, the implementation is based on JCP, 141, 184108
        leak_idx = r > r_vdw_type[symb]

        pol = sph.multipoles(x_nj, lmax)
        fak_pol = []
        for l in range(lmax+1):
            # x_{<}^{l} / x_{>}^{l+1} Y_l^m  in JCP, 141, 184108, Eq (31)
            #:Ys = sph.real_sph_vec(x_nj/r.reshape(-1,1), lmax, True)
            #:rr = numpy.zeros_like(r)
            #:rr[r<=r_vdw[ia]] = r[r<=r_vdw[ia]]**l / r_vdw[ia]**(l+1)
            #:rr[r> r_vdw[ia]] = r_vdw[ia]**l / r[r>r_vdw[ia]]**(l+1)
            #:xx_ylm = numpy.einsum('n,mn->mn', rr, Ys[l])
            xx_ylm = pol[l] * (1./r_vdw_type[symb]**(l+1))
            # The line below is not needed for JCTC, 9, 3637
            xx_ylm[:,leak_idx] *= (r_vdw_type[symb]/r[leak_idx])**(2*l+1)
            fak_pol.append(xx_ylm)
        cached_pol[symb] = (fak_pol, leak_idx)
    return cached_pol

def _vstack_factor_fak_pol(fak_pol, lmax):
    fac_pol = []
    for l in range(lmax+1):
        fac = 4*numpy.pi/(l*2+1)
        fac_pol.append(fac * fak_pol[l])
    return numpy.vstack(fac_pol)

def atoms_with_vdw_overlap(atm_id, atom_coords, r_vdw):
    atm_dist = atom_coords - atom_coords[atm_id]
    atm_dist = numpy.einsum('pi,pi->p', atm_dist, atm_dist)
    atm_dist[atm_id] = 1e200
    vdw_sum = r_vdw + r_vdw[atm_id]
    atoms_nearby = numpy.where(atm_dist < vdw_sum**2)[0]
    return atoms_nearby

class DDCOSMO(lib.StreamObject):
    def __init__(self, mol):
        self.mol = mol
        self.stdout = mol.stdout
        self.verbose = mol.verbose
        self.max_memory = mol.max_memory

        #self.radii_table = radii.VDW
        self.radii_table = radii.UFF*1.1
        #self.radii_table = radii.MM3
        self.atom_radii = None
        self.lebedev_order = 17
        self.lmax = 6  # max angular momentum of spherical harmonics basis
        self.eta = .1  # regularization parameter
        self.eps = 78.3553
        self.grids = gen_grid.Grids(mol)

        # The maximum iterations and convergence tolerance to update solvent
        # effects in CASCI, CC, MP, CI, ... methods
        self.max_cycle = 20
        self.conv_tol = 1e-7
        self.state_id = 0

        # Set frozen to enable/disable the frozen ddCOSMO solvent potential.
        # If frozen is set, _dm (density matrix) needs to be specified to
        # generate the potential.
        self.frozen = False

        # In the rapid process (such as vertical excitation), solvent does not
        # follow the fast change of electronic structure of solutes. A
        # calculation of non-equilibrium solvation should be used. For slow
        # process (like geometry optimization), solvent has enough time to
        # respond to the changes in electronic structure or geometry of
        # solutes. Equilibrium solvation should be enabled in the calculation.
        # See for example JPCA, 104, 5631 (2000)
        #
        # Note this attribute has no effects if .frozen is enabled.
        #
        self.equilibrium_solvation = False

##################################################
# don't modify the following attributes, they are not input options
        # e (the dielectric correction) and v (the additional potential) are
        # updated during the SCF iterations
        self.e = None
        self.v = None
        self._dm = None

        self._intermediates = None
        self._keys = set(self.__dict__.keys())

    @property
    def dm(self):
        '''Density matrix to generate the frozen ddCOSMO solvent potential.'''
        return self._dm
    @dm.setter
    def dm(self, dm):
        '''Set dm to enable/disable the frozen ddCOSMO solvent potential.
        Setting dm to None will disable the frozen potental, i.e. the
        potential will respond to the change of the density during SCF
        iterations.
        '''
        if isinstance(dm, numpy.ndarray):
            self._dm = dm
            self.e, self.v = self.kernel(dm)
        else:
            self.e = self.v = self._dm = None

    # define epcm and vpcm for backward compatibility
    @property
    def epcm(self):
        return self.e_solvent
    @epcm.setter
    def epcm(self, val):
        self.e_solvent = val

    @property
    def vpcm(self):
        return self.v_solvent
    @vpcm.setter
    def vpcm(self, val):
        self.v_solvent = val

    def __setattr__(self, key, val):
        if key in ('radii_table', 'atom_radii', 'lebedev_order', 'lmax',
                   'eta', 'eps', 'grids'):
            self.reset()
        super(DDCOSMO, self).__setattr__(key, val)

    def dump_flags(self, verbose=None):
        logger.info(self, '******** %s ********', self.__class__)
        logger.info(self, 'lebedev_order = %s (%d grids per sphere)',
                    self.lebedev_order, gen_grid.LEBEDEV_ORDER[self.lebedev_order])
        logger.info(self, 'lmax = %s'         , self.lmax)
        logger.info(self, 'eta = %s'          , self.eta)
        logger.info(self, 'eps = %s'          , self.eps)
        logger.info(self, 'frozen = %s'       , self.frozen)
        logger.info(self, 'equilibrium_solvation = %s', self.equilibrium_solvation)
        logger.debug2(self, 'radii_table %s', self.radii_table)
        if self.atom_radii:
            logger.info(self, 'User specified atomic radii %s', str(self.atom_radii))
        self.grids.dump_flags(verbose)
        return self

    # TODO: Testing the value of psi (make_psi_vmat).  All intermediates except
    # psi are tested against ddPCM implementation on github. Psi needs to be
    # computed by the host program. It requires the numerical integration code.
    def build(self):
        if self.grids.coords is None:
            self.grids.build(with_non0tab=True)

        mol = self.mol
        natm = mol.natm
        lmax = self.lmax

        r_vdw = self.get_atomic_radii()
        coords_1sph, weights_1sph = make_grids_one_sphere(self.lebedev_order)
        ylm_1sph = numpy.vstack(sph.real_sph_vec(coords_1sph, lmax, True))

        fi = make_fi(self, r_vdw)
        ui = 1 - fi
        ui[ui<0] = 0
        nexposed = numpy.count_nonzero(ui==1)
        nbury = numpy.count_nonzero(ui==0)
        on_shell = numpy.count_nonzero(ui>0) - nexposed
        logger.debug(self, 'Num points exposed %d', nexposed)
        logger.debug(self, 'Num points buried %d', nbury)
        logger.debug(self, 'Num points on shell %d', on_shell)

        nlm = (lmax+1)**2
        Lmat = make_L(self, r_vdw, ylm_1sph, fi)
        Lmat = Lmat.reshape(natm*nlm,-1)

        cached_pol = cache_fake_multipoles(self.grids, r_vdw, lmax)

        self._intermediates = {
            'r_vdw': r_vdw,
            'ylm_1sph': ylm_1sph,
            'ui': ui,
            'Lmat': Lmat,
            'cached_pol': cached_pol,
        }

    def kernel(self, dm):
        '''A single shot solvent effects for given density matrix.
        '''
        self._dm = dm
        self.e, self.v = self._get_vind(dm)
        logger.info(self, '%s E_diel = %.15g', self.__class__, self.e)
        return self.e, self.v

    def reset(self, mol=None):
        '''Reset mol and clean up relevant attributes for scanner mode'''
        if mol is not None:
            self.mol = mol
            self.grids.reset(mol)
        self._intermediates = None
        return self

    def _get_vind(self, dm):
        '''A single shot solvent effects for given density matrix.
        '''
        if not self._intermediates or self.grids.coords is None:
            self.build()

        mol = self.mol
        r_vdw      = self._intermediates['r_vdw'     ]
        ylm_1sph   = self._intermediates['ylm_1sph'  ]
        ui         = self._intermediates['ui'        ]
        Lmat       = self._intermediates['Lmat'      ]
        cached_pol = self._intermediates['cached_pol']

        if not (isinstance(dm, numpy.ndarray) and dm.ndim == 2):
            # spin-traced DM for UHF or ROHF
            dm = dm[0] + dm[1]

        phi = make_phi(self, dm, r_vdw, ui, ylm_1sph)
        Xvec = numpy.linalg.solve(Lmat, phi.ravel()).reshape(mol.natm,-1)
        psi, vmat = make_psi_vmat(self, dm, r_vdw, ui, ylm_1sph,
                                  cached_pol, Xvec, Lmat)[:2]
        dielectric = self.eps
        if dielectric > 0:
            f_epsilon = (dielectric-1.)/dielectric
        else:
            f_epsilon = 1
        epcm = .5 * f_epsilon * numpy.einsum('jx,jx', psi, Xvec)
        vpcm = .5 * f_epsilon * vmat
        return epcm, vpcm

    def _B_dot_x(self, dm):
        '''
        Compute the matrix-vector product B * x. The B matrix, as defined in
        the paper R. Cammi, JPCA, 104, 5631 (2000), is the second order
        derivatives of E_solvation wrt density matrices.

        Note: In ddCOSMO, strictly, B is not symmetric. To make it compatible
        with the CIS framework, it is symmetrized in current implementation.
        '''
        if not self._intermediates or self.grids.coords is None:
            self.build()

        mol = self.mol
        r_vdw      = self._intermediates['r_vdw'     ]
        ylm_1sph   = self._intermediates['ylm_1sph'  ]
        ui         = self._intermediates['ui'        ]
        Lmat       = self._intermediates['Lmat'      ]
        cached_pol = self._intermediates['cached_pol']
        natm = mol.natm
        nlm = (self.lmax+1)**2

        dms = numpy.asarray(dm)
        dm_shape = dms.shape
        nao = dm_shape[-1]
        dms = dms.reshape(-1,nao,nao)

        phi = make_phi(self, dms, r_vdw, ui, ylm_1sph, with_nuc=False)
        Xvec = numpy.linalg.solve(Lmat, phi.reshape(-1,natm*nlm).T)
        Xvec = Xvec.reshape(natm,nlm,-1).transpose(2,0,1)
        vmat = make_psi_vmat(self, dms, r_vdw, ui, ylm_1sph,
                             cached_pol, Xvec, Lmat, with_nuc=False)[1]
        dielectric = self.eps
        if dielectric > 0:
            f_epsilon = (dielectric-1.)/dielectric
        else:
            f_epsilon = 1
        return .5 * f_epsilon * vmat.reshape(dm_shape)

    energy = energy
    gen_solver = as_solver = gen_ddcosmo_solver
    get_atomic_radii = get_atomic_radii

    def regularize_xt(self, t, eta, scale=1):
        # scale = eta*scale, is it correct?
        return regularize_xt(t, eta*scale)

    def nuc_grad_method(self, grad_method):
        '''For grad_method in vacuum, add nuclear gradients of solvent
        '''
        from pyscf import tdscf
        from pyscf.solvent import ddcosmo_grad, _ddcosmo_tdscf_grad
        if self.frozen:
            raise RuntimeError('Frozen solvent model is not supported for '
                               'energy gradients')
        if isinstance(grad_method.base, tdscf.rhf.TDMixin):
            return _ddcosmo_tdscf_grad.make_grad_object(grad_method)
        else:
            return ddcosmo_grad.make_grad_object(grad_method)


if __name__ == '__main__':
    from pyscf import scf
    from pyscf import mcscf
    from pyscf import cc
    mol = gto.M(atom='H 0 0 0; H 0 1 1.2; H 1. .1 0; H .5 .5 1')
    natm = mol.natm
    r_vdw = [radii.VDW[gto.charge(mol.atom_symbol(i))]
             for i in range(natm)]
    r_vdw = numpy.asarray(r_vdw)
    pcmobj = DDCOSMO(mol)
    pcmobj.regularize_xt = lambda t, eta, scale: regularize_xt(t, eta)
    pcmobj.lebedev_order = 7
    pcmobj.lmax = 6
    pcmobj.eta = 0.1
    nlm = (pcmobj.lmax+1)**2
    coords_1sph, weights_1sph = make_grids_one_sphere(pcmobj.lebedev_order)
    fi = make_fi(pcmobj, r_vdw)
    ylm_1sph = numpy.vstack(sph.real_sph_vec(coords_1sph, pcmobj.lmax, True))
    L = make_L(pcmobj, r_vdw, ylm_1sph, fi)
    print(lib.fp(L) - 6.2823493771037473)

    mol = gto.Mole()
    mol.atom = ''' O                  0.00000000    0.00000000   -0.11081188
                   H                 -0.00000000   -0.84695236    0.59109389
                   H                 -0.00000000    0.89830571    0.52404783 '''
    mol.basis = '3-21g' #cc-pvdz'
    mol.build()
    cm = DDCOSMO(mol)
    cm.verbose = 4
    mf = ddcosmo_for_scf(scf.RHF(mol), cm)#.newton()
    mf.verbose = 4
    print(mf.kernel() - -75.570364368059)
    cm.verbose = 3
    e = ddcosmo_for_casci(mcscf.CASCI(mf, 4, 4)).kernel()[0]
    print(e - -75.5743583693215)
    cc_cosmo = ddcosmo_for_post_scf(cc.CCSD(mf)).run()
    print(cc_cosmo.e_tot - -75.70961637250134)

    mol = gto.Mole()
    mol.atom = ''' Fe                 0.00000000    0.00000000   -0.11081188
                   H                 -0.00000000   -0.84695236    0.59109389
                   H                 -0.00000000    0.89830571    0.52404783 '''
    mol.basis = '3-21g' #cc-pvdz'
    mol.build()
    cm = DDCOSMO(mol)
    cm.eps = -1
    cm.verbose = 4
    mf = ddcosmo_for_scf(scf.ROHF(mol), cm).newton()
    mf.verbose=4
    mf.kernel()


# pyscf/solvent/ddcosmo_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.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Analytical nuclear gradients for domain decomposition COSMO

See also

[1] Fast Domain Decomposition Algorithm for Continuum Solvation Models: Energy and First Derivatives.
F. Lipparini, B. Stamm, E. Cances, Y. Maday, B. Mennucci
J. Chem. Theory Comput., 9, 3637-3648 (2013)
http://dx.doi.org/10.1021/ct400280b

[2] Quantum, classical, and hybrid QM/MM calculations in solution: General implementation of the ddCOSMO linear scaling strategy.
F. Lipparini, G. Scalmani, L. Lagardere, B. Stamm, E. Cances, Y. Maday, J.-P.Piquemal, M. J. Frisch, B. Mennucci
J. Chem. Phys., 141, 184108 (2014)
http://dx.doi.org/10.1063/1.4901304
'''  # noqa: E501

import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf import gto
from pyscf import df
from pyscf.dft import gen_grid, numint
from pyscf.symm import sph
from pyscf.solvent import ddcosmo
from pyscf.solvent._attach_solvent import _Solvation
from pyscf.grad import rhf as rhf_grad
from pyscf.grad import rks as rks_grad
from pyscf.grad import tdrhf as tdrhf_grad  # noqa


# TODO: Define attribute grad_method.base to point to the class of the 0th
# order calculation for all gradients class. Then this function can be
# extended and used as the general interface to initialize solvent gradients.
def make_grad_object(grad_method):
    '''For grad_method in vacuum, add nuclear gradients of solvent pcmobj'''

    # Zeroth order method object must be a solvation-enabled method
    assert isinstance(grad_method.base, _Solvation)
    if grad_method.base.with_solvent.frozen:
        raise RuntimeError('Frozen solvent model is not avialbe for energy gradients')

    grad_method_class = grad_method.__class__
    class WithSolventGrad(grad_method_class):
        def __init__(self, grad_method):
            self.__dict__.update(grad_method.__dict__)
            self.de_solvent = None
            self.de_solute = None
            self._keys = self._keys.union(['de_solvent', 'de_solute'])

        # TODO: if moving to python3, change signature to
        # def kernel(self, *args, dm=None, atmlst=None, **kwargs):
        def kernel(self, *args, **kwargs):
            dm = kwargs.pop('dm', None)
            if dm is None:
                dm = self.base.make_rdm1(ao_repr=True)

            self.de_solvent = kernel(self.base.with_solvent, dm)
            self.de_solute = grad_method_class.kernel(self, *args, **kwargs)
            self.de = self.de_solute + self.de_solvent

            if self.verbose >= logger.NOTE:
                logger.note(self, '--------------- %s (+%s) gradients ---------------',
                            self.base.__class__.__name__,
                            self.base.with_solvent.__class__.__name__)
                rhf_grad._write(self, self.mol, self.de, self.atmlst)
                logger.note(self, '----------------------------------------------')
            return self.de

        def _finalize(self):
            # disable _finalize. It is called in grad_method.kernel method
            # where self.de was not yet initialized.
            pass

    return WithSolventGrad(grad_method)


def kernel(pcmobj, dm, verbose=None):
    mol = pcmobj.mol
    natm = mol.natm
    lmax = pcmobj.lmax
    if pcmobj.grids.coords is None:
        pcmobj.grids.build(with_non0tab=True)

    if not (isinstance(dm, numpy.ndarray) and dm.ndim == 2):
        # UHF density matrix
        dm = dm[0] + dm[1]

    r_vdw = ddcosmo.get_atomic_radii(pcmobj)
    coords_1sph, weights_1sph = ddcosmo.make_grids_one_sphere(pcmobj.lebedev_order)
    ylm_1sph = numpy.vstack(sph.real_sph_vec(coords_1sph, lmax, True))

    fi = ddcosmo.make_fi(pcmobj, r_vdw)
    ui = 1 - fi
    ui[ui<0] = 0

    cached_pol = ddcosmo.cache_fake_multipoles(pcmobj.grids, r_vdw, lmax)

    nlm = (lmax+1)**2
    L0 = ddcosmo.make_L(pcmobj, r_vdw, ylm_1sph, fi)
    L0 = L0.reshape(natm*nlm,-1)
    L1 = make_L1(pcmobj, r_vdw, ylm_1sph, fi)

    phi0 = ddcosmo.make_phi(pcmobj, dm, r_vdw, ui, ylm_1sph)
    phi1 = make_phi1(pcmobj, dm, r_vdw, ui, ylm_1sph)
    L0_X = numpy.linalg.solve(L0, phi0.ravel()).reshape(natm,-1)
    psi0, vmat, L0_S = \
            ddcosmo.make_psi_vmat(pcmobj, dm, r_vdw, ui, ylm_1sph,
                                  cached_pol, L0_X, L0)
    e_psi1 = make_e_psi1(pcmobj, dm, r_vdw, ui, ylm_1sph,
                         cached_pol, L0_X, L0)
    dielectric = pcmobj.eps
    if dielectric > 0:
        f_epsilon = (dielectric-1.)/dielectric
    else:
        f_epsilon = 1
    de = .5 * f_epsilon * e_psi1
    de+= .5 * f_epsilon * numpy.einsum('jx,azjx->az', L0_S, phi1)
    de-= .5 * f_epsilon * numpy.einsum('aziljm,il,jm->az', L1, L0_S, L0_X)
    return de

def make_L1(pcmobj, r_vdw, ylm_1sph, fi):
    # See JCTC, 9, 3637, Eq (18)
    mol = pcmobj.mol
    natm = mol.natm
    lmax = pcmobj.lmax
    eta = pcmobj.eta
    nlm = (lmax+1)**2

    coords_1sph, weights_1sph = ddcosmo.make_grids_one_sphere(pcmobj.lebedev_order)
    ngrid_1sph = weights_1sph.size
    atom_coords = mol.atom_coords()
    ylm_1sph = ylm_1sph.reshape(nlm,ngrid_1sph)

    Lmat = numpy.zeros((natm,3,natm,nlm,natm,nlm))
    fi1 = make_fi1(pcmobj, pcmobj.get_atomic_radii())

    for ja in range(natm):
        part_weights = weights_1sph.copy()
        part_weights[fi[ja]>1] /= fi[ja,fi[ja]>1]

        part_weights1 = numpy.zeros((natm,3,ngrid_1sph))
        tmp = part_weights[fi[ja]>1] / fi[ja,fi[ja]>1]
        part_weights1[:,:,fi[ja]>1] = -tmp * fi1[:,:,ja,fi[ja]>1]

        for ka in ddcosmo.atoms_with_vdw_overlap(ja, atom_coords, r_vdw):
            vjk = r_vdw[ja] * coords_1sph + atom_coords[ja] - atom_coords[ka]
            rv = lib.norm(vjk, axis=1)
            tjk = rv / r_vdw[ka]
            wjk0 = pcmobj.regularize_xt(tjk, eta, r_vdw[ka])
            wjk1 = regularize_xt1(tjk, eta*r_vdw[ka])
            sjk = vjk.T / rv
            wjk1 = 1./r_vdw[ka] * wjk1 * sjk

            wjk01 = wjk0 * part_weights1
            wjk0 *= part_weights
            wjk1 *= part_weights

            pol0 = sph.multipoles(vjk, lmax)
            pol1 = multipoles1(vjk, lmax)
            p1 = 0
            for l in range(lmax+1):
                fac = 4*numpy.pi/(l*2+1) / r_vdw[ka]**(l+1)
                p0, p1 = p1, p1 + (l*2+1)
                a = numpy.einsum('xn,zn,mn->zxm', ylm_1sph, wjk1, pol0[l])
                a+= numpy.einsum('xn,n,zmn->zxm', ylm_1sph, wjk0, pol1[l])
                Lmat[ja,:,ja,:,ka,p0:p1] += -fac * a
                Lmat[ka,:,ja,:,ka,p0:p1] -= -fac * a
                a = numpy.einsum('xn,azn,mn->azxm', ylm_1sph, wjk01, pol0[l])
                Lmat[:,:,ja,:,ka,p0:p1] += -fac * a
    return Lmat


def multipoles1(r, lmax, reorder_dipole=True):
    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,3,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,0] = lx * xs[lx-1] * ys[ly] * zs[lz]
                c[k,1] = ly * xs[lx] * ys[ly-1] * zs[lz]
                c[k,2] = lz * xs[lx] * ys[ly] * zs[lz-1]
                k += 1
        ylm = gto.cart2sph(l, c.reshape(nd,3*ngrid).T)
        ylm = ylm.reshape(3,ngrid,l*2+1).transpose(0,2,1)
        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 regularize_xt1(t, eta):
    xt = numpy.zeros_like(t)
    # no response if grids are inside the cavity
    # inner = t <= 1-eta
    # xt[inner] = 0
    on_shell = (1-eta < t) & (t < 1)
    ti = t[on_shell]
    xt[on_shell] = -30./eta**5 * (1-ti)**2 * (1-eta-ti)**2
    return xt

def make_fi1(pcmobj, r_vdw):
    coords_1sph, weights_1sph = ddcosmo.make_grids_one_sphere(pcmobj.lebedev_order)
    mol = pcmobj.mol
    eta = pcmobj.eta
    natm = mol.natm
    atom_coords = mol.atom_coords()
    ngrid_1sph = coords_1sph.shape[0]
    fi1 = numpy.zeros((natm,3,natm,ngrid_1sph))
    for ia in range(natm):
        for ja in ddcosmo.atoms_with_vdw_overlap(ia, atom_coords, r_vdw):
            v = r_vdw[ia]*coords_1sph + atom_coords[ia] - atom_coords[ja]
            rv = lib.norm(v, axis=1)
            t = rv / r_vdw[ja]
            xt1 = regularize_xt1(t, eta*r_vdw[ja])
            s_ij = v.T / rv
            xt1 = 1./r_vdw[ja] * xt1 * s_ij
            fi1[ia,:,ia] += xt1
            fi1[ja,:,ia] -= xt1

    fi = ddcosmo.make_fi(pcmobj, r_vdw)
    fi1[:,:,fi<1e-20] = 0
    return fi1

def make_phi1(pcmobj, dm, r_vdw, ui, ylm_1sph):
    mol = pcmobj.mol
    natm = mol.natm

    if not (isinstance(dm, numpy.ndarray) and dm.ndim == 2):
        dm = dm[0] + dm[1]
    tril_dm = lib.pack_tril(dm+dm.T)
    nao = dm.shape[0]
    diagidx = numpy.arange(nao)
    diagidx = diagidx*(diagidx+1)//2 + diagidx
    tril_dm[diagidx] *= .5

    atom_coords = mol.atom_coords()
    atom_charges = mol.atom_charges()

    coords_1sph, weights_1sph = ddcosmo.make_grids_one_sphere(pcmobj.lebedev_order)
    #extern_point_idx = ui > 0

    fi1 = make_fi1(pcmobj, pcmobj.get_atomic_radii())
    fi1[:,:,ui==0] = 0
    ui1 = -fi1

    ngrid_1sph = weights_1sph.size
    v_phi0 = numpy.empty((natm,ngrid_1sph))
    for ia in range(natm):
        cav_coords = atom_coords[ia] + r_vdw[ia] * coords_1sph
        d_rs = atom_coords.reshape(-1,1,3) - cav_coords
        v_phi0[ia] = numpy.einsum('z,zp->p', atom_charges, 1./lib.norm(d_rs,axis=2))
    phi1 = -numpy.einsum('n,ln,azjn,jn->azjl', weights_1sph, ylm_1sph, ui1, v_phi0)

    for ia in range(natm):
        cav_coords = atom_coords[ia] + r_vdw[ia] * coords_1sph
        for ja in range(natm):
            rs = atom_coords[ja] - cav_coords
            d_rs = lib.norm(rs, axis=1)
            v_phi = atom_charges[ja] * numpy.einsum('px,p->px', rs, 1./d_rs**3)
            tmp = numpy.einsum('n,ln,n,nx->xl', weights_1sph, ylm_1sph, ui[ia], v_phi)
            phi1[ja,:,ia] += tmp  # response of the other atoms
            phi1[ia,:,ia] -= tmp  # response of cavity grids

    int3c2e = mol._add_suffix('int3c2e')
    int3c2e_ip1 = mol._add_suffix('int3c2e_ip1')
    aoslices = mol.aoslice_by_atom()
    for ia in range(natm):
        cav_coords = atom_coords[ia] + r_vdw[ia] * coords_1sph
        #fakemol = gto.fakemol_for_charges(cav_coords[ui[ia]>0])
        fakemol = gto.fakemol_for_charges(cav_coords)
        v_nj = df.incore.aux_e2(mol, fakemol, intor=int3c2e, aosym='s1')
        v_phi = numpy.einsum('ij,ijk->k', dm, v_nj)
        phi1[:,:,ia] += numpy.einsum('n,ln,azn,n->azl', weights_1sph, ylm_1sph, ui1[:,:,ia], v_phi)

        v_e1_nj = df.incore.aux_e2(mol, fakemol, intor=int3c2e_ip1, comp=3, aosym='s1')
        phi1_e2_nj  = numpy.einsum('ij,xijr->xr', dm, v_e1_nj)
        phi1_e2_nj += numpy.einsum('ji,xijr->xr', dm, v_e1_nj)
        phi1[ia,:,ia] += numpy.einsum('n,ln,n,xn->xl', weights_1sph, ylm_1sph, ui[ia], phi1_e2_nj)

        for ja in range(natm):
            shl0, shl1, p0, p1 = aoslices[ja]
            phi1_nj  = numpy.einsum('ij,xijr->xr', dm[p0:p1  ], v_e1_nj[:,p0:p1])
            phi1_nj += numpy.einsum('ji,xijr->xr', dm[:,p0:p1], v_e1_nj[:,p0:p1])
            phi1[ja,:,ia] -= numpy.einsum('n,ln,n,xn->xl', weights_1sph, ylm_1sph, ui[ia], phi1_nj)
    return phi1

def make_e_psi1(pcmobj, dm, r_vdw, ui, ylm_1sph, cached_pol, Xvec, L):
    mol = pcmobj.mol
    natm = mol.natm
    lmax = pcmobj.lmax
    grids = pcmobj.grids

    if not (isinstance(dm, numpy.ndarray) and dm.ndim == 2):
        dm = dm[0] + dm[1]
    ni = numint.NumInt()
    make_rho, nset, nao = ni._gen_rho_evaluator(mol, dm)
    den = numpy.empty((4,grids.weights.size))

    ao_loc = mol.ao_loc_nr()
    vmat = numpy.zeros((3,nao,nao))
    psi1 = numpy.zeros((natm,3))
    i1 = 0
    for ia, (coords, weight, weight1) in enumerate(rks_grad.grids_response_cc(grids)):
        i0, i1 = i1, i1 + weight.size
        ao = ni.eval_ao(mol, coords, deriv=1)
        mask = gen_grid.make_mask(mol, coords)
        den[:,i0:i1] = make_rho(0, ao, mask, 'GGA')

        fak_pol, leak_idx = cached_pol[mol.atom_symbol(ia)]
        eta_nj = 0
        p1 = 0
        for l in range(lmax+1):
            fac = 4*numpy.pi/(l*2+1)
            p0, p1 = p1, p1 + (l*2+1)
            eta_nj += fac * numpy.einsum('mn,m->n', fak_pol[l], Xvec[ia,p0:p1])
        psi1 -= numpy.einsum('n,n,zxn->zx', den[0,i0:i1], eta_nj, weight1)
        psi1[ia] -= numpy.einsum('xn,n,n->x', den[1:4,i0:i1], eta_nj, weight)

        vtmp = numpy.zeros((3,nao,nao))
        aow = numpy.einsum('pi,p->pi', ao[0], weight*eta_nj)
        rks_grad._d1_dot_(vtmp, mol, ao[1:4], aow, mask, ao_loc, True)
        vmat += vtmp

    aoslices = mol.aoslice_by_atom()
    for ia in range(natm):
        shl0, shl1, p0, p1 = aoslices[ia]
        psi1[ia] += numpy.einsum('xij,ij->x', vmat[:,p0:p1], dm[p0:p1]) * 2
    return psi1


if __name__ == '__main__':
    from pyscf import scf
    mol = gto.M(atom='H 0 0 0; H 0 1 1.2; H 1. .1 0; H .5 .5 1', unit='B')
    mf = ddcosmo.ddcosmo_for_scf(scf.RHF(mol))
    mf.kernel()
    de = mf.nuc_grad_method().kernel()
    de_cosmo = kernel(mf.with_solvent, mf.make_rdm1())
    dm1 = mf.make_rdm1()

    mol = gto.M(atom='H 0 0 -0.001; H 0 1 1.2; H 1. .1 0; H .5 .5 1', unit='B')
    mf = ddcosmo.ddcosmo_for_scf(scf.RHF(mol))
    e1 = mf.kernel()
    e1_cosmo = mf.with_solvent.energy(dm1)

    mol = gto.M(atom='H 0 0 0.001; H 0 1 1.2; H 1. .1 0; H .5 .5 1', unit='B')
    mf = ddcosmo.ddcosmo_for_scf(scf.RHF(mol))
    e2 = mf.kernel()
    e2_cosmo = mf.with_solvent.energy(dm1)
    print(abs((e2-e1)/0.002 - de[0,2]).max())
    print(abs((e2_cosmo-e1_cosmo)/0.002 - de_cosmo[0,2]).max())


# pyscf/solvent/ddpcm.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>
#

'''
domain decomposition PCM (In testing)

See also
JCP, 144, 054101
JCP, 144, 160901
'''

import warnings
import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf import gto
from pyscf.dft import gen_grid
from pyscf.data import radii
from pyscf.solvent import ddcosmo
from pyscf.symm import sph
from pyscf.solvent import _attach_solvent

warnings.warn('Module ddPCM is under testing')


@lib.with_doc(_attach_solvent._for_scf.__doc__)
def ddpcm_for_scf(mf, solvent_obj=None, dm=None):
    if solvent_obj is None:
        solvent_obj = DDPCM(mf.mol)
    return _attach_solvent._for_scf(mf, solvent_obj, dm)

@lib.with_doc(_attach_solvent._for_casscf.__doc__)
def ddpcm_for_casscf(mc, solvent_obj=None, dm=None):
    if solvent_obj is None:
        if isinstance(getattr(mc._scf, 'with_solvent', None), DDPCM):
            solvent_obj = mc._scf.with_solvent
        else:
            solvent_obj = DDPCM(mc.mol)
    return _attach_solvent._for_casscf(mc, solvent_obj, dm)

@lib.with_doc(_attach_solvent._for_casci.__doc__)
def ddpcm_for_casci(mc, solvent_obj=None, dm=None):
    if solvent_obj is None:
        if isinstance(getattr(mc._scf, 'with_solvent', None), DDPCM):
            solvent_obj = mc._scf.with_solvent
        else:
            solvent_obj = DDPCM(mc.mol)
    return _attach_solvent._for_casci(mc, solvent_obj, dm)

@lib.with_doc(_attach_solvent._for_post_scf.__doc__)
def ddpcm_for_post_scf(method, solvent_obj=None, dm=None):
    if solvent_obj is None:
        if isinstance(getattr(method._scf, 'with_solvent', None), DDPCM):
            solvent_obj = method._scf.with_solvent
        else:
            solvent_obj = DDPCM(method.mol)
    return _attach_solvent._for_post_scf(method, solvent_obj, dm)

@lib.with_doc(_attach_solvent._for_tdscf.__doc__)
def ddpcm_for_tdscf(method, solvent_obj=None, dm=None):
    scf_solvent = getattr(method._scf, 'with_solvent', None)
    assert scf_solvent is None or isinstance(scf_solvent, DDPCM)

    if solvent_obj is None:
        solvent_obj = DDPCM(method.mol)
    return _attach_solvent._for_tdscf(method, solvent_obj, dm)


# Inject ddPCM to other methods
from pyscf import scf
from pyscf import mcscf
from pyscf import mp, ci, cc
from pyscf import tdscf
scf.hf.SCF.ddPCM    = scf.hf.SCF.DDPCM    = ddpcm_for_scf
mp.mp2.MP2.ddPCM    = mp.mp2.MP2.DDPCM    = ddpcm_for_post_scf
ci.cisd.CISD.ddPCM  = ci.cisd.CISD.DDPCM  = ddpcm_for_post_scf
cc.ccsd.CCSD.ddPCM  = cc.ccsd.CCSD.DDPCM  = ddpcm_for_post_scf
tdscf.rhf.TDMixin.ddPCM = tdscf.rhf.TDMixin.DDPCM = ddpcm_for_tdscf
mcscf.casci.CASCI.ddPCM = mcscf.casci.CASCI.DDPCM = ddpcm_for_casci
mcscf.mc1step.CASSCF.ddPCM = mcscf.mc1step.CASSCF.DDPCM = ddpcm_for_casscf

def gen_ddpcm_solver(pcmobj, verbose=None):
    mol = pcmobj.mol
    if pcmobj.grids.coords is None:
        pcmobj.grids.build(with_non0tab=True)

    natm = mol.natm
    lmax = pcmobj.lmax

    r_vdw = ddcosmo.get_atomic_radii(pcmobj)
    coords_1sph, weights_1sph = ddcosmo.make_grids_one_sphere(pcmobj.lebedev_order)
    ylm_1sph = numpy.vstack(sph.real_sph_vec(coords_1sph, lmax, True))

    fi = ddcosmo.make_fi(pcmobj, r_vdw)
    ui = 1 - fi
    ui[ui<0] = 0
    nexposed = numpy.count_nonzero(ui==1)
    nbury = numpy.count_nonzero(ui==0)
    on_shell = numpy.count_nonzero(ui>0) - nexposed
    logger.debug(pcmobj, 'Num points exposed %d', nexposed)
    logger.debug(pcmobj, 'Num points buried %d', nbury)
    logger.debug(pcmobj, 'Num points on shell %d', on_shell)

    nlm = (lmax+1)**2
    Lmat = ddcosmo.make_L(pcmobj, r_vdw, ylm_1sph, fi)
    Lmat = Lmat.reshape(natm*nlm,-1)

    Amat = make_A(pcmobj, r_vdw, ylm_1sph, ui).reshape(natm*nlm,-1)
    fac = 2*numpy.pi * (pcmobj.eps+1) / (pcmobj.eps-1)
    A_diele = Amat + fac * numpy.eye(natm*nlm)
    A_inf = Amat + 2*numpy.pi * numpy.eye(natm*nlm)

    cached_pol = ddcosmo.cache_fake_multipoles(pcmobj.grids, r_vdw, lmax)

    def gen_vind(dm):
        phi = ddcosmo.make_phi(pcmobj, dm, r_vdw, ui)
        phi = numpy.linalg.solve(A_diele, A_inf.dot(phi.ravel()))

        Xvec = numpy.linalg.solve(Lmat, phi.ravel()).reshape(natm,-1)
        psi, vmat = ddcosmo.make_psi_vmat(pcmobj, dm, r_vdw, ui, pcmobj.grids,
                                          ylm_1sph, cached_pol, Xvec, Lmat)[:2]
        dielectric = pcmobj.eps
        f_epsilon = (dielectric-1.)/dielectric
        epcm = .5 * f_epsilon * numpy.einsum('jx,jx', psi, Xvec)
        vpcm = .5 * f_epsilon * vmat
        return epcm, vpcm
    return gen_vind

def energy(pcmobj, dm):
    r'''
    ddPCM energy
    Es = 1/2 f(eps) \int rho(r) W(r) dr
    '''
    epcm = gen_ddpcm_solver(pcmobj, pcmobj.verbose)(dm)[0]
    return epcm

def regularize_xt(t, eta, scale=1):
    eta *= scale
    xt = numpy.zeros_like(t)
    inner = t <= 1-eta
    on_shell = (1-eta < t) & (t < 1)
    xt[inner] = 1
    ti = t[on_shell] - eta*.5
    # JCP, 144, 054101
    xt[on_shell] = 1./eta**4 * (1-ti)**2 * (ti-1+2*eta)**2
    return xt

def make_A(pcmobj, r_vdw, ylm_1sph, ui):
    # Part of A matrix defined in JCP, 144, 054101, Eq (43), (44)
    mol = pcmobj.mol
    natm = mol.natm
    lmax = pcmobj.lmax
    # eta = pcmobj.eta
    nlm = (lmax+1)**2

    coords_1sph, weights_1sph = ddcosmo.make_grids_one_sphere(pcmobj.lebedev_order)
    ngrid_1sph = weights_1sph.size
    atom_coords = mol.atom_coords()
    ylm_1sph = ylm_1sph.reshape(nlm,ngrid_1sph)
    Amat = numpy.zeros((natm,nlm,natm,nlm))

    for ja in range(natm):
        # w_u = precontract w_n U_j
        w_u = weights_1sph * ui[ja]
        p1 = 0
        for l in range(lmax+1):
            fac = 2*numpy.pi/(l*2+1)
            p0, p1 = p1, p1 + (l*2+1)
            a = numpy.einsum('xn,n,mn->xm', ylm_1sph, w_u, ylm_1sph[p0:p1])
            Amat[ja,:,ja,p0:p1] += -fac * a

        for ka in ddcosmo.atoms_with_vdw_overlap(ja, atom_coords, r_vdw):
            vjk = r_vdw[ja] * coords_1sph + atom_coords[ja] - atom_coords[ka]
            rjk = lib.norm(vjk, axis=1)
            pol = sph.multipoles(vjk, lmax)
            p1 = 0
            weights = w_u / rjk**(l*2+1)
            for l in range(lmax+1):
                fac = 4*numpy.pi*l/(l*2+1) * r_vdw[ka]**(l+1)
                p0, p1 = p1, p1 + (l*2+1)
                a = numpy.einsum('xn,n,mn->xm', ylm_1sph, weights, pol[l])
                Amat[ja,:,ka,p0:p1] += -fac * a
    return Amat

class DDPCM(ddcosmo.DDCOSMO):
    def __init__(self, mol):
        ddcosmo.DDCOSMO.__init__(self, mol)

    def dump_flags(self, verbose=None):
        logger.info(self, '******** %s (In testing) ********', self.__class__)
        logger.warn(self, 'ddPCM is an experimental feature. It is '
                    'still in testing.\nFeatures and APIs may be changed '
                    'in the future.')
        logger.info(self, 'lebedev_order = %s (%d grids per sphere)',
                    self.lebedev_order, gen_grid.LEBEDEV_ORDER[self.lebedev_order])
        logger.info(self, 'lmax = %s'         , self.lmax)
        logger.info(self, 'eta = %s'          , self.eta)
        logger.info(self, 'eps = %s'          , self.eps)
        logger.info(self, 'frozen = %s'       , self.frozen)
        logger.info(self, 'equilibrium_solvation = %s', self.equilibrium_solvation)
        logger.debug2(self, 'radii_table %s', self.radii_table)
        if self.atom_radii:
            logger.info(self, 'User specified atomic radii %s', str(self.atom_radii))
        self.grids.dump_flags(verbose)
        return self

    def build(self):
        if self.grids.coords is None:
            self.grids.build(with_non0tab=True)

        mol = self.mol
        natm = mol.natm
        lmax = self.lmax

        r_vdw = ddcosmo.get_atomic_radii(self)
        coords_1sph, weights_1sph = ddcosmo.make_grids_one_sphere(self.lebedev_order)
        ylm_1sph = numpy.vstack(sph.real_sph_vec(coords_1sph, lmax, True))

        fi = ddcosmo.make_fi(self, r_vdw)
        ui = 1 - fi
        ui[ui<0] = 0
        nexposed = numpy.count_nonzero(ui==1)
        nbury = numpy.count_nonzero(ui==0)
        on_shell = numpy.count_nonzero(ui>0) - nexposed
        logger.debug(self, 'Num points exposed %d', nexposed)
        logger.debug(self, 'Num points buried %d', nbury)
        logger.debug(self, 'Num points on shell %d', on_shell)

        nlm = (lmax+1)**2
        Lmat = ddcosmo.make_L(self, r_vdw, ylm_1sph, fi)
        Lmat = Lmat.reshape(natm*nlm,-1)

        Amat = make_A(self, r_vdw, ylm_1sph, ui).reshape(natm*nlm,-1)
        fac = 2*numpy.pi * (self.eps+1) / (self.eps-1)
        A_diele = Amat + fac * numpy.eye(natm*nlm)
        A_inf = Amat + 2*numpy.pi * numpy.eye(natm*nlm)

        cached_pol = ddcosmo.cache_fake_multipoles(self.grids, r_vdw, lmax)

        self._intermediates = {
            'r_vdw': r_vdw,
            'ylm_1sph': ylm_1sph,
            'ui': ui,
            'Lmat': Lmat,
            'A_diele': A_diele,
            'A_inf': A_inf,
            'cached_pol': cached_pol,
        }

    def _get_vind(self, dm):
        '''A single shot solvent effects for given density matrix.
        '''
        if not self._intermediates or self.grids.coords is None:
            self.build()

        mol = self.mol
        r_vdw      = self._intermediates['r_vdw'     ]
        ylm_1sph   = self._intermediates['ylm_1sph'  ]
        ui         = self._intermediates['ui'        ]
        Lmat       = self._intermediates['Lmat'      ]
        A_diele    = self._intermediates['A_diele'   ]
        A_inf      = self._intermediates['A_inf'     ]
        cached_pol = self._intermediates['cached_pol']

        if not (isinstance(dm, numpy.ndarray) and dm.ndim == 2):
            # spin-traced DM for UHF or ROHF
            dm = dm[0] + dm[1]

        phi = ddcosmo.make_phi(self, dm, r_vdw, ui, ylm_1sph)
        phi = numpy.linalg.solve(A_diele, A_inf.dot(phi.ravel()))

        Xvec = numpy.linalg.solve(Lmat, phi.ravel()).reshape(mol.natm,-1)
        psi, vmat = ddcosmo.make_psi_vmat(self, dm, r_vdw, ui, ylm_1sph,
                                          cached_pol, Xvec, Lmat)[:2]
        dielectric = self.eps
        f_epsilon = (dielectric-1.)/dielectric
        epcm = .5 * f_epsilon * numpy.einsum('jx,jx', psi, Xvec)
        vpcm = .5 * f_epsilon * vmat
        return epcm, vpcm

    def _B_dot_x(self, dm):
        '''
        Compute the matrix-vector product B * x. The B matrix, as defined in
        the paper R. Cammi, JPCA, 104, 5631 (2000), is the second order
        derivatives of E_solvation wrt density matrices.

        Note: In ddCOSMO, strictly, B is not symmetric. To make it compatible
        with the CIS framework, it is symmetrized in current implementation.
        '''
        if not self._intermediates or self.grids.coords is None:
            self.build()

        mol = self.mol
        r_vdw      = self._intermediates['r_vdw'     ]
        ylm_1sph   = self._intermediates['ylm_1sph'  ]
        ui         = self._intermediates['ui'        ]
        Lmat       = self._intermediates['Lmat'      ]
        A_diele    = self._intermediates['A_diele'   ]
        A_inf      = self._intermediates['A_inf'     ]
        cached_pol = self._intermediates['cached_pol']
        natm = mol.natm
        nlm = (self.lmax+1)**2

        phi = ddcosmo.make_phi(self, dm, r_vdw, ui, ylm_1sph, with_nuc=False)
        phi = numpy.linalg.solve(A_diele, A_inf.dot(phi.reshape(-1,natm*nlm).T))

        Xvec = numpy.linalg.solve(Lmat, phi)
        Xvec = Xvec.reshape(natm,nlm,-1).transpose(2,0,1)
        vmat = ddcosmo.make_psi_vmat(self, dm, r_vdw, ui, ylm_1sph,
                                     cached_pol, Xvec, Lmat, with_nuc=False)[1]
        dielectric = self.eps
        f_epsilon = (dielectric-1.)/dielectric
        return .5 * f_epsilon * vmat

    gen_solver = as_solver = gen_ddpcm_solver

    def regularize_xt(self, t, eta, scale=1):
        return regularize_xt(t, eta, scale)

    def nuc_grad_method(self, grad_method):
        raise NotImplementedError


if __name__ == '__main__':
    from pyscf import scf
    mol = gto.M(atom='H 0 0 0; H 0 1 1.2; H 1. .1 0; H .5 .5 1')
    numpy.random.seed(1)

    nao = mol.nao_nr()
    dm = numpy.random.random((nao,nao))
    dm = dm + dm.T
    #dm = scf.RHF(mol).run().make_rdm1()
    e, vmat = DDPCM(mol).kernel(dm)
    print(e + 1.2446306643473923)
    print(lib.fp(vmat) - 0.77873361914445294)

    mol = gto.Mole()
    mol.atom = ''' O                  0.00000000    0.00000000   -0.11081188
                   H                 -0.00000000   -0.84695236    0.59109389
                   H                 -0.00000000    0.89830571    0.52404783 '''
    mol.basis = '3-21g' #cc-pvdz'
    mol.build()
    cm = DDPCM(mol)
    cm.verbose = 4
    mf = ddpcm_for_scf(scf.RHF(mol), cm)#.newton()
    mf.verbose = 4
    mf.kernel()  # -75.5697645601958


# pyscf/solvent/pol_embed.py

#!/usr/bin/env python
# Copyright 2019-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>
#

'''
Interface to CPPE, a library of polarizable embedding solvent model.
This interface requires the cppe library

GitHub:      https://github.com/maxscheurer/cppe
Code:        10.5281/zenodo.3345696
Publication: https://doi.org/10.1021/acs.jctc.9b00758

The CPPE library can be installed via:
        pip install cppe
or
        pip install git+https://github.com/maxscheurer/cppe.git

The potential file required by CPPE library needs to be generated from the
PyFraME library  https://gitlab.com/FraME-projects/PyFraME

References:

  [1] Olsen, J. M., Aidas, K., & Kongsted, J. (2010). Excited States in Solution
  through Polarizable Embedding. J. Chem. Theory Comput., 6 (12), 3721-3734.
  https://doi.org/10.1021/ct1003803

  [2] Olsen, J. M. H., & Kongsted, J. (2011). Molecular Properties through
  Polarizable Embedding. Advances in Quantum Chemistry (Vol. 61).
  https://doi.org/10.1016/B978-0-12-386013-2.00003-6
'''

import sys
import numpy
from pkg_resources import parse_version

try:
    import cppe
except ModuleNotFoundError:
    sys.stderr.write('cppe library was not found\n')
    sys.stderr.write(__doc__)
    raise

from pyscf import lib
from pyscf.lib import logger
from pyscf import gto
from pyscf import df
from pyscf.solvent import _attach_solvent
from pyscf.data import elements


@lib.with_doc(_attach_solvent._for_scf.__doc__)
def pe_for_scf(mf, solvent_obj, dm=None):
    if not isinstance(solvent_obj, PolEmbed):
        solvent_obj = PolEmbed(mf.mol, solvent_obj)
    return _attach_solvent._for_scf(mf, solvent_obj, dm)

@lib.with_doc(_attach_solvent._for_casscf.__doc__)
def pe_for_casscf(mc, solvent_obj, dm=None):
    if not isinstance(solvent_obj, PolEmbed):
        solvent_obj = PolEmbed(mc.mol, solvent_obj)
    return _attach_solvent._for_casscf(mc, solvent_obj, dm)

@lib.with_doc(_attach_solvent._for_casci.__doc__)
def pe_for_casci(mc, solvent_obj, dm=None):
    if not isinstance(solvent_obj, PolEmbed):
        solvent_obj = PolEmbed(mc.mol, solvent_obj)
    return _attach_solvent._for_casci(mc, solvent_obj, dm)

@lib.with_doc(_attach_solvent._for_post_scf.__doc__)
def pe_for_post_scf(method, solvent_obj, dm=None):
    if not isinstance(solvent_obj, PolEmbed):
        solvent_obj = PolEmbed(method.mol, solvent_obj)
    return _attach_solvent._for_post_scf(method, solvent_obj, dm)

@lib.with_doc(_attach_solvent._for_tdscf.__doc__)
def pe_for_tdscf(method, solvent_obj, dm=None):
    scf_solvent = getattr(method._scf, 'with_solvent', None)
    assert scf_solvent is None or isinstance(scf_solvent, PolEmbed)

    if not isinstance(solvent_obj, PolEmbed):
        solvent_obj = PolEmbed(method.mol, solvent_obj)
    return _attach_solvent._for_tdscf(method, solvent_obj, dm)


# data from https://doi.org/10.1021/acs.jctc.9b01162
_pe_ecps = [
    ("X1", gto.parse_ecp(
        """
        X1 nelec 0
        X1 ul
        2      1.000000000000      0.000000000000
        X1 S
        2      0.509800000000      2.420000000000
        X1 P
        2      0.491650000000     -0.435900000000
        """)),
    ("X2", gto.parse_ecp(
        """
        X2 nelec 0
        X2 ul
        2      1.000000000000      0.000000000000
        X2 S
        2      2.047500000000     54.510000000000
        X2 P
        2      0.448150000000      1.465000000000
        X2 D
        2      0.492050000000     -0.838000000000
        """)),
    ("X3", gto.parse_ecp(
        """
        X3 nelec 0
        X3 ul
        2      1.000000000000      0.000000000000
        X3 S
        2      1.641000000000    275.000000000000
        X3 P
        2      0.273300000000      1.900000000000
        X3 D
        2      0.440000000000     -3.400000000000
        """))
]


def _get_element_row(symbol):
    """
    Helper function to determine the row of an element
    for choosing the correct ECP for PE(ECP)
    """
    nucchg = elements.charge(symbol)
    if nucchg <= 2:
        element_row = 0
    elif nucchg <= 10:
        element_row = 1
    elif nucchg <= 18:
        element_row = 2
    else:
        raise NotImplementedError("PE(ECP) only implemented for first, "
                                  "second, and third row elements")
    return element_row


class PolEmbed(lib.StreamObject):
    def __init__(self, mol, options_or_potfile):
        self.mol = mol
        self.stdout = mol.stdout
        self.verbose = mol.verbose
        self.max_memory = mol.max_memory

        self.frozen = False
        # FIXME: Should the solvent in PE model by default has the character
        # of rapid process?
        self.equilibrium_solvation = False

##################################################
# don't modify the following attributes, they are not input options
        if isinstance(options_or_potfile, str):
            options = {"potfile": options_or_potfile}
        else:
            options = options_or_potfile

        min_version = "0.3.1"
        if parse_version(cppe.__version__) < parse_version(min_version):
            raise ModuleNotFoundError("cppe version {} is required at least. "
                                      "Version {}"
                                      " was found.".format(min_version,
                                                           cppe.__version__))

        if not isinstance(options, dict):
            raise TypeError("Options should be a dictionary.")

        self.options = options
        # use PE(ECP) repulsive potentials
        self.do_ecp = self.options.pop("ecp", False)
        # use effective external field (EEF)
        self.eef = self.options.pop("eef", False)
        self.cppe_state = self._create_cppe_state(mol)
        self.potentials = self.cppe_state.potentials
        self.V_es = None

        if self.do_ecp:
            # Use ECPs in the environment as repulsive potentials
            # with parameters from https://doi.org/10.1021/acs.jctc.9b01162
            ecpatoms = []
            # one set of parameters for each row of elements (first 3 rows supported)
            for p in self.potentials:
                if p.element == "X":
                    continue
                element_row = _get_element_row(p.element)
                ecp_label, _ = _pe_ecps[element_row]
                ecpatoms.append([ecp_label, p.x, p.y, p.z])
            self.ecpmol = gto.M(atom=ecpatoms, ecp={l: k for (l, k) in _pe_ecps},
                                basis={}, unit="Bohr")
            # add the normal mol to compute integrals
            self.ecpmol += self.mol

        # e (the electrostatic and induction energy)
        # and v (the additional potential) are
        # updated during the SCF iterations
        self.e = None
        self.v = None
        self._dm = None

        self._keys = set(self.__dict__.keys())

    def dump_flags(self, verbose=None):
        logger.info(self, '******** %s flags ********', self.__class__)
        options = self.cppe_state.options
        option_keys = cppe.valid_option_keys
        logger.info(self, 'frozen = %s'       , self.frozen)
        logger.info(self, 'equilibrium_solvation = %s', self.equilibrium_solvation)
        logger.info(self, 'pe(ecp) repulsive potentials = %s', self.do_ecp)
        for key in option_keys:
            logger.info(self, "cppe.%s = %s", key, options[key])
        return self

    def _create_cppe_state(self, mol):
        cppe_mol = cppe.Molecule()
        for z, coord in zip(mol.atom_charges(), mol.atom_coords()):
            cppe_mol.append(cppe.Atom(z, *coord))

        def callback(output):
            logger.info(self, output)
        cppe_state = cppe.CppeState(self.options, cppe_mol, callback)
        cppe_state.calculate_static_energies_and_fields()
        # logger.info(self, "Static energies and fields computed")
        return cppe_state

    def reset(self, mol=None):
        '''Reset mol and clean up relevant attributes for scanner mode'''
        if mol is not None:
            self.mol = mol
        self.cppe_state = self._create_cppe_state(mol)
        self.potentials = self.cppe_state.potentials
        self.V_es = None
        return self

    def kernel(self, dm, elec_only=False):
        '''
        '''
        if not (isinstance(dm, numpy.ndarray) and dm.ndim == 2):
            # spin-traced DM for UHF or ROHF
            dm = dm[0] + dm[1]

        e, v = self._exec_cppe(dm, elec_only)
        logger.info(self, 'Polarizable embedding energy = %.15g', e)

        self.e = e
        self.v = v
        return e, v

    def effective_dipole_operator(self):
        """
        Compute the derivatives of induced moments wrt each coordinate
        and form integrals for effective dipole operator (EEF)
        """
        dips = self.mol.intor_symmetric('int1e_r', comp=3)
        if self.eef:
            logger.info(self, "Computing effective dipole operator for EEF.")
            positions = self.cppe_state.positions_polarizable
            n_sites = positions.shape[0]
            induced_moments = self.cppe_state.induced_moments_eef()
            induced_moments = induced_moments.reshape(n_sites, 3, 3)
            fakemol = gto.fakemol_for_charges(positions)
            j3c = df.incore.aux_e2(self.mol, fakemol, intor='int3c2e_ip1')
            V_ind = numpy.einsum('aijg,gax->xij', j3c, -induced_moments)
            dips += V_ind + V_ind.transpose(0, 2, 1)
        return list(dips)

    def _exec_cppe(self, dm, elec_only=False):
        dms = numpy.asarray(dm)
        is_single_dm = dms.ndim == 2

        nao = dms.shape[-1]
        dms = dms.reshape(-1,nao,nao)
        n_dm = dms.shape[0]

        max_memory = self.max_memory

        if self.V_es is None:
            # very conservative estimate (based on multipole potential integrals)
            # when all sites have a charge, dipole, and quadrupole moment
            max_memreq = 10 * len(self.potentials) * nao**2 * 8.0/1e6
            n_chunks_el = 1
            if max_memreq >= max_memory:
                n_chunks_el = int(max_memreq // max_memory + 1)
            positions = numpy.array([p.position for p in self.potentials])
            moments = []
            orders = []
            for p in self.potentials:
                p_moments = []
                for m in p.multipoles:
                    m.remove_trace()
                    p_moments.append(m.values)
                orders.append(m.k)
                moments.append(p_moments)
            self.V_es = self._compute_multipole_potential_integrals(
                positions, orders, moments, n_chunks_el
            )
            if self.do_ecp:
                self.V_ecp = self.ecpmol.intor("ECPscalar")

        e_static = numpy.einsum('ij,xij->x', self.V_es, dms)
        self.cppe_state.energies["Electrostatic"]["Electronic"] = (
            e_static[0]
        )

        e_ecp = 0.0
        if self.do_ecp:
            e_ecp = numpy.einsum('ij,xij->x', self.V_ecp, dms)[0]

        positions = self.cppe_state.positions_polarizable
        n_sites = positions.shape[0]
        V_ind = numpy.zeros((n_dm, nao, nao))

        e_tot = []
        e_pol = []
        if n_sites > 0:
            max_memreq = 6 * n_sites * nao**2 * 8.0/1e6
            n_chunks_ind = 1
            if max_memreq >= max_memory:
                n_chunks_ind = int(max_memreq // max_memory + 1)

            chunks = numpy.array_split(positions, n_chunks_ind)
            elec_fields_chunk = []
            for chunk in chunks:
                fakemol = gto.fakemol_for_charges(chunk)
                j3c = df.incore.aux_e2(self.mol, fakemol, intor='int3c2e_ip1')
                elf = (numpy.einsum('aijg,nij->nga', j3c, dms) +
                       numpy.einsum('aijg,nji->nga', j3c, dms))
                elec_fields_chunk.append(elf)
            elec_fields = numpy.concatenate(elec_fields_chunk, axis=1)

            induced_moments = numpy.empty((n_dm, n_sites * 3))
            for i_dm in range(n_dm):
                self.cppe_state.update_induced_moments(elec_fields[i_dm].ravel(), elec_only)
                induced_moments[i_dm] = numpy.array(self.cppe_state.get_induced_moments())

                e_tot.append(self.cppe_state.total_energy + e_ecp)
                e_pol.append(self.cppe_state.energies["Polarization"]["Electronic"])

            induced_moments = induced_moments.reshape(n_dm, n_sites, 3)
            induced_moments_chunked = numpy.array_split(induced_moments, n_chunks_ind, axis=1)
            for pos_chunk, ind_chunk in zip(chunks, induced_moments_chunked):
                fakemol = gto.fakemol_for_charges(pos_chunk)
                j3c = df.incore.aux_e2(self.mol, fakemol, intor='int3c2e_ip1')
                V_ind += numpy.einsum('aijg,nga->nij', j3c, -ind_chunk)
            V_ind = V_ind + V_ind.transpose(0, 2, 1)
        else:
            for i_dm in range(n_dm):
                e_tot.append(self.cppe_state.total_energy + e_ecp)
                e_pol.append(0.0)

        if not elec_only:
            vmat = self.V_es + V_ind
            if self.do_ecp:
                vmat += self.V_ecp
            e = numpy.array(e_tot)
        else:
            vmat = V_ind
            e = numpy.array(e_pol)

        if is_single_dm:
            e = e[0]
            vmat = vmat[0]
        return e, vmat

    def _compute_multipole_potential_integrals(self, all_sites, all_orders, all_moments, n_chunks=1):
        all_orders = numpy.asarray(all_orders)
        if numpy.any(all_orders > 2):
            raise NotImplementedError("""Multipole potential integrals not
                                      implemented for order > 2.""")

        chunks = numpy.array_split(all_sites, n_chunks)
        chunks_o = numpy.array_split(all_orders, n_chunks)
        chunks_m = numpy.array_split(all_moments, n_chunks)
        op = 0
        for (sites, orders, moments) in zip(chunks, chunks_o, chunks_m):
            # order 0
            fakemol = gto.fakemol_for_charges(sites)
            integral0 = df.incore.aux_e2(self.mol, fakemol, intor='int3c2e')
            moments_0 = numpy.array([m[0] for m in moments])
            op += numpy.einsum('ijg,ga->ij', integral0, moments_0 * cppe.prefactors(0))

            # order 1
            if numpy.any(orders >= 1):
                idx = numpy.where(orders >= 1)[0]
                fakemol = gto.fakemol_for_charges(sites[idx])
                integral1 = df.incore.aux_e2(self.mol, fakemol, intor='int3c2e_ip1')
                moments_1 = numpy.array([moments[i][1] for i in idx])
                v = numpy.einsum('aijg,ga,a->ij', integral1, moments_1, cppe.prefactors(1))
                op += v + v.T

            if numpy.any(orders >= 2):
                idx = numpy.where(orders >= 2)[0]
                fakemol = gto.fakemol_for_charges(sites[idx])
                n_sites = idx.size
                # moments_2 is the lower triangle of
                # [[XX, XY, XZ], [YX, YY, YZ], [ZX, ZY, ZZ]] i.e.
                # XX, XY, XZ, YY, YZ, ZZ = 0,1,2,4,5,8
                # symmetrize it to the upper triangle part
                # XX, YX, ZX, YY, ZY, ZZ = 0,3,6,4,7,8
                m2 = numpy.einsum('ga,a->ga', [moments[i][2] for i in idx],
                                  cppe.prefactors(2))
                moments_2 = numpy.zeros((n_sites, 9))
                moments_2[:, [0, 1, 2, 4, 5, 8]]  = m2
                moments_2[:, [0, 3, 6, 4, 7, 8]] += m2
                moments_2 *= .5

                integral2 = df.incore.aux_e2(self.mol, fakemol, intor='int3c2e_ipip1')
                v = numpy.einsum('aijg,ga->ij', integral2, moments_2)
                op += v + v.T
                integral2 = df.incore.aux_e2(self.mol, fakemol, intor='int3c2e_ipvip1')
                op += numpy.einsum('aijg,ga->ij', integral2, moments_2) * 2

        return op

    def _B_dot_x(self, dm):
        dms = numpy.asarray(dm)
        dm_shape = dms.shape
        nao = dm_shape[-1]
        dms = dms.reshape(-1,nao,nao)
        v_pe_ao = [self._exec_cppe(x, elec_only=True)[1] for x in dms]
        return numpy.asarray(v_pe_ao).reshape(dm_shape)

    def nuc_grad_method(self, grad_method):
        raise NotImplementedError("Nuclear gradients not implemented for PE.")

if __name__ == '__main__':
    import tempfile
    from pyscf.solvent import PE
    mol = gto.M(atom='''
           6        0.000000    0.000000   -0.542500
           8        0.000000    0.000000    0.677500
           1        0.000000    0.935307   -1.082500
           1        0.000000   -0.935307   -1.082500
                ''', basis='sto3g')
    mf = mol.RHF()
    with tempfile.NamedTemporaryFile() as f:
        f.write(b'''!
@COORDINATES
3
AA
O     3.53300000    2.99600000    0.88700000      1
H     4.11100000    3.13200000    1.63800000      2
H     4.10500000    2.64200000    0.20600000      3
@MULTIPOLES
ORDER 0
3
1     -0.67444000
2      0.33722000
3      0.33722000
@POLARIZABILITIES
ORDER 1 1
3
1      5.73935000     0.00000000     0.00000000     5.73935000     0.00000000     5.73935000
2      2.30839000     0.00000000     0.00000000     2.30839000     0.00000000     2.30839000
3      2.30839000     0.00000000     0.00000000     2.30839000     0.00000000     2.30839000
EXCLISTS
3 3
1   2  3
2   1  3
3   1  2''')
        f.flush()
        pe_options = {"potfile": f.name}
        # pe = pol_embed.PolEmbed(mol, pe_options)
        #mf = PE(mf, pe).run()
        mf = PE(mf, pe_options).run()
        print(mf.e_tot - -112.35232445743728)
        print(mf.with_solvent.e - 0.00020182314249546455)


# pyscf/soscf/__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.


# pyscf/soscf/ciah.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>
#

import sys

import numpy
import scipy.linalg
from pyscf import lib
from pyscf.lib import logger
from pyscf import __config__


def expmat(a):
    return scipy.linalg.expm(a)

class CIAHOptimizer(lib.StreamObject):

    conv_tol_grad = getattr(__config__, 'soscf_ciah_CIAHOptimizer_conv_tol_grad', 1e-4)
    max_stepsize = getattr(__config__, 'soscf_ciah_CIAHOptimizer_max_stepsize', .05)
    max_iters = getattr(__config__, 'soscf_ciah_CIAHOptimizer_max_iters', 10)
    kf_interval = getattr(__config__, 'soscf_ciah_CIAHOptimizer_kf_interval', 5)
    kf_trust_region = getattr(__config__, 'soscf_ciah_CIAHOptimizer_kf_trust_region', 5)
    ah_start_tol = getattr(__config__, 'soscf_ciah_CIAHOptimizer_ah_start_tol', 5.)
    ah_start_cycle = getattr(__config__, 'soscf_ciah_CIAHOptimizer_ah_start_cycle', 1)
    ah_level_shift = getattr(__config__, 'soscf_ciah_CIAHOptimizer_ah_level_shift', 0)
    ah_conv_tol = getattr(__config__, 'soscf_ciah_CIAHOptimizer_ah_conv_tol', 1e-12)
    ah_lindep = getattr(__config__, 'soscf_ciah_CIAHOptimizer_ah_lindep', 1e-14)
    ah_max_cycle = getattr(__config__, 'soscf_ciah_CIAHOptimizer_ah_max_cycle', 30)
    ah_trust_region = getattr(__config__, 'soscf_ciah_CIAHOptimizer_ah_trust_region', 3.)

    def __init__(self):
        self._keys = set(('conv_tol_grad', 'max_stepsize', 'max_iters',
                          'kf_interval', 'kf_trust_region', 'ah_start_tol',
                          'ah_start_cycle', 'ah_level_shift', 'ah_conv_tol',
                          'ah_lindep', 'ah_max_cycle', 'ah_trust_region'))

    def gen_g_hop(self, u):
        pass

    def pack_uniq_var(self, mat):
        nmo = mat.shape[0]
        idx = numpy.tril_indices(nmo, -1)
        return mat[idx]

    def unpack_uniq_var(self, v):
        nmo = int(numpy.sqrt(v.size*2)) + 1
        idx = numpy.tril_indices(nmo, -1)
        mat = numpy.zeros((nmo,nmo))
        mat[idx] = v
        return mat - mat.conj().T

    def extract_rotation(self, dr, u0=1):
        dr = self.unpack_uniq_var(dr)
        return numpy.dot(u0, expmat(dr))

    def get_grad(self, u):
        pass

    def cost_function(self, u):
        pass


def rotate_orb_cc(iah, u0, conv_tol_grad=None, verbose=logger.NOTE):
    t2m = (logger.process_clock(), logger.perf_counter())
    if isinstance(verbose, logger.Logger):
        log = verbose
    else:
        log = logger.Logger(sys.stdout, verbose)

    if conv_tol_grad is None:
        conv_tol_grad = iah.conv_tol_grad

    g_orb, h_op, h_diag = iah.gen_g_hop(u0)
    g_kf = g_orb
    norm_gkf = norm_gorb = numpy.linalg.norm(g_orb)
    log.debug('    |g|= %4.3g (keyframe)', norm_gorb)
    t3m = log.timer('gen h_op', *t2m)

    if h_diag is None:
        def precond(x, e):
            return x
    else:
        def precond(x, e):
            hdiagd = h_diag-(e-iah.ah_level_shift)
            hdiagd[abs(hdiagd)<1e-8] = 1e-8
            x = x/hdiagd
            return x

    def scale_down_step(dxi, hdxi, norm_gorb):
        dxmax = abs(dxi).max()
        if dxmax > iah.max_stepsize:
            scale = iah.max_stepsize / dxmax
            log.debug1('Scale rotation by %g', scale)
            dxi *= scale
            hdxi *= scale
        return dxi, hdxi

    class Statistic:
        def __init__(self):
            self.imic = 0
            self.tot_hop = 0
            self.tot_kf = 0

    kf_trust_region = iah.kf_trust_region
    g_op = lambda: g_orb
    x0_guess = g_orb
    while True:
        stat = Statistic()
        dr = 0
        ikf = 0
        ukf = 1

        for ah_conv, ihop, w, dxi, hdxi, residual, seig \
                in davidson_cc(h_op, g_op, precond, x0_guess,
                               tol=iah.ah_conv_tol, max_cycle=iah.ah_max_cycle,
                               lindep=iah.ah_lindep, verbose=log):
            stat.tot_hop = ihop
            norm_residual = numpy.linalg.norm(residual)
            if (ah_conv or ihop == iah.ah_max_cycle or # make sure to use the last step
                ((norm_residual < iah.ah_start_tol) and (ihop >= iah.ah_start_cycle)) or
                (seig < iah.ah_lindep)):
                stat.imic += 1
                dxmax = abs(dxi).max()
                dxi, hdxi = scale_down_step(dxi, hdxi, norm_gorb)

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

                max_cycle = max(iah.max_iters,
                                iah.max_iters-int(numpy.log(norm_gkf+1e-9)*2))
                log.debug1('Set max_cycle %d', max_cycle)
                ikf += 1
                if stat.imic > 3 and norm_gorb > norm_gkf*iah.ah_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 (stat.imic >= max_cycle or norm_gorb < conv_tol_grad*.2):
                    break

                elif (ikf > 2 and # avoid frequent keyframe
                      (ikf >= max(iah.kf_interval, iah.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
                    ukf = iah.extract_rotation(dr, ukf)
                    dr[:] = 0
                    g_kf1 = iah.get_grad(u0.dot(ukf))
                    stat.tot_kf += 1
                    norm_gkf1 = numpy.linalg.norm(g_kf1)
                    norm_dg = numpy.linalg.norm(g_kf1-g_orb)
                    log.debug('Adjust keyframe g_orb to |g|= %4.3g  '
                              '|g-correction|= %4.3g', norm_gkf1, norm_dg)

                    if (norm_dg < norm_gorb*iah.ah_trust_region  # kf not too diff
                        #or norm_gkf1 < norm_gkf  # grad is decaying
                        # close to solution
                        or norm_gkf1 < conv_tol_grad*iah.ah_trust_region):
                        kf_trust_region = min(max(norm_gorb/(norm_dg+1e-9), iah.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 = iah.extract_rotation(dr, ukf)
        log.debug('    tot inner=%d  |g|= %4.3g  |u-1|= %4.3g',
                  stat.imic, norm_gorb, numpy.linalg.norm(numpy.tril(u,-1)))
        h_op = h_diag = None
        t3m = log.timer('aug_hess in %d inner iters' % stat.imic, *t3m)
        u0 = (yield u, g_kf, stat)

        g_kf, h_op, h_diag = iah.gen_g_hop(u0)
        norm_gkf = numpy.linalg.norm(g_kf)
        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), iah.kf_trust_region), 10)
        log.debug1('Set  kf_trust_region = %g', kf_trust_region)
        g_orb = g_kf
        norm_gorb = norm_gkf
        x0_guess = dxi

def davidson_cc(h_op, g_op, precond, x0, tol=1e-10, xs=[], ax=[],
                max_cycle=30, lindep=1e-14, dot=numpy.dot, verbose=logger.WARN):

    if isinstance(verbose, logger.Logger):
        log = verbose
    else:
        log = logger.Logger(sys.stdout, verbose)

    toloose = numpy.sqrt(tol)
    # the first trial vector is (1,0,0,...), which is not included in xs
    xs = list(xs)
    ax = list(ax)
    nx = len(xs)

    problem_size = x0.size
    max_cycle = min(max_cycle, problem_size)
    heff = numpy.zeros((max_cycle+nx+1,max_cycle+nx+1), dtype=x0.dtype)
    ovlp = numpy.eye(max_cycle+nx+1, dtype=x0.dtype)
    if nx == 0:
        xs.append(x0)
        ax.append(h_op(x0))
    else:
        for i in range(1, nx+1):
            for j in range(1, i+1):
                heff[i,j] = dot(xs[i-1].conj(), ax[j-1])
                ovlp[i,j] = dot(xs[i-1].conj(), xs[j-1])
            heff[1:i,i] = heff[i,1:i].conj()
            ovlp[1:i,i] = ovlp[i,1:i].conj()

    w_t = 0
    for istep in range(max_cycle):
        g = g_op()
        nx = len(xs)
        for i in range(nx):
            heff[i+1,0] = dot(xs[i].conj(), g)
            heff[nx,i+1] = dot(xs[nx-1].conj(), ax[i])
            ovlp[nx,i+1] = dot(xs[nx-1].conj(), xs[i])
        heff[0,:nx+1] = heff[:nx+1,0].conj()
        heff[1:nx,nx] = heff[nx,1:nx].conj()
        ovlp[1:nx,nx] = ovlp[nx,1:nx].conj()
        nvec = nx + 1
        #s0 = scipy.linalg.eigh(ovlp[:nvec,:nvec])[0][0]
        #if s0 < lindep:
        #    yield True, istep, w_t, xtrial, hx, dx, s0
        #    break
        wlast = w_t
        xtrial, w_t, v_t, index, seig = \
                _regular_step(heff[:nvec,:nvec], ovlp[:nvec,:nvec], xs,
                              lindep, log)
        s0 = seig[0]
        hx = _dgemv(v_t[1:], ax)
        # note g*v_t[0], as the first trial vector is (1,0,0,...)
        dx = hx + g*v_t[0] - w_t * v_t[0]*xtrial
        norm_dx = numpy.linalg.norm(dx)
        log.debug1('... AH step %d  index= %d  |dx|= %.5g  eig= %.5g  v[0]= %.5g  lindep= %.5g',
                   istep+1, index, norm_dx, w_t, v_t[0].real, s0)
        hx *= 1/v_t[0] # == h_op(xtrial)
        if ((abs(w_t-wlast) < tol and norm_dx < toloose) or
            s0 < lindep or
            istep+1 == problem_size):
            # Avoid adding more trial vectors if hessian converged
            yield True, istep+1, w_t, xtrial, hx, dx, s0
            if s0 < lindep or norm_dx < lindep:# or numpy.linalg.norm(xtrial) < lindep:
                # stop the iteration because eigenvectors would be barely updated
                break
        else:
            yield False, istep+1, w_t, xtrial, hx, dx, s0
            x0 = precond(dx, w_t)
            xs.append(x0)
            ax.append(h_op(x0))


def _regular_step(heff, ovlp, xs, lindep, log):
    try:
        e, c = scipy.linalg.eigh(heff[1:,1:], ovlp[1:,1:])
    except scipy.linalg.LinAlgError:
        e, c = lib.safe_eigh(heff[1:,1:], ovlp[1:,1:], lindep)[:2]
    if numpy.any(e < -1e-5):
        log.debug('Negative hessians found %s', e[e<0])

    w, v, seig = lib.safe_eigh(heff, ovlp, lindep)
    if log.verbose >= logger.DEBUG3:
        numpy.set_printoptions(3, linewidth=1000)
        log.debug3('v[0] %s', v[0])
        log.debug3('AH eigs %s', w)
        numpy.set_printoptions(8, linewidth=75)

    #if e[0] < -.1:
    #    sel = 0
    #else:
    # There exists systems that the first eigenvalue of AH is -inf.
    # Dynamically choosing the eigenvectors may be better.
    idx = numpy.where(abs(v[0]) > 0.1)[0]
    sel = idx[0]
    log.debug1('CIAH eigen-sel %s', sel)

    w_t = w[sel]
    xtrial = _dgemv(v[1:,sel]/v[0,sel], xs)
    return xtrial, w_t, v[:,sel], sel, seig

def _dgemv(v, m):
    vm = v[0] * m[0]
    for i,vi in enumerate(v[1:]):
        vm += vi * m[i+1]
    return vm




# pyscf/soscf/newton_ah.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>
#

'''
Co-iterative augmented hessian second order SCF solver (CIAH-SOSCF)
'''

import sys

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 chkfile
from pyscf.scf import addons
from pyscf.scf import hf_symm, uhf_symm, ghf_symm
from pyscf.scf import hf, rohf, uhf, dhf
# import _response_functions to load gen_response methods in SCF class
from pyscf.scf import _response_functions  # noqa
from pyscf.soscf import ciah
from pyscf import __config__

WITH_EX_EY_DEGENERACY = getattr(__config__, 'soscf_newton_ah_Ex_Ey_degeneracy', True)


# http://scicomp.stackexchange.com/questions/1234/matrix-exponential-of-a-skew-hermitian-matrix-with-fortran-95-and-lapack
def expmat(a):
    return scipy.linalg.expm(a)

# kwarg with_symmetry is required the stability analysis. It controls whether
# the stability analysis can break the point group symmetry.
def gen_g_hop_rhf(mf, mo_coeff, mo_occ, fock_ao=None, h1e=None,
                  with_symmetry=True):
    mo_coeff0 = mo_coeff
    mol = mf.mol
    if getattr(mf, '_scf', None) and mf._scf.mol != mol:
        #TODO: construct vind with dual-basis treatment, (see also JCP, 118, 9497)
        # To project Hessians from another basis if different basis sets are used
        # in newton solver and underlying mean-filed solver.
        mo_coeff = addons.project_mo_nr2nr(mf._scf.mol, mo_coeff, mol)

    occidx = numpy.where(mo_occ==2)[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 = hf_symm.get_orbsym(mol, mo_coeff)
        sym_forbid = orbsym[viridx,None] != orbsym[occidx]

    if fock_ao is None:
        # dm0 is the density matrix in projected basis. Computing fock in
        # projected basis.
        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)
        fock = reduce(numpy.dot, (mo_coeff.conj().T, fock_ao, mo_coeff))
    else:
        # If fock is given, it corresponds to main basis. It needs to be
        # diagonalized with the mo_coeff of the main basis.
        fock = reduce(numpy.dot, (mo_coeff0.conj().T, fock_ao, mo_coeff0))

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

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

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

    if with_symmetry and mol.symmetry:
        g[sym_forbid] = 0
        h_diag[sym_forbid] = 0
    vind = mf.gen_response(mo_coeff, mo_occ, singlet=None, 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-= .5*numpy.einsum('ps,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
        problem_size = g_orb.size

        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 ihop == problem_size:
                    log.debug1('... Hx=g fully converged for small systems')
                elif dxmax > mf.max_stepsize:
                    scale = mf.max_stepsize / dxmax
                    log.debug1('... scale rotation size %g', scale)
                    dxi *= scale
                    hdxi *= scale

                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)

    homo = lumo = None
    if isinstance(mf, dhf.DHF):
        n2c = mo_energy.size // 2
        pos_mo_energy = mo_energy[n2c:]
        pos_mo_occ = mo_occ[n2c:]
        if numpy.any(pos_mo_occ==0):
            homo = pos_mo_energy[pos_mo_occ>0].max()
            lumo = pos_mo_energy[pos_mo_occ==0].min()
    else:
        if numpy.any(mo_occ==0):
            homo = mo_energy[mo_occ>0].max()
            lumo = mo_energy[mo_occ==0].min()
    if lumo is not None and homo > lumo:
        log.warn('HOMO %s > LUMO %s was found in the canonicalized orbitals.',
                 homo, lumo)
    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)

def remove_soscf(mf):
    '''Remove the SOSCF decorator'''
    if not isinstance(mf, _CIAH_SOSCF):
        return mf

    for cls in mf.__class__.__mro__:
        if not issubclass(cls, _CIAH_SOSCF):
            break
    mf = mf.view(cls)
    if hasattr(mf, '_scf'):
        delattr(mf, '_scf')
    return 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.dft import KohnShamDFT
    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)
    mf = mf.remove_soscf()
    if isinstance(mf, scf.rohf.ROHF):
        # Is it correct to call TDUHF for ROHF?
        mf = mf.to_uhf()
    return mf.TDHF()

def TDA(mf):
    mf = mf.remove_soscf()
    if isinstance(mf, scf.rohf.ROHF):
        if isinstance(mf, KohnShamDFT):
            mf = mf.to_uks()
        else:
            mf = mf.to_uhf()
    return mf.TDA()

def TDDFT(mf):
    if isinstance(mf, KohnShamDFT):
        mf = mf.remove_soscf()
        if isinstance(mf, scf.rohf.ROHF):
            mf = mf.to_uks()
        return mf.TDDFT()
    else:
        return TDHF(mf)

TD = TDDFT


def RPA(mf):
    return TDDFT(mf)

def dRPA(mf):
    mf = mf.remove_soscf()
    if isinstance(mf, scf.rohf.ROHF):
        if isinstance(mf, KohnShamDFT):
            mf = mf.to_uks()
        else:
            mf = mf.to_uhf()
    return mf.dRPA()

def dTDA(mf):
    mf = mf.remove_soscf()
    if isinstance(mf, scf.rohf.ROHF):
        if isinstance(mf, KohnShamDFT):
            mf = mf.to_uks()
        else:
            mf = mf.to_uhf()
    return mf.dTDA()


# 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)

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,po,qv->xov', v1ao, orbo.conj(), orbv)
        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

    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))
    c1 = .5 / lib.param.LIGHT_SPEED
    moL = numpy.asarray(mo[:n2c], order='F')
    moS = numpy.asarray(mo[n2c:], order='F') * c1
    orboL = moL[:,:nocc]
    orboS = 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_spsp1spsp2_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 = a + numpy.einsum('iabj->iajb', eri_mo[:nocc,nocc:,nocc:,:nocc])
        a = a - numpy.einsum('ijba->iajb', eri_mo[:nocc,:nocc,nocc:,nocc:]) * hyb
        b = b + numpy.einsum('iajb->iajb', eri_mo[:nocc,nocc:,:nocc,nocc:])
        b = b - numpy.einsum('jaib->iajb', eri_mo[:nocc,nocc:,:nocc,nocc:]) * hyb
        return a, b

    if isinstance(mf, dft.KohnShamDFT):
        from pyscf.dft import xc_deriv
        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)

        a, b = add_hf_(a, b, hyb)

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

        def get_mo_value(ao):
            aoLa, aoLb, aoSa, aoSb = ao
            if aoLa.ndim == 2:
                moLa = lib.einsum('rp,pi->ri', aoLa, moL)
                moLb = lib.einsum('rp,pi->ri', aoLb, moL)
                moSa = lib.einsum('rp,pi->ri', aoSa, moS)
                moSb = lib.einsum('rp,pi->ri', aoSb, moS)
                return (moLa[:,:nocc], moLa[:,nocc:], moLb[:,:nocc], moLb[:,nocc:],
                        moSa[:,:nocc], moSa[:,nocc:], moSb[:,:nocc], moSb[:,nocc:])
            else:
                moLa = lib.einsum('xrp,pi->xri', aoLa, moL)
                moLb = lib.einsum('xrp,pi->xri', aoLb, moL)
                moSa = lib.einsum('xrp,pi->xri', aoSa, moS)
                moSb = lib.einsum('xrp,pi->xri', aoSb, moS)
                return (moLa[:,:,:nocc], moLa[:,:,nocc:], moLb[:,:,:nocc], moLb[:,:,nocc:],
                        moSa[:,:,:nocc], moSa[:,:,nocc:], moSb[:,:,:nocc], moSb[:,:,nocc:])

        def ud2tm(aa, ab, ba, bb):
            return numpy.stack([aa + bb,        # rho
                                ba + ab,        # mx
                                (ba - ab) * 1j, # my
                                aa - bb])       # mz
        def addLS(rhoL, rhoS):
            rhoS[1:4] *= -1  # beta * Sigma
            return rhoL + rhoS

        if xctype == 'LDA':
            ao_deriv = 0
            for ao, mask, weight, coords \
                    in ni.block_loop(mol, mf.grids, nao, ao_deriv, max_memory,
                                     with_s=True):
                if ni.collinear[0] == 'm':
                    rho = ni.eval_rho(mol, ao, dm0, mask, xctype, hermi=1, with_lapl=False)
                    eval_xc = ni.mcfun_eval_xc_adapter(mf.xc)
                    fxc = eval_xc(mf.xc, rho, deriv=2, xctype=xctype)[2]
                    wfxc = weight * fxc.reshape(4,4,-1)
                    moLoa, moLva, moLob, moLvb, moSoa, moSva, moSob, moSvb = get_mo_value(ao)
                    rhoL_ov_aa = numpy.einsum('ri,ra->ria', moLoa.conj(), moLva)
                    rhoL_ov_ab = numpy.einsum('ri,ra->ria', moLoa.conj(), moLvb)
                    rhoL_ov_ba = numpy.einsum('ri,ra->ria', moLob.conj(), moLva)
                    rhoL_ov_bb = numpy.einsum('ri,ra->ria', moLob.conj(), moLvb)
                    rhoS_ov_aa = numpy.einsum('ri,ra->ria', moSoa.conj(), moSva)
                    rhoS_ov_ab = numpy.einsum('ri,ra->ria', moSoa.conj(), moSvb)
                    rhoS_ov_ba = numpy.einsum('ri,ra->ria', moSob.conj(), moSva)
                    rhoS_ov_bb = numpy.einsum('ri,ra->ria', moSob.conj(), moSvb)
                    rhoL_ov = ud2tm(rhoL_ov_aa, rhoL_ov_ab, rhoL_ov_ba, rhoL_ov_bb)
                    rhoS_ov = ud2tm(rhoS_ov_aa, rhoS_ov_ab, rhoS_ov_ba, rhoS_ov_bb)
                    rho_ov = addLS(rhoL_ov, rhoS_ov)
                    rho_vo = rho_ov.conj()
                    w_ov = numpy.einsum('tsr,tria->sria', wfxc, rho_ov)
                    a += lib.einsum('sria,srjb->iajb', w_ov, rho_vo)
                    b += lib.einsum('sria,srjb->iajb', w_ov, rho_ov)
                elif ni.collinear[0] == 'c':
                    rho = ni.eval_rho(mol, ao, dm0, mask, xctype, hermi=1, with_lapl=False)
                    fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2)[2]
                    wv_a, wv_b = weight * fxc.reshape(2,2,-1)
                    moLoa, moLva, moLob, moLvb, moSoa, moSva, moSob, moSvb = get_mo_value(ao)
                    rhoL_ov_a = numpy.einsum('ri,ra->ria', moLoa.conj(), moLva)
                    rhoL_ov_b = numpy.einsum('ri,ra->ria', moLob.conj(), moLvb)
                    rhoS_ov_a = numpy.einsum('ri,ra->ria', moSoa.conj(), moSva)
                    rhoS_ov_b = numpy.einsum('ri,ra->ria', moSob.conj(), moSvb)
                    rhoL_vo_a = rhoL_ov_a.conj()
                    rhoL_vo_b = rhoL_ov_b.conj()
                    rhoS_vo_a = rhoS_ov_a.conj()
                    rhoS_vo_b = rhoS_ov_b.conj()
                    w_ov  = wv_a[:,:,None,None] * rhoL_ov_a
                    w_ov += wv_b[:,:,None,None] * rhoL_ov_b
                    w_ov += wv_b[:,:,None,None] * rhoS_ov_a  # for beta*Sigma
                    w_ov += wv_a[:,:,None,None] * rhoS_ov_b
                    wa_ov, wb_ov = w_ov
                    a += lib.einsum('ria,rjb->iajb', wa_ov, rhoL_vo_a)
                    a += lib.einsum('ria,rjb->iajb', wb_ov, rhoL_vo_b)
                    a += lib.einsum('ria,rjb->iajb', wb_ov, rhoS_vo_a)
                    a += lib.einsum('ria,rjb->iajb', wa_ov, rhoS_vo_b)
                    b += lib.einsum('ria,rjb->iajb', wa_ov, rhoL_ov_a)
                    b += lib.einsum('ria,rjb->iajb', wb_ov, rhoL_ov_b)
                    b += lib.einsum('ria,rjb->iajb', wb_ov, rhoS_ov_a)
                    b += lib.einsum('ria,rjb->iajb', wa_ov, rhoS_ov_b)
                else:
                    raise NotImplementedError(ni.collinear)

        elif xctype == 'GGA':
            ao_deriv = 1
            for ao, mask, weight, coords \
                    in ni.block_loop(mol, mf.grids, nao, ao_deriv, max_memory,
                                     with_s=True):
                if ni.collinear[0] == 'm':
                    rho = ni.eval_rho(mol, ao, dm0, mask, xctype, hermi=1, with_lapl=False)
                    eval_xc = ni.mcfun_eval_xc_adapter(mf.xc)
                    fxc = eval_xc(mf.xc, rho, deriv=2, xctype=xctype)[2]
                    wfxc = weight * fxc
                    moLoa, moLva, moLob, moLvb, moSoa, moSva, moSob, moSvb = get_mo_value(ao)
                    rhoL_ov_aa = numpy.einsum('ri,xra->xria', moLoa[0].conj(), moLva)
                    rhoL_ov_ab = numpy.einsum('ri,xra->xria', moLoa[0].conj(), moLvb)
                    rhoL_ov_ba = numpy.einsum('ri,xra->xria', moLob[0].conj(), moLva)
                    rhoL_ov_bb = numpy.einsum('ri,xra->xria', moLob[0].conj(), moLvb)
                    rhoS_ov_aa = numpy.einsum('ri,xra->xria', moSoa[0].conj(), moSva)
                    rhoS_ov_ab = numpy.einsum('ri,xra->xria', moSoa[0].conj(), moSvb)
                    rhoS_ov_ba = numpy.einsum('ri,xra->xria', moSob[0].conj(), moSva)
                    rhoS_ov_bb = numpy.einsum('ri,xra->xria', moSob[0].conj(), moSvb)
                    rhoL_ov_aa[1:4] += numpy.einsum('xri,ra->xria', moLoa[1:4].conj(), moLva[0])
                    rhoL_ov_ab[1:4] += numpy.einsum('xri,ra->xria', moLoa[1:4].conj(), moLvb[0])
                    rhoL_ov_ba[1:4] += numpy.einsum('xri,ra->xria', moLob[1:4].conj(), moLva[0])
                    rhoL_ov_bb[1:4] += numpy.einsum('xri,ra->xria', moLob[1:4].conj(), moLvb[0])
                    rhoS_ov_aa[1:4] += numpy.einsum('xri,ra->xria', moSoa[1:4].conj(), moSva[0])
                    rhoS_ov_ab[1:4] += numpy.einsum('xri,ra->xria', moSoa[1:4].conj(), moSvb[0])
                    rhoS_ov_ba[1:4] += numpy.einsum('xri,ra->xria', moSob[1:4].conj(), moSva[0])
                    rhoS_ov_bb[1:4] += numpy.einsum('xri,ra->xria', moSob[1:4].conj(), moSvb[0])
                    rhoL_ov = ud2tm(rhoL_ov_aa, rhoL_ov_ab, rhoL_ov_ba, rhoL_ov_bb)
                    rhoS_ov = ud2tm(rhoS_ov_aa, rhoS_ov_ab, rhoS_ov_ba, rhoS_ov_bb)
                    rho_ov = addLS(rhoL_ov, rhoS_ov)
                    rho_vo = rho_ov.conj()
                    w_ov = numpy.einsum('txsyr,txria->syria', wfxc, rho_ov)
                    a += lib.einsum('syria,syrjb->iajb', w_ov, rho_vo)
                    b += lib.einsum('syria,syrjb->iajb', w_ov, rho_ov)
                elif ni.collinear[0] == 'c':
                    rho = ni.eval_rho(mol, ao, dm0, mask, xctype, hermi=1, with_lapl=False)
                    fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2)[2]
                    fxc = xc_deriv.ud2ts(fxc)
                    wfxc = weight * fxc
                    moLoa, moLva, moLob, moLvb, moSoa, moSva, moSob, moSvb = get_mo_value(ao)
                    rhoL_ov_a = numpy.einsum('xri,ra->xria', moLoa.conj(), moLva[0])
                    rhoL_ov_b = numpy.einsum('xri,ra->xria', moLob.conj(), moLvb[0])
                    rhoL_ov_a[1:4] += numpy.einsum('ri,xra->xria', moLoa[0].conj(), moLva[1:4])
                    rhoL_ov_b[1:4] += numpy.einsum('ri,xra->xria', moLob[0].conj(), moLvb[1:4])
                    rhoS_ov_a = numpy.einsum('xri,ra->xria', moSoa.conj(), moSva[0])
                    rhoS_ov_b = numpy.einsum('xri,ra->xria', moSob.conj(), moSvb[0])
                    rhoS_ov_a[1:4] += numpy.einsum('ri,xra->xria', moSoa[0].conj(), moSva[1:4])
                    rhoS_ov_b[1:4] += numpy.einsum('ri,xra->xria', moSob[0].conj(), moSvb[1:4])
                    rhoL_ov = numpy.stack((rhoL_ov_a+rhoL_ov_b, rhoL_ov_a-rhoL_ov_b))
                    rhoS_ov = numpy.stack((rhoS_ov_a+rhoS_ov_b, rhoS_ov_a-rhoS_ov_b))
                    rhoS_ov[1] *= -1
                    rho_ov = rhoL_ov + rhoS_ov
                    rho_vo = rho_ov.conj()
                    w_ov = numpy.einsum('txsyr,txria->syria', wfxc, rho_ov)
                    a += lib.einsum('syria,syrjb->iajb', w_ov, rho_vo)
                    b += lib.einsum('syria,syrjb->iajb', w_ov, rho_ov)
                else:
                    raise NotImplementedError(ni.collinear)

        elif xctype == 'HF':
            pass

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

        elif xctype == 'MGGA':
            ao_deriv = 1
            for ao, mask, weight, coords \
                    in ni.block_loop(mol, mf.grids, nao, ao_deriv, max_memory,
                                     with_s=True):
                if ni.collinear[0] == 'm':
                    rho = ni.eval_rho(mol, ao, dm0, mask, xctype, hermi=1, with_lapl=False)
                    eval_xc = ni.mcfun_eval_xc_adapter(mf.xc)
                    fxc = eval_xc(mf.xc, rho, deriv=2, xctype=xctype)[2]
                    wfxc = weight * fxc
                    moLoa, moLva, moLob, moLvb, moSoa, moSva, moSob, moSvb = get_mo_value(ao)
                    rhoL_ov_aa = numpy.einsum('ri,xra->xria', moLoa[0].conj(), moLva)
                    rhoL_ov_ab = numpy.einsum('ri,xra->xria', moLoa[0].conj(), moLvb)
                    rhoL_ov_ba = numpy.einsum('ri,xra->xria', moLob[0].conj(), moLva)
                    rhoL_ov_bb = numpy.einsum('ri,xra->xria', moLob[0].conj(), moLvb)
                    rhoS_ov_aa = numpy.einsum('ri,xra->xria', moSoa[0].conj(), moSva)
                    rhoS_ov_ab = numpy.einsum('ri,xra->xria', moSoa[0].conj(), moSvb)
                    rhoS_ov_ba = numpy.einsum('ri,xra->xria', moSob[0].conj(), moSva)
                    rhoS_ov_bb = numpy.einsum('ri,xra->xria', moSob[0].conj(), moSvb)
                    rhoL_ov_aa[1:4] += numpy.einsum('xri,ra->xria', moLoa[1:4].conj(), moLva[0])
                    rhoL_ov_ab[1:4] += numpy.einsum('xri,ra->xria', moLoa[1:4].conj(), moLvb[0])
                    rhoL_ov_ba[1:4] += numpy.einsum('xri,ra->xria', moLob[1:4].conj(), moLva[0])
                    rhoL_ov_bb[1:4] += numpy.einsum('xri,ra->xria', moLob[1:4].conj(), moLvb[0])
                    rhoS_ov_aa[1:4] += numpy.einsum('xri,ra->xria', moSoa[1:4].conj(), moSva[0])
                    rhoS_ov_ab[1:4] += numpy.einsum('xri,ra->xria', moSoa[1:4].conj(), moSvb[0])
                    rhoS_ov_ba[1:4] += numpy.einsum('xri,ra->xria', moSob[1:4].conj(), moSva[0])
                    rhoS_ov_bb[1:4] += numpy.einsum('xri,ra->xria', moSob[1:4].conj(), moSvb[0])
                    tauL_ov_aa = numpy.einsum('xri,xra->ria', moLoa[1:4].conj(), moLva[1:4]) * .5
                    tauL_ov_ab = numpy.einsum('xri,xra->ria', moLoa[1:4].conj(), moLvb[1:4]) * .5
                    tauL_ov_ba = numpy.einsum('xri,xra->ria', moLob[1:4].conj(), moLva[1:4]) * .5
                    tauL_ov_bb = numpy.einsum('xri,xra->ria', moLob[1:4].conj(), moLvb[1:4]) * .5
                    tauS_ov_aa = numpy.einsum('xri,xra->ria', moSoa[1:4].conj(), moSva[1:4]) * .5
                    tauS_ov_ab = numpy.einsum('xri,xra->ria', moSoa[1:4].conj(), moSvb[1:4]) * .5
                    tauS_ov_ba = numpy.einsum('xri,xra->ria', moSob[1:4].conj(), moSva[1:4]) * .5
                    tauS_ov_bb = numpy.einsum('xri,xra->ria', moSob[1:4].conj(), moSvb[1:4]) * .5
                    rhoL_ov_aa = numpy.vstack([rhoL_ov_aa, tauL_ov_aa[numpy.newaxis]])
                    rhoL_ov_ab = numpy.vstack([rhoL_ov_ab, tauL_ov_ab[numpy.newaxis]])
                    rhoL_ov_ba = numpy.vstack([rhoL_ov_ba, tauL_ov_ba[numpy.newaxis]])
                    rhoL_ov_bb = numpy.vstack([rhoL_ov_bb, tauL_ov_bb[numpy.newaxis]])
                    rhoS_ov_aa = numpy.vstack([rhoS_ov_aa, tauS_ov_aa[numpy.newaxis]])
                    rhoS_ov_ab = numpy.vstack([rhoS_ov_ab, tauS_ov_ab[numpy.newaxis]])
                    rhoS_ov_ba = numpy.vstack([rhoS_ov_ba, tauS_ov_ba[numpy.newaxis]])
                    rhoS_ov_bb = numpy.vstack([rhoS_ov_bb, tauS_ov_bb[numpy.newaxis]])
                    rhoL_ov = ud2tm(rhoL_ov_aa, rhoL_ov_ab, rhoL_ov_ba, rhoL_ov_bb)
                    rhoS_ov = ud2tm(rhoS_ov_aa, rhoS_ov_ab, rhoS_ov_ba, rhoS_ov_bb)
                    rho_ov = addLS(rhoL_ov, rhoS_ov)
                    rho_vo = rho_ov.conj()
                    w_ov = numpy.einsum('txsyr,txria->syria', wfxc, rho_ov)
                    a += lib.einsum('syria,syrjb->iajb', w_ov, rho_vo)
                    b += lib.einsum('syria,syrjb->iajb', w_ov, rho_ov)
                elif ni.collinear[0] == 'c':
                    rho = ni.eval_rho(mol, ao, dm0, mask, xctype, hermi=1, with_lapl=False)
                    fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2)[2]
                    fxc = xc_deriv.ud2ts(fxc)
                    wfxc = weight * fxc
                    moLoa, moLva, moLob, moLvb, moSoa, moSva, moSob, moSvb = get_mo_value(ao)
                    rhoL_ov_a = numpy.einsum('xri,ra->xria', moLoa.conj(), moLva[0])
                    rhoL_ov_b = numpy.einsum('xri,ra->xria', moLob.conj(), moLvb[0])
                    rhoS_ov_a = numpy.einsum('xri,ra->xria', moSoa.conj(), moSva[0])
                    rhoS_ov_b = numpy.einsum('xri,ra->xria', moSob.conj(), moSvb[0])
                    rhoL_ov_a[1:4] += numpy.einsum('ri,xra->xria', moLoa[0].conj(), moLva[1:4])
                    rhoL_ov_b[1:4] += numpy.einsum('ri,xra->xria', moLob[0].conj(), moLvb[1:4])
                    rhoS_ov_a[1:4] += numpy.einsum('ri,xra->xria', moSoa[0].conj(), moSva[1:4])
                    rhoS_ov_b[1:4] += numpy.einsum('ri,xra->xria', moSob[0].conj(), moSvb[1:4])
                    tauL_ov_a = numpy.einsum('xri,xra->ria', moLoa[1:4].conj(), moLva[1:4]) * .5
                    tauL_ov_b = numpy.einsum('xri,xra->ria', moLob[1:4].conj(), moLvb[1:4]) * .5
                    tauS_ov_a = numpy.einsum('xri,xra->ria', moSoa[1:4].conj(), moSva[1:4]) * .5
                    tauS_ov_b = numpy.einsum('xri,xra->ria', moSob[1:4].conj(), moSvb[1:4]) * .5
                    rhoL_ov_a = numpy.vstack([rhoL_ov_a, tauL_ov_a[numpy.newaxis]])
                    rhoL_ov_b = numpy.vstack([rhoL_ov_b, tauL_ov_b[numpy.newaxis]])
                    rhoS_ov_a = numpy.vstack([rhoS_ov_a, tauS_ov_a[numpy.newaxis]])
                    rhoS_ov_b = numpy.vstack([rhoS_ov_b, tauS_ov_b[numpy.newaxis]])
                    rhoL_ov = numpy.stack((rhoL_ov_a+rhoL_ov_b, rhoL_ov_a-rhoL_ov_b))
                    rhoS_ov = numpy.stack((rhoS_ov_a+rhoS_ov_b, rhoS_ov_a-rhoS_ov_b))
                    rhoS_ov[1] *= -1
                    rho_ov = rhoL_ov + rhoS_ov
                    rho_vo = rho_ov.conj()
                    w_ov = numpy.einsum('txsyr,txria->syria', wfxc, rho_ov)
                    a += lib.einsum('syria,syrjb->iajb', w_ov, rho_vo)
                    b += lib.einsum('syria,syrjb->iajb', w_ov, rho_ov)
                else:
                    raise NotImplementedError(ni.collinear)

    else:
        a, b = 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=None):
        '''Generate function to compute Ax'''
        if mf is None:
            mf = self._scf
        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 > self.positive_eig_threshold)[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.conj())  # (A*)Y
        v1vo -= lib.einsum('xpr,sp->xsr', ys, foo.conj())  # (A*)Y

        # (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=None):
        if mf is None:
            mf = self._scf
        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.asarray(numpy.block([[x0, y0], [y0, x0.conj()]]))

    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 > self.positive_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  = scf.dhf.RDHF.TDA  = lib.class_as_method(TDA)
scf.dhf.DHF.TDHF = scf.dhf.RDHF.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__


class TDA(dhf.TDA):
    pass

class TDDFT(dhf.TDHF):
    pass

RPA = TDDKS = TDDFT

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


# 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)


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,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

    mol = mf.mol
    nmo = mo_occ.size
    nao = mol.nao
    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).astype(mo_coeff.dtype)
    b = numpy.zeros_like(a)

    def add_hf_(a, b, hyb=1):
        if mo_coeff.dtype == numpy.double:
            eri_mo  = ao2mo.general(mol, [orboa,moa,moa,moa], compact=False)
            eri_mo += ao2mo.general(mol, [orbob,mob,mob,mob], compact=False)
            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)
        else:
            eri_ao = mol.intor('int2e').reshape([nao]*4)
            eri_mo_a = lib.einsum('pqrs,pi,qj->ijrs', eri_ao, orboa.conj(), moa)
            eri_mo_a+= lib.einsum('pqrs,pi,qj->ijrs', eri_ao, orbob.conj(), mob)
            eri_mo = lib.einsum('ijrs,rk,sl->ijkl', eri_mo_a, moa.conj(), moa)
            eri_mo+= lib.einsum('ijrs,rk,sl->ijkl', eri_mo_a, mob.conj(), mob)
        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
        return a, b

    if isinstance(mf, dft.KohnShamDFT):
        from pyscf.dft import xc_deriv
        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)

        a, b = add_hf_(a, b, hyb)

        if ni.collinear[0] == 'm':  # mcol
            a = a.astype(numpy.complex128)
            b = b.astype(numpy.complex128)

        xctype = ni._xc_type(mf.xc)
        dm0 = mf.make_rdm1(mo_coeff, mo_occ)
        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:]

        def ud2tm(aa, ab, ba, bb):
            return numpy.stack([aa + bb,        # rho
                                ba + ab,        # mx
                                (ba - ab) * 1j, # my
                                aa - bb])       # mz

        if xctype == 'LDA':
            ao_deriv = 0
            for ao, mask, weight, coords \
                    in ni.block_loop(mol, mf.grids, nao, ao_deriv, max_memory):
                if ni.collinear[0] == 'm':
                    rho = ni.eval_rho(mol, ao, dm0, mask, xctype, hermi=1, with_lapl=False)
                    eval_xc = ni.mcfun_eval_xc_adapter(mf.xc)
                    fxc = eval_xc(mf.xc, rho, deriv=2, xctype=xctype)[2]
                    wfxc = weight * fxc.reshape(4,4,-1)
                    wr, wmx, wmy, wmz = weight * fxc.reshape(4,4,-1)
                    mo_oa, mo_va, mo_ob, mo_vb = get_mo_value(ao)
                    rho_ov_aa = numpy.einsum('ri,ra->ria', mo_oa.conj(), mo_va)
                    rho_ov_ab = numpy.einsum('ri,ra->ria', mo_oa.conj(), mo_vb)
                    rho_ov_ba = numpy.einsum('ri,ra->ria', mo_ob.conj(), mo_va)
                    rho_ov_bb = numpy.einsum('ri,ra->ria', mo_ob.conj(), mo_vb)
                    rho_ov = ud2tm(rho_ov_aa, rho_ov_ab, rho_ov_ba, rho_ov_bb)
                    rho_vo = rho_ov.conj()
                    w_ov = numpy.einsum('tsr,tria->sria', wfxc, rho_ov)
                    a += lib.einsum('sria,srjb->iajb', w_ov, rho_vo)
                    b += lib.einsum('sria,srjb->iajb', w_ov, rho_ov)
                elif ni.collinear[0] == 'c':
                    rho = ni.eval_rho(mol, ao, dm0, mask, xctype, hermi=1, with_lapl=False)
                    fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2)[2]
                    wv_a, wv_b = weight * fxc.reshape(2,2,-1)
                    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  = wv_a[:,:,None,None] * rho_ov_a
                    w_ov += wv_b[:,:,None,None] * rho_ov_b
                    wa_ov, wb_ov = w_ov
                    a += lib.einsum('ria,rjb->iajb', wa_ov, rho_vo_a)
                    a += lib.einsum('ria,rjb->iajb', wb_ov, rho_vo_b)
                    b += lib.einsum('ria,rjb->iajb', wa_ov, rho_ov_a)
                    b += lib.einsum('ria,rjb->iajb', wb_ov, rho_ov_b)
                else:
                    raise NotImplementedError(ni.collinear)

        elif xctype == 'GGA':
            ao_deriv = 1
            for ao, mask, weight, coords \
                    in ni.block_loop(mol, mf.grids, nao, ao_deriv, max_memory):
                if ni.collinear[0] == 'm':
                    rho = ni.eval_rho(mol, ao, dm0, mask, xctype, hermi=1, with_lapl=False)
                    eval_xc = ni.mcfun_eval_xc_adapter(mf.xc)
                    fxc = eval_xc(mf.xc, rho, deriv=2, xctype=xctype)[2]
                    wfxc = weight * fxc
                    wr, wmx, wmy, wmz = weight * fxc
                    mo_oa, mo_va, mo_ob, mo_vb = get_mo_value(ao)
                    rho_ov_aa = numpy.einsum('ri,xra->xria', mo_oa[0].conj(), mo_va)
                    rho_ov_ab = numpy.einsum('ri,xra->xria', mo_oa[0].conj(), mo_vb)
                    rho_ov_ba = numpy.einsum('ri,xra->xria', mo_ob[0].conj(), mo_va)
                    rho_ov_bb = numpy.einsum('ri,xra->xria', mo_ob[0].conj(), mo_vb)
                    rho_ov_aa[1:4] += numpy.einsum('xri,ra->xria', mo_oa[1:4].conj(), mo_va[0])
                    rho_ov_ab[1:4] += numpy.einsum('xri,ra->xria', mo_oa[1:4].conj(), mo_vb[0])
                    rho_ov_ba[1:4] += numpy.einsum('xri,ra->xria', mo_ob[1:4].conj(), mo_va[0])
                    rho_ov_bb[1:4] += numpy.einsum('xri,ra->xria', mo_ob[1:4].conj(), mo_vb[0])
                    rho_ov = ud2tm(rho_ov_aa, rho_ov_ab, rho_ov_ba, rho_ov_bb)
                    rho_vo = rho_ov.conj()
                    w_ov = numpy.einsum('txsyr,txria->syria', wfxc, rho_ov)
                    a += lib.einsum('syria,syrjb->iajb', w_ov, rho_vo)
                    b += lib.einsum('syria,syrjb->iajb', w_ov, rho_ov)
                elif ni.collinear[0] == 'c':
                    rho = ni.eval_rho(mol, ao, dm0, mask, xctype, hermi=1, with_lapl=False)
                    fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2)[2]
                    wv_a, wv_b = weight * fxc
                    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_ov_b.conj()
                    w_ov  = numpy.einsum('xsyr,xria->syria', wv_a, rho_ov_a)
                    w_ov += numpy.einsum('xsyr,xria->syria', wv_b, rho_ov_b)
                    wa_ov, wb_ov = w_ov
                    a += lib.einsum('xria,xrjb->iajb', wa_ov, rho_vo_a)
                    a += lib.einsum('xria,xrjb->iajb', wb_ov, rho_vo_b)
                    b += lib.einsum('xria,xrjb->iajb', wa_ov, rho_ov_a)
                    b += lib.einsum('xria,xrjb->iajb', wb_ov, rho_ov_b)
                else:
                    raise NotImplementedError(ni.collinear)

        elif xctype == 'HF':
            pass

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

        elif xctype == 'MGGA':
            ao_deriv = 1
            for ao, mask, weight, coords \
                    in ni.block_loop(mol, mf.grids, nao, ao_deriv, max_memory):
                if ni.collinear[0] == 'm':
                    rho = ni.eval_rho(mol, ao, dm0, mask, xctype, hermi=1, with_lapl=False)
                    eval_xc = ni.mcfun_eval_xc_adapter(mf.xc)
                    fxc = eval_xc(mf.xc, rho, deriv=2, xctype=xctype)[2]
                    wfxc = weight * fxc
                    wr, wmx, wmy, wmz = weight * fxc
                    mo_oa, mo_va, mo_ob, mo_vb = get_mo_value(ao)
                    rho_ov_aa = numpy.einsum('ri,xra->xria', mo_oa[0].conj(), mo_va)
                    rho_ov_ab = numpy.einsum('ri,xra->xria', mo_oa[0].conj(), mo_vb)
                    rho_ov_ba = numpy.einsum('ri,xra->xria', mo_ob[0].conj(), mo_va)
                    rho_ov_bb = numpy.einsum('ri,xra->xria', mo_ob[0].conj(), mo_vb)
                    rho_ov_aa[1:4] += numpy.einsum('xri,ra->xria', mo_oa[1:4].conj(), mo_va[0])
                    rho_ov_ab[1:4] += numpy.einsum('xri,ra->xria', mo_oa[1:4].conj(), mo_vb[0])
                    rho_ov_ba[1:4] += numpy.einsum('xri,ra->xria', mo_ob[1:4].conj(), mo_va[0])
                    rho_ov_bb[1:4] += numpy.einsum('xri,ra->xria', mo_ob[1:4].conj(), mo_vb[0])
                    tau_ov_aa = numpy.einsum('xri,xra->ria', mo_oa[1:4].conj(), mo_va[1:4]) * .5
                    tau_ov_ab = numpy.einsum('xri,xra->ria', mo_oa[1:4].conj(), mo_vb[1:4]) * .5
                    tau_ov_ba = numpy.einsum('xri,xra->ria', mo_ob[1:4].conj(), mo_va[1:4]) * .5
                    tau_ov_bb = numpy.einsum('xri,xra->ria', mo_ob[1:4].conj(), mo_vb[1:4]) * .5
                    rho_ov_aa = numpy.vstack([rho_ov_aa, tau_ov_aa[numpy.newaxis]])
                    rho_ov_ab = numpy.vstack([rho_ov_ab, tau_ov_ab[numpy.newaxis]])
                    rho_ov_ba = numpy.vstack([rho_ov_ba, tau_ov_ba[numpy.newaxis]])
                    rho_ov_bb = numpy.vstack([rho_ov_bb, tau_ov_bb[numpy.newaxis]])
                    rho_ov = ud2tm(rho_ov_aa, rho_ov_ab, rho_ov_ba, rho_ov_bb)
                    rho_vo = rho_ov.conj()
                    w_ov = numpy.einsum('txsyr,txria->syria', wfxc, rho_ov)
                    a += lib.einsum('syria,syrjb->iajb', w_ov, rho_vo)
                    b += lib.einsum('syria,syrjb->iajb', w_ov, rho_ov)
                elif ni.collinear[0] == 'c':
                    rho = ni.eval_rho(mol, ao, dm0, mask, xctype, hermi=1, with_lapl=False)
                    fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2)[2]
                    wv_a, wv_b = weight * fxc
                    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])
                    tau_ov_a = numpy.einsum('xri,xra->ria', mo_oa[1:4].conj(), mo_va[1:4]) * .5
                    tau_ov_b = numpy.einsum('xri,xra->ria', mo_ob[1:4].conj(), mo_vb[1:4]) * .5
                    rho_ov_a = numpy.vstack([rho_ov_a, tau_ov_a[numpy.newaxis]])
                    rho_ov_b = numpy.vstack([rho_ov_b, tau_ov_b[numpy.newaxis]])
                    rho_vo_a = rho_ov_a.conj()
                    rho_vo_b = rho_ov_b.conj()
                    w_ov  = numpy.einsum('xsyr,xria->syria', wv_a, rho_ov_a)
                    w_ov += numpy.einsum('xsyr,xria->syria', wv_b, rho_ov_b)
                    wa_ov, wb_ov = w_ov
                    a += lib.einsum('xria,xrjb->iajb', wa_ov, rho_vo_a)
                    a += lib.einsum('xria,xrjb->iajb', wb_ov, rho_vo_b)
                    b += lib.einsum('xria,xrjb->iajb', wa_ov, rho_ov_a)
                    b += lib.einsum('xria,xrjb->iajb', wb_ov, rho_ov_b)
                else:
                    raise NotImplementedError(ni.collinear)

    else:
        a, b = 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=None):
        '''Generate function to compute Ax'''
        if mf is None:
            mf = self._scf
        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==1)[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 > self.positive_eig_threshold)[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.conj())  # (A*)Y
        v1vo -= lib.einsum('xpr,sp->xsr', ys, foo.conj())  # (A*)Y

        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=None):
        if mf is None:
            mf = self._scf
        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.asarray(numpy.block([[x0, y0], [y0, x0.conj()]]))

    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)

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

        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__


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=None):
        if mf is None:
            mf = self._scf
        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,po,qv->xpq', zs*d_ia, orbo, orbv)
            # +cc for A+B because 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 > self.positive_eig_threshold)[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 > self.positive_eig_threshold)[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 and mf.collinear[0] != 'm'):
        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)


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,qv,po->xpq', zs*2, orbv.conj(), orbo)
        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):
        from pyscf.dft import xc_deriv
        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]
                fxc = xc_deriv.transform_fxc(rho, vxc, fxc, xctype, spin=0)
                wfxc = fxc * weight
                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])
                w_ov = numpy.einsum('xyr,xria->yria', wfxc, rho_ov)
                iajb = lib.einsum('xria,xrjb->iajb', w_ov, rho_ov) * 2
                a += iajb
                b += iajb

        elif xctype == 'HF':
            pass

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

        elif xctype == 'MGGA':
            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]
                fxc = xc_deriv.transform_fxc(rho, vxc, fxc, xctype, spin=0)
                wfxc = fxc * weight
                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])
                tau_ov = numpy.einsum('xri,xra->ria', rho_o[1:4], rho_v[1:4]) * .5
                rho_ov = numpy.vstack([rho_ov, tau_ov[numpy.newaxis]])
                w_ov = numpy.einsum('xyr,xria->yria', wfxc, rho_ov)
                iajb = lib.einsum('xria,xrjb->iajb', w_ov, rho_ov) * 2
                a += iajb
                b += iajb

    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)
    # Low excitation filter to avoid numerical instability
    positive_eig_threshold = getattr(__config__, 'tdscf_rhf_TDDFT_positive_eig_threshold', 1e-3)

    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=None):
        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=None):
        '''Generate function to compute Ax'''
        if mf is None:
            mf = self._scf
        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 > self.positive_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)

        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,qv,po->xpq', xs*2, orbv.conj(), orbo)
        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,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, -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=None):
        if mf is None:
            mf = self._scf
        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.asarray(numpy.block([[x0, y0], [y0, x0.conj()]]))

    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 > self.positive_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__


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=None):
        if mf is None:
            mf = self._scf
        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 > self.positive_eig_threshold)[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 > self.positive_eig_threshold)[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_rks()
        # 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_rks()
        # 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)


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,qv,po->xpq', za, orbva.conj(), orboa)
        dmovb = lib.einsum('xov,qv,po->xpq', zb, orbvb.conj(), orbob)

        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):
        from pyscf.dft import xc_deriv
        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, xctype)
                rho0b = make_rho(1, ao, mask, xctype)
                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, xctype)
                rho0b = make_rho(1, ao, mask, xctype)
                rho = (rho0a, rho0b)
                vxc, fxc = ni.eval_xc(mf.xc, rho, 1, deriv=2)[1:3]
                fxc = xc_deriv.transform_fxc(rho, vxc, fxc, xctype, spin=1)
                wfxc = fxc * weight
                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])
                w_ov_aa = numpy.einsum('xyr,xria->yria', wfxc[0,:,0], rho_ov_a)
                w_ov_ab = numpy.einsum('xyr,xria->yria', wfxc[0,:,1], rho_ov_a)
                w_ov_bb = numpy.einsum('xyr,xria->yria', wfxc[1,:,1], rho_ov_b)

                iajb = lib.einsum('xria,xrjb->iajb', w_ov_aa, rho_ov_a)
                a_aa += iajb
                b_aa += iajb

                iajb = lib.einsum('xria,xrjb->iajb', w_ov_bb, rho_ov_b)
                a_bb += iajb
                b_bb += iajb

                iajb = lib.einsum('xria,xrjb->iajb', w_ov_ab, rho_ov_b)
                a_ab += iajb
                b_ab += iajb

        elif xctype == 'HF':
            pass

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

        elif xctype == 'MGGA':
            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, xctype)
                rho0b = make_rho(1, ao, mask, xctype)
                rho = (rho0a, rho0b)
                vxc, fxc = ni.eval_xc(mf.xc, rho, 1, deriv=2)[1:3]
                fxc = xc_deriv.transform_fxc(rho, vxc, fxc, xctype, spin=1)
                wfxc = fxc * weight
                rho_oa = lib.einsum('xrp,pi->xri', ao, orbo_a)
                rho_ob = lib.einsum('xrp,pi->xri', ao, orbo_b)
                rho_va = lib.einsum('xrp,pi->xri', ao, orbv_a)
                rho_vb = lib.einsum('xrp,pi->xri', ao, orbv_b)
                rho_ov_a = numpy.einsum('xri,ra->xria', rho_oa, rho_va[0])
                rho_ov_b = numpy.einsum('xri,ra->xria', rho_ob, rho_vb[0])
                rho_ov_a[1:4] += numpy.einsum('ri,xra->xria', rho_oa[0], rho_va[1:4])
                rho_ov_b[1:4] += numpy.einsum('ri,xra->xria', rho_ob[0], rho_vb[1:4])
                tau_ov_a = numpy.einsum('xri,xra->ria', rho_oa[1:4], rho_va[1:4]) * .5
                tau_ov_b = numpy.einsum('xri,xra->ria', rho_ob[1:4], rho_vb[1:4]) * .5
                rho_ov_a = numpy.vstack([rho_ov_a, tau_ov_a[numpy.newaxis]])
                rho_ov_b = numpy.vstack([rho_ov_b, tau_ov_b[numpy.newaxis]])
                w_ov_aa = numpy.einsum('xyr,xria->yria', wfxc[0,:,0], rho_ov_a)
                w_ov_ab = numpy.einsum('xyr,xria->yria', wfxc[0,:,1], rho_ov_a)
                w_ov_bb = numpy.einsum('xyr,xria->yria', wfxc[1,:,1], rho_ov_b)

                iajb = lib.einsum('xria,xrjb->iajb', w_ov_aa, rho_ov_a)
                a_aa += iajb
                b_aa += iajb

                iajb = lib.einsum('xria,xrjb->iajb', w_ov_bb, rho_ov_b)
                a_bb += iajb
                b_bb += iajb

                iajb = lib.einsum('xria,xrjb->iajb', w_ov_ab, rho_ov_b)
                a_ab += iajb
                b_ab += iajb

    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=None):
        '''Generate function to compute Ax'''
        if mf is None:
            mf = self._scf
        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 > self.positive_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,qv,po->xpq', xa, orbva.conj(), orboa)
        dmsb  = lib.einsum('xov,qv,po->xpq', xb, orbvb.conj(), orbob)
        dmsa += lib.einsum('xov,pv,qo->xpq', ya, orbva, orboa.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)
        v1avo = lib.einsum('xpq,qo,pv->xov', v1ao[0], orboa, orbva.conj())
        v1bov = lib.einsum('xpq,po,qv->xov', v1ao[1], orbob.conj(), orbvb)
        v1bvo = lib.einsum('xpq,qo,pv->xov', v1ao[1], orbob, orbvb.conj())

        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=None):
        if mf is None:
            mf = self._scf
        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.asarray(numpy.block([[x0, y0], [y0, x0.conj()]]))

    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 > self.positive_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__


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 gen_vind(self, mf=None):
        if mf is None:
            mf = self._scf
        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.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 > self.positive_eig_threshold)[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] < self.positive_eig_threshold:
                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_uks()
        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_uks()
        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.scf._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)

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

    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]
    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 = a + numpy.einsum('iabj->iajb', eri_mo[:nocc,nocc:,nocc:,:nocc])
        a = a - numpy.einsum('ijba->iajb', eri_mo[:nocc,:nocc,nocc:,nocc:]) * hyb
        b = b + numpy.einsum('iajb->iajb', eri_mo[:nocc,nocc:,:nocc,nocc:])
        b = b - numpy.einsum('jaib->iajb', eri_mo[:nocc,nocc:,:nocc,nocc:]) * hyb
        return a, b

    if isinstance(mf, dft.KohnShamDFT):
        from pyscf.dft import xc_deriv
        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)

        a, b = add_hf_(a, b, hyb)

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

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

        def ud2tm(aa, ab, ba, bb):
            return numpy.stack([aa + bb,        # rho
                                ba + ab,        # mx
                                (ba - ab) * 1j, # my
                                aa - bb])       # mz

        if xctype == 'LDA':
            ao_deriv = 0
            for ao, mask, weight, coords \
                    in ni.block_loop(mol, mf.grids, nao, ao_deriv, max_memory,
                                     with_s=False):
                if ni.collinear[0] == 'm':
                    rho = ni.eval_rho(mol, ao, dm0, mask, xctype, hermi=1, with_lapl=False)
                    eval_xc = ni.mcfun_eval_xc_adapter(mf.xc)
                    fxc = eval_xc(mf.xc, rho, deriv=2, xctype=xctype)[2]
                    wfxc = weight * fxc.reshape(4,4,-1)
                    mo_oa, mo_va, mo_ob, mo_vb = get_mo_value(ao)
                    rho_ov_aa = numpy.einsum('ri,ra->ria', mo_oa.conj(), mo_va)
                    rho_ov_ab = numpy.einsum('ri,ra->ria', mo_oa.conj(), mo_vb)
                    rho_ov_ba = numpy.einsum('ri,ra->ria', mo_ob.conj(), mo_va)
                    rho_ov_bb = numpy.einsum('ri,ra->ria', mo_ob.conj(), mo_vb)
                    rho_ov = ud2tm(rho_ov_aa, rho_ov_ab, rho_ov_ba, rho_ov_bb)
                    rho_vo = rho_ov.conj()
                    w_ov = numpy.einsum('tsr,tria->sria', wfxc, rho_ov)
                    a += lib.einsum('sria,srjb->iajb', w_ov, rho_vo)
                    b += lib.einsum('sria,srjb->iajb', w_ov, rho_ov)
                elif ni.collinear[0] == 'c':
                    rho = ni.eval_rho(mol, ao, dm0, mask, xctype, hermi=1)
                    fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2)[2]
                    wv_a, wv_b = weight * fxc.reshape(2,2,-1)
                    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  = wv_a[:,:,None,None] * rho_ov_a
                    w_ov += wv_b[:,:,None,None] * rho_ov_b
                    wa_ov, wb_ov = w_ov
                    a += lib.einsum('ria,rjb->iajb', wa_ov, rho_vo_a)
                    a += lib.einsum('ria,rjb->iajb', wb_ov, rho_vo_b)
                    b += lib.einsum('ria,rjb->iajb', wa_ov, rho_ov_a)
                    b += lib.einsum('ria,rjb->iajb', wb_ov, rho_ov_b)
                else:
                    raise NotImplementedError(ni.collinear)

        elif xctype == 'GGA':
            ao_deriv = 1
            for ao, mask, weight, coords \
                    in ni.block_loop(mol, mf.grids, nao, ao_deriv, max_memory,
                                     with_s=False):
                if ni.collinear[0] == 'm':
                    rho = ni.eval_rho(mol, ao, dm0, mask, xctype, hermi=1, with_lapl=False)
                    eval_xc = ni.mcfun_eval_xc_adapter(mf.xc)
                    fxc = eval_xc(mf.xc, rho, deriv=2, xctype=xctype)[2]
                    wfxc = weight * fxc
                    mo_oa, mo_va, mo_ob, mo_vb = get_mo_value(ao)
                    rho_ov_aa = numpy.einsum('ri,xra->xria', mo_oa[0].conj(), mo_va)
                    rho_ov_ab = numpy.einsum('ri,xra->xria', mo_oa[0].conj(), mo_vb)
                    rho_ov_ba = numpy.einsum('ri,xra->xria', mo_ob[0].conj(), mo_va)
                    rho_ov_bb = numpy.einsum('ri,xra->xria', mo_ob[0].conj(), mo_vb)
                    rho_ov_aa[1:4] += numpy.einsum('xri,ra->xria', mo_oa[1:4].conj(), mo_va[0])
                    rho_ov_ab[1:4] += numpy.einsum('xri,ra->xria', mo_oa[1:4].conj(), mo_vb[0])
                    rho_ov_ba[1:4] += numpy.einsum('xri,ra->xria', mo_ob[1:4].conj(), mo_va[0])
                    rho_ov_bb[1:4] += numpy.einsum('xri,ra->xria', mo_ob[1:4].conj(), mo_vb[0])
                    rho_ov = ud2tm(rho_ov_aa, rho_ov_ab, rho_ov_ba, rho_ov_bb)
                    rho_vo = rho_ov.conj()
                    w_ov = numpy.einsum('txsyr,txria->syria', wfxc, rho_ov)
                    a += lib.einsum('syria,syrjb->iajb', w_ov, rho_vo)
                    b += lib.einsum('syria,syrjb->iajb', w_ov, rho_ov)
                elif ni.collinear[0] == 'c':
                    rho = ni.eval_rho(mol, ao, dm0, mask, xctype, hermi=1)
                    fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2)[2]
                    wv_a, wv_b = weight * fxc
                    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_ov_b.conj()
                    w_ov  = numpy.einsum('xsyr,xria->syria', wv_a, rho_ov_a)
                    w_ov += numpy.einsum('xsyr,xria->syria', wv_b, rho_ov_b)
                    wa_ov, wb_ov = w_ov
                    a += lib.einsum('xria,xrjb->iajb', wa_ov, rho_vo_a)
                    a += lib.einsum('xria,xrjb->iajb', wb_ov, rho_vo_b)
                    b += lib.einsum('xria,xrjb->iajb', wa_ov, rho_ov_a)
                    b += lib.einsum('xria,xrjb->iajb', wb_ov, rho_ov_b)
                else:
                    raise NotImplementedError(ni.collinear)

        elif xctype == 'HF':
            pass

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

        elif xctype == 'MGGA':
            ao_deriv = 1
            for ao, mask, weight, coords \
                    in ni.block_loop(mol, mf.grids, nao, ao_deriv, max_memory,
                                     with_s=False):
                if ni.collinear[0] == 'm':
                    rho = ni.eval_rho(mol, ao, dm0, mask, xctype, hermi=1, with_lapl=False)
                    eval_xc = ni.mcfun_eval_xc_adapter(mf.xc)
                    fxc = eval_xc(mf.xc, rho, deriv=2, xctype=xctype)[2]
                    wfxc = weight * fxc
                    mo_oa, mo_va, mo_ob, mo_vb = get_mo_value(ao)
                    rho_ov_aa = numpy.einsum('ri,xra->xria', mo_oa[0].conj(), mo_va)
                    rho_ov_ab = numpy.einsum('ri,xra->xria', mo_oa[0].conj(), mo_vb)
                    rho_ov_ba = numpy.einsum('ri,xra->xria', mo_ob[0].conj(), mo_va)
                    rho_ov_bb = numpy.einsum('ri,xra->xria', mo_ob[0].conj(), mo_vb)
                    rho_ov_aa[1:4] += numpy.einsum('xri,ra->xria', mo_oa[1:4].conj(), mo_va[0])
                    rho_ov_ab[1:4] += numpy.einsum('xri,ra->xria', mo_oa[1:4].conj(), mo_vb[0])
                    rho_ov_ba[1:4] += numpy.einsum('xri,ra->xria', mo_ob[1:4].conj(), mo_va[0])
                    rho_ov_bb[1:4] += numpy.einsum('xri,ra->xria', mo_ob[1:4].conj(), mo_vb[0])
                    tau_ov_aa = numpy.einsum('xri,xra->ria', mo_oa[1:4].conj(), mo_va[1:4]) * .5
                    tau_ov_ab = numpy.einsum('xri,xra->ria', mo_oa[1:4].conj(), mo_vb[1:4]) * .5
                    tau_ov_ba = numpy.einsum('xri,xra->ria', mo_ob[1:4].conj(), mo_va[1:4]) * .5
                    tau_ov_bb = numpy.einsum('xri,xra->ria', mo_ob[1:4].conj(), mo_vb[1:4]) * .5
                    rho_ov_aa = numpy.vstack([rho_ov_aa, tau_ov_aa[numpy.newaxis]])
                    rho_ov_ab = numpy.vstack([rho_ov_ab, tau_ov_ab[numpy.newaxis]])
                    rho_ov_ba = numpy.vstack([rho_ov_ba, tau_ov_ba[numpy.newaxis]])
                    rho_ov_bb = numpy.vstack([rho_ov_bb, tau_ov_bb[numpy.newaxis]])
                    rho_ov = ud2tm(rho_ov_aa, rho_ov_ab, rho_ov_ba, rho_ov_bb)
                    rho_vo = rho_ov.conj()
                    w_ov = numpy.einsum('txsyr,txria->syria', wfxc, rho_ov)
                    a += lib.einsum('syria,syrjb->iajb', w_ov, rho_vo)
                    b += lib.einsum('syria,syrjb->iajb', w_ov, rho_ov)
                elif ni.collinear[0] == 'c':
                    rho = ni.eval_rho(mol, ao, dm0, mask, xctype, hermi=1, with_lapl=False)
                    fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2)[2]
                    wv_a, wv_b = weight * fxc
                    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])
                    tau_ov_a = numpy.einsum('xri,xra->ria', mo_oa[1:4].conj(), mo_va[1:4]) * .5
                    tau_ov_b = numpy.einsum('xri,xra->ria', mo_ob[1:4].conj(), mo_vb[1:4]) * .5
                    rho_ov_a = numpy.vstack([rho_ov_a, tau_ov_a[numpy.newaxis]])
                    rho_ov_b = numpy.vstack([rho_ov_b, tau_ov_b[numpy.newaxis]])
                    rho_vo_a = rho_ov_a.conj()
                    rho_vo_b = rho_ov_b.conj()
                    w_ov  = numpy.einsum('xsyr,xria->syria', wv_a, rho_ov_a)
                    w_ov += numpy.einsum('xsyr,xria->syria', wv_b, rho_ov_b)
                    wa_ov, wb_ov = w_ov
                    a += lib.einsum('xria,xrjb->iajb', wa_ov, rho_vo_a)
                    a += lib.einsum('xria,xrjb->iajb', wb_ov, rho_vo_b)
                    b += lib.einsum('xria,xrjb->iajb', wa_ov, rho_ov_a)
                    b += lib.einsum('xria,xrjb->iajb', wb_ov, rho_ov_b)
                else:
                    raise NotImplementedError(ni.collinear)

    else:
        a, b = 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=None):
        '''Generate function to compute Ax'''
        if mf is None:
            mf = self._scf
        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=None):
        if mf is None:
            mf = self._scf
        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.
    '''
    vj, vk = _vhf.rdirect_mapdm('int2e_spinor', 's8',
                                ('ji->s2kl', 'jk->s1il'), dm, 1,
                                mol._atm, mol._bas, mol._env, mf_opt)
    vj = vj.reshape(dm.shape)
    vk = vk.reshape(dm.shape)
    return dhf._jk_triu_(mol, 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)
        nocc = mol.nelectron
        mo_occ[:nocc] = 1
        if nocc < len(mo_energy):
            if mo_energy[nocc-1]+1e-3 > mo_energy[nocc]:
                logger.warn(self, 'HOMO %.15g == LUMO %.15g',
                            mo_energy[nocc-1], mo_energy[nocc])
            else:
                logger.info(self, 'nocc = %d  HOMO = %.12g  LUMO = %.12g',
                            nocc, mo_energy[nocc-1], mo_energy[nocc])
        else:
            logger.info(self, 'nocc = %d  HOMO = %.12g  no LUMO',
                        nocc, mo_energy[nocc-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_ah 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'],
    'mcfun': ['mcfun>=0.2.1'],
}
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/gto/test/test_ecp.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>
#

import unittest
import numpy
from pyscf import gto
from pyscf import scf
from pyscf import lib


def setUpModule():
    global mol, mol1, mol2, cu1_basis
    cu1_basis = gto.basis.parse('''
     H    S
           1.8000000              1.0000000
     H    S
           2.8000000              0.0210870             -0.0045400              0.0000000
           1.3190000              0.3461290             -0.1703520              0.0000000
           0.9059000              0.0393780              0.1403820              1.0000000
     H    P
           2.1330000              0.0868660              0.0000000
           1.2000000              0.0000000              0.5000000
           0.3827000              0.5010080              1.0000000
     H    D
           0.3827000              1.0000000
     H    F
           2.1330000              0.1868660              0.0000000
           0.3827000              0.2010080              1.0000000
                                   ''')

    mol = gto.M(atom='''
    Cu1 0. 0. 0.
    Cu 0. 1. 0.
    He 1. 0. 0.
    ''',
                basis={'Cu':'lanl2dz', 'Cu1': cu1_basis, 'He':'sto3g'},
                ecp = {'cu':'lanl2dz'})

    mol1 = gto.M(atom='''
    Cu1 0.  0.  0.
    Cu 0. 1. 0.
    He 1. 0. 0.
    Ghost-Cu1 0.  0.  0.0001
    ''',
                 basis={'Cu':'lanl2dz', 'Cu1': cu1_basis, 'He':'sto3g'},
                 ecp = {'cu':'lanl2dz'})

    mol2 = gto.M(atom='''
    Cu1 0.  0.  0.
    Cu 0. 1. 0.
    He 1. 0. 0.
    Ghost-Cu1 0.  0. -0.0001
    ''',
                 basis={'Cu':'lanl2dz', 'Cu1': cu1_basis, 'He':'sto3g'},
                 ecp = {'cu':'lanl2dz'})

def tearDownModule():
    global mol, mol1, mol2, cu1_basis
    del mol, mol1, mol2, cu1_basis

class KnownValues(unittest.TestCase):
    def test_ecp_by_shell(self):
        for i in (0,2,3,6,9):
            for j in (1,2,3,5,6):
                ref = mol.intor_by_shell('ECPscalar_sph', (i,j))
                dat = gto.ecp.type1_by_shell(mol, (i, j))
                dat+= gto.ecp.type2_by_shell(mol, (i, j))
                self.assertAlmostEqual(abs(ref-dat).max(), 0, 12)

                ref = mol.intor_by_shell('ECPscalar_cart', (i,j))
                dat = gto.ecp.type1_by_shell(mol, (i, j), cart=True)
                dat+= gto.ecp.type2_by_shell(mol, (i, j), cart=True)
                self.assertAlmostEqual(abs(ref-dat).max(), 0, 12)

    def test_nr_rhf(self):
        mol = gto.M(atom='Na 0. 0. 0.;  H  0.  0.  1.',
                    basis={'Na':'lanl2dz', 'H':'sto3g'},
                    ecp = {'Na':'lanl2dz'},
                    verbose=0)
        self.assertAlmostEqual(lib.fp(mol.intor('ECPscalar')), -0.19922134780248762, 9)
        mf = scf.RHF(mol)
        self.assertAlmostEqual(mf.kernel(), -0.45002315563472206, 10)

    def test_bfd(self):
        mol = gto.M(atom='H 0. 0. 0.',
                    basis={'H':'bfd-vdz'},
                    ecp = {'H':'bfd-pp'},
                    spin = 1,
                    verbose=0)
        mf = scf.RHF(mol)
        self.assertAlmostEqual(mf.kernel(), -0.499045, 6)

        mol = gto.M(atom='Na 0. 0. 0.',
                    basis={'Na':'bfd-vtz'},
                    ecp = {'Na':'bfd-pp'},
                    spin = 1,
                    verbose=0)
        mf = scf.RHF(mol)
        self.assertAlmostEqual(mf.kernel(), -0.181799, 6)

        mol = gto.M(atom='Mg 0. 0. 0.',
                    basis={'Mg':'bfd-vtz'},
                    ecp = {'Mg':'bfd-pp'},
                    spin = 0,
                    verbose=0)
        mf = scf.RHF(mol)
        self.assertAlmostEqual(mf.kernel(), -0.784579, 6)

#        mol = gto.M(atom='Ne 0. 0. 0.',
#                    basis={'Ne':'bfd-vdz'},
#                    ecp = {'Ne':'bfd-pp'},
#                    verbose=0)
#        mf = scf.RHF(mol)
#        self.assertAlmostEqual(mf.kernel(), -34.709059, 6)

    def test_ecp_grad(self):
        aoslices = mol.aoslice_nr_by_atom()
        ish0, ish1 = aoslices[0][:2]
        for i in range(ish0, ish1):
            for j in range(mol.nbas):
                shls = (i,j)
                shls1 = (shls[0] + mol.nbas, shls[1])
                ref = (mol1.intor_by_shell('ECPscalar_cart', shls1) -
                       mol2.intor_by_shell('ECPscalar_cart', shls1)) / 0.0002 * lib.param.BOHR
                dat = mol.intor_by_shell('ECPscalar_ipnuc_cart', shls, comp=3)
                self.assertAlmostEqual(abs(-dat[2]-ref).max(), 0, 4)

    def test_ecp_iprinv(self):
        mol = gto.M(atom='''
        Cu 0. 0. 0.
        H  1. 0. 0.
        ''',
                    basis={'Cu':'lanl2dz', 'H':'ccpvdz'},
                    ecp = {'cu':'lanl2dz'})
        mol1 = gto.M(atom='''
        Cu 0. 0. 0.
        H  1. 0. 0.
        Ghost-Cu 0.  0.  0.0001
        ''',
                    basis={'Cu':'lanl2dz', 'H':'ccpvdz'},
                    ecp = {'cu':'lanl2dz'})
        mol2 = gto.M(atom='''
        Cu 0. 0. 0.
        H  1. 0. 0.
        Ghost-Cu 0.  0. -0.0001
        ''',
                    basis={'Cu':'lanl2dz', 'H':'ccpvdz'},
                    ecp = {'cu':'lanl2dz'})
        aoslices = mol.aoslice_nr_by_atom()
        ish0, ish1 = aoslices[0][:2]
        for i in range(ish0, ish1):
            for j in range(mol.nbas):
                shls = (i,j)
                shls1 = (shls[0] + mol.nbas, shls[1])
                ref = (mol1.intor_by_shell('ECPscalar_cart', shls1) -
                       mol2.intor_by_shell('ECPscalar_cart', shls1)) / 0.0002 * lib.param.BOHR
                with mol.with_rinv_at_nucleus(0):
                    dat = mol.intor_by_shell('ECPscalar_iprinv_cart', shls, comp=3)
                self.assertAlmostEqual(abs(-dat[2]-ref).max(), 0, 4)

    def test_ecp_hessian(self):
        aoslices = mol.aoslice_nr_by_atom()
        ish0, ish1 = aoslices[0][:2]
        for i in range(ish0, ish1):
            for j in range(mol.nbas):
                shls = (i,j)
                shls1 = (shls[0] + mol.nbas, shls[1])
                ref =-(mol1.intor_by_shell('ECPscalar_ipnuc_cart', shls1, comp=3) -
                       mol2.intor_by_shell('ECPscalar_ipnuc_cart', shls1, comp=3)) / 0.0002 * lib.param.BOHR
                dat = mol.intor_by_shell('ECPscalar_ipipnuc_cart', shls, comp=9)
                di, dj = dat.shape[1:]
                dat = dat.reshape(3,3,di,dj)
                self.assertAlmostEqual(abs(dat[2]-ref).max(), 0, 3)

        for i in range(mol.nbas):
            for j in range(ish0, ish1):
                shls = (i,j)
                shls1 = (shls[0], shls[1] + mol.nbas)
                ref =-(mol1.intor_by_shell('ECPscalar_ipnuc_cart', shls1, comp=3) -
                       mol2.intor_by_shell('ECPscalar_ipnuc_cart', shls1, comp=3)) / 0.0002 * lib.param.BOHR
                dat = mol.intor_by_shell('ECPscalar_ipnucip_cart', shls, comp=9)
                di, dj = dat.shape[1:]
                dat = dat.reshape(3,3,di,dj)
                self.assertAlmostEqual(abs(dat[:,2]-ref).max(), 0, 3)

    def test_pp_int(self):
        from pyscf import gto, scf
        from pyscf.pbc import gto as pbcgto
        from pyscf.pbc import scf as pbcscf
        from pyscf.pbc import df
        cell = pbcgto.Cell()
        cell.atom = 'He 1. .5 .5; C .1 1.3 2.1'
        cell.basis = {'He': [(0, (2.5, 1)), (0, (1., 1))],
                      'C' :'gth-szv',}
        cell.pseudo = {'C':'gth-pade',
                       'He': pbcgto.pseudo.parse('''He
        2
         0.40000000    3    -1.98934751    -0.75604821    0.95604821
        2
         0.29482550    3     1.23870466    .855         .3
                                           .71         -1.1
                                                        .9
         0.32235865    2     2.25670239    -0.39677748
                                            0.93894690
                                                     ''')}
        cell.a = numpy.eye(3)
        cell.dimension = 0
        cell.build()
        mol = cell.to_mol()

        hcore = scf.RHF(mol).get_hcore()
        mydf = df.AFTDF(cell)
        ref = mydf.get_pp() + mol.intor('int1e_kin')
        self.assertAlmostEqual(abs(hcore-ref).max(), 0, 2)

        mf = pbcscf.RHF(cell)
        mf.with_df = mydf
        mf.run()
        e_ref = mf.e_tot

        e_tot = scf.RHF(mol).run().e_tot
        self.assertAlmostEqual(abs(e_ref-e_tot).max(), 0, 6)

    def test_scalar_vs_int1e_rinv(self):
        mol = gto.M(atom='''
                    Na 0.5 0.5 0.
                    H  1.0 0.  0.2
                    ''',
                    basis={'Na': [(0, (1, 1)), (1, (4, 1)), (2, (1, 1))],
                           'H': 'ccpvtz'},
                    ecp = {'Na': gto.basis.parse_ecp('''
Na nelec 8
Na ul
1      0.    -3.
''')})
        mat = mol.intor('ECPscalar')
        with mol.with_rinv_orig(mol.atom_coord(0)):
            ref = mol.intor('int1e_rinv')*-3
        self.assertAlmostEqual(abs(mat-ref).max(), 0, 9)

    def test_so_vs_int1e_rinv(self):
        mol = gto.M(atom='''
                    Na 0.5 0.5 0.
                    ''',
                    charge=1,
                    basis={'Na': [(0, (1, 1)), (1, (4, 1)), (1, (1, 1)), (2, (1, 1))]},
                    ecp = {'Na': gto.basis.parse_ecp('''
Na nelec 8
Na S
0      0.     0     0
1      0.    -3.    -3.
Na P
1      0.    -3.    -3.
Na D
1      0.    -3.    -3.
Na F
1      0.    -3.    -3.
''')})
        u = mol.sph2spinor_coeff()
        ref = numpy.einsum('sxy,spq,xpi,yqj->ij', lib.PauliMatrices,
                           mol.intor('int1e_inuc_rxp'), u.conj(), u)

        mat = mol.intor('ECPso_spinor')
        self.assertAlmostEqual(abs(ref-mat).max(), 0, 11)

        mat = numpy.einsum('sxy,spq,xpi,yqj->ij', lib.PauliMatrices,
                           mol.intor('ECPso'), u.conj(), u)
        self.assertAlmostEqual(abs(ref-mat).max(), 0, 11)

    def test_ecp_grad1(self):
        mol = gto.M(atom='Na, 0.00, 0.00, 0.00; Cl, 0.00, 0.00, 2.050',
                    basis='lanl2dz', ecp = 'lanl2dz', unit='B')
        with mol.with_rinv_at_nucleus(0):
            iprinv0 = mol.intor('ECPscalar_iprinv')
        with mol.with_rinv_at_nucleus(1):
            iprinv1 = mol.intor('ECPscalar_iprinv')
        ipnuc = mol.intor('ECPscalar_ipnuc')
        self.assertAlmostEqual(abs(iprinv0 + iprinv1 - ipnuc).max(), 0, 12)

        aoslices = mol.aoslice_by_atom()
        atm_id = 1
        shl0, shl1, p0, p1 = aoslices[atm_id]
        mat0 = iprinv1
        mat0[:,p0:p1] -= ipnuc[:,p0:p1]
        mat0 = mat0 + mat0.transpose(0,2,1)

        mat1 = mol.set_geom_('Na, 0.00, 0.00, 0.00; Cl, 0.00, 0.00, 2.049').intor('ECPscalar')
        mat2 = mol.set_geom_('Na, 0.00, 0.00, 0.00; Cl, 0.00, 0.00, 2.051').intor('ECPscalar')
        self.assertAlmostEqual(abs(mat0[2] - (mat2 - mat1) / 0.002).max(), 0, 6)

    def test_ecp_hessian1(self):
        mol = gto.M(atom='Na, 0.00, 0.00, 0.00; Cl, 0.00, 0.00, 2.050',
                    basis='lanl2dz', ecp = 'lanl2dz', unit='B')
        with mol.with_rinv_at_nucleus(1):
            rinv1 = mol.intor('ECPscalar_ipiprinv', comp=9)
            rinv1+= mol.intor('ECPscalar_iprinvip', comp=9)
        ipipnuc = mol.intor('ECPscalar_ipipnuc', comp=9)
        ipnucip = mol.intor('ECPscalar_ipnucip', comp=9)

        aoslices = mol.aoslice_by_atom()
        atm_id = 1
        shl0, shl1, p0, p1 = aoslices[atm_id]
        mat0 = rinv1
        mat0[:,p0:p1] -= ipipnuc[:,p0:p1]
        mat0[:,:,p0:p1] -= ipnucip[:,:,p0:p1]

        nao = mol.nao
        mat1 = mol.set_geom_('Na, 0.00, 0.00, 0.00; Cl, 0.00, 0.00, 2.049').intor('ECPscalar_ipnuc')
        mat2 = mol.set_geom_('Na, 0.00, 0.00, 0.00; Cl, 0.00, 0.00, 2.051').intor('ECPscalar_ipnuc')
        self.assertAlmostEqual(abs(mat0.reshape(3,3,nao,nao)[:,2] - (mat2 - mat1) / 0.002).max(), 0, 5)


if __name__ == '__main__':
    print("Full Tests for ECP")
    unittest.main()



# pyscf/gto/test/test_eval_gto.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 scf
from pyscf import lib


def setUpModule():
    global mol, mol1, r, r1
    mol = gto.M(atom='''
    H 0. 0. 0.
    H 8. 0. 0.
    ''', basis='ccpvqz')
    numpy.random.seed(1)
    r = numpy.random.random((100,3)) * 2

    mol1 = gto.M(atom='''
    H 0. 0. 0.
    H 8. 0. 0.
    ''', basis=[[0, (1, 1)],
                [1, (.9, 1)],
                [2, (.8, 1)],
                [3, (.6, 1)],
                [4, (.5, 1)],
                [5, (.4, 1)]])
    r1 = r[:70]

def tearDownModule():
    global mol, mol1, r, r1
    del mol, mol1, r, r1

class KnownValues(unittest.TestCase):
    def test_cart(self):
        val = mol.eval_gto('GTOval', r)
        self.assertEqual(val.shape, (100,60))
        self.assertAlmostEqual(lib.fp(val), -3.0283379087553808, 9)

        val = mol.eval_gto('GTOval_spinor', r)
        self.assertEqual(val.shape, (2,100,120))
        self.assertTrue(val.dtype == numpy.complex128)
        self.assertAlmostEqual(lib.fp(val), -2.029604328041736+0.80571570114202995j, 9)

        val = mol.eval_gto('GTOval_ip', r)
        self.assertEqual(val.shape, (3,100,60))
        self.assertAlmostEqual(lib.fp(val), -14.526634330008513, 9)

        val = mol.eval_gto('GTOval_ig', r)
        self.assertEqual(val.shape, (3,100,60))
        self.assertAlmostEqual(lib.fp(val), -9.8708097471236808e-08, 9)

        val = mol.eval_gto('GTOval_sph_deriv3', r)
        self.assertEqual(val.shape, (20,100,60))
        self.assertAlmostEqual(lib.fp(val), 207.32371064250577, 9)

        val = mol.eval_gto('GTOval_cart_deriv3', r)
        self.assertEqual(val.shape, (20,100,70))
        self.assertAlmostEqual(lib.fp(val), -93.310347533061787, 9)

        val = mol.eval_gto('GTOval_spinor_deriv3', r)
        self.assertEqual(val.shape, (2,20,100,120))
        self.assertTrue(val.dtype == numpy.complex128)
        self.assertAlmostEqual(lib.fp(val), -2470.5090409822333+64.324145522087917j, 8)

        mol1 = mol.copy()
        mol1.cart = True
        val = mol1.eval_gto('GTOval_ip', r)
        self.assertEqual(val.shape, (3,100,70))
        self.assertAlmostEqual(lib.fp(val), -15.411590075004403, 9)

    def test_deriv(self):
        val = mol1.eval_gto('GTOval_sph_deriv4', r1)
        self.assertEqual(val.shape, (35,70,72))
        self.assertAlmostEqual(lib.fp(val), 9.0043160890191309, 9)

        ref = val
        val = mol1.eval_gto('GTOval_sph_deriv3', r1)
        self.assertEqual(val.shape, (20,70,72))
        self.assertAlmostEqual(abs(val - ref[:20]).max(), 0, 12)

        val = mol1.eval_gto('GTOval_sph_deriv2', r1)
        self.assertEqual(val.shape, (10,70,72))
        self.assertAlmostEqual(abs(val - ref[:10]).max(), 0, 12)

        val = mol1.eval_gto('GTOval_sph_deriv1', r1)
        self.assertEqual(val.shape, (4,70,72))
        self.assertAlmostEqual(abs(val - ref[:4]).max(), 0, 12)

        val = mol1.eval_gto('GTOval_ip_sph', r1)
        self.assertAlmostEqual(abs(val - ref[1:4]).max(), 0, 12)

        val = mol1.eval_gto('GTOval_cart_deriv4', r1)
        self.assertEqual(val.shape, (35,70,112))
        self.assertAlmostEqual(lib.fp(val), -210.58019667772351, 9)

        ref = val
        val = mol1.eval_gto('GTOval_cart_deriv3', r1)
        self.assertAlmostEqual(abs(val - ref[:20]).max(), 0, 12)

        val = mol1.eval_gto('GTOval_cart_deriv2', r1)
        self.assertAlmostEqual(abs(val - ref[:10]).max(), 0, 12)

        val = mol1.eval_gto('GTOval_cart_deriv1', r1)
        self.assertAlmostEqual(abs(val - ref[:4]).max(), 0, 12)

        val = mol1.eval_gto('GTOval_ip_cart', r1)
        self.assertAlmostEqual(abs(val - ref[1:4]).max(), 0, 12)

        val = mol1.eval_gto('GTOval_spinor_deriv4', r1)
        self.assertEqual(val.shape, (2,35,70,144))
        self.assertTrue(val.dtype == numpy.complex128)
        self.assertAlmostEqual(lib.fp(val), 91.706259060820059+147.16582428195531j, 9)

        ref = val
        val = mol1.eval_gto('GTOval_spinor_deriv3', r1)
        self.assertAlmostEqual(abs(val - ref[:,:20]).max(), 0, 12)

        val = mol1.eval_gto('GTOval_spinor_deriv2', r1)
        self.assertAlmostEqual(abs(val - ref[:,:10]).max(), 0, 12)

        val = mol1.eval_gto('GTOval_spinor_deriv1', r1)
        self.assertAlmostEqual(abs(val - ref[:,:4]).max(), 0, 12)

        val = mol1.eval_gto('GTOval_ip_spinor', r1)
        self.assertAlmostEqual(abs(val - ref[:,1:4]).max(), 0, 12)


if __name__ == '__main__':
    print("Full Tests for eval_gto")
    unittest.main()



# pyscf/gto/test/test_ft_ao.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 ctypes
import unittest
import numpy
from pyscf import lib
from pyscf import gto
from pyscf.gto import ft_ao

libpbc = lib.load_library('libpbc')

def setUpModule():
    global mol, b, Gvbase, Gv, gxyz
    mol = gto.Mole()
    mol.atom = '''
    C    1.3    .2       .3
    C     .1   -.1      1.1 '''
    mol.basis = 'ccpvdz'
    mol.build()
    mesh = (7,9,11)

    numpy.random.seed(12)
    invh = numpy.diag(numpy.random.random(3))
    b = 2*numpy.pi * invh
    Gvbase = (numpy.fft.fftfreq(mesh[0], 1./mesh[0]),
              numpy.fft.fftfreq(mesh[1], 1./mesh[1]),
              numpy.fft.fftfreq(mesh[2], 1./mesh[2]))
    Gv = numpy.dot(lib.cartesian_prod(Gvbase), b)
    gxyz = lib.cartesian_prod([numpy.arange(len(x)) for x in Gvbase])

def tearDownModule():
    global mol, Gvbase, Gv, gxyz
    del mol, Gvbase, Gv, gxyz


def ft_ao_o0(mol, Gv):
    nao = mol.nao_nr()
    ngrids = Gv.shape[0]
    aoG = numpy.zeros((nao,ngrids), dtype=numpy.complex128)
    gx = numpy.empty((12,ngrids), dtype=numpy.complex128)
    gy = numpy.empty((12,ngrids), dtype=numpy.complex128)
    gz = numpy.empty((12,ngrids), dtype=numpy.complex128)
    buf = numpy.empty((64,ngrids), dtype=numpy.complex128)
    kk = numpy.einsum('ki,ki->k', Gv, Gv)

    i0 = 0
    for ib in range(mol.nbas):
        ci = mol._libcint_ctr_coeff(ib)
        ei = mol.bas_exp(ib)
        li = mol.bas_angular(ib)
        ri = mol.bas_coord(ib)
        ni = ci.shape[1]
        di = (li*2+1) * ni
        nfi = (li+1)*(li+2)//2
        kr = numpy.dot(Gv,ri)
        cs = numpy.exp(-1j*kr)

        buf[:nfi*ni] = 0
        for ip in range(ci.shape[0]):
            ai = ei[ip]
            fac = (numpy.pi/ai)**1.5 * numpy.exp(-.25/ai*kk)
            gx[0] = 1
            gy[0] = 1
            gz[0] = cs * fac
            if li > 0:
                gx[1] = -1j*Gv[:,0]/(2*ai) * gx[0]
                gy[1] = -1j*Gv[:,1]/(2*ai) * gy[0]
                gz[1] = -1j*Gv[:,2]/(2*ai) * gz[0]
                for m in range(1, li):
                    gx[m+1] = m/(2*ai) * gx[m-1] - 1j*Gv[:,0]/(2*ai) * gx[m]
                    gy[m+1] = m/(2*ai) * gy[m-1] - 1j*Gv[:,1]/(2*ai) * gy[m]
                    gz[m+1] = m/(2*ai) * gz[m-1] - 1j*Gv[:,2]/(2*ai) * gz[m]

            for m,(ix,iy,iz) in enumerate(loop_cart(li)):
                val = gx[ix] * gy[iy] * gz[iz]
                for i, cip in enumerate(ci[ip]):
                    buf[i*nfi+m] += cip*val

        ti = c2s_bra(li, numpy.eye(nfi)).T
        tmp1 = numpy.empty((di,ngrids), dtype=numpy.complex128)
        for i in range(ni):
            tmp1[i*(li*2+1):(i+1)*(li*2+1)] = \
                    numpy.einsum('pi,px->ix', ti, buf[i*nfi:(i+1)*nfi])
        aoG[i0:i0+di] += tmp1
        i0 += di
    return aoG.T

def loop_cart(l):
    for ix in reversed(range(l+1)):
        for iy in reversed(range(l-ix+1)):
            iz = l - ix - iy
            yield ix, iy, iz

def c2s_bra(l, gcart):
    if l == 0:
        return gcart * 0.282094791773878143
    elif l == 1:
        return gcart * 0.488602511902919921
    else:
        m = gcart.shape[1]
        gsph = numpy.empty((l*2+1,m))
        fc2s = gto.moleintor.libcgto.CINTc2s_ket_sph
        fc2s(gsph.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(m),
             gcart.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(l))
        return gsph

class KnownValues(unittest.TestCase):
    def test_ft_ao1(self):
        ref = ft_ao_o0(mol, Gv)
        dat = ft_ao.ft_ao(mol, Gv)
        self.assertTrue(numpy.allclose(ref, dat))

        dat = ft_ao.ft_ao(mol, Gv, b=b, gxyz=gxyz, Gvbase=Gvbase)
        self.assertTrue(numpy.allclose(ref, dat))

    def test_ft_ao2(self):
        numpy.random.seed(12)
        invh = numpy.random.random(3) + numpy.eye(3) * 2.5
        b = 2*numpy.pi * invh
        Gv = numpy.dot(lib.cartesian_prod(Gvbase), b)
        ref = ft_ao_o0(mol, Gv)
        dat = ft_ao.ft_ao(mol, Gv)
        self.assertTrue(numpy.allclose(ref, dat))

        mol1 = mol.copy()
        mol1.cart = True
        ref = ft_ao.ft_ao(mol1, Gv)
        dat = ft_ao.ft_ao(mol1, Gv, b=b, gxyz=gxyz, Gvbase=Gvbase)
        self.assertTrue(numpy.allclose(ref, dat))

    def test_ft_aopair1(self):
        dat = ft_ao.ft_aopair(mol, Gv)
        self.assertAlmostEqual(lib.fp(dat), (-5.9794759129252348+8.07254562525371j), 9)

        dat_s2 = ft_ao.ft_aopair(mol, Gv, aosym='s2')
        nao = dat.shape[-1]
        for i in range(nao):
            for j in range(i+1):
                dat[:,i,j] = dat[:,j,i] = dat_s2[:,i*(i+1)//2+j]
        self.assertAlmostEqual(lib.fp(dat), (-5.9794759129252348+8.07254562525371j), 9)

        dat1 = ft_ao.ft_aopair(mol, Gv, b=b, gxyz=gxyz, Gvbase=Gvbase)
        self.assertAlmostEqual(lib.fp(dat1), (-5.9794759129252348+8.07254562525371j), 9)

    def test_ft_aopair2(self):
        numpy.random.seed(12)
        invh = numpy.random.random(3) + numpy.eye(3) * 2.5
        b = 2*numpy.pi * invh
        Gv = numpy.dot(lib.cartesian_prod(Gvbase), b)
        dat = ft_ao.ft_aopair(mol, Gv)
        self.assertAlmostEqual(lib.fp(dat), (-3.1468496579780125-0.019209667673850885j), 9)

        dat1 = ft_ao.ft_aopair(mol, Gv, b=b, gxyz=gxyz, Gvbase=Gvbase)
        self.assertAlmostEqual(lib.fp(dat1), (-3.1468496579780125-0.019209667673850885j), 9)

    def test_ft_aopair_pdotp(self):
        dat = ft_ao.ft_aopair(mol, Gv, intor='GTO_ft_pdotp_sph')
        self.assertAlmostEqual(lib.fp(dat), (-80.69687735727976+69.239798150854909j), 9)

    def test_ft_aopair_pxp(self):
        dat = ft_ao.ft_aopair(mol, Gv, intor='GTO_ft_pxp_sph', comp=3)
        self.assertAlmostEqual(lib.fp(dat), (3.7490985032017079+43.665863070814687j), 8)

    def test_ft_aopair_overlap0(self):
        G = numpy.asarray([[-1.679872,  1.679872,  2.937055],
                           [-1.425679,  1.425679 , 2.492629],
                           [-1.187609 , 1.187609 , 2.076392]])
        mol = gto.M(atom='Ne 7 0.0 0.0; Ne 7 0.0 0.0', basis='3-21g')
        dat = ft_ao.ft_aopair(mol, G)
        self.assertAlmostEqual(lib.fp(dat), (-1.4150713647161861-0.8020058716859948j), 12)

if __name__ == '__main__':
    print('Full Tests for ft_ao')
    unittest.main()


# pyscf/gto/test/test_mole.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 copy
import unittest
import tempfile
from functools import reduce
import numpy
import scipy.linalg
from pyscf import gto
from pyscf import lib
import pyscf.lib.parameters as param
from pyscf.lib.exceptions import BasisNotFoundError, PointGroupSymmetryError

def setUpModule():
    global mol0, ftmp
    mol0 = gto.Mole()
    mol0.atom = [
        [1  , (0.,1.,1.)],
        ["O1", (0.,0.,0.)],
        [1  , (1.,1.,0.)], ]
    mol0.nucmod = { "O":'gaussian', 3:'g' }
    mol0.unit = 'ang'
    mol0.basis = {
        "O": [(0, 0, (15, 1)), ] + gto.etbs(((0, 4, 1, 1.8),
                                             (1, 3, 2, 1.8),
                                             (2, 2, 1, 1.8),)),
        "H": [(0, 0, (1, 1, 0), (3, 3, 1), (5, 1, 0)),
              (1, -2, (1, 1)), ]}
    mol0.symmetry = 1
    mol0.charge = 1
    mol0.spin = 1
    mol0.verbose = 7
    mol0.ecp = {'O1': 'lanl2dz'}
    ftmp = tempfile.NamedTemporaryFile()
    mol0.output = ftmp.name
    mol0.build()

def tearDownModule():
    global mol0, ftmp
    mol0.stdout.close()
    del mol0, ftmp

class KnownValues(unittest.TestCase):
    def test_intor_cross(self):
        mol1 = mol0.unpack(mol0.pack())
        mol1.symmetry = True
        mol1.unit = 'Ang'
        mol1.atom = '''
                1    0  1  .5*2
                O    0  0  0*np.exp(0)
                h    1  1  0'''
        mol1.basis = {'O': gto.basis.parse('''
C    S
   3047.5249000              0.0018347*1.0
    457.3695100              0.0140373*1.0
    103.9486900              0.0688426*1.0
     29.2101550              0.2321844*1.0
      9.2866630              0.4679413*1.0
      3.1639270              0.3623120*1.0
#     1.                     0.1
C    SP
      7.8682724             -0.1193324*1.0          0.0689991        
      1.8812885             -0.1608542*1.0          0.3164240        
      0.5442493              1.1434564*1.0          0.7443083        
C    SP
      0.1687144              1.0000000              1.0000000'''),
                      'H': '6-31g'}
        mol1.build()
        v = gto.mole.intor_cross('cint1e_ovlp_sph', mol0, mol1)
        self.assertAlmostEqual(numpy.linalg.norm(v), 3.6489423434168562, 1)

    def test_num_basis(self):
        self.assertEqual(mol0.nao_nr(), 34)
        self.assertEqual(mol0.nao_2c(), 64)

    def test_time_reversal_map(self):
        tao = [ -2,  1, -4,  3,  8, -7,  6, -5,-10,  9,-12, 11,-14, 13,-16, 15,-18, 17,
                20,-19, 24,-23, 22,-21, 26,-25, 30,-29, 28,-27, 32,-31, 36,-35, 34,-33,
               -40, 39,-38, 37,-46, 45,-44, 43,-42, 41,-50, 49,-48, 47,-56, 55,-54, 53,
               -52, 51,-58, 57,-60, 59, 64,-63, 62,-61]
        self.assertEqual(list(mol0.time_reversal_map()), tao)

    def test_check_sanity(self):
        mol1 = mol0.copy()
        mol1.x = None
        mol1.copy = None
        mol1.check_sanity()

    def test_nao_range(self):
        self.assertEqual(mol0.nao_nr_range(1,4), (2, 7))
        self.assertEqual(mol0.nao_2c_range(1,4), (4, 12))
        self.assertEqual(numpy.dot(range(mol0.nbas+1), mol0.ao_loc_nr()), 2151)
        self.assertEqual(numpy.dot(range(mol0.nbas+1), mol0.ao_loc_2c()), 4066)

    def test_search_bas(self):
        self.assertEqual(mol0.search_shell_id(1, 1), 7)
        self.assertRaises(RuntimeError, mol0.search_ao_nr, 1, 1, -1, 5)
        self.assertEqual(mol0.search_ao_nr(1, 1, -1, 4), 16)
        mol0.cart = True
        self.assertEqual(mol0.search_ao_nr(2, 1, -1, 1), 30)
        mol0.cart = False

    def test_atom_types(self):
        atoms = [['H0', ( 0, 0, 0)],
                 ['H1', ( 0, 0, 0)],
                 ['H',  ( 0, 0, 0)],
                 ['H3', ( 0, 0, 0)]]
        basis = {'H':'sto3g', 'H1': '6-31g'}
        atmgroup = gto.mole.atom_types(atoms, basis)
        self.assertEqual(atmgroup, {'H': [0, 2, 3], 'H1': [1]})
        atoms = [['H0', ( 0, 0, 0)],
                 ['H1', ( 0, 0, 0)],
                 ['H2', ( 0, 0, 0)],
                 ['H3', ( 0, 0, 0)]]
        basis = {'H2':'sto3g', 'H3':'6-31g', 'H0':'sto3g', 'H1': '6-31g'}
        atmgroup = gto.mole.atom_types(atoms, basis)
        self.assertEqual(atmgroup, {'H2': [2], 'H3': [3], 'H0': [0], 'H1': [1]})

    def test_input_symmetry(self):
        mol = gto.M(atom='H 0 0 -1; H 0 0 1', symmetry='D2h')
        self.assertEqual(mol.irrep_id, [0, 5])
        mol = gto.M(atom='H 0 0 -1; H 0 0 1', symmetry='D2')
        self.assertEqual(mol.irrep_id, [0, 1])
        mol = gto.M(atom='H 0 0 -1; H 0 0 1', symmetry='C2v')
        self.assertEqual(mol.irrep_id, [0])

    def test_dumps_loads(self):
        import warnings
        mol1 = gto.M()
        mol1.x = lambda *args: None
        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter("always")
            d = mol1.dumps()
            self.assertTrue(w[0].category, UserWarning)
        mol1.loads(mol0.dumps())

    def test_symm_orb_serialization(self):
        '''Handle the complex symmetry-adapted orbitals'''
        mol = gto.M(atom='He', basis='ccpvdz', symmetry=True)
        mol.loads(mol.dumps())

        lz_minus = numpy.sqrt(.5) * (mol.symm_orb[3] - mol.symm_orb[2] * 1j)
        lz_plus = -numpy.sqrt(.5) * (mol.symm_orb[3] + mol.symm_orb[2] * 1j)
        mol.symm_orb[2] = lz_minus
        mol.symm_orb[3] = lz_plus
        mol.loads(mol.dumps())
        self.assertTrue(mol.symm_orb[0].dtype == numpy.double)
        self.assertTrue(mol.symm_orb[2].dtype == numpy.complex128)
        self.assertTrue(mol.symm_orb[3].dtype == numpy.complex128)

    def test_same_mol1(self):
        self.assertTrue(gto.same_mol(mol0, mol0))
        mol1 = gto.M(atom='h   0  1  1; O1  0  0  0;  h   1  1  0')
        self.assertTrue(not gto.same_mol(mol0, mol1))
        self.assertTrue(gto.same_mol(mol0, mol1, cmp_basis=False))

        mol1 = gto.M(atom='h   0  1  1; O1  0  0  0;  h   1  1  0.01')
        self.assertTrue(not gto.same_mol(mol0, mol1, cmp_basis=False))
        self.assertTrue(gto.same_mol(mol0, mol1, tol=.02, cmp_basis=False))

        mol1 = gto.M(atom='''H 0.0052917700 0.0000000000 -0.8746076326
                             F 0.0000000000 0.0000000000 0.0516931447''')
        mol2 = gto.M(atom='''H 0.0000000000 0.0000000000 -0.8746076326
                             F 0.0000000000 0.0000000000 0.0516931447''')
        self.assertTrue(gto.same_mol(mol1, mol2))
        self.assertTrue(not gto.same_mol(mol1, mol2, tol=1e-6))
        mol3 = gto.M(atom='''H 0.0000000000 0.0000000000 -0.8746076326
                             H 0.0000000000 0.0000000000 0.0516931447''')
        self.assertTrue(not gto.same_mol(mol3, mol2))

    def test_same_mol2(self):
        mol1 = gto.M(atom='H 0.0052917700 0.0000000000 -0.8746076326; F 0.0000000000 0.0000000000 0.0464013747')
        mol2 = gto.M(atom='H 0.0000000000 0.0000000000 -0.8746076326; F 0.0052917700 0.0000000000 0.0464013747')
        self.assertTrue(gto.same_mol(mol1, mol2))

        mol1 = gto.M(atom='H 0.0052917700 0.0000000000 -0.8693158626; F 0.0000000000 0.0000000000 0.0464013747')
        mol2 = gto.M(atom='H 0.0000000000 0.0052917700 -0.8693158626; F 0.0000000000 0.0000000000 0.0464013747')
        mol3 = gto.M(atom='H 0.0000000000 0.0000000000 -0.8693158626; F 0.0052917700 0.0000000000 0.0464013747')
        mol4 = gto.M(atom='H -0.0052917700 0.0000000000 -0.8746076326; F 0.0000000000 0.0000000000 0.0411096047')
        mols = (mol1, mol2, mol3, mol4)
        for i,mi in enumerate(mols):
            for j in range(i):
                self.assertTrue(gto.same_mol(mols[i], mols[j]))

        mol1 = gto.M(atom='''H 0.0000000000 0.0000000000 0.0000000000
          H 0.9497795800 1.3265673200 0.0000000000
          H 0.9444878100 -1.3265673200 0.0000000000
          H1 -0.9444878100 0.0000000000 1.3265673200
          H1 -0.9444878100 0.0000000000 -1.3265673200''', basis={'H':'sto3g', 'H1':'sto3g'}, charge=1)
        mol2 = gto.M(atom='''H 0.0000000000 0.0000000000 0.0000000000
          H 0.9444878100 1.3265673200 0.0000000000
          H 0.9497795800 -1.3265673200 0.0000000000
          H1 -0.9444878100 0.0000000000 1.3265673200
          H1 -0.9444878100 0.0000000000 -1.3265673200''', basis={'H':'sto3g', 'H1':'sto3g'}, charge=1)
        self.assertTrue(gto.same_mol(mol1, mol2))
        self.assertEqual(len(gto.atom_types(mol1._atom)), 2)
        mol3 = gto.M(atom='''H 0.0000000000 0.0000000000 0.0000000000
          H1 0.9497795800 1.3265673200 0.0000000000
          H1 0.9444878100 -1.3265673200 0.0000000000
          H1 -0.9444878100 0.0000000000 1.3265673200
          H1 -0.9444878100 0.0000000000 -1.3265673200''', basis={'H':'sto3g', 'H1':'321g'}, charge=1)
        self.assertTrue(not gto.same_mol(mol3, mol2))

    def test_inertia_momentum(self):
        mol1 = gto.Mole()
        mol1.atom = mol0.atom
        mol1.nucmod = 'G'
        mol1.verbose = 5
        mol1.nucprop = {'H': {'mass': 3}}
        mol1.output = '/dev/null'
        mol1.build(False, False)
        self.assertAlmostEqual(lib.fp(gto.inertia_moment(mol1)),
                               5.340587366981696, 9)

        mass = mol0.atom_mass_list(isotope_avg=True)
        self.assertAlmostEqual(lib.fp(gto.inertia_moment(mol1, mass)),
                               2.1549269955776205, 9)

    def test_chiral_mol(self):
        mol1 = gto.M(atom='C 0 0 0; H 1 1 1; He -1 -1 1; Li -1 1 -1; Be 1 -1 -1')
        mol2 = gto.M(atom='C 0 0 0; H 1 1 1; He -1 -1 1; Be -1 1 -1; Li 1 -1 -1')
        self.assertTrue(gto.chiral_mol(mol1, mol2))
        self.assertTrue(gto.chiral_mol(mol1))

        mol1 = gto.M(atom='''H 0.9444878100 1.3265673200 0.0052917700
                            H 0.9444878100 -1.3265673200 0.0000000000
                            H -0.9444878100 0.0000000000 1.3265673200
                            H -0.9444878100 0.0000000000 -1.3265673200''')
        mol2 = gto.M(atom='''H 0.9444878100 1.3265673200 0.0000000000
                            H 0.9444878100 -1.3265673200 0.0052917700
                            H -0.9444878100 0.0000000000 1.3265673200
                            H -0.9444878100 0.0000000000 -1.3265673200''')
        self.assertTrue(gto.chiral_mol(mol1, mol2))

        mol1 = gto.M(atom='''H 0.9444878100 1.3265673200 0.0052917700
                            H 0.9444878100 -1.3265673200 0.0000000000
                            H -0.9444878100 0.0000000000 1.3265673200
                            H -0.9444878100 0.0000000000 -1.3265673200''')
        self.assertTrue(gto.chiral_mol(mol1))

    def test_first_argument(self):
        mol1 = gto.Mole()
        mol1.build('He')
        self.assertEqual(mol1.natm, 1)

    def test_atom_as_file(self):
        ftmp = tempfile.NamedTemporaryFile('w')
        # file in xyz format
        ftmp.write('He 0 0 0\nHe 0 0 1\n')
        ftmp.flush()
        mol1 = gto.M(atom=ftmp.name)
        self.assertEqual(mol1.natm, 2)

        # file in zmatrix format
        ftmp = tempfile.NamedTemporaryFile('w')
        ftmp.write('He\nHe 1 1.5\n')
        ftmp.flush()
        mol1 = gto.M(atom=ftmp.name)
        self.assertEqual(mol1.natm, 2)

    def test_format_atom(self):
        atoms = [['h' , 0,1,1], "O1  0. 0. 0.", [1, 1.,1.,0.],]
        self.assertTrue(numpy.allclose(gto.mole.format_atom(atoms, unit='Ang')[0][1],
                                       [0.0, 1.8897261245650618, 1.8897261245650618]))
        atoms = '''h 0 1 1
        O1 0 0 0; 1 1 1 0; #H 0 0 3'''
        self.assertTrue(numpy.allclose(gto.mole.format_atom(atoms, unit=1)[0][1],
                                       [0.0, 1., 1.]))
        atoms = 'O1; h 1 1; 1 1 1 2 90'
        atoms = gto.mole.format_atom(atoms, unit=1)[2]
        self.assertEqual(atoms[0], 'H')
        self.assertTrue(numpy.allclose(atoms[1], [0, 0, 1.]))

    def test_format_basis(self):
        mol = gto.M(atom = '''O 0 0 0; 1 0 1 0; H 0 0 1''',
                    basis = {8: 'ccpvdz'})
        self.assertEqual(mol.nao_nr(), 14)

        mol = gto.M(atom = '''O 0 0 0; H:1 0 1 0; H@2 0 0 1''',
                    basis = {'O': 'ccpvdz', 'H:1': 'sto3g', 'H': 'unc-iglo3'})
        self.assertEqual(mol.nao_nr(), 32)

        mol = gto.M(
            atom = '''O 0 0 0; H1 0 1 0; H2 0 0 1''',
            basis = {'default': ('6-31g', [[0, [.05, 1.]], []]), 'H2': 'sto3g'}
        )
        self.assertEqual(mol.nao_nr(), 14)

        mol = gto.M(
            atom = '''O 0 0 0; H1 0 1 0; H2 0 0 1''',
            basis = {'H1': gto.parse('''
# Parse NWChem format basis string (see https://bse.pnl.gov/bse/portal).
# Comment lines are ignored
#BASIS SET: (6s,3p) -> [2s,1p]
        H    S
              2.9412494             -0.09996723
              0.6834831              0.39951283
              0.2222899              0.70011547
        H    S
              2.9412494             0.15591627
              0.6834831             0.60768372
              0.2222899             0.39195739
                                    ''', optimize=True),
                     'O': 'unc-ccpvdz',
                     'H2': gto.load('sto-3g', 'He')  # or use basis of another atom
                    }
        )
        self.assertEqual(mol.nao_nr(), 29)

        mol = gto.M(
            atom = '''O 0 0 0; H1 0 1 0; H2 0 0 1''',
            basis = {'H': ['sto3g', '''unc
        C    S
             71.6168370              0.15432897
             13.0450960              0.53532814
              3.5305122              0.44463454
        C    SP
              2.9412494             -0.09996723             0.15591627
              0.6834831              0.39951283             0.60768372
              0.2222899              0.70011547             0.39195739
        '''],
                     'O': mol.expand_etbs([(0, 4, 1.5, 2.2),  # s-function
                                           (1, 2, 0.5, 2.2)]) # p-function
                    }
        )
        self.assertEqual(mol.nao_nr(), 42)

        mol = gto.M(
            atom = '''O 0 0 0; H1 0 1 0; H2 0 0 1''',
            basis = ('sto3g', 'ccpvdz', '3-21g',
                     gto.etbs([(0, 4, 1.5, 2.2), (1, 2, 0.5, 2.2)]),
                    [[0, numpy.array([1e3, 1.])]])
        )
        self.assertEqual(mol.nao_nr(), 77)

        mol.atom = 'Hg'
        mol.basis = 'ccpvdz'
        self.assertRaises(RuntimeError, mol.build)

    def test_default_basis(self):
        mol = gto.M(atom=[['h' , 0,1,1], ["O1", (0.,0.,0.)], [1, 1.,1.,0.],],
                    basis={'default':'321g', 'O1': 'sto3g'})
        self.assertEqual(sorted(mol._basis.keys()), ['H', 'O1'])

    def test_parse_pople_basis(self):
        self.assertEqual(len(gto.basis.load('6-31G(d)'      , 'H')), 2)
        self.assertEqual(len(gto.basis.load('6-31G(d)'      , 'C')), 6)
        self.assertEqual(len(gto.basis.load('6-31Gs'        , 'C')), 6)
        self.assertEqual(len(gto.basis.load('6-31G*'        , 'C')), 6)
        self.assertEqual(len(gto.basis.load('6-31G(d,p)'    , 'H')), 3)
        self.assertEqual(len(gto.basis.load('6-31G(d,p)'    , 'C')), 6)
        self.assertEqual(len(gto.basis.load('6-31G(2d,2p)'  , 'H')), 4)
        self.assertEqual(len(gto.basis.load('6-31G(2d,2p)'  , 'C')), 7)
        self.assertEqual(len(gto.basis.load('6-31G(3df,3pd)', 'H')), 6)
        self.assertEqual(len(gto.basis.load('6-31G(3df,3pd)', 'C')), 9)

    def test_parse_basis(self):
        mol = gto.M(atom='''
                    6        0    0   -0.5
                    8        0    0    0.5
                    1        1    0   -1.0
                    1       -1    0   -1.0''',
                    basis='''
#BASIS SET: (3s) -> [2s]
H    S
      5.4471780              0.1562849787        
      0.82454724             0.9046908767        
H    S
      0.18319158             1.0000000        
#BASIS SET: (6s,3p) -> [3s,2p]
C    S
    172.2560000              0.0617669        
     25.9109000              0.3587940        
      5.5333500              0.7007130        
C    SP
      3.6649800             -0.3958970              0.2364600        
      0.7705450              1.2158400              0.8606190        
C    SP
      0.1958570              1.0000000              1.0000000        
#BASIS SET: (6s,3p) -> [3s,2p]
O    S
    322.0370000              0.0592394        
     48.4308000              0.3515000        
     10.4206000              0.7076580        
O    SP
      7.4029400             -0.4044530              0.2445860        
      1.5762000              1.2215600              0.8539550        
O    SP
      0.3736840              1.0000000              1.0000000        
''')
        self.assertTrue(mol.nao_nr() == 22)

    def test_ghost(self):
        mol = gto.M(
            atom = 'C 0 0 0; ghost 0 0 2',
            basis = {'C': 'sto3g', 'ghost': gto.basis.load('sto3g', 'H')}
        )
        self.assertEqual(mol.nao_nr(), 6)

        mol = gto.M(atom='''
        ghost-O     0.000000000     0.000000000     2.500000000
        X_H        -0.663641000    -0.383071000     3.095377000
        ghost.H     0.663588000     0.383072000     3.095377000
        O     1.000000000     0.000000000     2.500000000
        H    -1.663641000    -0.383071000     3.095377000
        H     1.663588000     0.383072000     3.095377000
        ''',
        basis='631g')
        self.assertEqual(mol.nao_nr(), 26)

        mol = gto.M(atom='''
        ghost-O     0.000000000     0.000000000     2.500000000
        X_H        -0.663641000    -0.383071000     3.095377000
        ghost.H     0.663588000     0.383072000     3.095377000
        O     1.000000000     0.000000000     2.500000000
        ''',
                basis={'H': '3-21g', 'o': '3-21g', 'ghost-O': 'sto3g'})
        self.assertEqual(mol.nao_nr(), 18) # 5 + 2 + 2 + 9

        mol = gto.M(atom='Zn 0 0 0; ghost-Fe 0 0 1',
                    basis='lanl2dz', ecp='lanl2dz')
        self.assertTrue(len(mol._ecp) == 1)  # only Zn ecp

        mol = gto.M(atom='Zn 0 0 0; ghost-Fe 0 0 1',
                    basis='lanl2dz', ecp={'Zn': 'lanl2dz', 'ghost-Fe': 'lanl2dz'})
        self.assertTrue(len(mol._ecp) == 2)  # Zn and ghost-Fe in ecp

    def test_nucmod(self):
        gto.filatov_nuc_mod(80)
        self.assertEqual(gto.mole._parse_nuc_mod(1), gto.NUC_GAUSS)
        self.assertEqual(gto.mole._parse_nuc_mod('Gaussian'), gto.NUC_GAUSS)
        mol1 = gto.Mole()
        mol1.atom = mol0.atom
        mol1.nucmod = 'G'
        mol1.verbose = 5
        mol1.nucprop = {'H': {'mass': 3}}
        mol1.output = '/dev/null'
        mol1.build(False, False)
        mol1.set_nuc_mod(0, 2)
        self.assertTrue(mol1._atm[1,gto.NUC_MOD_OF] == gto.NUC_GAUSS)
        self.assertAlmostEqual(mol1._env[mol1._atm[0,gto.PTR_ZETA]], 2, 9)
        self.assertAlmostEqual(mol1._env[mol1._atm[1,gto.PTR_ZETA]], 586314366.54656982, 4)

        mol1.set_nuc_mod(1, 0)
        self.assertTrue(mol1._atm[1,gto.NUC_MOD_OF] == gto.NUC_POINT)

        mol1.nucmod = None
        mol1.build(False, False)
        self.assertTrue(mol1._atm[1,gto.NUC_MOD_OF] == gto.NUC_POINT)

        mol1.nucmod = {'H': gto.filatov_nuc_mod}
        mol1.build(False, False)
        self.assertTrue(mol1._atm[0,gto.NUC_MOD_OF] == gto.NUC_GAUSS)
        self.assertTrue(mol1._atm[1,gto.NUC_MOD_OF] == gto.NUC_POINT)
        self.assertTrue(mol1._atm[2,gto.NUC_MOD_OF] == gto.NUC_GAUSS)

    def test_zmat(self):
        coord = numpy.array((
            (0.200000000000, -1.889726124565,  0.000000000000),
            (1.300000000000, -1.889726124565,  0.000000000000),
            (2.400000000000, -1.889726124565,  0.000000000000),
            (3.500000000000, -1.889726124565,  0.000000000000),
            (0.000000000000,  0.000000000000, -1.889726124565),
            (0.000000000000,  1.889726124565,  0.000000000000),
            (0.200000000000, -0.800000000000,  0.000000000000),
            (1.889726124565,  0.000000000000,  1.133835674739)))
        zstr0 = gto.cart2zmat(coord)
        zstr = '\n'.join(['H '+x for x in zstr0.splitlines()])
        atoms = gto.zmat2cart(zstr)
        zstr1 = gto.cart2zmat([x[1] for x in atoms])
        self.assertTrue(zstr0 == zstr1)

        numpy.random.seed(1)
        coord = numpy.random.random((6,3))
        zstr0 = gto.cart2zmat(coord)
        zstr = '\n'.join(['H '+x for x in zstr0.splitlines()])
        atoms = gto.zmat2cart(zstr)
        zstr1 = gto.cart2zmat([x[1] for x in atoms])
        self.assertTrue(zstr0 == zstr1)

    def test_c2s(self):  # Transformation of cart <-> sph, sph <-> spinor
        c = mol0.sph2spinor_coeff()
        s0 = mol0.intor('int1e_ovlp_spinor')
        s1 = mol0.intor('int1e_ovlp_sph')
        sa = reduce(numpy.dot, (c[0].T.conj(), s1, c[0]))
        sa+= reduce(numpy.dot, (c[1].T.conj(), s1, c[1]))
        mol0.cart = True
        s2 = mol0.intor('int1e_ovlp')
        mol0.cart = False
        self.assertAlmostEqual(abs(s0 - sa).max(), 0, 12)
        c = mol0.cart2sph_coeff()
        sa = reduce(numpy.dot, (c.T.conj(), s2, c))
        self.assertAlmostEqual(abs(s1 - sa).max(), 0, 12)

        c0 = gto.mole.cart2sph(1)
        ca, cb = gto.mole.cart2spinor_l(1)
        ua, ub = gto.mole.sph2spinor_l(1)
        self.assertAlmostEqual(abs(c0.dot(ua)-ca).max(), 0, 9)
        self.assertAlmostEqual(abs(c0.dot(ub)-cb).max(), 0, 9)

        c0 = gto.mole.cart2sph(0, normalized='sp')
        ca, cb = gto.mole.cart2spinor_kappa(-1, 0, normalized='sp')
        ua, ub = gto.mole.sph2spinor_kappa(-1, 0)
        self.assertAlmostEqual(abs(c0.dot(ua)-ca).max(), 0, 9)
        self.assertAlmostEqual(abs(c0.dot(ub)-cb).max(), 0, 9)

        c1 = gto.mole.cart2sph(0, numpy.eye(1))
        self.assertAlmostEqual(abs(c0*0.282094791773878143-c1).max(), 0, 12)

        c0 = gto.mole.cart2sph(1, normalized='sp')
        ca, cb = gto.mole.cart2spinor_kappa(1, 1, normalized='sp')
        ua, ub = gto.mole.sph2spinor_kappa(1, 1)
        self.assertAlmostEqual(abs(c0.dot(ua)-ca).max(), 0, 9)
        self.assertAlmostEqual(abs(c0.dot(ub)-cb).max(), 0, 9)

        c1 = gto.mole.cart2sph(1, numpy.eye(3).T)
        self.assertAlmostEqual(abs(c0*0.488602511902919921-c1).max(), 0, 12)

    def test_bas_method(self):
        self.assertEqual([mol0.bas_len_cart(x) for x in range(mol0.nbas)],
                         [1, 3, 1, 1, 1, 1, 1, 3, 3, 3, 6, 6, 1, 3])
        self.assertEqual([mol0.bas_len_spinor(x) for x in range(mol0.nbas)],
                         [2, 4, 2, 2, 2, 2, 2, 6, 6, 6, 10, 10, 2, 4])
        c0 = mol0.bas_ctr_coeff(0)
        self.assertAlmostEqual(abs(c0[:,0]/c0[0,0] - (1,3,1)).max(), 0, 9)
        self.assertAlmostEqual(abs(c0[:,1] - (0,1,0)).max(), 0, 9)

        self.assertRaises(ValueError, mol0.gto_norm, -1, 1.)

    def test_nelectron(self):
        mol = gto.Mole()
        mol.atom = [
            [1  , (0.,1.,1.)],
            ["O1", (0.,0.,0.)],
            [1  , (1.,1.,0.)], ]
        mol.charge = 1
        self.assertEqual(mol.nelectron, 9)

        mol0.nelectron = mol0.nelectron
        mol0.nelectron = mol0.nelectron
        mol0.spin = 2
        self.assertRaises(RuntimeError, lambda *args: mol0.nelec)
        mol0.spin = 1

        mol1 = copy.copy(mol0)
        self.assertEqual(mol1.nelec, (5, 4))
        mol1.nelec = (3, 6)
        self.assertEqual(mol1.nelec, (3, 6))

    def test_multiplicity(self):
        mol1 = copy.copy(mol0)
        self.assertEqual(mol1.multiplicity, 2)
        mol1.multiplicity = 5
        self.assertEqual(mol1.multiplicity, 5)
        self.assertEqual(mol1.spin, 4)
        self.assertRaises(RuntimeError, lambda:mol1.nelec)

    def test_ms(self):
        mol1 = copy.copy(mol0)
        self.assertEqual(mol1.ms, 0.5)
        mol1.ms = 1
        self.assertEqual(mol1.multiplicity, 3)
        self.assertEqual(mol1.spin, 2)
        self.assertRaises(RuntimeError, lambda:mol1.nelec)

    def test_basis_not_found(self):
        mol = gto.M(atom='''
        H     -0.663641000    -0.383071000     3.095377000
        H     0.663588000     0.383072000     3.095377000
        O     0.000000000     0.000000000     2.500000000
        H     -0.663641000    -0.383071000     3.095377000
        H     0.663588000     0.383072000     3.095377000
        O     1.000000000     0.000000000     2.500000000
        H     -0.663641000    -0.383071000     3.095377000
        H     0.663588000     0.383072000     3.095377000
        ''', basis={'O': '3-21g'})
        #TODO: assert the warning "Warn: Basis not found for atom 1 H"
        self.assertEqual(mol.nao_nr(), 18)

        aoslice = mol.aoslice_by_atom()
        self.assertEqual(aoslice[:,0].tolist(), [0, 0, 0, 5, 5, 5,10,10])
        self.assertEqual(aoslice[:,1].tolist(), [0, 0, 5, 5, 5,10,10,10])

    def test_atom_method(self):
        aoslice = mol0.aoslice_by_atom()
        for i in range(mol0.natm):
            symb = mol0.atom_pure_symbol(i)
            shls = mol0.atom_shell_ids(i)
            nshls = aoslice[i][1] - aoslice[i][0]
            self.assertEqual(shls[0], aoslice[i][0])
            self.assertEqual(len(shls), nshls)
            self.assertEqual(mol0.atom_nshells(i), nshls)
        aoslice = mol0.aoslice_2c_by_atom()
        mol0.elements  # test property(elements) in Mole
        self.assertEqual([x[2] for x in aoslice], [0, 8, 56])
        self.assertEqual([x[3] for x in aoslice], [8, 56, 64])

    def test_dump_loads_skip(self):
        import json
        with tempfile.NamedTemporaryFile() as tmpfile:
            lib.chkfile.save_mol(mol0, tmpfile.name)
            mol1 = gto.Mole()
            mol1.update(tmpfile.name)
            # dumps() may produce different orders in different runs
            self.assertEqual(json.loads(mol1.dumps()), json.loads(mol0.dumps()))
        mol1.loads(mol1.dumps())
        mol1.loads_(mol0.dumps())
        mol1.unpack(mol1.pack())
        mol1.unpack_(mol0.pack())

    def test_set_geom(self):
        mol1 = gto.Mole()
        mol1.verbose = 5
        mol1.set_geom_(mol0._atom, 'B', symmetry=True)
        mol1.set_geom_(mol0.atom_coords(), 'B', inplace=False)

        mol1.symmetry = False
        mol1.set_geom_(mol0.atom_coords(), 'B')
        mol1.set_geom_(mol0.atom_coords(), inplace=False)
        mol1.set_geom_(mol0.atom_coords(), unit=1.)
        mol1.set_geom_(mol0.atom_coords(), unit='Ang', inplace=False)

    def test_apply(self):
        from pyscf import scf, mp
        self.assertTrue(isinstance(mol0.apply('RHF'), scf.rohf.ROHF))
        self.assertTrue(isinstance(mol0.apply('MP2'), mp.ump2.UMP2))
        self.assertTrue(isinstance(mol0.apply(scf.RHF), scf.rohf.ROHF))
        self.assertTrue(isinstance(mol0.apply(scf.uhf.UHF), scf.uhf.UHF))

    def test_with_MoleContext(self):
        mol1 = mol0.copy()
        with mol1.with_rinv_at_nucleus(1):
            self.assertTrue(mol1._env[gto.PTR_RINV_ZETA] != 0)
            self.assertAlmostEqual(abs(mol1._env[gto.PTR_RINV_ORIG+2]), 0, 9)
        self.assertAlmostEqual(mol1._env[gto.PTR_RINV_ZETA], 0, 9)
        self.assertAlmostEqual(mol1._env[gto.PTR_RINV_ORIG+2], 0, 9)
        with mol1.with_rinv_at_nucleus(0):
            self.assertAlmostEqual(abs(mol1._env[gto.PTR_RINV_ORIG+2]), 1.8897261245650618, 9)
        self.assertAlmostEqual(mol1._env[gto.PTR_RINV_ORIG+2], 0, 9)

        with mol1.with_rinv_zeta(20):
            self.assertAlmostEqual(mol1._env[gto.PTR_RINV_ZETA], 20, 9)
            mol1.set_rinv_zeta(3.)
        self.assertAlmostEqual(mol1._env[gto.PTR_RINV_ZETA], 0, 9)

        with mol1.with_rinv_origin((1,2,3)):
            self.assertAlmostEqual(mol1._env[gto.PTR_RINV_ORIG+2], 3, 9)
        self.assertAlmostEqual(mol1._env[gto.PTR_RINV_ORIG+2], 0, 9)

        with mol1.with_range_coulomb(20):
            self.assertAlmostEqual(mol1._env[gto.PTR_RANGE_OMEGA], 20, 9)
            mol1.set_range_coulomb(2.)
        self.assertAlmostEqual(mol1._env[gto.PTR_RANGE_OMEGA], 0, 9)

        with mol1.with_common_origin((1,2,3)):
            self.assertAlmostEqual(mol1._env[gto.PTR_COMMON_ORIG+2], 3, 9)
        self.assertAlmostEqual(mol1._env[gto.PTR_COMMON_ORIG+2], 0, 9)

        mol1.set_f12_zeta(2.)

    def test_input_symmetry1(self):
        mol1 = gto.Mole()
        mol1.atom = 'H 1 1 1; H -1 -1 1; H 1 -1 -1; H -1 1 -1'
        mol1.unit = 'B'
        mol1.symmetry = True
        mol1.verbose = 5
        mol1.output = '/dev/null'
        mol1.build()
        self.assertAlmostEqual(lib.fp(mol1.atom_coords()), 3.4708548731841296, 9)

        mol1 = gto.Mole()
        mol1.atom = 'H 0 0 -1; H 0 0 1'
        mol1.cart = True
        mol1.unit = 'B'
        mol1.symmetry = 'Dooh'
        mol1.verbose = 5
        mol1.output = '/dev/null'
        mol1.build()
        self.assertAlmostEqual(lib.fp(mol1.atom_coords()), 0.69980902201036865, 9)

        mol1 = gto.Mole()
        mol1.atom = 'H 0 -1 0; H 0 1 0'
        mol1.unit = 'B'
        mol1.symmetry = True
        mol1.symmetry_subgroup = 'D2h'
        mol1.build()
        self.assertAlmostEqual(lib.fp(mol1.atom_coords()), -1.1939459267317516, 9)

        mol1.atom = 'H 0 0 -1; H 0 0 1'
        mol1.unit = 'B'
        mol1.symmetry = 'Coov'
        mol1.symmetry_subgroup = 'C2'
        mol1.build()
        self.assertAlmostEqual(lib.fp(mol1.atom_coords()), 0.69980902201036865, 9)

        mol1.atom = 'H 1 0 -1; H 0 0 1; He 0 0 2'
        mol1.symmetry = 'Coov'
        self.assertRaises(PointGroupSymmetryError, mol1.build)

        mol1.atom = '''
        C 0. 0. 0.7264
        C 0. 0. -.7264
        H 0.92419 0. 1.29252
        H -.92419 0. 1.29252
        H 0. 0.92419 -1.29252
        H 0. -.92419 -1.29252'''
        mol1.symmetry = True
        mol1.symmetry_subgroup = 'C2v'
        mol1.build()
        self.assertAlmostEqual(lib.fp(mol1.atom_coords()), 2.9413856643164618, 9)

        mol1 = gto.Mole()
        mol1.atom = [
            ["O" , (0. , 0.     , 0.)],
            ["H" , (0. , -0.757 , 0.587)],
            ["H" , (0. , 0.757  , 0.587)]]
        mol1.symmetry = "C3"
        self.assertRaises(PointGroupSymmetryError, mol1.build)

        mol1 = gto.Mole()
        mol1.atom = 'H 0 0 0; H 1 0 0'
        mol1.basis = 'sto-3g'
        mol1.symmetry = 'Dooh'
        mol1.build()
        self.assertAlmostEqual(abs(mol1._symm_axes - numpy.eye(3)[[1,2,0]]).max(), 0, 9)

    def test_symm_orb(self):
        rs = numpy.array([[.1, -.3, -.2],
                          [.3,  .1,  .8]])
        mol = gto.M(atom=[('H', c) for c in rs], unit='Bohr',
                    basis={'H': [[0, (1, 1)], [1, (.9, 1)], [2, (.8, 1)], [3, (.7, 1)]]})

        numpy.random.seed(1)
        u, w, vh = numpy.linalg.svd(numpy.random.random((3,3)))
        rs1 = rs.dot(u) + numpy.array([-.5, -.3, .9])
        mol1 = gto.M(atom=[('H', c) for c in rs1], unit='Bohr',
                     basis={'H': [[0, (1, 1)], [1, (.9, 1)], [2, (.8, 1)], [3, (.7, 1)]]})

        mol.symmetry = 1
        mol.build()
        mol1.symmetry = 1
        mol1.build()

        s0 = mol.intor('int1e_ovlp')
        s0 = [abs(c.T.dot(s0).dot(c)) for c in mol.symm_orb]
        s1 = mol1.intor('int1e_ovlp')
        s1 = [abs(c.T.dot(s1).dot(c)) for c in mol1.symm_orb]
        self.assertTrue(all(abs(s0[i]-s1[i]).max()<1e-12 for i in range(len(mol.symm_orb))))

        mol.cart = True
        mol.symmetry = 1
        mol.build()
        mol1.cart = True
        mol1.symmetry = 1
        mol1.build()

        s0 = mol.intor('int1e_ovlp')
        s0 = [abs(c.T.dot(s0).dot(c)) for c in mol.symm_orb]
        s1 = mol1.intor('int1e_ovlp')
        s1 = [abs(c.T.dot(s1).dot(c)) for c in mol1.symm_orb]
        self.assertTrue(all(abs(s0[i]-s1[i]).max()<1e-12 for i in range(len(mol.symm_orb))))

    def test_search_ao_label(self):
        mol1 = mol0.copy()
        mol1.atom = mol0.atom + ['Mg 1,1,1']
        mol1.ecp['Mg'] = 'lanl2dz'
        mol1.basis['Mg'] = 'lanl2dz'
        mol1.build(0, 0)
        self.assertEqual(list(mol1.search_ao_label('O.*2p')), [10,11,12])
        self.assertEqual(list(mol1.search_ao_label('O1 2p')), [10,11,12])
        self.assertEqual(list(mol1.search_ao_label(['O.*2p','0 H 1s'])), [0, 10,11,12])
        self.assertEqual(list(mol1.search_ao_label([10,11,12])), [10,11,12])
        self.assertEqual(list(mol1.search_ao_label(lambda x: '4d' in x)), [24,25,26,27,28])
        mol1.ao_labels(fmt='%s%s%s%s')
        mol1.sph_labels(fmt=None)
        mol1.cart = True
        self.assertEqual(list(mol1.search_ao_label('4d')), [25,26,27,28,29,30])
        mol1.ao_labels(fmt='%s%s%s%s')
        mol1.ao_labels(fmt=None)
        mol1.cart = False
        mol1.spinor_labels()
        mol1.spinor_labels(fmt='%s%s%s%s')
        mol1.spinor_labels(fmt=None)

    def test_input_ecp(self):
        mol1 = gto.Mole()
        mol1.atom = mol0.atom
        mol1.ecp = 'lanl2dz'
        mol1.build(False, False)
        gto.basis.load_ecp('lanl08', 'O')
        gto.format_ecp({'O':'lanl08', 1:'lanl2dz'})
        self.assertRaises(BasisNotFoundError, gto.format_ecp, {'H':'lan2ldz'})

    def test_condense_to_shell(self):
        mol1 = mol0.copy()
        mol1.symmetry = False
        mol1.build(False, False)
        v = gto.condense_to_shell(mol1, mol1.intor('int1e_ovlp'), numpy.max)
        self.assertAlmostEqual(lib.fp(v), 5.7342530154117846, 9)

    def test_input_ghost_atom(self):
        mol = gto.M(
            atom = 'C 0 0 0; ghost 0 0 2',
            basis = {'C': 'sto3g', 'ghost': gto.basis.load('sto3g', 'H')}
        )

        mol = gto.M(atom='''
        ghost1     0.000000000     0.000000000     2.500000000
        ghost2    -0.663641000    -0.383071000     3.095377000
        ghost2     0.663588000     0.383072000     3.095377000
        O     1.000000000     0.000000000     2.500000000
        H    -1.663641000    -0.383071000     3.095377000
        H     1.663588000     0.383072000     3.095377000
        ''',
                    basis={'ghost1':gto.basis.load('sto3g', 'O'),
                           'ghost2':gto.basis.load('631g', 'H'),
                           'O':'631g', 'H':'631g'}
        )

        mol = gto.M(atom='''
        ghost-O     0.000000000     0.000000000     2.500000000
        ghost_H    -0.663641000    -0.383071000     3.095377000
        ghost:H     0.663588000     0.383072000     3.095377000
        O     1.000000000     0.000000000     2.500000000
        H    -1.663641000    -0.383071000     3.095377000
        H     1.663588000     0.383072000     3.095377000
        ''', basis='631g')

        mol = gto.M(atom='''
        X1     0.000000000     0.000000000     2.500000000
        X2    -0.663641000    -0.383071000     3.095377000
        X2     0.663588000     0.383072000     3.095377000
        O     1.000000000     0.000000000     2.500000000
        H    -1.663641000    -0.383071000     3.095377000
        H     1.663588000     0.383072000     3.095377000
        ''',
                    basis={'X1':gto.basis.load('sto3g', 'O'),
                           'X2':gto.basis.load('631g', 'H'),
                           'O':'631g', 'H':'631g'}
        )

        mol = gto.M(atom='''
        X-O     0.000000000     0.000000000     2.500000000
        X_H1   -0.663641000    -0.383071000     3.095377000
        X:H     0.663588000     0.383072000     3.095377000
        O     1.000000000     0.000000000     2.500000000
        H    -1.663641000    -0.383071000     3.095377000
        H     1.663588000     0.383072000     3.095377000
        ''', basis='631g')

    def test_conc_mole(self):
        mol1 = gto.M(atom='Mg', ecp='LANL2DZ', basis='lanl2dz')
        mol2 = mol1 + mol0
        self.assertEqual(mol2.natm, 4)
        self.assertEqual(mol2.nbas, 18)
        self.assertEqual(mol2.nao_nr(), 42)
        mol2 = mol0 + mol1
        self.assertEqual(mol2.natm, 4)
        self.assertEqual(mol2.nbas, 18)
        self.assertEqual(mol2.nao_nr(), 42)
        n0 = mol0.npgto_nr()
        n1 = mol1.npgto_nr()
        self.assertEqual(mol2.npgto_nr(), n0+n1)
        mol2 = mol2 + mol2
        mol2.cart = True
        self.assertEqual(mol2.npgto_nr(), 100)
        mol3 = gto.M(atom='Cu', basis='lanl2dz', ecp='lanl2dz', spin=None)
        mol4 = mol1 + mol3
        self.assertEqual(len(mol4._ecpbas), 16)

    def test_intor_cross_cart(self):
        mol1 = gto.M(atom='He', basis={'He': [(2,(1.,1))]}, cart=True)
        s0 = gto.intor_cross('int1e_ovlp', mol1, mol0)
        self.assertEqual(s0.shape, (6, 34))
        s0 = gto.intor_cross('int1e_ovlp', mol0, mol1)
        self.assertEqual(s0.shape, (34, 6))
        s0 = gto.intor_cross('int1e_ovlp_cart', mol0, mol1)
        self.assertEqual(s0.shape, (36, 6))

    def test_energy_nuc(self):
        self.assertAlmostEqual(mol0.get_enuc(), 6.3611415029455705, 9)
        self.assertAlmostEqual(gto.M().energy_nuc(), 0, 9)

    def test_fakemol(self):
        numpy.random.seed(1)
        coords = numpy.random.random((6,3))*4
        vref = 0
        mol = mol0.copy()
        for c in coords:
            mol.set_rinv_origin(c)
            vref += mol.intor('int1e_rinv')

        fakemol = gto.fakemol_for_charges(coords)
        pmol = mol + fakemol
        shls_slice = (0, mol.nbas, 0, mol.nbas, mol.nbas, pmol.nbas)
        v = pmol.intor('int3c2e', comp=1, shls_slice=shls_slice)
        v = numpy.einsum('pqk->pq', v)
        self.assertAlmostEqual(abs(vref-v).max(), 0, 12)

    def test_to_uncontracted_cartesian_basis(self):
        pmol, ctr_coeff = mol0.to_uncontracted_cartesian_basis()
        c = scipy.linalg.block_diag(*ctr_coeff)
        s = reduce(numpy.dot, (c.T, pmol.intor('int1e_ovlp'), c))
        self.assertAlmostEqual(abs(s-mol0.intor('int1e_ovlp')).max(), 0, 9)

        mol0.cart = True
        pmol, ctr_coeff = mol0.to_uncontracted_cartesian_basis()
        c = scipy.linalg.block_diag(*ctr_coeff)
        s = reduce(numpy.dot, (c.T, pmol.intor('int1e_ovlp'), c))
        self.assertAlmostEqual(abs(s-mol0.intor('int1e_ovlp')).max(), 0, 9)
        mol0.cart = False

    def test_getattr(self):
        from pyscf import scf, dft, ci, tdscf
        mol = gto.M(atom='He')
        self.assertEqual(mol.HF().__class__, scf.HF(mol).__class__)
        self.assertEqual(mol.KS().__class__, dft.KS(mol).__class__)
        self.assertEqual(mol.UKS().__class__, dft.UKS(mol).__class__)
        self.assertEqual(mol.CISD().__class__, ci.cisd.RCISD)
        self.assertEqual(mol.TDA().__class__, tdscf.rhf.TDA)
        self.assertEqual(mol.dTDA().__class__, tdscf.rks.dTDA)
        self.assertEqual(mol.TDBP86().__class__, tdscf.rks.TDDFTNoHybrid)
        self.assertEqual(mol.TDB3LYP().__class__, tdscf.rks.TDDFT)
        self.assertRaises(AttributeError, lambda: mol.xyz)
        self.assertRaises(AttributeError, lambda: mol.TDxyz)

    def test_ao2mo(self):
        mol = gto.M(atom='He')
        nao = mol.nao
        eri = mol.ao2mo(numpy.eye(nao))
        self.assertAlmostEqual(eri[0,0], 1.0557129427350722, 12)

    def test_tofile(self):
        tmpfile = tempfile.NamedTemporaryFile()
        mol = gto.M(atom=[[1  , (0.,1.,1.)],
                          ["O1", (0.,0.,0.)],
                          [1  , (1.,1.,0.)], ])
        out1 = mol.tofile(tmpfile.name, format='xyz')
        ref = '''3
XYZ from PySCF
H           0.00000        1.00000        1.00000
O           0.00000        0.00000        0.00000
H           1.00000        1.00000        0.00000
'''
        with open(tmpfile.name, 'r') as f:
            self.assertEqual(f.read(), ref)
        self.assertEqual(out1, ref[:-1])

        tmpfile = tempfile.NamedTemporaryFile(suffix='.zmat')
        str1 = mol.tofile(tmpfile.name, format='zmat')
        #FIXME:self.assertEqual(mol._atom, mol.fromfile(tmpfile.name))

    def test_frac_particles(self):
        mol = gto.M(atom=[['h', (0.,1.,1.)],
                          ['O', (0.,0.,0.)],
                          ['h', (1.,1.,0.)],],
                     basis='sto3g')
        mol._atm[1, gto.NUC_MOD_OF] = gto.NUC_FRAC_CHARGE
        mol._env[mol._atm[1, gto.PTR_FRAC_CHARGE]] = 2.5
        self.assertAlmostEqual(mol.atom_charges().sum(), 4.5, 12)
        self.assertAlmostEqual(mol.atom_charge(1), 2.5, 12)

        # Add test after updating cint
        ref = 0
        for ia in range(mol.natm):
            with mol.with_rinv_origin(mol.atom_coord(ia)):
                ref -= mol.intor('int1e_rinv') * mol.atom_charge(ia)
        v = mol.intor('int1e_nuc')
        self.assertAlmostEqual(abs(ref-v).max(), 0, 12)

    def test_fromstring(self):
        mol = gto.Mole()
        mol.fromstring('2\n\nH 0 0 1\nH 0 -1 0')
        print(mol._atom == [('H', [0.0, 0.0, 1.8897261245650618]), ('H', [0.0, -1.8897261245650618, 0.0])])
        print(mol.atom == [('H', [0.0, 0.0, 1.0]), ('H', [0.0, -1.0, 0.0])])
        print(mol.unit == 'Angstrom')

    def test_fromfile(self):
        with tempfile.NamedTemporaryFile(mode='w+', suffix='.xyz') as f:
            f.write('2\n\nH 0 0 1; H 0 -1 0')
            f.flush()
            mol = gto.Mole()
            mol.fromfile(f.name)
            print(mol._atom == [('H', [0.0, 0.0, 1.8897261245650618]), ('H', [0.0, -1.8897261245650618, 0.0])])
            print(mol.atom == [('H', [0.0, 0.0, 1.0]), ('H', [0.0, -1.0, 0.0])])
            print(mol.unit == 'Angstrom')

    def test_uncontract(self):
        basis = gto.basis.parse('''
H    S
0.9  0.8  0
0.5  0.5  0.6
0.3  0.5  0.8
H    S
0.3  1
H    P
0.9  0.6
0.5  0.6
0.3  0.6
''')
        self.assertEqual(gto.uncontract(basis),
                         [[0, [0.9, 1]], [0, [0.5, 1]], [0, [0.3, 1]],
                          [1, [0.9, 1]], [1, [0.5, 1]], [1, [0.3, 1]]])

        basis = [[1, 0, [0.9, .7], [0.5, .7]], [1, [0.5, .8], [0.3, .6]], [1, [0.3, 1]]]
        self.assertEqual(gto.uncontract(basis),
                         [[1, [0.9, 1]], [1, [0.5, 1]], [1, [0.3, 1]]])

        basis = [[1, -2, [0.9, .7], [0.5, .7]], [1, [0.5, .8], [0.3, .6]], [1, [0.3, 1]]]
        self.assertEqual(gto.uncontract(basis),
                         [[1, -2, [0.9, 1]], [1, -2, [0.5, 1]], [1, [0.3, 1]]])

        # FIXME:
        #basis = [[1, [0.9, .7], [0.5, .7]], [1, -2, [0.5, .8], [0.3, .6]], [1, [0.3, 1]]]
        #serl.assertEqual(gto.uncontract(basis),
        #                 [[1, [0.9, 1]], [1, [0.5, 1]], [1, [0.3, 1]]])

if __name__ == "__main__":
    print("test mole.py")
    unittest.main()


# pyscf/gto/test/test_moleintor.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, lib

def setUpModule():
    global mol
    mol = gto.Mole()
    mol.verbose = 0
    mol.output = None#"out_h2o"
    mol.atom = [
        ["C", (-0.65830719,  0.61123287, -0.00800148)],
        ["C1", ( 0.73685281,  0.61123287, -0.00800148)],
        ["C2", ( 1.43439081,  1.81898387, -0.00800148)],
        ["C3", ( 0.73673681,  3.02749287, -0.00920048)],
        ["C4", (-0.65808819,  3.02741487, -0.00967948)],
        ["C5", (-1.35568919,  1.81920887, -0.00868348)],
        ["H", (-1.20806619, -0.34108413, -0.00755148)],
        ["H", ( 1.28636081, -0.34128013, -0.00668648)],
        ["H", ( 2.53407081,  1.81906387, -0.00736748)],
        ["H", ( 1.28693681,  3.97963587, -0.00925948)],
        ["H", (-1.20821019,  3.97969587, -0.01063248)],
        ["H", (-2.45529319,  1.81939187, -0.00886348)],]

    mol.basis = {'H': 'cc-pvdz',
                 'C1': 'CC PVDZ',
                 'C2': 'CC PVDZ',
                 'C3': 'cc-pVDZ',
                 'C4': gto.basis.parse('''
#BASIS SET: (9s,4p,1d) -> [3s,2p,1d]
    C    S
       6665.0000000              0.0006920             -0.0001460
       1000.0000000              0.0053290             -0.0011540
        228.0000000              0.0270770             -0.0057250
         64.7100000              0.1017180             -0.0233120
         21.0600000              0.2747400             -0.0639550
          7.4950000              0.4485640             -0.1499810
          2.7970000              0.2850740             -0.1272620
          0.5215000              0.0152040              0.5445290
    C    S
          0.1596000              1.0000000
    C    P
          9.4390000              0.0381090
          2.0020000              0.2094800
          0.5456000              0.5085570
    C    P
          0.1517000              1.0000000
    C    D
          0.5500000              1.0000000        '''),
                 'C': 'CC PVDZ',}
    mol.ecp = {'C1': 'LANL2DZ'}
    mol.build()

def tearDownModule():
    global mol
    del mol


def fp(mat):
    return abs(mat).sum()


class KnownValues(unittest.TestCase):
    def test_intor_nr(self):
        nao = mol.nao_nr()
        s0 = numpy.empty((3,nao,nao))
        ip = 0
        for i in range(mol.nbas):
            jp = 0
            for j in range(mol.nbas):
                buf = mol.intor_by_shell('int1e_ipovlp_sph', (i,j), comp=3)
                di,dj = buf.shape[1:]
                s0[:,ip:ip+di,jp:jp+dj] = buf
                jp += dj
            ip += di
        s = mol.intor('int1e_ipovlp_sph', comp=3)
        self.assertTrue(numpy.allclose(s,s0))
        self.assertAlmostEqual(fp(s), 960.67081839920604, 11)

    def test_intor_nr0(self):
        s = mol.intor('int1e_ovlp_sph')
        self.assertAlmostEqual(fp(s), 622.29059965181796, 11)

    def test_intor_nr1(self):
        s = mol.intor_symmetric('int1e_ovlp_sph')
        self.assertAlmostEqual(fp(s), 622.29059965181796, 11)

    def test_intor_nr2(self):
        s = mol.intor_asymmetric('int1e_ovlp_sph')
        self.assertAlmostEqual(fp(s), 622.29059965181796, 11)

    def test_intor_nr_cross(self):
        shls_slice = (0, mol.nbas//4, mol.nbas//4, mol.nbas)
        s = mol.intor('int1e_ovlp_sph', shls_slice=shls_slice)
        self.assertAlmostEqual(fp(s), 99.38188078749701, 11)

    def test_intor_r(self):
        s = mol.intor('int1e_ovlp_spinor')
        self.assertAlmostEqual(fp(s), 1592.2297864313475, 11)

    def test_intor_r1(self):
        s0 = mol.intor('int1e_ovlp_spinor')
        s1 = mol.intor_symmetric('int1e_ovlp_spinor')
        self.assertTrue(numpy.allclose(s1,s0))

    def test_intor_r2(self):
        s = mol.intor_asymmetric('int1e_ovlp_spinor')
        self.assertAlmostEqual(fp(s), 1592.2297864313475, 11)

    def test_intor_r_comp(self):
        s = mol.intor('int1e_ipkin_spinor')
        self.assertAlmostEqual(fp(s), 4409.86758420756, 10)
        s1 = mol.intor_asymmetric('int1e_ipkin_spinor')
        self.assertTrue(numpy.allclose(s1,s))

    def test_intor_nr2e(self):
        mol1 = gto.M(atom=[["O" , (0. , 0.     , 0.)],
                           [1   , (0. , -0.757 , 0.587)],
                           [1   , (0. , 0.757  , 0.587)]],
                     basis = '631g')
        nao = mol1.nao_nr()
        eri0 = numpy.empty((3,nao,nao,nao,nao))
        ip = 0
        for i in range(mol1.nbas):
            jp = 0
            for j in range(mol1.nbas):
                kp = 0
                for k in range(mol1.nbas):
                    lp = 0
                    for l in range(mol1.nbas):
                        buf = mol1.intor_by_shell('int2e_ip1_sph', (i,j,k,l), 3)
                        di,dj,dk,dl = buf.shape[1:]
                        eri0[:,ip:ip+di,jp:jp+dj,kp:kp+dk,lp:lp+dl] = buf
                        lp += dl
                    kp += dk
                jp += dj
            ip += di
        buf = mol1.intor_by_shell('int2e_sph', (i,j,k,l))
        self.assertEqual(buf.ndim, 4)

        eri1 = mol1.intor('int2e_ip1_sph').reshape(3,13,13,13,13)
        self.assertTrue(numpy.allclose(eri0, eri1))

        idx = numpy.tril_indices(13)
        naopair = nao * (nao+1) // 2
        ref = eri0[:,idx[0],idx[1]]
        eri1 = mol1.intor('int2e_ip1_sph', comp=3, aosym='s2ij')
        self.assertTrue(numpy.allclose(ref, eri1))

        idx = numpy.tril_indices(13)
        ref = eri0[:,:,:,idx[0],idx[1]]
        eri1 = mol1.intor('int2e_ip1_sph', comp=3, aosym='s2kl')
        self.assertTrue(numpy.allclose(ref, eri1))

        idx = numpy.tril_indices(13)
        ref = eri0[:,idx[0],idx[1]][:,:,idx[0],idx[1]]
        eri1 = mol1.intor('int2e_ip1_sph', comp=3, aosym='4')
        self.assertTrue(numpy.allclose(ref, eri1))


        shls_slice = [0, 4, 2, 6, 1, 5, 3, 5]
        i0,i1,j0,j1,k0,k1,l0,l1 = mol1.ao_loc_nr()[shls_slice]
        eri1 = mol1.intor('int2e_ip1_sph', comp=3, shls_slice=shls_slice)
        eri1 = eri1.reshape(3,i1-i0,j1-j0,k1-k0,l1-l0)
        self.assertTrue(numpy.allclose(eri0[:,i0:i1,j0:j1,k0:k1,l0:l1], eri1))

        shls_slice = [5, 7, 7, 9, 1, 5, 4, 6]
        i0,i1,j0,j1,k0,k1,l0,l1 = mol1.ao_loc_nr()[shls_slice]
        idx = numpy.tril_indices(2)
        ref = eri0[:,i0:i1,j0:j1,k0:k1,l0:l1][:,idx[0],idx[1]]
        eri1 = mol1.intor('int2e_ip1_sph', comp=3, aosym='s2ij', shls_slice=shls_slice)
        eri1 = eri1.reshape(3,-1,k1-k0,l1-l0)
        self.assertTrue(numpy.allclose(ref, eri1))

        shls_slice = [1, 5, 4, 6, 5, 7, 7, 9]
        i0,i1,j0,j1,k0,k1,l0,l1 = mol1.ao_loc_nr()[shls_slice]
        idx = numpy.tril_indices(2)
        ref = eri0[:,i0:i1,j0:j1,k0:k1,l0:l1][:,:,:,idx[0],idx[1]]
        eri1 = mol1.intor('int2e_ip1_sph', comp=3, aosym='s2kl', shls_slice=shls_slice)
        eri1 = eri1.reshape(3,i1-i0,j1-j0,-1)
        self.assertTrue(numpy.allclose(ref, eri1))

        shls_slice = [5, 7, 7, 9, 7, 9, 5, 7]
        i0,i1,j0,j1,k0,k1,l0,l1 = mol1.ao_loc_nr()[shls_slice]
        idx = numpy.tril_indices(2)
        ref = eri0[:,i0:i1,j0:j1,k0:k1,l0:l1][:,idx[0],idx[1]][:,:,idx[0],idx[1]]
        eri1 = mol1.intor('int2e_ip1_sph', comp=3, aosym='s4', shls_slice=shls_slice)
        self.assertTrue(numpy.allclose(ref, eri1))

        eri1 = mol1.intor('int2e_ip1_sph', comp=3, aosym='s1', shls_slice=(2,4,1,3))
        self.assertEqual(eri1.shape, (3,4,2,13,13))
        self.assertTrue(numpy.allclose(eri0[:,2:6,1:3], eri1))

    def test_intor_r2e(self):
        mol1 = gto.M(atom=[[1   , (0. , -0.7 , 0.)],
                           [1   , (0. ,  0.7 , 0.)]],
                     basis = '631g')
        nao = mol1.nao_2c()
        eri0 = numpy.empty((3,nao,nao,nao,nao), dtype=numpy.complex128)
        ip = 0
        for i in range(mol1.nbas):
            jp = 0
            for j in range(mol1.nbas):
                kp = 0
                for k in range(mol1.nbas):
                    lp = 0
                    for l in range(mol1.nbas):
                        buf = mol1.intor_by_shell('int2e_ip1_spinor', (i,j,k,l), comp=3)
                        di,dj,dk,dl = buf.shape[1:]
                        eri0[:,ip:ip+di,jp:jp+dj,kp:kp+dk,lp:lp+dl] = buf
                        lp += dl
                    kp += dk
                jp += dj
            ip += di

    def test_input_cint(self):
        '''Compatibility to old cint functions
        '''
        self.assertEqual(gto.moleintor.ascint3('cint2e'), 'int2e_spinor')
        self.assertEqual(gto.moleintor.ascint3('cint2e_sph'), 'int2e_sph')

    def test_rinv_with_zeta(self):
        with mol.with_rinv_orig((.2,.3,.4)), mol.with_rinv_zeta(2.2):
            v1 = mol.intor('int1e_rinv_sph')
        pmol = gto.M(atom='Ghost .2 .3 .4', unit='b', basis={'Ghost':[[0,(2.2*.5, 1)]]})
        pmol._atm, pmol._bas, pmol._env = \
            gto.conc_env(mol._atm, mol._bas, mol._env,
                         pmol._atm, pmol._bas, pmol._env)
        shls_slice = (pmol.nbas-1,pmol.nbas, pmol.nbas-1,pmol.nbas,
                      0, pmol.nbas, 0, pmol.nbas)
        v0 = pmol.intor('int2e_sph', shls_slice=shls_slice)
        nao = pmol.nao_nr()
        v0 = v0.reshape(nao,nao)[:-1,:-1]
        self.assertTrue(numpy.allclose(v0, v1))

    def test_intor_nr3c(self):
        mol = gto.M(verbose = 0, output = None,
            atom = '''H 0  1 .5
                      H 1 .8 1.1
                      H .2 1.8 0
                      H .5 .8 .1''',
            basis = 'cc-pvdz')
        nao = mol.nao_nr()
        eri0 = numpy.empty((nao,nao,nao))
        ip = 0
        for i in range(mol.nbas):
            jp = 0
            for j in range(mol.nbas):
                kp = 0
                for k in range(mol.nbas):
                    buf = mol.intor_by_shell('int3c1e_sph', (i,j,k))
                    di,dj,dk = buf.shape
                    eri0[ip:ip+di,jp:jp+dj,kp:kp+dk] = buf
                    kp += dk
                jp += dj
            ip += di
        buf = mol.intor_by_shell('int3c1e_sph', (i,j,k))
        self.assertEqual(buf.ndim, 3)

        eri1 = mol.intor('int3c1e_sph')
        self.assertTrue(numpy.allclose(eri0, eri1))

        idx = numpy.tril_indices(nao)
        eri1 = mol.intor('int3c1e_sph', aosym='s2ij')
        self.assertTrue(numpy.allclose(eri0[idx], eri1))

        eri1 = mol.intor('int3c1e_sph', aosym='s2ij',
                         shls_slice=(2,5,0,5,0,mol.nbas))
        self.assertTrue(numpy.allclose(eri0[idx][3:28], eri1))

        eri1 = mol.intor('int3c1e_sph', shls_slice=(2,5,4,9,0,mol.nbas))
        self.assertTrue(numpy.allclose(eri0[2:7,6:15], eri1))

        eri1 = mol.intor('int3c2e_ip1_sph', comp=3, shls_slice=(2,5,4,9,0,mol.nbas))
        self.assertEqual(eri1.shape, (3,5,9,nao))
        self.assertAlmostEqual(fp(eri1), 642.70512922279079, 11)

        eri1 = mol.intor('int3c2e_spinor')
        self.assertEqual(eri1.shape, (nao*2,nao*2,nao))
        self.assertAlmostEqual(fp(eri1), 9462.9659834308495, 11)

#        eri1 = mol.intor('int3c2e_spinor_ssc')
#        self.assertEqual(eri1.shape, (nao*2,nao*2,20))
#        self.assertAlmostEqual(fp(eri1), 1227.1971656655824, 11)

    def test_nr_s8(self):
        mol = gto.M(atom="He 0 0 0; Ne 3 0 0", basis='ccpvdz')
        eri0 = mol.intor('int2e_sph', aosym='s8')
        self.assertAlmostEqual(lib.fp(eri0), -10.685918926843847, 9)

    # FIXME
    #def test_nr_s8_skip(self):
    #    eri1 = mol.intor('int2e_yp_sph', aosym=8)
    #    self.assertAlmostEqual(lib.fp(eri1), -10.685918926843847, 9)
    #    self.assertAlmostEqual(abs(eri0-eri1).max(), 0, 9)

    def test_unknonw(self):
        self.assertRaises(KeyError, mol.intor, 'int4c3e')

    def test_nr_2c2e(self):
        mat = mol.intor('int2c2e')
        self.assertAlmostEqual(lib.fp(mat), -460.83033192375615, 9)


if __name__ == "__main__":
    unittest.main()


# pyscf/gto/test/test_tao.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
from pyscf import gto

# . . . /
# . . ./
#     /. .
#    / . .
#   /  . .
#
def rotatesub(smat, i0, j0, tao):
    di = abs(tao[i0]) - i0
    dj = abs(tao[j0]) - j0
    rmat = numpy.empty((dj,di),dtype=smat.dtype)
    for i in range(di):
        for j in range(dj):
            rmat[dj-1-j,di-1-i] = smat[i,j]
    if tao[j0] < 0: # |j,-j> => -|j,j>
        rmat[1::2] = -rmat[1::2]
    else:
        rmat[::2] = -rmat[::2]
    if tao[i0] < 0:
        rmat[:,1::2] = -rmat[:,1::2]
    else:
        rmat[:,::2] = -rmat[:,::2]
    return rmat

def rotatesub1(smat, i0, j0, tao):
    di = abs(tao[i0]) - i0
    dj = abs(tao[j0]) - j0
    rmat = numpy.empty((dj,di),dtype=smat.dtype)
    for j in range(dj):
        for i in range(di):
            rmat[j,i] = smat[di-1-i,dj-1-j]
    if tao[j0] < 0: # |j,-j> => -|j,j>
        rmat[1::2] = -rmat[1::2]
    else:
        rmat[::2] = -rmat[::2]
    if tao[i0] < 0:
        rmat[:,1::2] = -rmat[:,1::2]
    else:
        rmat[:,::2] = -rmat[:,::2]
    return rmat

def test_tao():
    mol = gto.Mole()
    mol.atom = [
        ['N', (0.,0.,0.)],
        ['H', (0.,1.,1.)],
        ['H', (1.,0.,1.)],
        ['H', (1.,1.,0.)], ]
    mol.basis = {
        "N": [(0, 0, (15, 1)), ],
        "H": [(0, 0, (1, 1, 0), (3, 3, 1), (5, 1, 0)),
              (1, 0, (1, 1)), (2, 0, (.8, 1)), ]}
    mol.basis['N'].extend(gto.mole.expand_etbs(((0, 4, 1, 1.8),
                                                (1, 3, 2, 1.8),
                                                (2, 2, 1, 1.8),)))

    mol.verbose = 0
    mol.output = None
    mol.build()

    tao = mol.time_reversal_map()

    ao_loc = mol.ao_loc_2c()
    s = mol.intor('cint1e_spnucsp')
    for ish in range(mol.nbas):
        for jsh in range(mol.nbas):
            i0 = ao_loc[ish]
            i1 = abs(tao[i0])
            j0 = ao_loc[jsh]
            j1 = abs(tao[j0])
            rmat = rotatesub(s[i0:i1,j0:j1], i0, j0, tao)
            dd = abs(s[j0:j1,i0:i1]-rmat).sum()
            assert(numpy.allclose(rmat,
                                  rotatesub1(s[i0:i1,j0:j1], i0, j0, tao)))
            assert dd < 1e-12, f'{ish}, {jsh}, {dd}'

    nao = mol.nao_2c()
    j = 0
    idx0 = []
    for i in range(nao):
        if abs(tao[i]) > j:
            idx0.append(i)
            j = abs(tao[i])
    j = 0
    idx1 = [0]
    while idx1[-1] < nao:
        j += 1
        idx1.append(abs(tao[idx1[-1]]))

    assert all(numpy.array(idx0[:j-1]) == numpy.array(idx1[:j-1]))
    print('pass')


# 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/gto/__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>
#

from pyscf.gto import mole
from pyscf.gto import basis
from pyscf.gto.basis import parse, load, parse_ecp, load_ecp
from pyscf.gto.mole import *
from pyscf.gto.moleintor import getints, getints_by_shell
from pyscf.gto.eval_gto import eval_gto
from pyscf.gto import ecp

parse = basis.parse
#import pyscf.gto.mole.cmd_args


# pyscf/gto/cmd_args.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.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

import sys
import pyscf.lib.logger

import argparse

def cmd_args():
    '''
    get input from cmdline
    '''
    parser = argparse.ArgumentParser()
    parser.add_argument('-v', '--verbose',
                        action='store_false', dest='verbose', default=0,
                        help='make lots of noise')
    parser.add_argument('-q', '--quiet',
                        action='store_false', dest='quite', default=False,
                        help='be very quiet')
    parser.add_argument('-o', '--output',
                        dest='output', metavar='FILE', help='write output to FILE')
    parser.add_argument('-m', '--max-memory',
                        action='store', dest='max_memory', metavar='NUM',
                        help='maximum memory to use (in MB)')

    (opts, args_left) = parser.parse_known_args()

    if opts.quite:
        opts.verbose = pyscf.lib.logger.QUIET

    if opts.verbose:
        opts.verbose = pyscf.lib.logger.DEBUG

    if opts.max_memory:
        opts.max_memory = float(opts.max_memory)

    return opts


if __name__ == '__main__':
    opts = cmd_args()
    print(opts.verbose, opts.output, opts.max_memory)


# pyscf/gto/ecp.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>
#

'''
Effective core potential (ECP)

This module exposes some ecp integration functions from the C implementation.

Reference for ecp integral computation
* Analytical integration
J. Chem. Phys. 65, 3826
J. Chem. Phys. 111, 8778
J. Comput. Phys. 44, 289

* Numerical integration
J. Comput. Chem. 27, 1009
Chem. Phys. Lett. 296, 445
'''


import ctypes
import numpy
from pyscf import lib
from pyscf.gto import moleintor

libecp = moleintor.libcgto
libecp.ECPscalar_cache_size.restype = ctypes.c_int

def type1_by_shell(mol, shls, cart=False):
    li = mol.bas_angular(shls[0])
    lj = mol.bas_angular(shls[1])
    if cart:
        fn = libecp.ECPtype1_cart
        di = (li+1)*(li+2)//2 * mol.bas_nctr(shls[0])
        dj = (lj+1)*(lj+2)//2 * mol.bas_nctr(shls[1])
    else:
        fn = libecp.ECPtype1_sph
        di = (li*2+1) * mol.bas_nctr(shls[0])
        dj = (lj*2+1) * mol.bas_nctr(shls[1])
    cache_size = libecp.ECPscalar_cache_size(
        ctypes.c_int(1), (ctypes.c_int*2)(*shls),
        mol._atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.natm),
        mol._bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.nbas),
        mol._env.ctypes.data_as(ctypes.c_void_p))
    cache = numpy.empty(cache_size)
    buf = numpy.empty((di,dj), order='F')

    fn(buf.ctypes.data_as(ctypes.c_void_p),
       (ctypes.c_int*2)(*shls),
       mol._ecpbas.ctypes.data_as(ctypes.c_void_p),
       ctypes.c_int(len(mol._ecpbas)),
       mol._atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.natm),
       mol._bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.nbas),
       mol._env.ctypes.data_as(ctypes.c_void_p), lib.c_null_ptr(),
       cache.ctypes.data_as(ctypes.c_void_p))
    return buf

def type2_by_shell(mol, shls, cart=False):
    li = mol.bas_angular(shls[0])
    lj = mol.bas_angular(shls[1])
    if cart:
        fn = libecp.ECPtype2_cart
        di = (li+1)*(li+2)//2 * mol.bas_nctr(shls[0])
        dj = (lj+1)*(lj+2)//2 * mol.bas_nctr(shls[1])
    else:
        fn = libecp.ECPtype2_sph
        di = (li*2+1) * mol.bas_nctr(shls[0])
        dj = (lj*2+1) * mol.bas_nctr(shls[1])
    cache_size = libecp.ECPscalar_cache_size(
        ctypes.c_int(1), (ctypes.c_int*2)(*shls),
        mol._atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.natm),
        mol._bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.nbas),
        mol._env.ctypes.data_as(ctypes.c_void_p))
    cache = numpy.empty(cache_size)
    buf = numpy.empty((di,dj), order='F')

    fn(buf.ctypes.data_as(ctypes.c_void_p),
       (ctypes.c_int*2)(*shls),
       mol._ecpbas.ctypes.data_as(ctypes.c_void_p),
       ctypes.c_int(len(mol._ecpbas)),
       mol._atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.natm),
       mol._bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.nbas),
       mol._env.ctypes.data_as(ctypes.c_void_p), lib.c_null_ptr(),
       cache.ctypes.data_as(ctypes.c_void_p))
    return buf

AS_ECPBAS_OFFSET= 18
AS_NECPBAS      = 19
def so_by_shell(mol, shls):
    '''Spin-orbit coupling ECP in spinor basis
    i/2 <Pauli_matrix dot l U(r)>
    '''
    li = mol.bas_angular(shls[0])
    lj = mol.bas_angular(shls[1])
    di = (li*4+2) * mol.bas_nctr(shls[0])
    dj = (lj*4+2) * mol.bas_nctr(shls[1])
    bas = numpy.vstack((mol._bas, mol._ecpbas))
    mol._env[AS_ECPBAS_OFFSET] = len(mol._bas)
    mol._env[AS_NECPBAS] = len(mol._ecpbas)
    buf = numpy.empty((di,dj), order='F', dtype=numpy.complex128)
    cache = numpy.empty(buf.size*48+100000)
    fn = libecp.ECPso_spinor
    fn(buf.ctypes.data_as(ctypes.c_void_p),
       (ctypes.c_int*2)(di, dj),
       (ctypes.c_int*2)(*shls),
       mol._atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.natm),
       bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.nbas),
       mol._env.ctypes.data_as(ctypes.c_void_p), lib.c_null_ptr(),
       cache.ctypes.data_as(ctypes.c_void_p))
    return buf

def core_configuration(nelec_core):
    conf_dic = {
        0 : '0s0p0d0f',
        2 : '1s0p0d0f',
        10: '2s1p0d0f',
        18: '3s2p0d0f',
        28: '3s2p1d0f',
        36: '4s3p1d0f',
        46: '4s3p2d0f',
        54: '5s4p2d0f',
        60: '4s3p2d1f',
        68: '5s4p2d1f',
        78: '5s4p3d1f',
        92: '5s4p3d2f',
    }
    if nelec_core not in conf_dic:
        raise RuntimeError('Core configuration for %d core electrons is not available.' % nelec_core)
    coreshell = [int(x) for x in conf_dic[nelec_core][::2]]
    return coreshell


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

    mol = gto.M(atom='''
 Cu 0. 0. 0.
 H  0.  0. -1.56
 H  0.  0.  1.56
''',
                basis={'Cu':'lanl2dz', 'H':'sto3g'},
                ecp = {'cu':'lanl2dz'},
                #basis={'Cu':'crenbs', 'H':'sto3g'},
                #ecp = {'cu':'crenbs'},
                charge=-1,
                verbose=4)
    mf = scf.RHF(mol)
    print(mf.kernel(), -196.09477546034623)

    mol = gto.M(atom='''
 Na 0. 0. 0.
 H  0.  0.  1.
''',
                basis={'Na':'lanl2dz', 'H':'sto3g'},
                ecp = {'Na':'lanl2dz'},
                verbose=0)
    mf = scf.RHF(mol)
    print(mf.kernel(), -0.45002315562861461)



# pyscf/gto/eval_gto.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>
#

import warnings
import ctypes
import numpy
from pyscf import lib
from pyscf.gto.moleintor import make_loc

BLKSIZE = 104  # must be equal to lib/gto/grid_ao_drv.h

libcgto = lib.load_library('libcgto')

def eval_gto(mol, eval_name, coords, comp=None, shls_slice=None, non0tab=None,
             ao_loc=None, cutoff=None, out=None):
    r'''Evaluate AO function value on the given grids,

    Args:
        eval_name : str

            ========================  ======  =======================
            Function                  comp    Expression
            ========================  ======  =======================
            "GTOval_sph"              1       |AO>
            "GTOval_ip_sph"           3       nabla |AO>
            "GTOval_ig_sph"           3       (#C(0 1) g) |AO>
            "GTOval_ipig_sph"         9       (#C(0 1) nabla g) |AO>
            "GTOval_cart"             1       |AO>
            "GTOval_ip_cart"          3       nabla |AO>
            "GTOval_ig_cart"          3       (#C(0 1) g)|AO>
            "GTOval_sph_deriv1"       4       GTO value and 1st order GTO values
            "GTOval_sph_deriv2"       10      All derivatives up to 2nd order
            "GTOval_sph_deriv3"       20      All derivatives up to 3rd order
            "GTOval_sph_deriv4"       35      All derivatives up to 4th order
            "GTOval_sp_spinor"        1       sigma dot p |AO> (spinor basis)
            "GTOval_ipsp_spinor"      3       nabla sigma dot p |AO> (spinor basis)
            "GTOval_ipipsp_spinor"    9       nabla nabla sigma dot p |AO> (spinor basis)
            ========================  ======  =======================

        atm : int32 ndarray
            libcint integral function argument
        bas : int32 ndarray
            libcint integral function argument
        env : float64 ndarray
            libcint integral function argument

        coords : 2D array, shape (N,3)
            The coordinates of the grids.

    Kwargs:
        comp : int
            Number of the components of the operator
        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:`dft.gen_grid.make_mask`
        cutoff : float
            AO values smaller than cutoff will be set to zero. The default
            cutoff threshold is ~1e-22 (defined in gto/grid_ao_drv.h)
        out : ndarray
            If provided, results are written into this array.

    Returns:
        2D array of shape (N,nao) Or 3D array of shape (\*,N,nao) to store AO
        values on grids.

    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 = mol.eval_gto("GTOval_sph", coords)
    >>> print(ao_value.shape)
    (100, 24)
    >>> ao_value = mol.eval_gto("GTOval_ig_sph", coords)
    >>> print(ao_value.shape)
    (3, 100, 24)
    '''
    eval_name, comp = _get_intor_and_comp(mol, eval_name, comp)

    atm = numpy.asarray(mol._atm, dtype=numpy.int32, order='C')
    bas = numpy.asarray(mol._bas, dtype=numpy.int32, order='C')
    env = numpy.asarray(mol._env, dtype=numpy.double, order='C')
    coords = numpy.asarray(coords, dtype=numpy.double, order='F')
    natm = atm.shape[0]
    nbas = bas.shape[0]
    ngrids = coords.shape[0]

    if ao_loc is None:
        ao_loc = make_loc(bas, eval_name)

    if shls_slice is None:
        shls_slice = (0, nbas)
    sh0, sh1 = shls_slice
    nao = ao_loc[sh1] - ao_loc[sh0]
    if 'spinor' in eval_name:
        ao = numpy.ndarray((2,comp,nao,ngrids), dtype=numpy.complex128, buffer=out)
    else:
        ao = numpy.ndarray((comp,nao,ngrids), buffer=out)

    if non0tab is None:
        non0tab = numpy.ones(((ngrids+BLKSIZE-1)//BLKSIZE,nbas),
                             dtype=numpy.uint8)

    with mol.with_integral_screen(cutoff):
        drv = getattr(libcgto, eval_name)
        drv(ctypes.c_int(ngrids),
            (ctypes.c_int*2)(*shls_slice), ao_loc.ctypes.data_as(ctypes.c_void_p),
            ao.ctypes.data_as(ctypes.c_void_p),
            coords.ctypes.data_as(ctypes.c_void_p),
            non0tab.ctypes.data_as(ctypes.c_void_p),
            atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(natm),
            bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nbas),
            env.ctypes.data_as(ctypes.c_void_p))

    ao = numpy.swapaxes(ao, -1, -2)
    if comp == 1:
        if 'spinor' in eval_name:
            ao = ao[:,0]
        else:
            ao = ao[0]
    return ao

def _get_intor_and_comp(mol, eval_name, comp=None):
    if not ('_sph' in eval_name or '_cart' in eval_name or
            '_spinor' in eval_name):
        if mol.cart:
            eval_name = eval_name + '_cart'
        else:
            eval_name = eval_name + '_sph'

    if comp is None:
        if '_spinor' in eval_name:
            fname = eval_name.replace('_spinor', '')
            comp = _GTO_EVAL_FUNCTIONS.get(fname, (None,None))[1]
        else:
            fname = eval_name.replace('_sph', '').replace('_cart', '')
            comp = _GTO_EVAL_FUNCTIONS.get(fname, (None,None))[0]
        if comp is None:
            warnings.warn('Function %s not found.  Set its comp to 1' % eval_name)
            comp = 1
    return eval_name, comp

_GTO_EVAL_FUNCTIONS = {
    # Functiona name        : (comp-for-scalar, comp-for-spinor)
    'GTOval'                : (1, 1 ),
    'GTOval_ip'             : (3, 3 ),
    'GTOval_ig'             : (3, 3 ),
    'GTOval_ipig'           : (9, 9 ),
    'GTOval_deriv0'         : (1, 1 ),
    'GTOval_deriv1'         : (4, 4 ),
    'GTOval_deriv2'         : (10,10),
    'GTOval_deriv3'         : (20,20),
    'GTOval_deriv4'         : (35,35),
    'GTOval_sp'             : (4, 1 ),
    'GTOval_ipsp'           : (12,3 ),
    'GTOval_ipipsp'         : (36,9 ),
}


if __name__ == '__main__':
    from pyscf import gto
    mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='ccpvdz')
    coords = numpy.random.random((100,3))
    ao_value = eval_gto(mol, "GTOval_sph", coords)
    print(ao_value.shape)


# pyscf/gto/ft_ao.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>
#

'''
Analytic Fourier transformation for AO and AO-pair value
'''

import ctypes
import numpy
import scipy.linalg
from pyscf import lib
from pyscf import gto
from pyscf.gto.moleintor import libcgto

# TODO: in C code, store complex data in two vectors for real and imag part

#
# \int mu*nu*exp(-ik*r) dr
#
# gxyz is the index for Gvbase
def ft_aopair(mol, Gv, shls_slice=None, aosym='s1', b=numpy.eye(3),
              gxyz=None, Gvbase=None, buf=None, intor='GTO_ft_ovlp',
              comp=1, verbose=None):
    r''' FT transform AO pair
    \int i(r) j(r) exp(-ikr) dr^3
    '''

    intor = mol._add_suffix(intor)

    if shls_slice is None:
        shls_slice = (0, mol.nbas, 0, mol.nbas)
    nGv = Gv.shape[0]
    if (gxyz is None or b is None or Gvbase is None
        # backward compatibility for pyscf-1.2, in which the argument Gvbase is gs
        or (Gvbase is not None and isinstance(Gvbase[0], (int, numpy.integer)))):
        GvT = numpy.asarray(Gv.T, order='C')
        p_gxyzT = lib.c_null_ptr()
        p_gs = (ctypes.c_int*3)(0,0,0)
        p_b = (ctypes.c_double*1)(0)
        eval_gz = 'GTO_Gv_general'
    else:
        if abs(b-numpy.diag(b.diagonal())).sum() < 1e-8:
            eval_gz = 'GTO_Gv_orth'
        else:
            eval_gz = 'GTO_Gv_nonorth'
        GvT = numpy.asarray(Gv.T, order='C')
        gxyzT = numpy.asarray(gxyz.T, order='C', dtype=numpy.int32)
        p_gxyzT = gxyzT.ctypes.data_as(ctypes.c_void_p)
        b = numpy.hstack((b.ravel(), numpy.zeros(3)) + Gvbase)
        p_b = b.ctypes.data_as(ctypes.c_void_p)
        p_gs = (ctypes.c_int*3)(*[len(x) for x in Gvbase])

    ao_loc = gto.moleintor.make_loc(mol._bas, intor)
    ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]]
    nj = ao_loc[shls_slice[3]] - ao_loc[shls_slice[2]]

    if aosym == 's1':
        if (shls_slice[:2] == shls_slice[2:4] and
            intor.startswith('GTO_ft_ovlp')):
            fill = getattr(libcgto, 'GTO_ft_fill_s1hermi')
        else:
            fill = getattr(libcgto, 'GTO_ft_fill_s1')
        shape = (nGv,ni,nj,comp)
    else:
        fill = getattr(libcgto, 'GTO_ft_fill_s2')
        i0 = ao_loc[shls_slice[0]]
        i1 = ao_loc[shls_slice[1]]
        nij = i1*(i1+1)//2 - i0*(i0+1)//2
        shape = (nGv,nij,comp)
    mat = numpy.ndarray(shape, order='F', dtype=numpy.complex128, buffer=buf)

    fn = libcgto.GTO_ft_fill_drv
    intor = getattr(libcgto, intor)
    eval_gz = getattr(libcgto, eval_gz)

    fn(intor, eval_gz, fill, mat.ctypes.data_as(ctypes.c_void_p),
       ctypes.c_int(comp), (ctypes.c_int*4)(*shls_slice),
       ao_loc.ctypes.data_as(ctypes.c_void_p), ctypes.c_double(0),
       GvT.ctypes.data_as(ctypes.c_void_p),
       p_b, p_gxyzT, p_gs, ctypes.c_int(nGv),
       mol._atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.natm),
       mol._bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(mol.nbas),
       mol._env.ctypes.data_as(ctypes.c_void_p))

    mat = numpy.rollaxis(mat, -1, 0)
    if comp == 1:
        mat = mat[0]
    return mat


# gxyz is the index for Gvbase
def ft_ao(mol, Gv, shls_slice=None, b=numpy.eye(3),
          gxyz=None, Gvbase=None, verbose=None):
    ''' FT transform AO
    '''
    if shls_slice is None:
        shls_slice = (0, mol.nbas)
    nGv = Gv.shape[0]
    if (gxyz is None or b is None or Gvbase is None
        # backward compatibility for pyscf-1.2, in which the argument Gvbase is gs
        or (Gvbase is not None and isinstance(Gvbase[0], (int, numpy.integer)))):
        GvT = numpy.asarray(Gv.T, order='C')
        p_gxyzT = lib.c_null_ptr()
        p_gs = (ctypes.c_int*3)(0,0,0)
        p_b = (ctypes.c_double*1)(0)
        eval_gz = 'GTO_Gv_general'
    else:
        if abs(b-numpy.diag(b.diagonal())).sum() < 1e-8:
            eval_gz = 'GTO_Gv_orth'
        else:
            eval_gz = 'GTO_Gv_nonorth'
        GvT = numpy.asarray(Gv.T, order='C')
        gxyzT = numpy.asarray(gxyz.T, order='C', dtype=numpy.int32)
        p_gxyzT = gxyzT.ctypes.data_as(ctypes.c_void_p)
        b = numpy.hstack((b.ravel(), numpy.zeros(3)) + Gvbase)
        p_b = b.ctypes.data_as(ctypes.c_void_p)
        p_gs = (ctypes.c_int*3)(*[len(x) for x in Gvbase])

    fn = libcgto.GTO_ft_fill_drv
    if mol.cart:
        intor = getattr(libcgto, 'GTO_ft_ovlp_cart')
    else:
        intor = getattr(libcgto, 'GTO_ft_ovlp_sph')
    eval_gz = getattr(libcgto, eval_gz)
    fill = getattr(libcgto, 'GTO_ft_fill_s1')

    ghost_atm = numpy.array([[0,0,0,0,0,0]], dtype=numpy.int32)
    ghost_bas = numpy.array([[0,0,1,1,0,0,3,0]], dtype=numpy.int32)
    ghost_env = numpy.zeros(4)
    ghost_env[3] = numpy.sqrt(4*numpy.pi)  # s function spherical norm
    atm, bas, env = gto.conc_env(mol._atm, mol._bas, mol._env,
                                 ghost_atm, ghost_bas, ghost_env)
    ao_loc = mol.ao_loc_nr()
    nao = ao_loc[mol.nbas]
    ao_loc = numpy.asarray(numpy.hstack((ao_loc, [nao+1])), dtype=numpy.int32)
    ni = ao_loc[shls_slice[1]] - ao_loc[shls_slice[0]]
    mat = numpy.zeros((nGv,ni), order='F', dtype=numpy.complex128)

    shls_slice = shls_slice + (mol.nbas, mol.nbas+1)
    fn(intor, eval_gz, fill, mat.ctypes.data_as(ctypes.c_void_p),
       ctypes.c_int(1), (ctypes.c_int*4)(*shls_slice),
       ao_loc.ctypes.data_as(ctypes.c_void_p),
       ctypes.c_double(0),
       GvT.ctypes.data_as(ctypes.c_void_p),
       p_b, p_gxyzT, p_gs, ctypes.c_int(nGv),
       atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(len(atm)),
       bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(len(bas)),
       env.ctypes.data_as(ctypes.c_void_p))
    return mat


if __name__ == '__main__':
    mol = gto.Mole()
    mol.atom = '''C    1.3    .2       .3
                  C     .1    .1      1.1
                  '''
    mol.basis = 'ccpvdz'
    #mol.basis = {'C': [[0, (2.4, .1, .6), (1.0,.8, .4)], [1, (1.1, 1)]]}
    #mol.basis = {'C': [[0, (2.4, 1)]]}
    mol.unit = 'B'
    mol.build(0,0)

    L = 5.
    n = 20
    a = numpy.diag([L,L,L])
    b = scipy.linalg.inv(a)
    gs = [n,n,n]
    gxrange = range(gs[0]+1)+range(-gs[0],0)
    gyrange = range(gs[1]+1)+range(-gs[1],0)
    gzrange = range(gs[2]+1)+range(-gs[2],0)
    gxyz = lib.cartesian_prod((gxrange, gyrange, gzrange))
    Gv = 2*numpy.pi * numpy.dot(gxyz, b)

    print(numpy.linalg.norm(ft_aopair(mol, Gv, None, 's1', b, gxyz, gs)) - 63.0239113778)
    print(numpy.linalg.norm(ft_ao(mol, Gv, None, b, gxyz, gs))-56.8273147065)


# pyscf/gto/mole.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>
#

'''
Mole class and helper functions to handle paramters and attributes for GTO
integrals. This module serves the interface to the integral library libcint.
'''

import os, sys
import types
import re
import platform
import gc
import time

import json
import ctypes
import numpy
import numpy as np
import h5py
import scipy.special
import scipy.linalg
import contextlib
import threading
from pyscf import lib
from pyscf.lib import param
from pyscf.data import elements
from pyscf.lib import logger
from pyscf.gto import cmd_args
from pyscf.gto import basis
from pyscf.gto import moleintor
from pyscf.gto.eval_gto import eval_gto
from pyscf.gto.ecp import core_configuration
from pyscf import __config__

from pyscf.data.elements import ELEMENTS, ELEMENTS_PROTON, \
        _rm_digit, charge, _symbol, _std_symbol, _atom_symbol, is_ghost_atom, \
        _std_symbol_without_ghost

from pyscf.lib.exceptions import BasisNotFoundError, PointGroupSymmetryError
import warnings

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


# for _atm, _bas, _env
CHARGE_OF  = 0
PTR_COORD  = 1
NUC_MOD_OF = 2
PTR_ZETA   = PTR_FRAC_CHARGE = 3
ATM_SLOTS  = 6
ATOM_OF    = 0
ANG_OF     = 1
NPRIM_OF   = 2
NCTR_OF    = 3
RADI_POWER = 3 # for ECP
KAPPA_OF   = 4
SO_TYPE_OF = 4 # for ECP
PTR_EXP    = 5
PTR_COEFF  = 6
BAS_SLOTS  = 8
# pointer to env
PTR_EXPCUTOFF   = 0
PTR_COMMON_ORIG = 1
PTR_RINV_ORIG   = 4
PTR_RINV_ZETA   = 7
PTR_RANGE_OMEGA = 8
PTR_F12_ZETA    = 9
PTR_GTG_ZETA    = 10
NGRIDS          = 11
PTR_GRIDS       = 12
AS_RINV_ORIG_ATOM = 17
AS_ECPBAS_OFFSET = 18
AS_NECPBAS      = 19
PTR_ENV_START   = 20
# parameters from libcint
NUC_POINT = 1
NUC_GAUSS = 2
# nucleus with fractional charges. It can be used to mimic MM particles
NUC_FRAC_CHARGE = 3
NUC_ECP = 4  # atoms with pseudo potential

BASE = getattr(__config__, 'BASE', 0)
NORMALIZE_GTO = getattr(__config__, 'NORMALIZE_GTO', True)
DISABLE_EVAL = getattr(__config__, 'DISABLE_EVAL', False)
# Whether to disable the explicit call to gc.collect(). gc.collect() may cause
# non-negligible overhead (https://github.com/pyscf/pyscf/issues/1038).
DISABLE_GC = getattr(__config__, 'DISABLE_GC', False)

def M(**kwargs):
    r'''This is a shortcut to build up Mole object.

    Args: Same to :func:`Mole.build`

    Examples:

    >>> from pyscf import gto
    >>> mol = gto.M(atom='H 0 0 0; F 0 0 1', basis='6-31g')
    '''
    mol = Mole()
    mol.build(**kwargs)
    return mol

def gaussian_int(n, alpha):
    r'''int_0^inf x^n exp(-alpha x^2) dx'''
    n1 = (n + 1) * .5
    return scipy.special.gamma(n1) / (2. * alpha**n1)

def gto_norm(l, expnt):
    r'''Normalized factor for GTO radial part   :math:`g=r^l e^{-\alpha r^2}`

    .. math::

        \frac{1}{\sqrt{\int g^2 r^2 dr}}
        = \sqrt{\frac{2^{2l+3} (l+1)! (2a)^{l+1.5}}{(2l+2)!\sqrt{\pi}}}

    Ref: H. B. Schlegel and M. J. Frisch, Int. J. Quant.  Chem., 54(1995), 83-87.

    Args:
        l (int):
            angular momentum
        expnt :
            exponent :math:`\alpha`

    Returns:
        normalization factor

    Examples:

    >>> print(gto_norm(0, 1))
    2.5264751109842591
    '''
    if l >= 0:
        #f = 2**(2*l+3) * math.factorial(l+1) * (2*expnt)**(l+1.5) \
        #        / (math.factorial(2*l+2) * math.sqrt(math.pi))
        #return math.sqrt(f)
        return 1/numpy.sqrt(gaussian_int(l*2+2, 2*expnt))
    else:
        raise ValueError('l should be > 0')

def cart2sph(l, c_tensor=None, normalized=None):
    '''
    Cartesian to real spherical transformation matrix

    Kwargs:
        normalized :
            How the Cartesian GTOs are normalized.  'sp' means the s and p
            functions are normalized (this is the convention used by libcint
            library).
    '''
    nf = (l+1)*(l+2)//2
    if c_tensor is None:
        c_tensor = numpy.eye(nf)
    else:
        c_tensor = numpy.asarray(c_tensor, order='F').reshape(-1,nf)
    if l == 0 or l == 1:
        if normalized == 'sp':
            return c_tensor
        elif l == 0:
            return c_tensor * 0.282094791773878143
        else:
            return c_tensor * 0.488602511902919921
    else:
        assert l <= 15
        nd = l * 2 + 1
        ngrid = c_tensor.shape[0]
        c2sph = numpy.zeros((ngrid,nd), order='F')
        fn = moleintor.libcgto.CINTc2s_ket_sph
        fn(c2sph.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(ngrid),
           c_tensor.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(l))
        return c2sph

def cart2spinor_kappa(kappa, l=None, normalized=None):
    '''Cartesian to spinor transformation matrix for kappa

    Kwargs:
        normalized :
            How the Cartesian GTOs are normalized.  'sp' means the s and p
            functions are normalized (this is the convention used by libcint
            library).
    '''
    if kappa < 0:
        l = -kappa - 1
        nd = l * 2 + 2
    elif kappa > 0:
        l = kappa
        nd = l * 2
    else:
        assert(l is not None)
        assert(l <= 12)
        nd = l * 4 + 2
    nf = (l+1)*(l+2)//2
    c2smat = numpy.zeros((nf*2,nd), order='F', dtype=numpy.complex128)
    cmat = numpy.eye(nf)
    fn = moleintor.libcgto.CINTc2s_ket_spinor_sf1
    fn(c2smat.ctypes.data_as(ctypes.c_void_p),
       c2smat[nf:].ctypes.data_as(ctypes.c_void_p),
       cmat.ctypes.data_as(ctypes.c_void_p),
       ctypes.c_int(nf*2), ctypes.c_int(nf),
       ctypes.c_int(1), ctypes.c_int(kappa), ctypes.c_int(l))
    if normalized != 'sp':
        if l == 0:
            c2smat *= 0.282094791773878143
        elif l == 1:
            c2smat *= 0.488602511902919921
    # c2smat[0] is the transformation for spin up
    # c2smat[1] is the transformation for spin down
    c2smat = c2smat.reshape(2,nf,nd)
    return c2smat
cart2j_kappa = cart2spinor_kappa

def cart2spinor_l(l, normalized=None):
    '''Cartesian to spinor transformation matrix for angular moment l

    Kwargs:
        normalized :
            How the Cartesian GTOs are normalized.  'sp' means the s and p
            functions are normalized (this is the convention used by libcint
            library).
    '''
    return cart2spinor_kappa(0, l, normalized)
cart2j_l = cart2spinor_l

def sph2spinor_kappa(kappa, l=None):
    '''Real spherical to spinor transformation matrix for kappa'''
    from pyscf.symm.sph import sph2spinor
    ua, ub = sph2spinor(l)
    if kappa < 0:
        l = -kappa - 1
        ua = ua[:,l*2:]
        ub = ub[:,l*2:]
    elif kappa > 0:
        l = kappa
        ua = ua[:,:l*2]
        ub = ub[:,:l*2]
    else:
        assert(l is not None)
        assert(l <= 12)
    return ua, ub

def sph2spinor_l(l):
    '''Real spherical to spinor transformation matrix for angular moment l'''
    return sph2spinor_kappa(0, l)

def atom_types(atoms, basis=None):
    '''symmetry inequivalent atoms'''
    atmgroup = {}
    for ia, a in enumerate(atoms):
        if 'GHOST' in a[0].upper():
            a = ['X'+a[0][5:]] + list(a[1:])
        if a[0] in atmgroup:
            atmgroup[a[0]].append(ia)
        elif basis is None:
            atmgroup[a[0]] = [ia]
        else:
            stdsymb = _std_symbol(a[0])
            if a[0] in basis:
                if stdsymb in basis and basis[a[0]] == basis[stdsymb]:
                    if stdsymb in atmgroup:
                        atmgroup[stdsymb].append(ia)
                    else:
                        atmgroup[stdsymb] = [ia]
                else:
                    atmgroup[a[0]] = [ia]
            elif stdsymb in atmgroup:
                atmgroup[stdsymb].append(ia)
            else:
                atmgroup[stdsymb] = [ia]
    return atmgroup


def format_atom(atoms, origin=0, axes=None,
                unit=getattr(__config__, 'UNIT', 'Ang')):
    '''Convert the input :attr:`Mole.atom` to the internal data format.
    Including, changing the nuclear charge to atom symbol, converting the
    coordinates to AU, rotate and shift the molecule.
    If the :attr:`~Mole.atom` is a string, it takes ";" and "\\n"
    for the mark to separate atoms;  "," and arbitrary length of blank space
    to spearate the individual terms for an atom.  Blank line will be ignored.

    Args:
        atoms : list or str
            the same to :attr:`Mole.atom`

    Kwargs:
        origin : ndarray
            new axis origin.
        axes : ndarray
            (new_x, new_y, new_z), new coordinates
        unit : str or number
            If unit is one of strings (B, b, Bohr, bohr, AU, au), the coordinates
            of the input atoms are the atomic unit;  If unit is one of strings
            (A, a, Angstrom, angstrom, Ang, ang), the coordinates are in the
            unit of angstrom;  If a number is given, the number are considered
            as the Bohr value (in angstrom), which should be around 0.53.
            Set unit=1 if wishing to preserve the unit of the coordinates.

    Returns:
        "atoms" in the internal format. The internal format is
            | atom = [[atom1, (x, y, z)],
            |         [atom2, (x, y, z)],
            |         ...
            |         [atomN, (x, y, z)]]

    Examples:

    >>> gto.format_atom('9,0,0,0; h@1 0 0 1', origin=(1,1,1))
    [['F', [-1.0, -1.0, -1.0]], ['H@1', [-1.0, -1.0, 0.0]]]
    >>> gto.format_atom(['9,0,0,0', (1, (0, 0, 1))], origin=(1,1,1))
    [['F', [-1.0, -1.0, -1.0]], ['H', [-1, -1, 0]]]
    '''
    def str2atm(line):
        dat = line.split()
        try:
            coords = [float(x) for x in dat[1:4]]
        except ValueError:
            if DISABLE_EVAL:
                raise ValueError('Failed to parse geometry %s' % line)
            else:
                coords = list(eval(','.join(dat[1:4])))
        if len(coords) != 3:
            raise ValueError('Coordinates error in %s' % line)
        return [_atom_symbol(dat[0]), coords]

    if isinstance(atoms, (str, unicode)):
        # The input atoms points to a geometry file
        if os.path.isfile(atoms):
            try:
                atoms = fromfile(atoms)
            except ValueError:
                sys.stderr.write('\nFailed to parse geometry file  %s\n\n' % atoms)
                raise

        atoms = atoms.replace(';','\n').replace(',',' ').replace('\t',' ')
        fmt_atoms = []
        for dat in atoms.split('\n'):
            dat = dat.strip()
            if dat and dat[0] != '#':
                fmt_atoms.append(dat)

        if len(fmt_atoms[0].split()) < 4:
            fmt_atoms = from_zmatrix('\n'.join(fmt_atoms))
        else:
            fmt_atoms = [str2atm(line) for line in fmt_atoms]
    else:
        fmt_atoms = []
        for atom in atoms:
            if isinstance(atom, (str, unicode)):
                if atom.lstrip()[0] != '#':
                    fmt_atoms.append(str2atm(atom.replace(',',' ')))
            else:
                if isinstance(atom[1], (int, float)):
                    fmt_atoms.append([_atom_symbol(atom[0]), atom[1:4]])
                else:
                    fmt_atoms.append([_atom_symbol(atom[0]), atom[1]])

    if len(fmt_atoms) == 0:
        return []

    if axes is None:
        axes = numpy.eye(3)

    if isinstance(unit, (str, unicode)):
        if unit.upper().startswith(('B', 'AU')):
            unit = 1.
        else: #unit[:3].upper() == 'ANG':
            unit = 1./param.BOHR
    else:
        unit = 1./unit

    c = numpy.array([a[1] for a in fmt_atoms], dtype=numpy.double)
    c = numpy.einsum('ix,kx->ki', axes * unit, c - origin)
    z = [a[0] for a in fmt_atoms]
    return list(zip(z, c.tolist()))

#TODO: sort exponents
def format_basis(basis_tab):
    '''Convert the input :attr:`Mole.basis` to the internal data format.

    ``{ atom: [(l, ((-exp, c_1, c_2, ..),
                    (-exp, c_1, c_2, ..))),
               (l, ((-exp, c_1, c_2, ..),
                    (-exp, c_1, c_2, ..)))], ... }``

    Args:
        basis_tab : dict
            Similar to :attr:`Mole.basis`, it **cannot** be a str

    Returns:
        Formated :attr:`~Mole.basis`

    Examples:

    >>> gto.format_basis({'H':'sto-3g', 'H^2': '3-21g'})
    {'H': [[0,
        [3.4252509099999999, 0.15432897000000001],
        [0.62391373000000006, 0.53532813999999995],
        [0.16885539999999999, 0.44463454000000002]]],
     'H^2': [[0,
        [5.4471780000000001, 0.15628500000000001],
        [0.82454700000000003, 0.90469100000000002]],
        [0, [0.18319199999999999, 1.0]]]}
    '''
    basis_converter = _generate_basis_converter()
    fmt_basis = {}
    for atom, atom_basis in basis_tab.items():
        symb = _atom_symbol(atom)
        fmt_basis[symb] = basis_converter(symb, atom_basis)

        if len(fmt_basis[symb]) == 0:
            raise BasisNotFoundError('Basis not found for  %s' % symb)
    return fmt_basis

def _generate_basis_converter():
    def nparray_to_list(item):
        val = []
        for x in item:
            if isinstance(x, (tuple, list)):
                val.append(nparray_to_list(x))
            elif isinstance(x, numpy.ndarray):
                val.append(x.tolist())
            else:
                val.append(x)
        return val

    def load(basis_name, symb):
        if basis_name.lower().startswith('unc'):
            return uncontract(basis.load(basis_name[3:], symb))
        else:
            return basis.load(basis_name, symb)

    def converter(symb, raw_basis):
        if isinstance(raw_basis, (str, unicode)):
            bset = load(str(raw_basis), _std_symbol_without_ghost(symb))
        elif (any(isinstance(x, (str, unicode)) for x in raw_basis)
              # The first element is the basis of internal format
              or not isinstance(raw_basis[0][0], int)):
            stdsymb = _std_symbol_without_ghost(symb)
            bset = []
            for rawb in raw_basis:
                if isinstance(rawb, (str, unicode)):
                    bset += load(str(rawb), stdsymb)
                else:
                    bset += nparray_to_list(rawb)
        else:
            bset = nparray_to_list(raw_basis)
        return bset
    return converter

def uncontracted_basis(_basis):
    '''Uncontract internal format _basis

    Examples:

    >>> gto.uncontract(gto.load('sto3g', 'He'))
    [[0, [6.36242139, 1]], [0, [1.158923, 1]], [0, [0.31364979, 1]]]
    '''
    MAXL = 10
    ubasis_raw = [[] for l in range(MAXL)]
    ubasis_exp = [[] for l in range(MAXL)]
    for b in _basis:
        angl = b[0]
        kappa = b[1]
        if isinstance(kappa, int):
            coeffs = b[2:]
        else:
            coeffs = b[1:]

        if isinstance(kappa, int) and kappa != 0:
            warnings.warn('For basis with kappa != 0, the uncontract basis might be wrong. '
                          'Please double check the resultant attribute mol._basis')
            for p in coeffs:
                ubasis_raw[angl].append([angl, kappa, [p[0], 1]])
                ubasis_exp[angl].append(p[0])
        else:
            for p in coeffs:
                ubasis_raw[angl].append([angl, [p[0], 1]])
                ubasis_exp[angl].append(p[0])

    # Check linear dependency
    ubasis = []
    for l in range(MAXL):
        basis_l = ubasis_raw[l]
        if basis_l:
            es = numpy.array(ubasis_exp[l])
            # Remove duplicated primitive basis functions
            es, e_idx = numpy.unique(es.round(9), True)
            # from large exponent to small exponent
            for i in reversed(e_idx):
                ubasis.append(basis_l[i])
    return ubasis
uncontract = uncontracted_basis
contract = contracted_basis = basis.to_general_contraction

def to_uncontracted_cartesian_basis(mol):
    '''Decontract the basis of a Mole or a Cell.  Returns a Mole (Cell) object
    with the uncontracted basis environment and a list of coefficients that
    transform the uncontracted cartesian basis to the original basis.  Each
    element in the list corresponds to one shell of the original Mole (Cell).

    Examples:

    >>> mol = gto.M(atom='Ne', basis='ccpvdz')
    >>> pmol, ctr_coeff = mol.to_uncontracted_cartesian_basis()
    >>> c = scipy.linalg.block_diag(*ctr_coeff)
    >>> s = reduce(numpy.dot, (c.T, pmol.intor('int1e_ovlp'), c))
    >>> abs(s-mol.intor('int1e_ovlp')).max()
    0.0
    '''
    import copy
    lmax = mol._bas[:,ANG_OF].max()
    if mol.cart:
        c2s = [numpy.eye((l+1)*(l+2)//2) for l in range(lmax+1)]
    else:
        c2s = [cart2sph(l, normalized='sp') for l in range(lmax+1)]

    pmol = copy.copy(mol)
    pmol.cart = True
    _bas = []
    _env = mol._env.copy()
    contr_coeff = []
    for ib in range(mol.nbas):
        l = mol._bas[ib,ANG_OF]
        ncart = (l+1)*(l+2)//2
        es = mol.bas_exp(ib)
        cs = mol._libcint_ctr_coeff(ib)
        nprim = cs.shape[0]
        norm = gto_norm(l, es)
        c = numpy.einsum('pi,p,xm->pxim', cs, 1./norm, c2s[l])
        contr_coeff.append(c.reshape(nprim*ncart,-1))

        pexp = mol._bas[ib,PTR_EXP]
        pc = mol._bas[ib,PTR_COEFF]
        bs = numpy.empty((nprim,8), dtype=numpy.int32)
        bs[:] = mol._bas[ib]
        bs[:,NCTR_OF] = bs[:,NPRIM_OF] = 1
        bs[:,PTR_EXP] = numpy.arange(pexp, pexp+nprim)
        bs[:,PTR_COEFF] = numpy.arange(pc, pc+nprim)
        _env[pc:pc+nprim] = norm
        _bas.append(bs)

    pmol._bas = numpy.asarray(numpy.vstack(_bas), dtype=numpy.int32)
    pmol._env = _env
    return pmol, contr_coeff

def format_ecp(ecp_tab):
    r'''Convert the input :attr:`ecp` (dict) to the internal data format::

      { atom: (nelec,  # core electrons
               ((l,  # l=-1 for UL, l>=0 for Ul to indicate |l><l|
                 (((exp_1, c_1),  # for r^0
                   (exp_2, c_2),
                   ...),
                  ((exp_1, c_1),  # for r^1
                   (exp_2, c_2),
                   ...),
                  ((exp_1, c_1),  # for r^2
                   ...))))),
       ...}
    '''
    fmt_ecp = {}
    for atom, atom_ecp in ecp_tab.items():
        symb = _atom_symbol(atom)

        if isinstance(atom_ecp, (str, unicode)):
            stdsymb = _std_symbol_without_ghost(symb)
            ecp_dat = basis.load_ecp(str(atom_ecp), stdsymb)
            if ecp_dat is None or len(ecp_dat) == 0:
                #raise BasisNotFoundError('ECP not found for  %s' % symb)
                sys.stderr.write('ECP %s not found for  %s\n' %
                                 (atom_ecp, symb))
            else:
                fmt_ecp[symb] = ecp_dat
        else:
            fmt_ecp[symb] = atom_ecp
    return fmt_ecp

# transform etb to basis format
def expand_etb(l, n, alpha, beta):
    r'''Generate the exponents of even tempered basis for :attr:`Mole.basis`.
    .. math::

        e = e^{-\alpha * \beta^{i-1}} for i = 1 .. n

    Args:
        l : int
            Angular momentum
        n : int
            Number of GTOs

    Returns:
        Formated :attr:`~Mole.basis`

    Examples:

    >>> gto.expand_etb(1, 3, 1.5, 2)
    [[1, [6.0, 1]], [1, [3.0, 1]], [1, [1.5, 1]]]
    '''
    return [[l, [alpha*beta**i, 1]] for i in reversed(range(n))]
def expand_etbs(etbs):
    r'''Generate even tempered basis.  See also :func:`expand_etb`

    Args:
        etbs = [(l, n, alpha, beta), (l, n, alpha, beta),...]

    Returns:
        Formated :attr:`~Mole.basis`

    Examples:

    >>> gto.expand_etbs([(0, 2, 1.5, 2.), (1, 2, 1, 2.)])
    [[0, [6.0, 1]], [0, [3.0, 1]], [1, [1., 1]], [1, [2., 1]]]
    '''
    return lib.flatten([expand_etb(*etb) for etb in etbs])
etbs = expand_etbs

# concatenate two mol
def conc_env(atm1, bas1, env1, atm2, bas2, env2):
    r'''Concatenate two Mole's integral parameters.  This function can be used
    to construct the environment for cross integrals like

    .. math::

        \langle \mu | \nu \rangle, \mu \in mol1, \nu \in mol2

    Returns:
        Concatenated atm, bas, env

    Examples:
        Compute the overlap between H2 molecule and O atom

    >>> mol1 = gto.M(atom='H 0 1 0; H 0 0 1', basis='sto3g')
    >>> mol2 = gto.M(atom='O 0 0 0', basis='sto3g')
    >>> atm3, bas3, env3 = gto.conc_env(mol1._atm, mol1._bas, mol1._env,
    ...                                 mol2._atm, mol2._bas, mol2._env)
    >>> gto.moleintor.getints('int1e_ovlp_sph', atm3, bas3, env3, range(2), range(2,5))
    [[ 0.04875181  0.44714688  0.          0.37820346  0.        ]
     [ 0.04875181  0.44714688  0.          0.          0.37820346]]
    '''
    off = len(env1)
    natm_off = len(atm1)
    atm2 = numpy.copy(atm2)
    bas2 = numpy.copy(bas2)
    atm2[:,PTR_COORD] += off
    atm2[:,PTR_ZETA ] += off
    bas2[:,ATOM_OF  ] += natm_off
    bas2[:,PTR_EXP  ] += off
    bas2[:,PTR_COEFF] += off
    return (numpy.asarray(numpy.vstack((atm1,atm2)), dtype=numpy.int32),
            numpy.asarray(numpy.vstack((bas1,bas2)), dtype=numpy.int32),
            numpy.hstack((env1,env2)))

def conc_mol(mol1, mol2):
    '''Concatenate two Mole objects.
    '''
    if not mol1._built:
        logger.warn(mol1, 'Warning: object %s not initialized. Initializing %s',
                    mol1, mol1)
        mol1.build()
    if not mol2._built:
        logger.warn(mol2, 'Warning: object %s not initialized. Initializing %s',
                    mol2, mol2)
        mol2.build()

    mol3 = Mole()
    mol3._built = True

    mol3._atm, mol3._bas, mol3._env = \
            conc_env(mol1._atm, mol1._bas, mol1._env,
                     mol2._atm, mol2._bas, mol2._env)
    off = len(mol1._env)
    natm_off = len(mol1._atm)
    if len(mol2._ecpbas) == 0:
        mol3._ecpbas = mol1._ecpbas
    else:
        ecpbas2 = numpy.copy(mol2._ecpbas)
        ecpbas2[:,ATOM_OF  ] += natm_off
        ecpbas2[:,PTR_EXP  ] += off
        ecpbas2[:,PTR_COEFF] += off
        if len(mol1._ecpbas) == 0:
            mol3._ecpbas = ecpbas2
        else:
            mol3._ecpbas = numpy.vstack((mol1._ecpbas, ecpbas2))

    mol3.verbose = mol1.verbose
    mol3.output = mol1.output
    mol3.max_memory = mol1.max_memory
    mol3.charge = mol1.charge + mol2.charge
    mol3.spin = abs(mol1.spin - mol2.spin)
    mol3.symmetry = False
    mol3.symmetry_subgroup = None
    mol3.cart = mol1.cart and mol2.cart

    mol3._atom = mol1._atom + mol2._atom
    mol3.atom = mol3._atom
    mol3.unit = 'Bohr'

    mol3._basis.update(mol2._basis)
    mol3._basis.update(mol1._basis)
    mol3._pseudo.update(mol2._pseudo)
    mol3._pseudo.update(mol1._pseudo)
    mol3._ecp.update(mol2._ecp)
    mol3._ecp.update(mol1._ecp)
    mol3.basis = mol3._basis
    mol3.ecp = mol3._ecp

    mol3.nucprop.update(mol1.nucprop)
    mol3.nucprop.update(mol2.nucprop)
    return mol3

# <bas-of-mol1|intor|bas-of-mol2>
def intor_cross(intor, mol1, mol2, comp=None, grids=None):
    r'''1-electron integrals from two molecules like

    .. math::

        \langle \mu | intor | \nu \rangle, \mu \in mol1, \nu \in mol2

    Args:
        intor : str
            Name of the 1-electron integral, such as int1e_ovlp_sph (spherical overlap),
            int1e_nuc_cart (cartesian nuclear attraction), int1e_ipovlp_spinor
            (spinor overlap gradients), etc.  Ref to :func:`getints` for the
            full list of available 1-electron integral names
        mol1, mol2:
            :class:`Mole` objects

    Kwargs:
        comp : int
            Components of the integrals, e.g. int1e_ipovlp_sph has 3 components
        grids : ndarray
            Coordinates of grids for the int1e_grids integrals

    Returns:
        ndarray of 1-electron integrals, can be either 2-dim or 3-dim, depending on comp

    Examples:
        Compute the overlap between H2 molecule and O atom

    >>> mol1 = gto.M(atom='H 0 1 0; H 0 0 1', basis='sto3g')
    >>> mol2 = gto.M(atom='O 0 0 0', basis='sto3g')
    >>> gto.intor_cross('int1e_ovlp_sph', mol1, mol2)
    [[ 0.04875181  0.44714688  0.          0.37820346  0.        ]
     [ 0.04875181  0.44714688  0.          0.          0.37820346]]
    '''
    nbas1 = len(mol1._bas)
    nbas2 = len(mol2._bas)
    atmc, basc, envc = conc_env(mol1._atm, mol1._bas, mol1._env,
                                mol2._atm, mol2._bas, mol2._env)
    if '_grids' in intor:
        assert grids is not None
        envc = numpy.append(envc, grids.ravel())
        envc[NGRIDS] = grids.shape[0]
        envc[PTR_GRIDS] = envc.size - grids.size

    shls_slice = (0, nbas1, nbas1, nbas1+nbas2)

    if (intor.endswith('_sph') or intor.startswith('cint') or
        intor.endswith('_spinor') or intor.endswith('_cart')):
        return moleintor.getints(intor, atmc, basc, envc, shls_slice, comp, 0)
    elif mol1.cart == mol2.cart:
        intor = mol1._add_suffix(intor)
        return moleintor.getints(intor, atmc, basc, envc, shls_slice, comp, 0)
    elif mol1.cart:
        mat = moleintor.getints(intor+'_cart', atmc, basc, envc, shls_slice, comp, 0)
        return numpy.dot(mat, mol2.cart2sph_coeff())
    else:
        mat = moleintor.getints(intor+'_cart', atmc, basc, envc, shls_slice, comp, 0)
        return numpy.dot(mol1.cart2sph_coeff().T, mat)

# append (charge, pointer to coordinates, nuc_mod) to _atm
def make_atm_env(atom, ptr=0, nuclear_model=NUC_POINT, nucprop={}):
    '''Convert the internal format :attr:`Mole._atom` to the format required
    by ``libcint`` integrals
    '''
    nuc_charge = charge(atom[0])
    if nuclear_model == NUC_POINT:
        zeta = 0
    elif nuclear_model == NUC_GAUSS:
        zeta = dyall_nuc_mod(nuc_charge, nucprop)
    else:  # callable(nuclear_model)
        zeta = nuclear_model(nuc_charge, nucprop)
        nuclear_model = NUC_GAUSS

    _env = numpy.hstack((atom[1], zeta))
    _atm = numpy.zeros(6, dtype=numpy.int32)
    _atm[CHARGE_OF] = nuc_charge
    _atm[PTR_COORD] = ptr
    _atm[NUC_MOD_OF] = nuclear_model
    _atm[PTR_ZETA ] = ptr + 3
    return _atm, _env

# append (atom, l, nprim, nctr, kappa, ptr_exp, ptr_coeff, 0) to bas
# absorb normalization into GTO contraction coefficients
def make_bas_env(basis_add, atom_id=0, ptr=0):
    '''Convert :attr:`Mole.basis` to the argument ``bas`` for ``libcint`` integrals
    '''
    _bas = []
    _env = []
    for b in basis_add:
        if not b:  # == []
            continue
        angl = b[0]
        if angl > 14:
            sys.stderr.write('Warning: integral library does not support basis '
                             'with angular momentum > 14\n')

        if isinstance(b[1], int):
            kappa = b[1]
            b_coeff = numpy.array(sorted(list(b[2:]), reverse=True))
        else:
            kappa = 0
            b_coeff = numpy.array(sorted(list(b[1:]), reverse=True))
        es = b_coeff[:,0]
        cs = b_coeff[:,1:]
        nprim, nctr = cs.shape
        cs = numpy.einsum('pi,p->pi', cs, gto_norm(angl, es))
        if NORMALIZE_GTO:
            cs = _nomalize_contracted_ao(angl, es, cs)

        _env.append(es)
        _env.append(cs.T.reshape(-1))
        ptr_exp = ptr
        ptr_coeff = ptr_exp + nprim
        ptr = ptr_coeff + nprim * nctr
        _bas.append([atom_id, angl, nprim, nctr, kappa, ptr_exp, ptr_coeff, 0])
    _env = lib.flatten(_env) # flatten nested lists
    return (numpy.array(_bas, numpy.int32).reshape(-1,BAS_SLOTS),
            numpy.array(_env, numpy.double))

def _nomalize_contracted_ao(l, es, cs):
    #ee = numpy.empty((nprim,nprim))
    #for i in range(nprim):
    #    for j in range(i+1):
    #        ee[i,j] = ee[j,i] = gaussian_int(angl*2+2, es[i]+es[j])
    #s1 = 1/numpy.sqrt(numpy.einsum('pi,pq,qi->i', cs, ee, cs))
    ee = es.reshape(-1,1) + es.reshape(1,-1)
    ee = gaussian_int(l*2+2, ee)
    s1 = 1. / numpy.sqrt(numpy.einsum('pi,pq,qi->i', cs, ee, cs))
    return numpy.einsum('pi,i->pi', cs, s1)

def make_env(atoms, basis, pre_env=[], nucmod={}, nucprop={}):
    '''Generate the input arguments for ``libcint`` library based on internal
    format :attr:`Mole._atom` and :attr:`Mole._basis`
    '''
    _atm = []
    _bas = []
    _env = []
    ptr_env = len(pre_env)

    for ia, atom in enumerate(atoms):
        symb = atom[0]
        stdsymb = _rm_digit(symb)

        if ia+1 in nucprop:
            prop = nucprop[ia+1]
        elif symb in nucprop:
            prop = nucprop[symb]
        else:
            prop = nucprop.get(stdsymb, {})

        nuclear_model = NUC_POINT
        if nucmod:
            if nucmod is None:
                nuclear_model = NUC_POINT
            elif isinstance(nucmod, (int, str, unicode, types.FunctionType)):
                nuclear_model = _parse_nuc_mod(nucmod)
            elif ia+1 in nucmod:
                nuclear_model = _parse_nuc_mod(nucmod[ia+1])
            elif symb in nucmod:
                nuclear_model = _parse_nuc_mod(nucmod[symb])
            elif stdsymb in nucmod:
                nuclear_model = _parse_nuc_mod(nucmod[stdsymb])
        atm0, env0 = make_atm_env(atom, ptr_env, nuclear_model, prop)
        ptr_env = ptr_env + len(env0)
        _atm.append(atm0)
        _env.append(env0)

    _basdic = {}
    for symb, basis_add in basis.items():
        bas0, env0 = make_bas_env(basis_add, 0, ptr_env)
        ptr_env = ptr_env + len(env0)
        _basdic[symb] = bas0
        _env.append(env0)

    for ia, atom in enumerate(atoms):
        symb = atom[0]
        puresymb = _rm_digit(symb)
        if symb in _basdic:
            b = _basdic[symb].copy()
        elif puresymb in _basdic:
            b = _basdic[puresymb].copy()
        else:
            if symb[:2].upper() == 'X-':
                symb = symb[2:]
            elif symb[:6].upper() == 'GHOST-':
                symb = symb[6:]
            puresymb = _rm_digit(symb)
            if symb in _basdic:
                b = _basdic[symb].copy()
            elif puresymb in _basdic:
                b = _basdic[puresymb].copy()
            else:
                sys.stderr.write('Warning: Basis not found for atom %d %s\n' % (ia, symb))
                continue
        b[:,ATOM_OF] = ia
        _bas.append(b)

    if _atm:
        _atm = numpy.asarray(numpy.vstack(_atm), numpy.int32).reshape(-1, ATM_SLOTS)
    else:
        _atm = numpy.zeros((0,ATM_SLOTS), numpy.int32)
    if _bas:
        _bas = numpy.asarray(numpy.vstack(_bas), numpy.int32).reshape(-1, BAS_SLOTS)
    else:
        _bas = numpy.zeros((0,BAS_SLOTS), numpy.int32)
    if _env:
        _env = numpy.hstack((pre_env,numpy.hstack(_env)))
    else:
        _env = numpy.array(pre_env, copy=False)
    return _atm, _bas, _env

def make_ecp_env(mol, _atm, ecp, pre_env=[]):
    '''Generate the input arguments _ecpbas for ECP integrals
    '''
    _env = []
    ptr_env = len(pre_env)

    _ecpdic = {}
    for symb, ecp_add in ecp.items():
        ecp0 = []
        nelec = ecp_add[0]
        for lb in ecp_add[1]:
            for rorder, bi in enumerate(lb[1]):
                if len(bi) > 0:
                    ec = numpy.array(sorted(bi, reverse=True))
                    nexp, ncol = ec.shape
                    _env.append(ec[:,0])
                    _env.append(ec[:,1])
                    ptr_exp, ptr_coeff = ptr_env, ptr_env + nexp
                    ecp0.append([0, lb[0], nexp, rorder, 0,
                                 ptr_exp, ptr_coeff, 0])
                    ptr_env += nexp * 2

                    if ncol == 3:  # Has SO-ECP
                        _env.append(ec[:,2])
                        ptr_coeff, ptr_env = ptr_env, ptr_env + nexp
                        ecp0.append([0, lb[0], nexp, rorder, 1,
                                     ptr_exp, ptr_coeff, 0])

        _ecpdic[symb] = (nelec, numpy.asarray(ecp0, dtype=numpy.int32))

    _ecpbas = []
    if _ecpdic:
        _atm = _atm.copy()
        for ia, atom in enumerate(mol._atom):
            symb = atom[0]
            if symb in _ecpdic:
                ecp0 = _ecpdic[symb]
            elif _rm_digit(symb) in _ecpdic:
                ecp0 = _ecpdic[_rm_digit(symb)]
            else:
                ecp0 = None
            if ecp0 is not None:
                _atm[ia,CHARGE_OF ] = charge(symb) - ecp0[0]
                _atm[ia,NUC_MOD_OF] = NUC_ECP
                b = ecp0[1].copy().reshape(-1,BAS_SLOTS)
                b[:,ATOM_OF] = ia
                _ecpbas.append(b)
    if _ecpbas:
        _ecpbas = numpy.asarray(numpy.vstack(_ecpbas), numpy.int32)
        _env = numpy.hstack([pre_env] + _env)
    else:
        _ecpbas = numpy.zeros((0,BAS_SLOTS), numpy.int32)
        _env = pre_env
    return _atm, _ecpbas, _env

def tot_electrons(mol):
    '''Total number of electrons for the given molecule

    Returns:
        electron number in integer

    Examples:

    >>> mol = gto.M(atom='H 0 1 0; C 0 0 1', charge=1)
    >>> mol.tot_electrons()
    6
    '''
    if mol._atm.size != 0:
        nelectron = mol.atom_charges().sum()
    elif mol._atom:
        nelectron = sum(charge(a[0]) for a in mol._atom)
    else:
        nelectron = sum(charge(a[0]) for a in format_atom(mol.atom))
    nelectron -= mol.charge
    nelectron_int = int(round(nelectron))

    if abs(nelectron - nelectron_int) > 1e-4:
        logger.warn(mol, 'Found fractional number of electrons %f. Round it to %d',
                    nelectron, nelectron_int)
    return nelectron_int

def copy(mol):
    '''Deepcopy of the given :class:`Mole` object
    '''
    import copy
    newmol = copy.copy(mol)
    newmol._atm    = numpy.copy(mol._atm)
    newmol._bas    = numpy.copy(mol._bas)
    newmol._env    = numpy.copy(mol._env)
    newmol._ecpbas = numpy.copy(mol._ecpbas)

    newmol.atom    = copy.deepcopy(mol.atom)
    newmol._atom   = copy.deepcopy(mol._atom)
    newmol.basis   = copy.deepcopy(mol.basis)
    newmol._basis  = copy.deepcopy(mol._basis)
    newmol.ecp     = copy.deepcopy(mol.ecp)
    newmol._ecp    = copy.deepcopy(mol._ecp)
    return newmol

def pack(mol):
    '''Pack the input args of :class:`Mole` to a dict.

    Note this function only pack the input arguments (not the entire Mole
    class).  Modifications to mol._atm, mol._bas, mol._env are not tracked.
    Use :func:`dumps` to serialize the entire Mole object.
    '''
    mdic = {'atom'    : mol.atom,
            'unit'    : mol.unit,
            'basis'   : mol.basis,
            'charge'  : mol.charge,
            'spin'    : mol.spin,
            'symmetry': mol.symmetry,
            'nucmod'  : mol.nucmod,
            'nucprop' : mol.nucprop,
            'ecp'     : mol.ecp,
            '_nelectron': mol._nelectron,
            'verbose' : mol.verbose}
    return mdic
def unpack(moldic):
    '''Unpack a dict which is packed by :func:`pack`, to generate the input
    arguments for :class:`Mole` object.
    '''
    mol = Mole()
    mol.__dict__.update(moldic)
    return mol


def dumps(mol):
    '''Serialize Mole object to a JSON formatted str.
    '''
    exclude_keys = set(('output', 'stdout', '_keys',
                        # Constructing in function loads
                        'symm_orb', 'irrep_id', 'irrep_name'))
    nparray_keys = set(('_atm', '_bas', '_env', '_ecpbas',
                        '_symm_orig', '_symm_axes'))

    moldic = dict(mol.__dict__)
    for k in exclude_keys:
        if k in moldic:
            del(moldic[k])
    for k in nparray_keys:
        if isinstance(moldic[k], numpy.ndarray):
            moldic[k] = moldic[k].tolist()
    moldic['atom'] = repr(mol.atom)
    moldic['basis']= repr(mol.basis)
    moldic['ecp' ] = repr(mol.ecp)

    try:
        return json.dumps(moldic)
    except TypeError:
        def skip_value(dic):
            dic1 = {}
            for k,v in dic.items():
                if (v is None or
                    isinstance(v, (str, unicode, bool, int, float))):
                    dic1[k] = v
                elif isinstance(v, (list, tuple)):
                    dic1[k] = v   # Should I recursively skip_vaule?
                elif isinstance(v, set):
                    dic1[k] = list(v)
                elif isinstance(v, dict):
                    dic1[k] = skip_value(v)
                else:
                    msg =('Function mol.dumps drops attribute %s because '
                          'it is not JSON-serializable' % k)
                    warnings.warn(msg)
            return dic1
        return json.dumps(skip_value(moldic), skipkeys=True)

def loads(molstr):
    '''Deserialize a str containing a JSON document to a Mole object.
    '''
    # the numpy function array is used by eval function
    from numpy import array  # noqa
    moldic = json.loads(molstr)
    if sys.version_info < (3,):
        # Convert to utf8 because JSON loads fucntion returns unicode.
        def byteify(inp):
            if isinstance(inp, dict):
                return dict([(byteify(k), byteify(v)) for k, v in inp.iteritems()])
            elif isinstance(inp, (tuple, list)):
                return [byteify(x) for x in inp]
            elif isinstance(inp, unicode):
                return inp.encode('utf-8')
            else:
                return inp
        moldic = byteify(moldic)
    mol = Mole()
    mol.__dict__.update(moldic)
    mol.atom = eval(mol.atom)
    mol.basis= eval(mol.basis)
    mol.ecp  = eval(mol.ecp)
    mol._atm = numpy.array(mol._atm, dtype=numpy.int32)
    mol._bas = numpy.array(mol._bas, dtype=numpy.int32)
    mol._env = numpy.array(mol._env, dtype=numpy.double)
    mol._ecpbas = numpy.array(mol._ecpbas, dtype=numpy.int32)

    if mol.symmetry and mol._symm_orig is not None:
        from pyscf import symm
        mol._symm_orig = numpy.array(mol._symm_orig)
        mol._symm_axes = numpy.array(mol._symm_axes)
        mol.symm_orb, mol.irrep_id = \
                symm.symm_adapted_basis(mol, mol.groupname,
                                        mol._symm_orig, mol._symm_axes)
        mol.irrep_name = [symm.irrep_id2name(mol.groupname, ir)
                           for ir in mol.irrep_id]

    elif mol.symmetry and mol.symm_orb is not None:
        # Backward compatibility. To load symm_orb from chkfile of pyscf-1.6
        # and earlier.
        symm_orb = []

        # decompress symm_orb
        for val, x, y, shape in mol.symm_orb:
            if isinstance(val[0], list):
                # backward compatibility for chkfile of pyscf-1.4 in which val
                # is an array of real floats. In pyscf-1.5, val can be a list
                # of list, to include the imaginary part
                val_real, val_imag = val
            else:
                val_real, val_imag = val, None
            if val_imag is None:
                c = numpy.zeros(shape)
                c[numpy.array(x),numpy.array(y)] = numpy.array(val_real)
            else:
                c = numpy.zeros(shape, dtype=numpy.complex128)
                val = numpy.array(val_real) + numpy.array(val_imag) * 1j
                c[numpy.array(x),numpy.array(y)] = val
            symm_orb.append(c)
        mol.symm_orb = symm_orb

    return mol


def len_spinor(l, kappa):
    '''The number of spinor associated with given angular momentum and kappa.  If kappa is 0,
    return 4l+2
    '''
    if kappa == 0:
        n = (l * 4 + 2)
    elif kappa < 0:
        n = (l * 2 + 2)
    else:
        n = (l * 2)
    return n

def len_cart(l):
    '''The number of Cartesian function associated with given angular momentum.
    '''
    return (l + 1) * (l + 2) // 2

def npgto_nr(mol, cart=None):
    '''Total number of primitive spherical GTOs for the given :class:`Mole` object'''
    if cart is None:
        cart = mol.cart
    l = mol._bas[:,ANG_OF]
    if cart:
        return ((l+1)*(l+2)//2 * mol._bas[:,NPRIM_OF]).sum()
    else:
        return ((l*2+1) * mol._bas[:,NPRIM_OF]).sum()
def nao_nr(mol, cart=None):
    '''Total number of contracted GTOs for the given :class:`Mole` object'''
    if cart is None:
        cart = mol.cart
    if cart:
        return nao_cart(mol)
    else:
        return ((mol._bas[:,ANG_OF]*2+1) * mol._bas[:,NCTR_OF]).sum()
def nao_cart(mol):
    '''Total number of contracted cartesian GTOs for the given :class:`Mole` object'''
    l = mol._bas[:,ANG_OF]
    return ((l+1)*(l+2)//2 * mol._bas[:,NCTR_OF]).sum()

# nao_id0:nao_id1 corresponding to bas_id0:bas_id1
def nao_nr_range(mol, bas_id0, bas_id1):
    '''Lower and upper boundary of contracted spherical basis functions associated
    with the given shell range

    Args:
        mol :
            :class:`Mole` object
        bas_id0 : int
            start shell id
        bas_id1 : int
            stop shell id

    Returns:
        tupel of start basis function id and the stop function id

    Examples:

    >>> mol = gto.M(atom='O 0 0 0; C 0 0 1', basis='6-31g')
    >>> gto.nao_nr_range(mol, 2, 4)
    (2, 6)
    '''
    ao_loc = moleintor.make_loc(mol._bas[:bas_id1], 'sph')
    nao_id0 = ao_loc[bas_id0]
    nao_id1 = ao_loc[-1]
    return nao_id0, nao_id1

def nao_2c(mol):
    '''Total number of contracted spinor GTOs for the given :class:`Mole` object'''
    l = mol._bas[:,ANG_OF]
    kappa = mol._bas[:,KAPPA_OF]
    dims = (l*4+2) * mol._bas[:,NCTR_OF]
    dims[kappa<0] = l[kappa<0] * 2 + 2
    dims[kappa>0] = l[kappa>0] * 2
    return dims.sum()

# nao_id0:nao_id1 corresponding to bas_id0:bas_id1
def nao_2c_range(mol, bas_id0, bas_id1):
    '''Lower and upper boundary of contracted spinor basis functions associated
    with the given shell range

    Args:
        mol :
            :class:`Mole` object
        bas_id0 : int
            start shell id, 0-based
        bas_id1 : int
            stop shell id, 0-based

    Returns:
        tupel of start basis function id and the stop function id

    Examples:

    >>> mol = gto.M(atom='O 0 0 0; C 0 0 1', basis='6-31g')
    >>> gto.nao_2c_range(mol, 2, 4)
    (4, 12)
    '''
    ao_loc = moleintor.make_loc(mol._bas[:bas_id1], '')
    nao_id0 = ao_loc[bas_id0]
    nao_id1 = ao_loc[-1]
    return nao_id0, nao_id1

def ao_loc_nr(mol, cart=None):
    '''Offset of every shell in the spherical basis function spectrum

    Returns:
        list, each entry is the corresponding start basis function id

    Examples:

    >>> mol = gto.M(atom='O 0 0 0; C 0 0 1', basis='6-31g')
    >>> gto.ao_loc_nr(mol)
    [0, 1, 2, 3, 6, 9, 10, 11, 12, 15, 18]
    '''
    if cart is None:
        cart = mol.cart
    if cart:
        return moleintor.make_loc(mol._bas, 'cart')
    else:
        return moleintor.make_loc(mol._bas, 'sph')

def ao_loc_2c(mol):
    '''Offset of every shell in the spinor basis function spectrum

    Returns:
        list, each entry is the corresponding start id of spinor function

    Examples:

    >>> mol = gto.M(atom='O 0 0 0; C 0 0 1', basis='6-31g')
    >>> gto.ao_loc_2c(mol)
    [0, 2, 4, 6, 12, 18, 20, 22, 24, 30, 36]
    '''
    return moleintor.make_loc(mol._bas, 'spinor')

def time_reversal_map(mol):
    r'''The index to map the spinor functions and its time reversal counterpart.
    The returned indices have postive or negative values.  For the i-th basis function,
    if the returned j = idx[i] < 0, it means :math:`T|i\rangle = -|j\rangle`,
    otherwise :math:`T|i\rangle = |j\rangle`
    '''
    tao = []
    i = 0
    for b in mol._bas:
        l = b[ANG_OF]
        if b[KAPPA_OF] == 0:
            djs = (l * 2, l * 2 + 2)
        elif b[KAPPA_OF] > 0:
            djs = (l * 2,)
        else:
            djs = (l * 2 + 2,)
        if l % 2 == 0:
            for n in range(b[NCTR_OF]):
                for dj in djs:
                    for m in range(0, dj, 2):
                        tao.append(-(i + dj - m))
                        tao.append(  i + dj - m - 1)
                    i += dj
        else:
            for n in range(b[NCTR_OF]):
                for dj in djs:
                    for m in range(0, dj, 2):
                        tao.append(  i + dj - m)
                        tao.append(-(i + dj - m - 1))
                    i += dj
    return numpy.asarray(tao, dtype=numpy.int32)


CHECK_GEOM = getattr(__config__, 'gto_mole_check_geom', True)

def classical_coulomb_energy(mol, charges=None, coords=None):
    '''Compute nuclear repulsion energy (AU) or static Coulomb energy

    Returns
        float
    '''
    if charges is None: charges = mol.atom_charges()
    if len(charges) <= 1:
        return 0
    rr = inter_distance(mol, coords)
    rr[numpy.diag_indices_from(rr)] = 1e200
    if CHECK_GEOM and numpy.any(rr < 1e-5):
        raise_err = False
        for atm_idx in numpy.argwhere(rr<1e-5):
            # Only raise error if atoms with charge != 0 have the same coordinates
            if charges[atm_idx[0]] * charges[atm_idx[1]] != 0:
                logger.warn(mol, 'Atoms %s have the same coordinates', atm_idx)
                raise_err = True
            # At least one of the atoms is a ghost atom; suppress divide by 0 warning
            else:
                rr[atm_idx[0], atm_idx[1]] = 1e200
        if raise_err: raise RuntimeError('Ill geometry')
    e = numpy.einsum('i,ij,j->', charges, 1./rr, charges) * .5
    return e

energy_nuc = classical_coulomb_energy

def inter_distance(mol, coords=None):
    '''
    Inter-particle distance array
    '''
    if coords is None: coords = mol.atom_coords()
    rr = numpy.linalg.norm(coords.reshape(-1,1,3) - coords, axis=2)
    rr[numpy.diag_indices_from(rr)] = 0
    return rr

def sph_labels(mol, fmt=True, base=BASE):
    '''Labels for spherical GTO functions

    Kwargs:
        fmt : str or bool
        if fmt is boolean, it controls whether to format the labels and the
        default format is "%d%3s %s%-4s".  if fmt is string, the string will
        be used as the print format.

    Returns:
        List of [(atom-id, symbol-str, nl-str, str-of-real-spherical-notation]
        or formatted strings based on the argument "fmt"

    Examples:

    >>> mol = gto.M(atom='H 0 0 0; Cl 0 0 1', basis='sto-3g')
    >>> gto.sph_labels(mol)
    [(0, 'H', '1s', ''), (1, 'Cl', '1s', ''), (1, 'Cl', '2s', ''), (1, 'Cl', '3s', ''),
     (1, 'Cl', '2p', 'x'), (1, 'Cl', '2p', 'y'), (1, 'Cl', '2p', 'z'), (1, 'Cl', '3p', 'x'),
     (1, 'Cl', '3p', 'y'), (1, 'Cl', '3p', 'z')]
    '''
    count = numpy.zeros((mol.natm, 9), dtype=int)
    label = []
    for ib in range(mol.nbas):
        ia = mol.bas_atom(ib)
        l = mol.bas_angular(ib)
        strl = param.ANGULAR[l]
        nc = mol.bas_nctr(ib)
        symb = mol.atom_symbol(ia)
        nelec_ecp = mol.atom_nelec_core(ia)
        if nelec_ecp == 0 or l > 3:
            shl_start = count[ia,l]+l+1
        else:
            coreshl = core_configuration(nelec_ecp)
            shl_start = coreshl[l]+count[ia,l]+l+1
        count[ia,l] += nc
        for n in range(shl_start, shl_start+nc):
            for m in range(-l, l+1):
                label.append((ia+base, symb, '%d%s' % (n, strl),
                              str(param.REAL_SPHERIC[l][l+m])))

    if isinstance(fmt, (str, unicode)):
        return [(fmt % x) for x in label]
    elif fmt:
        return ['%d %s %s%-4s' % x for x in label]
    else:
        return label
spheric_labels = spherical_labels = sph_labels

def cart_labels(mol, fmt=True, base=BASE):
    '''Labels of Cartesian GTO functions

    Kwargs:
        fmt : str or bool
        if fmt is boolean, it controls whether to format the labels and the
        default format is "%d%3s %s%-4s".  if fmt is string, the string will
        be used as the print format.

    Returns:
        List of [(atom-id, symbol-str, nl-str, str-of-xyz-notation)]
        or formatted strings based on the argument "fmt"
    '''
    cartxyz = []
    for l in range(max(mol._bas[:,ANG_OF])+1):
        xyz = []
        for x in range(l, -1, -1):
            for y in range(l-x, -1, -1):
                z = l-x-y
                xyz.append('x'*x + 'y'*y + 'z'*z)
        cartxyz.append(xyz)

    count = numpy.zeros((mol.natm, 9), dtype=int)
    label = []
    for ib in range(len(mol._bas)):
        ia = mol.bas_atom(ib)
        l = mol.bas_angular(ib)
        strl = param.ANGULAR[l]
        nc = mol.bas_nctr(ib)
        symb = mol.atom_symbol(ia)
        nelec_ecp = mol.atom_nelec_core(ia)
        if nelec_ecp == 0 or l > 3:
            shl_start = count[ia,l]+l+1
        else:
            coreshl = core_configuration(nelec_ecp)
            shl_start = coreshl[l]+count[ia,l]+l+1
        count[ia,l] += nc
        ncart = (l + 1) * (l + 2) // 2
        for n in range(shl_start, shl_start+nc):
            for m in range(ncart):
                label.append((ia+base, symb, '%d%s' % (n, strl), cartxyz[l][m]))

    if isinstance(fmt, (str, unicode)):
        return [(fmt % x) for x in label]
    elif fmt:
        return ['%d%3s %s%-4s' % x for x in label]
    else:
        return label

def ao_labels(mol, fmt=True, base=BASE):
    '''Labels of AO basis functions

    Kwargs:
        fmt : str or bool
            if fmt is boolean, it controls whether to format the labels and the
            default format is "%d%3s %s%-4s".  if fmt is string, the string will
            be used as the print format.

    Returns:
        List of [(atom-id, symbol-str, nl-str, str-of-AO-notation)]
        or formatted strings based on the argument "fmt"
    '''
    if mol.cart:
        return mol.cart_labels(fmt, base)
    else:
        return mol.sph_labels(fmt, base)

def spinor_labels(mol, fmt=True, base=BASE):
    '''
    Labels of spinor GTO functions
    '''
    count = numpy.zeros((mol.natm, 9), dtype=int)
    label = []
    for ib in range(mol.nbas):
        ia = mol.bas_atom(ib)
        l = mol.bas_angular(ib)
        kappa = mol.bas_kappa(ib)
        strl = param.ANGULAR[l]
        nc = mol.bas_nctr(ib)
        symb = mol.atom_symbol(ia)
        nelec_ecp = mol.atom_nelec_core(ia)
        if nelec_ecp == 0 or l > 3:
            shl_start = count[ia,l]+l+1
        else:
            coreshl = core_configuration(nelec_ecp)
            shl_start = coreshl[l]+count[ia,l]+l+1
        count[ia,l] += nc
        for n in range(shl_start, shl_start+nc):
            if kappa >= 0:
                for m in range(-l*2+1, l*2, 2):
                    label.append((ia+base, symb, '%d%s%d/2' % (n, strl, l*2-1),
                                  '%d/2'%m))
            if kappa <= 0:
                for m in range(-l*2-1, l*2+2, 2):
                    label.append((ia+base, symb, '%d%s%d/2' % (n, strl, l*2+1),
                                  '%d/2'%m))

    if isinstance(fmt, (str, unicode)):
        return [(fmt % x) for x in label]
    elif fmt:
        return ['%d %s %s,%-5s' % x for x in label]
    else:
        return label

def search_ao_label(mol, label):
    '''Find the index of the AO basis function based on the given ao_label

    Args:
        ao_label : string or a list of strings
            The regular expression pattern to match the orbital labels
            returned by mol.ao_labels()

    Returns:
        A list of index for the AOs that matches the given ao_label RE pattern

    Examples:

    >>> mol = gto.M(atom='H 0 0 0; Cl 0 0 1', basis='ccpvtz')
    >>> mol.search_ao_label('Cl.*p')
    [19 20 21 22 23 24 25 26 27 28 29 30]
    >>> mol.search_ao_label('Cl 2p')
    [19 20 21]
    >>> mol.search_ao_label(['Cl.*d', 'Cl 4p'])
    [25 26 27 31 32 33 34 35 36 37 38 39 40]
    '''
    return _aolabels2baslst(mol, label)

def _aolabels2baslst(mol, aolabels_or_baslst, base=BASE):
    if callable(aolabels_or_baslst):
        baslst = [i for i,x in enumerate(mol.ao_labels(base=base))
                  if aolabels_or_baslst(x)]
    elif isinstance(aolabels_or_baslst, str):
        aolabels = re.sub(' +', ' ', aolabels_or_baslst.strip(), count=1)
        aolabels = re.compile(aolabels)
        baslst = [i for i,s in enumerate(mol.ao_labels(base=base))
                  if re.search(aolabels, s)]
    elif len(aolabels_or_baslst) > 0 and isinstance(aolabels_or_baslst[0], str):
        aolabels = [re.compile(re.sub(' +', ' ', x.strip(), count=1))
                    for x in aolabels_or_baslst]
        baslst = [i for i,t in enumerate(mol.ao_labels(base=base))
                  if any(re.search(x, t) for x in aolabels)]
    else:
        baslst = [i-base for i in aolabels_or_baslst]
    return numpy.asarray(baslst, dtype=int)

def search_shell_id(mol, atm_id, l):
    '''Search the first basis/shell id (**not** the basis function id) which
    matches the given atom-id and angular momentum

    Args:
        atm_id : int
            atom id, 0-based
        l : int
            angular momentum

    Returns:
        basis id, 0-based.  If not found, return None

    Examples:

    >>> mol = gto.M(atom='H 0 0 0; Cl 0 0 1', basis='sto-3g')
    >>> mol.search_shell_id(1, 1) # Cl p shell
    4
    >>> mol.search_shell_id(1, 2) # Cl d shell
    None
    '''
    for ib in range(len(mol._bas)):
        ia = mol.bas_atom(ib)
        l1 = mol.bas_angular(ib)
        if ia == atm_id and l1 == l:
            return ib

def search_ao_nr(mol, atm_id, l, m, atmshell):
    '''Search the first basis function id (**not** the shell id) which matches
    the given atom-id, angular momentum magnetic angular momentum, principal shell.

    Args:
        atm_id : int
            atom id, 0-based
        l : int
            angular momentum
        m : int
            magnetic angular momentum
        atmshell : int
            principal quantum number

    Returns:
        basis function id, 0-based.  If not found, return None

    Examples:

    >>> mol = gto.M(atom='H 0 0 0; Cl 0 0 1', basis='sto-3g')
    >>> mol.search_ao_nr(1, 1, -1, 3) # Cl 3px
    7
    '''
    ibf = 0
    for ib in range(len(mol._bas)):
        ia = mol.bas_atom(ib)
        l1 = mol.bas_angular(ib)
        if mol.cart:
            degen = (l1 + 1) * (l1 + 2) // 2
        else:
            degen = l1 * 2 + 1
        nc = mol.bas_nctr(ib)
        if ia == atm_id and l1 == l:
            if atmshell > nc+l1:
                atmshell = atmshell - nc
            else:
                return ibf + (atmshell-l1-1)*degen + (l1+m)
        ibf += degen * nc
    raise RuntimeError('Required AO not found')

def search_ao_r(mol, atm_id, l, j, m, atmshell):
    raise RuntimeError('TODO')
#TODO:    ibf = 0
#TODO:    for ib in range(len(mol._bas)):
#TODO:        ia = mol.bas_atom(ib)
#TODO:        l1 = mol.bas_angular(ib)
#TODO:        nc = mol.bas_nctr(ib)
#TODO:        k = mol.bas_kappa(bas_id)
#TODO:        degen = len_spinor(l1, k)
#TODO:        if ia == atm_id and l1 == l and k == kappa:
#TODO:            if atmshell > nc+l1:
#TODO:                atmshell = atmshell - nc
#TODO:            else:
#TODO:                return ibf + (atmshell-l1-1)*degen + (degen+m)
#TODO:        ibf += degen

def offset_2c_by_atom(mol):
    '''2-component AO offset for each atom.  Return a list, each item
    of the list gives (start-shell-id, stop-shell-id, start-AO-id, stop-AO-id)
    '''
    return aoslice_by_atom(mol, mol.ao_loc_2c())

def aoslice_by_atom(mol, ao_loc=None):
    '''AO offsets for each atom.  Return a list, each item of the list gives
    (start-shell-id, stop-shell-id, start-AO-id, stop-AO-id)
    '''
    if ao_loc is None:
        ao_loc = mol.ao_loc_nr()

    aorange = numpy.empty((mol.natm,4), dtype=int)

    if mol.natm == 0:
        return aorange

    bas_atom = mol._bas[:,ATOM_OF]
    delimiter = numpy.where(bas_atom[0:-1] != bas_atom[1:])[0] + 1

    if mol.natm == len(delimiter) + 1:
        aorange[:,0] = shell_start = numpy.append(0, delimiter)
        aorange[:,1] = shell_end = numpy.append(delimiter, mol.nbas)

    else:  # Some atoms miss basis
        shell_start = numpy.empty(mol.natm, dtype=int)
        shell_start[:] = -1
        shell_start[0] = 0
        shell_start[bas_atom[0]] = 0
        shell_start[bas_atom[delimiter]] = delimiter

        shell_end = numpy.empty(mol.natm, dtype=int)
        shell_end[0] = 0
        shell_end[bas_atom[delimiter-1]] = delimiter
        shell_end[bas_atom[-1]] = mol.nbas

        for i in range(1, mol.natm):
            if shell_start[i] == -1:
                shell_start[i] = shell_end[i] = shell_end[i-1]

    aorange[:,0] = shell_start
    aorange[:,1] = shell_end
    aorange[:,2] = ao_loc[shell_start]
    aorange[:,3] = ao_loc[shell_end]
    return aorange
offset_nr_by_atom = aoslice_by_atom

def same_basis_set(mol1, mol2):
    '''Check whether two molecules use the same basis sets.
    The two molecules can have different geometry.
    '''
    atomtypes1 = atom_types(mol1._atom, mol1._basis)
    atomtypes2 = atom_types(mol2._atom, mol2._basis)
    if set(atomtypes1.keys()) != set(atomtypes2.keys()):
        return False
    for k in atomtypes1:
        if len(atomtypes1[k]) != len(atomtypes2[k]):
            return False
        elif mol1._basis.get(k, None) != mol2._basis.get(k, None):
            return False
    return True

def same_mol(mol1, mol2, tol=1e-5, cmp_basis=True, ignore_chiral=False):
    '''Compare the two molecules whether they have the same structure.

    Kwargs:
        tol : float
            In Bohr
        cmp_basis : bool
            Whether to compare basis functions for the two molecules
    '''
    from pyscf import symm

    if mol1._atom.__len__() != mol2._atom.__len__():
        return False

    chg1 = mol1._atm[:,CHARGE_OF]
    chg2 = mol2._atm[:,CHARGE_OF]
    if not numpy.all(numpy.sort(chg1) == numpy.sort(chg2)):
        return False

    if cmp_basis and not same_basis_set(mol1, mol2):
        return False

    def finger(mol, chgs, coord):
        center = numpy.einsum('z,zr->r', chgs, coord) / chgs.sum()
        im = inertia_moment(mol, chgs, coord)
        # Divid im by chgs.sum(), to normalize im. Otherwise the input tol may
        # not reflect the actual deviation.
        im /= chgs.sum()

        w, v = scipy.linalg.eigh(im)
        axes = v.T
        if numpy.linalg.det(axes) < 0:
            axes *= -1
        r = numpy.dot(coord-center, axes.T)
        return w, r

    coord1 = mol1.atom_coords()
    coord2 = mol2.atom_coords()
    w1, r1 = finger(mol1, chg1, coord1)
    w2, r2 = finger(mol2, chg2, coord2)
    if not (numpy.allclose(w1, w2, atol=tol)):
        return False

    rotate_xy  = numpy.array([[-1., 0., 0.],
                              [ 0.,-1., 0.],
                              [ 0., 0., 1.]])
    rotate_yz  = numpy.array([[ 1., 0., 0.],
                              [ 0.,-1., 0.],
                              [ 0., 0.,-1.]])
    rotate_zx  = numpy.array([[-1., 0., 0.],
                              [ 0., 1., 0.],
                              [ 0., 0.,-1.]])

    def inspect(z1, r1, z2, r2):
        place = int(-numpy.log10(tol)) - 1
        idx = symm.argsort_coords(r2, place)
        z2 = z2[idx]
        r2 = r2[idx]
        for rot in (1, rotate_xy, rotate_yz, rotate_zx):
            r1new = numpy.dot(r1, rot)
            idx = symm.argsort_coords(r1new, place)
            if (numpy.all(z1[idx] == z2) and
                numpy.allclose(r1new[idx], r2, atol=tol)):
                return True
        return False

    return (inspect(chg1, r1, chg2, r2) or
            (ignore_chiral and inspect(chg1, r1, chg2, -r2)))
is_same_mol = same_mol

def chiral_mol(mol1, mol2=None):
    '''Detect whether the given molelcule is chiral molecule or two molecules
    are chiral isomers.
    '''
    if mol2 is None:
        mol2 = mol1.copy()
        ptr_coord = mol2._atm[:,PTR_COORD]
        mol2._env[ptr_coord  ] *= -1
        mol2._env[ptr_coord+1] *= -1
        mol2._env[ptr_coord+2] *= -1
    return (not same_mol(mol1, mol2, ignore_chiral=False) and
            same_mol(mol1, mol2, ignore_chiral=True))

def inertia_moment(mol, mass=None, coords=None):
    if mass is None:
        mass = mol.atom_mass_list()
    if coords is None:
        coords = mol.atom_coords()
    mass_center = numpy.einsum('i,ij->j', mass, coords)/mass.sum()
    coords = coords - mass_center
    im = numpy.einsum('i,ij,ik->jk', mass, coords, coords)
    im = numpy.eye(3) * im.trace() - im
    return im

def atom_mass_list(mol, isotope_avg=False):
    '''A list of mass for all atoms in the molecule

    Kwargs:
        isotope_avg : boolean
            Whether to use the isotope average mass as the atomic mass
    '''
    if isotope_avg:
        mass_table = elements.MASSES
    else:
        mass_table = elements.ISOTOPE_MAIN

    nucprop = mol.nucprop
    if nucprop:
        mass = []
        for ia in range(mol.natm):
            z = mol.atom_charge(ia)
            symb = mol.atom_symbol(ia)
            stdsymb = _std_symbol(symb)
            if ia+1 in nucprop:
                prop = nucprop[ia+1]
            elif symb in nucprop:
                prop = nucprop[symb]
            else:
                prop = nucprop.get(stdsymb, {})

            mass.append(prop.get('mass', mass_table[z]))
    else:
        #mass = [mass_table[z] for z in mol.atom_charges()]
        mass = []
        for ia in range(mol.natm):
            z = charge(mol.atom_symbol(ia))
            mass.append(mass_table[z])

    return numpy.array(mass)

def condense_to_shell(mol, mat, compressor='max'):
    '''The given matrix is first partitioned to blocks, based on AO shell as
    delimiter. Then call compressor function to abstract each block.

    Args:
        compressor: string or function
            if compressor is a string, its value can be  sum, max, min, abssum,
            absmax, absmin, norm
    '''
    ao_loc = mol.ao_loc_nr()
    if callable(compressor):
        abstract = numpy.empty((mol.nbas, mol.nbas))
        for i, i0 in enumerate(ao_loc[:mol.nbas]):
            for j, j0 in enumerate(ao_loc[:mol.nbas]):
                abstract[i,j] = compressor(mat[i0:ao_loc[i+1],j0:ao_loc[j+1]])
    else:
        abstract = lib.condense(compressor, mat, ao_loc)
    return abstract


def tostring(mol, format='raw'):
    '''Convert molecular geometry to a string of the required format.

    Supported output formats:
        | raw: Each line is  <symobl> <x> <y> <z>
        | xyz: XYZ cartesian coordinates format
        | zmat: Z-matrix format
    '''
    format = format.lower()
    if format == 'xyz' or format == 'raw':
        coords = mol.atom_coords() * param.BOHR
        output = []
        if format == 'xyz':
            output.append('%d' % mol.natm)
            output.append('XYZ from PySCF')

        for i in range(mol.natm):
            symb = mol.atom_pure_symbol(i)
            x, y, z = coords[i]
            output.append('%-4s %14.5f %14.5f %14.5f' %
                          (symb, x, y, z))
        return '\n'.join(output)
    elif format == 'zmat':
        coords = mol.atom_coords() * param.BOHR
        zmat = cart2zmat(coords).splitlines()
        output = []
        for i, line in enumerate(zmat):
            symb = mol.atom_pure_symbol(i)
            output.append('%-4s   %s' % (symb, line))
        return '\n'.join(output)
    else:
        raise NotImplementedError

def tofile(mol, filename, format=None):
    '''Write molecular geometry to a file of the required format.

    Supported output formats:
        | raw: Each line is  <symobl> <x> <y> <z>
        | xyz: XYZ cartesian coordinates format
        | zmat: Z-matrix format
    '''
    if format is None:  # Guess format based on filename
        format = os.path.splitext(filename)[1][1:]
    string = tostring(mol, format)
    with open(filename, 'w') as f:
        f.write(string)
        f.write('\n')
    return string

def fromfile(filename, format=None):
    '''Read molecular geometry from a file
    (in testing)

    Supported formats:
        | raw: Each line is  <symobl> <x> <y> <z>
        | xyz: XYZ cartesian coordinates format
        | zmat: Z-matrix format
    '''
    if format is None:  # Guess format based on filename
        format = os.path.splitext(filename)[1][1:].lower()
        if format not in ('xyz', 'zmat', 'sdf'):
            format = 'raw'
    with open(filename, 'r') as f:
        return fromstring(f.read(), format)


def fromstring(string, format='xyz'):
    '''Convert the string of the specified format to internal format
    (in testing)

    Supported formats:
        | raw: Each line is  <symobl> <x> <y> <z>
        | xyz: XYZ cartesian coordinates format
        | zmat: Z-matrix format
    '''
    format = format.lower()
    if format == 'zmat':
        return from_zmatrix(string)
    elif format == 'xyz':
        dat = string.splitlines()
        natm = int(dat[0])
        return '\n'.join(dat[2:natm+2])
    elif format == 'sdf':
        raw = string.splitlines()
        natoms, nbonds = raw[3].split()[:2]
        atoms = []
        for line in raw[4:4+int(natoms)]:
            d = line.split()
            atoms.append('%s %s %s %s' % (d[3], d[0], d[1], d[2]))
        return '\n'.join(atoms)
    elif format == 'raw':
        return string
    else:
        raise NotImplementedError


#
# Mole class handles three layers: input, internal format, libcint arguments.
# The relationship of the three layers are, eg
#    .atom (input) <=>  ._atom (for python) <=> ._atm (for libcint)
#   .basis (input) <=> ._basis (for python) <=> ._bas (for libcint)
# input layer does not talk to libcint directly.  Data are held in python
# internal fomrat layer.  Most of methods defined in this class only operates
# on the internal format.  Exceptions are make_env, make_atm_env, make_bas_env,
# set_common_orig_, set_rinv_orig_ which are used to manipulate the libcint arguments.
#
class Mole(lib.StreamObject):
    '''Basic class to hold molecular structure and global options

    Attributes:
        verbose : int
            Print level
        output : str or None
            Output file, default is None which dumps msg to sys.stdout
        max_memory : int, float
            Allowed memory in MB
        charge : int
            Charge of molecule. It affects the electron numbers
        spin : int or None
            2S, num. alpha electrons - num. beta electrons to control
            multiplicity. If spin = None is set, multiplicity will be guessed
            based on the neutral molecule.
        symmetry : bool or str
            Whether to use symmetry.  When this variable is set to True, the
            molecule will be rotated and the highest rotation axis will be
            placed z-axis.
            If a string is given as the name of point group, the given point
            group symmetry will be used.  Note that the input molecular
            coordinates will not be changed in this case.
        symmetry_subgroup : str
            subgroup

        atom : list or str
            To define molecluar structure.  The internal format is

            | atom = [[atom1, (x, y, z)],
            |         [atom2, (x, y, z)],
            |         ...
            |         [atomN, (x, y, z)]]

        unit : str
            Angstrom or Bohr
        basis : dict or str
            To define basis set.
        nucmod : dict or str or [function(nuc_charge, nucprop) => zeta]
            Nuclear model.  0 or None means point nuclear model.  Other
            values will enable Gaussian nuclear model.  If a function is
            assigned to this attribute, the function will be called to
            generate the nuclear charge distribution value "zeta" and the
            relevant nuclear model will be set to Gaussian model.
            Default is point nuclear model.
        nucprop : dict
            Nuclear properties (like g-factor 'g', quadrupole moments 'Q').
            It is needed by pyscf.prop module and submodules.
        cart : boolean
            Using Cartesian GTO basis and integrals (6d,10f,15g)

        ** Following attributes are generated by :func:`Mole.build` **

        stdout : file object
            Default is sys.stdout if :attr:`Mole.output` is not set
        topgroup : str
            Point group of the system.
        groupname : str
            The supported subgroup of the point group. It can be one of SO3,
            Dooh, Coov, D2h, C2h, C2v, D2, Cs, Ci, C2, C1
        nelectron : int
            sum of nuclear charges - :attr:`Mole.charge`
        symm_orb : a list of numpy.ndarray
            Symmetry adapted basis.  Each element is a set of symm-adapted orbitals
            for one irreducible representation.  The list index does **not** correspond
            to the id of irreducible representation.
        irrep_id : a list of int
            Each element is one irreducible representation id associated with the basis
            stored in symm_orb.  One irrep id stands for one irreducible representation
            symbol.  The irrep symbol and the relevant id are defined in
            :attr:`symm.param.IRREP_ID_TABLE`
        irrep_name : a list of str
            Each element is one irreducible representation symbol associated with the basis
            stored in symm_orb.  The irrep symbols are defined in
            :attr:`symm.param.IRREP_ID_TABLE`
        _built : bool
            To label whether :func:`Mole.build` has been called.  It is to
            ensure certain functions being initialized only once.
        _basis : dict
            like :attr:`Mole.basis`, the internal format which is returned from the
            parser :func:`format_basis`
        _keys : a set of str
            Store the keys appeared in the module.  It is used to check misinput attributes

        ** Following attributes are arguments used by ``libcint`` library **

        _atm :
            :code:`[[charge, ptr-of-coord, nuc-model, ptr-zeta, 0, 0], [...]]`
            each element reperesents one atom
        natm :
            number of atoms
        _bas :
            :code:`[[atom-id, angular-momentum, num-primitive-GTO, num-contracted-GTO, 0, ptr-of-exps, ptr-of-contract-coeff, 0], [...]]`
            each element reperesents one shell
        nbas :
            number of shells
        _env :
            list of floats to store the coordinates, GTO exponents, contract-coefficients

    Examples:

    >>> mol = Mole(atom='H^2 0 0 0; H 0 0 1.1', basis='sto3g').build()
    >>> print(mol.atom_symbol(0))
    H^2
    >>> print(mol.atom_pure_symbol(0))
    H
    >>> print(mol.nao_nr())
    2
    >>> print(mol.intor('int1e_ovlp_sph'))
    [[ 0.99999999  0.43958641]
     [ 0.43958641  0.99999999]]
    >>> mol.charge = 1
    >>> mol.build()
    <class 'pyscf.gto.mole.Mole'> has no attributes Charge
    '''  # noqa: E501

    verbose = getattr(__config__, 'VERBOSE', logger.NOTE)

    # the unit (angstrom/bohr) of the coordinates defined by the input self.atom
    unit = getattr(__config__, 'UNIT', 'angstrom')

    # Whether to hold everything in memory
    incore_anyway = getattr(__config__, 'INCORE_ANYWAY', False)

    # Using cartesian GTO (6d,10f,15g)
    cart = getattr(__config__, 'gto_mole_Mole_cart', False)

    def __init__(self, **kwargs):
        self.output = None
        self.max_memory = param.MAX_MEMORY

        self.charge = 0
        self.spin = 0 # 2j == nelec_alpha - nelec_beta
        self.symmetry = False
        self.symmetry_subgroup = None
        self.cart = False

# Save inputs
# self.atom = [(symb/nuc_charge, (coord(Angstrom):0.,0.,0.)), ...]
        self.atom = []
# self.basis = {atom_type/nuc_charge: [l, kappa, (expnt, c_1, c_2,..),..]}
        self.basis = 'sto-3g'
# self.nucmod = {atom_symbol: nuclear_model, atom_id: nuc_mod}, atom_id is 1-based
        self.nucmod = {}
# self.ecp = {atom_symbol: [[l, (r_order, expnt, c),...]]}
        self.ecp = {}
# Nuclear property. self.nucprop = {atom_symbol: {key: value}}
        self.nucprop = {}
##################################################
# don't modify the following private variables, they are not input options
        self._atm = numpy.zeros((0,6), dtype=numpy.int32)
        self._bas = numpy.zeros((0,8), dtype=numpy.int32)
        self._env = numpy.zeros(PTR_ENV_START)
        self._ecpbas = numpy.zeros((0,8), dtype=numpy.int32)

        self.stdout = sys.stdout
        self.groupname = 'C1'
        self.topgroup = 'C1'
        self.symm_orb = None
        self.irrep_id = None
        self.irrep_name = None
        self._symm_orig = None
        self._symm_axes = None
        self._nelectron = None
        self._nao = None
        self._enuc = None
        self._atom = []
        self._basis = {}
        self._ecp = {}
        self._built = False

        # _pseudo is created to make the mol object consistenet with the mol
        # object converted from Cell.to_mol(). It is initialized in the
        # Cell.build() method only. Assigning _pseudo to mol object basically
        # has no effects. Mole.build() method does not have code to access the
        # contents of _pseudo.
        self._pseudo = {}

        # Some methods modify ._env. These method are executed in the context
        # _TemporaryMoleContext which is protected by the _ctx_lock.
        self._ctx_lock = None
        keys = set(('verbose', 'unit', 'cart', 'incore_anyway'))
        self._keys = set(self.__dict__.keys()).union(keys)
        self.__dict__.update(kwargs)

    @property
    def natm(self):
        return len(self._atm)
    @property
    def nbas(self):
        return len(self._bas)

    @property
    def nelec(self):
        ne = self.nelectron
        nalpha = (ne + self.spin) // 2
        nbeta = nalpha - self.spin
        assert(nalpha >= 0 and nbeta >= 0)
        if nalpha + nbeta != ne:
            raise RuntimeError('Electron number %d and spin %d are not consistent\n'
                               'Note mol.spin = 2S = Nalpha - Nbeta, not 2S+1' %
                               (ne, self.spin))
        return nalpha, nbeta
    @nelec.setter
    def nelec(self, neleca_nelecb):
        neleca, nelecb = neleca_nelecb
        self._nelectron = neleca + nelecb
        self.spin = neleca - nelecb

    @property
    def nelectron(self):
        if self._nelectron is None:
            return self.tot_electrons()
        else:
            return self._nelectron
    @nelectron.setter
    def nelectron(self, n):
        self._nelectron = n

    @property
    def multiplicity(self):
        return self.spin + 1
    @multiplicity.setter
    def multiplicity(self, x):
        if x is None:
            self.spin = None
        else:
            self.spin = x - 1

    @property
    def ms(self):
        '''Spin quantum number. multiplicity = ms*2+1'''
        if self.spin % 2 == 0:
            return self.spin // 2
        else:
            return self.spin * .5
    @ms.setter
    def ms(self, x):
        if x is None:
            self.spin = None
        else:
            self.spin = int(round(2*x, 4))

    def __getattr__(self, key):
        '''To support accessing methods (mol.HF, mol.KS, mol.CCSD, mol.CASSCF, ...)
        from Mole object.
        '''
        if key[0] == '_':  # Skip private attributes and Python builtins
            raise AttributeError('Mole object does not have attribute %s' % key)
        elif key in ('_ipython_canary_method_should_not_exist_',
                     '_repr_mimebundle_'):
            # https://github.com/mewwts/addict/issues/26
            # https://github.com/jupyter/notebook/issues/2014
            raise AttributeError

        # Import all available modules. Some methods are registered to other
        # classes/modules when importing modules in __all__.
        from pyscf import __all__  # noqa
        from pyscf import scf, dft
        for mod in (scf, dft):
            method = getattr(mod, key, None)
            if callable(method):
                return method(self)

        if 'TD' in key[:3]:
            if key in ('TDHF', 'TDA'):
                mf = scf.HF(self)
            else:
                mf = dft.KS(self)
                xc = key.split('TD', 1)[1]
                if xc in dft.XC:
                    mf.xc = xc
                    key = 'TDDFT'
        else:
            mf = scf.HF(self)

        method = getattr(mf, key, None)
        if method is None:
            raise AttributeError('Mole object has no attribute %s' % key)

        # Initialize SCF object for post-SCF methods if applicable
        if self.nelectron != 0:
            mf.run()
        return method

# need "deepcopy" here because in shallow copy, _env may get new elements but
# with ptr_env unchanged
# def __copy__(self):
#        cls = self.__class__
#        newmol = cls.__new__(cls)
#        newmol = ...
# do not use __copy__ to aovid iteratively call copy.copy
    copy = copy

    pack = pack

    @classmethod
    @lib.with_doc(unpack.__doc__)
    def unpack(cls, moldic):
        return unpack(moldic)

    @lib.with_doc(unpack.__doc__)
    def unpack_(self, moldic):
        self.__dict__.update(moldic)
        return self

    dumps = dumps

    @classmethod
    @lib.with_doc(loads.__doc__)
    def loads(cls, molstr):
        return loads(molstr)

    @lib.with_doc(loads.__doc__)
    def loads_(self, molstr):
        self.__dict__.update(loads(molstr).__dict__)
        return self

    def build(self, dump_input=True, parse_arg=True,
              verbose=None, output=None, max_memory=None,
              atom=None, basis=None, unit=None, nucmod=None, ecp=None,
              charge=None, spin=0, symmetry=None, symmetry_subgroup=None,
              cart=None):
        '''Setup moleclue and initialize some control parameters.  Whenever you
        change the value of the attributes of :class:`Mole`, you need call
        this function to refresh the internal data of Mole.

        Kwargs:
            dump_input : bool
                whether to dump the contents of input file in the output file
            parse_arg : bool
                whether to read the sys.argv and overwrite the relevant parameters
            verbose : int
                Print level.  If given, overwrite :attr:`Mole.verbose`
            output : str or None
                Output file.  If given, overwrite :attr:`Mole.output`
            max_memory : int, float
                Allowd memory in MB.  If given, overwrite :attr:`Mole.max_memory`
            atom : list or str
                To define molecluar structure.
            basis : dict or str
                To define basis set.
            nucmod : dict or str
                Nuclear model.  If given, overwrite :attr:`Mole.nucmod`
            charge : int
                Charge of molecule. It affects the electron numbers
                If given, overwrite :attr:`Mole.charge`
            spin : int
                2S, num. alpha electrons - num. beta electrons to control
                multiplicity. If setting spin = None , multiplicity will be
                guessed based on the neutral molecule.
                If given, overwrite :attr:`Mole.spin`
            symmetry : bool or str
                Whether to use symmetry.  If given a string of point group
                name, the given point group symmetry will be used.

        '''
        if not DISABLE_GC:
            gc.collect()  # To release circular referred objects

        if isinstance(dump_input, (str, unicode)):
            sys.stderr.write('Assigning the first argument %s to mol.atom\n' %
                             dump_input)
            dump_input, atom = True, dump_input

        if verbose is not None: self.verbose = verbose
        if output is not None: self.output = output
        if max_memory is not None: self.max_memory = max_memory
        if atom is not None: self.atom = atom
        if basis is not None: self.basis = basis
        if unit is not None: self.unit = unit
        if nucmod is not None: self.nucmod = nucmod
        if ecp is not None: self.ecp = ecp
        if charge is not None: self.charge = charge
        if spin != 0: self.spin = spin
        if symmetry is not None: self.symmetry = symmetry
        if symmetry_subgroup is not None: self.symmetry_subgroup = symmetry_subgroup
        if cart is not None: self.cart = cart

        if parse_arg:
            _update_from_cmdargs_(self)

        # avoid opening output file twice
        if (self.output is not None
            # StringIO() does not have attribute 'name'
            and getattr(self.stdout, 'name', None) != self.output):

            if self.verbose > logger.QUIET:
                if os.path.isfile(self.output):
                    print('overwrite output file: %s' % self.output)
                else:
                    print('output file: %s' % self.output)

            if self.output == '/dev/null':
                self.stdout = open(os.devnull, 'w')
            else:
                self.stdout = open(self.output, 'w')

        if self.verbose >= logger.WARN:
            self.check_sanity()

        self._atom = self.format_atom(self.atom, unit=self.unit)
        uniq_atoms = set([a[0] for a in self._atom])

        if isinstance(self.basis, (str, unicode, tuple, list)):
            # specify global basis for whole molecule
            _basis = dict(((a, self.basis) for a in uniq_atoms))
        elif 'default' in self.basis:
            default_basis = self.basis['default']
            _basis = dict(((a, default_basis) for a in uniq_atoms))
            _basis.update(self.basis)
            del(_basis['default'])
        else:
            _basis = self.basis
        self._basis = self.format_basis(_basis)

# TODO: Consider ECP info in point group symmetry initialization
        if self.ecp:
            # Unless explicitly input, ECP should not be assigned to ghost atoms
            if isinstance(self.ecp, (str, unicode)):
                _ecp = dict([(a, str(self.ecp))
                             for a in uniq_atoms if not is_ghost_atom(a)])
            elif 'default' in self.ecp:
                default_ecp = self.ecp['default']
                _ecp = dict(((a, default_ecp)
                             for a in uniq_atoms if not is_ghost_atom(a)))
                _ecp.update(self.ecp)
                del(_ecp['default'])
            else:
                _ecp = self.ecp
            self._ecp = self.format_ecp(_ecp)

        env = self._env[:PTR_ENV_START]
        self._atm, self._bas, self._env = \
                self.make_env(self._atom, self._basis, env, self.nucmod,
                              self.nucprop)
        self._atm, self._ecpbas, self._env = \
                self.make_ecp_env(self._atm, self._ecp, self._env)

        if self.spin is None:
            self.spin = self.nelectron % 2
        else:
            # Access self.nelec in which the code checks whether the spin and
            # number of electrons are consistent.
            self.nelec

        if self.symmetry:
            self._build_symmetry()

        if dump_input and not self._built and self.verbose > logger.NOTE:
            self.dump_input()

        if self.verbose >= logger.DEBUG3:
            logger.debug3(self, 'arg.atm = %s', self._atm)
            logger.debug3(self, 'arg.bas = %s', self._bas)
            logger.debug3(self, 'arg.env = %s', self._env)
            logger.debug3(self, 'ecpbas  = %s', self._ecpbas)

        self._built = True
        return self
    kernel = build

    def _build_symmetry(self):
        '''
        Update symmetry related attributes: topgroup, groupname, _symm_orig,
        _symm_axes, irrep_id, irrep_name, symm_orb
        '''
        from pyscf import symm
        self.topgroup, orig, axes = symm.detect_symm(self._atom, self._basis)

        if isinstance(self.symmetry, (str, unicode)):
            self.symmetry = str(symm.std_symb(self.symmetry))
            groupname = None
            if abs(axes - np.eye(3)).max() < symm.TOLERANCE:
                if symm.check_symm(self.symmetry, self._atom, self._basis):
                    # Try to use original axes (issue #1209)
                    groupname = self.symmetry
                    axes = np.eye(3)
                else:
                    logger.warn(self, 'Unable to to identify input symmetry using original axes.\n'
                                'Different symmetry axes will be used.')
            if groupname is None:
                try:
                    groupname, axes = symm.as_subgroup(self.topgroup, axes,
                                                       self.symmetry)
                except PointGroupSymmetryError as e:
                    raise PointGroupSymmetryError(
                        'Unable to identify input symmetry %s. Try symmetry="%s"' %
                        (self.symmetry, self.topgroup)) from e
        else:
            groupname, axes = symm.as_subgroup(self.topgroup, axes,
                                               self.symmetry_subgroup)
        self._symm_orig = orig
        self._symm_axes = axes

        if self.cart and groupname in ('Dooh', 'Coov', 'SO3'):
            if groupname == 'Coov':
                groupname, lgroup = 'C2v', groupname
            else:
                groupname, lgroup = 'D2h', groupname
            logger.warn(self, 'This version does not support symmetry %s '
                        'for cartesian GTO basis. Its subgroup %s is used',
                        lgroup, groupname)
        self.groupname = groupname

        self.symm_orb, self.irrep_id = \
                symm.symm_adapted_basis(self, groupname, orig, axes)
        self.irrep_name = [symm.irrep_id2name(groupname, ir)
                           for ir in self.irrep_id]
        return self

    @lib.with_doc(format_atom.__doc__)
    def format_atom(self, atom, origin=0, axes=None, unit='Ang'):
        return format_atom(atom, origin, axes, unit)

    @lib.with_doc(format_basis.__doc__)
    def format_basis(self, basis_tab):
        return format_basis(basis_tab)

    @lib.with_doc(format_ecp.__doc__)
    def format_ecp(self, ecp_tab):
        return format_ecp(ecp_tab)

    @lib.with_doc(expand_etb.__doc__)
    def expand_etb(self, l, n, alpha, beta):
        return expand_etb(l, n, alpha, beta)

    @lib.with_doc(expand_etbs.__doc__)
    def expand_etbs(self, etbs):
        return expand_etbs(etbs)
    etbs = expand_etbs

    @lib.with_doc(make_env.__doc__)
    def make_env(self, atoms, basis, pre_env=[], nucmod={}, nucprop=None):
        if nucprop is None:
            nucprop = self.nucprop
        return make_env(atoms, basis, pre_env, nucmod, nucprop)

    @lib.with_doc(make_atm_env.__doc__)
    def make_atm_env(self, atom, ptr=0, nucmod=NUC_POINT, nucprop=None):
        if nucprop is None:
            nucprop = self.nucprop.get(atom[0], {})
        return make_atm_env(atom, ptr, nucmod, nucprop)

    @lib.with_doc(make_bas_env.__doc__)
    def make_bas_env(self, basis_add, atom_id=0, ptr=0):
        return make_bas_env(basis_add, atom_id, ptr)

    @lib.with_doc(make_ecp_env.__doc__)
    def make_ecp_env(self, _atm, _ecp, pre_env=[]):
        if _ecp:
            _atm, _ecpbas, _env = make_ecp_env(self, _atm, _ecp, pre_env)
        else:
            _atm, _ecpbas, _env = _atm, numpy.zeros((0,BAS_SLOTS)), pre_env
        return _atm, _ecpbas, _env

    tot_electrons = tot_electrons

    @lib.with_doc(gto_norm.__doc__)
    def gto_norm(self, l, expnt):
        return gto_norm(l, expnt)


    def dump_input(self):
        import __main__
        if hasattr(__main__, '__file__'):
            try:
                filename = os.path.abspath(__main__.__file__)
                finput = open(filename, 'r')
                self.stdout.write('#INFO: **** input file is %s ****\n' % filename)
                self.stdout.write(finput.read())
                self.stdout.write('#INFO: ******************** input file end ********************\n')
                self.stdout.write('\n')
                self.stdout.write('\n')
                finput.close()
            except IOError:
                logger.warn(self, 'input file does not exist')

        self.stdout.write('System: %s  Threads %s\n' %
                          (str(platform.uname()), lib.num_threads()))
        self.stdout.write('Python %s\n' % sys.version)
        self.stdout.write('numpy %s  scipy %s\n' %
                          (numpy.__version__, scipy.__version__))
        self.stdout.write('Date: %s\n' % time.ctime())
        import pyscf
        self.stdout.write('PySCF version %s\n' % pyscf.__version__)
        info = lib.repo_info(os.path.join(__file__, '..', '..'))
        self.stdout.write('PySCF path  %s\n' % info['path'])
        if 'git' in info:
            self.stdout.write(info['git'] + '\n')

        self.stdout.write('\n')
        for key in os.environ:
            if 'PYSCF' in key:
                self.stdout.write('[ENV] %s %s\n' % (key, os.environ[key]))
        if self.verbose >= logger.DEBUG2:
            for key in dir(__config__):
                if key[:2] != '__':
                    self.stdout.write('[CONFIG] %s = %s\n' %
                                      (key, getattr(__config__, key)))
        else:
            conf_file = getattr(__config__, 'conf_file', None)
            self.stdout.write('[CONFIG] conf_file %s\n' % conf_file)

        self.stdout.write('[INPUT] verbose = %d\n' % self.verbose)
        if self.verbose >= logger.DEBUG:
            self.stdout.write('[INPUT] max_memory = %s \n' % self.max_memory)
        self.stdout.write('[INPUT] num. atoms = %d\n' % self.natm)
        self.stdout.write('[INPUT] num. electrons = %d\n' % self.nelectron)
        self.stdout.write('[INPUT] charge = %d\n' % self.charge)
        self.stdout.write('[INPUT] spin (= nelec alpha-beta = 2S) = %d\n' % self.spin)
        self.stdout.write('[INPUT] symmetry %s subgroup %s\n' %
                          (self.symmetry, self.symmetry_subgroup))
        self.stdout.write('[INPUT] Mole.unit = %s\n' % self.unit)
        if self.cart:
            self.stdout.write('[INPUT] Cartesian GTO integrals (6d 10f)\n')

        for ia,atom in enumerate(self._atom):
            coorda = tuple([x * param.BOHR for x in atom[1]])
            coordb = tuple([x for x in atom[1]])
            self.stdout.write('[INPUT]%3d %-4s %16.12f %16.12f %16.12f AA  '
                              '%16.12f %16.12f %16.12f Bohr\n'
                              % ((ia+1, _symbol(atom[0])) + coorda + coordb))
        if self.nucmod:
            if isinstance(self.nucmod, (int, str, unicode,
                                        types.FunctionType)):
                nucatms = [_symbol(atom[0]) for atom in self._atom]
            else:
                nucatms = self.nucmod.keys()
            self.stdout.write('[INPUT] Gaussian nuclear model for atoms %s\n' %
                              nucatms)

        if self.nucprop:
            self.stdout.write('[INPUT] nucprop %s\n' % self.nucprop)

        if self.verbose >= logger.DEBUG:
            self.stdout.write('[INPUT] ---------------- BASIS SET ---------------- \n')
            self.stdout.write('[INPUT] l, kappa, [nprim/nctr], '
                              'expnt,             c_1 c_2 ...\n')
            for atom, basis_set in self._basis.items():
                self.stdout.write('[INPUT] %s\n' % atom)
                for b in basis_set:
                    if isinstance(b[1], int):
                        kappa = b[1]
                        b_coeff = b[2:]
                    else:
                        kappa = 0
                        b_coeff = b[1:]
                    nprim = len(b_coeff)
                    nctr = len(b_coeff[0])-1
                    if nprim < nctr:
                        logger.warn(self, 'num. primitives smaller than num. contracted basis')
                    self.stdout.write('[INPUT] %d   %2d    [%-5d/%-4d]  '
                                      % (b[0], kappa, nprim, nctr))
                    for k, x in enumerate(b_coeff):
                        if k == 0:
                            self.stdout.write('%-15.12g  ' % x[0])
                        else:
                            self.stdout.write(' '*32+'%-15.12g  ' % x[0])
                        for c in x[1:]:
                            self.stdout.write(' %4.12g' % c)
                        self.stdout.write('\n')

        if self.verbose >= logger.INFO:
            self.stdout.write('\n')
            logger.info(self, 'nuclear repulsion = %.15g', self.energy_nuc())
            if self.symmetry:
                if self.topgroup == self.groupname:
                    logger.info(self, 'point group symmetry = %s', self.topgroup)
                else:
                    logger.info(self, 'point group symmetry = %s, use subgroup %s',
                                self.topgroup, self.groupname)
                logger.info(self, "symmetry origin: %s", self._symm_orig)
                logger.info(self, "symmetry axis x: %s", self._symm_axes[0])
                logger.info(self, "symmetry axis y: %s", self._symm_axes[1])
                logger.info(self, "symmetry axis z: %s", self._symm_axes[2])
                for ir in range(self.symm_orb.__len__()):
                    logger.info(self, 'num. orbitals of irrep %s = %d',
                                self.irrep_name[ir], self.symm_orb[ir].shape[1])
            logger.info(self, 'number of shells = %d', self.nbas)
            logger.info(self, 'number of NR pGTOs = %d', self.npgto_nr())
            logger.info(self, 'number of NR cGTOs = %d', self.nao_nr())
            logger.info(self, 'basis = %s', self.basis)
            logger.info(self, 'ecp = %s', self.ecp)
        if self.verbose >= logger.DEBUG2:
            for i in range(len(self._bas)):
                exps = self.bas_exp(i)
                logger.debug1(self, 'bas %d, expnt(s) = %s', i, str(exps))

        logger.info(self, 'CPU time: %12.2f', logger.process_clock())
        return self

    def set_common_origin(self, coord):
        '''Update common origin for integrals of dipole, rxp etc.
        **Note** the unit of the coordinates needs to be Bohr

        Examples:

        >>> mol.set_common_origin(0)
        >>> mol.set_common_origin((1,0,0))
        '''
        self._env[PTR_COMMON_ORIG:PTR_COMMON_ORIG+3] = coord
        return self
    set_common_orig = set_common_origin
    set_common_orig_ = set_common_orig    # for backward compatibility
    set_common_origin_ = set_common_orig  # for backward compatibility

    def with_common_origin(self, coord):
        '''Return a temporary mol context which has the rquired common origin.
        The required common origin has no effects out of the temporary context.
        See also :func:`mol.set_common_origin`

        Examples:

        >>> with mol.with_common_origin((1,0,0)):
        ...     mol.intor('int1e_r', comp=3)
        '''
        coord0 = self._env[PTR_COMMON_ORIG:PTR_COMMON_ORIG+3].copy()
        return self._TemporaryMoleContext(self.set_common_origin, (coord,), (coord0,))
    with_common_orig = with_common_origin

    def set_rinv_origin(self, coord):
        r'''Update origin for operator :math:`\frac{1}{|r-R_O|}`.
        **Note** the unit is Bohr

        Examples:

        >>> mol.set_rinv_origin(0)
        >>> mol.set_rinv_origin((0,1,0))
        '''
        self._env[PTR_RINV_ORIG:PTR_RINV_ORIG+3] = coord[:3]
        return self
    set_rinv_orig = set_rinv_origin
    set_rinv_orig_ = set_rinv_orig    # for backward compatibility
    set_rinv_origin_ = set_rinv_orig  # for backward compatibility

    def with_rinv_origin(self, coord):
        '''Return a temporary mol context which has the rquired origin of 1/r
        operator.  The required origin has no effects out of the temporary
        context.  See also :func:`mol.set_rinv_origin`

        Examples:

        >>> with mol.with_rinv_origin((1,0,0)):
        ...     mol.intor('int1e_rinv')
        '''
        coord0 = self._env[PTR_RINV_ORIG:PTR_RINV_ORIG+3].copy()
        return self._TemporaryMoleContext(self.set_rinv_origin, (coord,), (coord0,))
    with_rinv_orig = with_rinv_origin

    def set_range_coulomb(self, omega):
        '''Switch on range-separated Coulomb operator for **all** 2e integrals

        Args:
            omega : double

                | = 0 : Regular electron repulsion integral
                | > 0 : Long-range operator  erf(omega r12) / r12
                | < 0 : Short-range operator  erfc(omega r12) /r12
        '''
        if omega is None:
            self._env[PTR_RANGE_OMEGA] = 0
        else:
            self._env[PTR_RANGE_OMEGA] = omega
    set_range_coulomb_ = set_range_coulomb  # for backward compatibility

    @property
    def omega(self):
        return self._env[PTR_RANGE_OMEGA]
    omega = omega.setter(set_range_coulomb)

    def with_range_coulomb(self, omega):
        '''Return a temporary mol context which sets the required parameter
        omega for range-separated Coulomb operator.
        If omega = None, return the context for regular Coulomb integrals.
        See also :func:`mol.set_range_coulomb`

        Examples:

        >>> with mol.with_range_coulomb(omega=1.5):
        ...     mol.intor('int2e')
        '''
        omega0 = self._env[PTR_RANGE_OMEGA].copy()
        return self._TemporaryMoleContext(self.set_range_coulomb, (omega,), (omega0,))

    def with_long_range_coulomb(self, omega):
        '''Return a temporary mol context for long-range part of
        range-separated Coulomb operator.
        '''
        return self.with_range_coulomb(abs(omega))

    def with_short_range_coulomb(self, omega):
        '''Return a temporary mol context for short-range part of
        range-separated Coulomb operator.
        '''
        return self.with_range_coulomb(-abs(omega))

    def set_f12_zeta(self, zeta):
        '''Set zeta for YP exp(-zeta r12)/r12 or STG exp(-zeta r12) type integrals
        '''
        self._env[PTR_F12_ZETA] = zeta

    def set_nuc_mod(self, atm_id, zeta):
        '''Change the nuclear charge distribution of the given atom ID.  The charge
        distribution is defined as: rho(r) = nuc_charge * Norm * exp(-zeta * r^2).
        This function can **only** be called after .build() method is executed.

        Examples:

        >>> for ia in range(mol.natm):
        ...     zeta = gto.filatov_nuc_mod(mol.atom_charge(ia))
        ...     mol.set_nuc_mod(ia, zeta)
        '''
        ptr = self._atm[atm_id,PTR_ZETA]
        self._env[ptr] = zeta
        if zeta == 0:
            self._atm[atm_id,NUC_MOD_OF] = NUC_POINT
        else:
            self._atm[atm_id,NUC_MOD_OF] = NUC_GAUSS
        return self
    set_nuc_mod_ = set_nuc_mod  # for backward compatibility

    def set_rinv_zeta(self, zeta):
        '''Assume the charge distribution on the "rinv_origin".  zeta is the parameter
        to control the charge distribution: rho(r) = Norm * exp(-zeta * r^2).
        **Be careful** when call this function. It affects the behavior of
        int1e_rinv_* functions.  Make sure to set it back to 0 after using it!
        '''
        self._env[PTR_RINV_ZETA] = zeta
        return self
    set_rinv_zeta_ = set_rinv_zeta  # for backward compatibility

    def with_rinv_zeta(self, zeta):
        '''Return a temporary mol context which has the rquired Gaussian charge
        distribution placed at "rinv_origin": rho(r) = Norm * exp(-zeta * r^2).
        See also :func:`mol.set_rinv_zeta`

        Examples:

        >>> with mol.with_rinv_zeta(zeta=1.5), mol.with_rinv_origin((1.,0,0)):
        ...     mol.intor('int1e_rinv')
        '''
        zeta0 = self._env[PTR_RINV_ZETA].copy()
        return self._TemporaryMoleContext(self.set_rinv_zeta, (zeta,), (zeta0,))

    def with_rinv_at_nucleus(self, atm_id):
        '''Return a temporary mol context in which the rinv operator (1/r) is
        treated like the Coulomb potential of a Gaussian charge distribution
        rho(r) = Norm * exp(-zeta * r^2) at the place of the input atm_id.

        Examples:

        >>> with mol.with_rinv_at_nucleus(3):
        ...     mol.intor('int1e_rinv')
        '''
        zeta = self._env[self._atm[atm_id,PTR_ZETA]]
        rinv = self.atom_coord(atm_id)
        if zeta == 0:
            self._env[AS_RINV_ORIG_ATOM] = atm_id  # required by ecp gradients
            return self.with_rinv_origin(rinv)
        else:
            self._env[AS_RINV_ORIG_ATOM] = atm_id  # required by ecp gradients
            rinv0 = self._env[PTR_RINV_ORIG:PTR_RINV_ORIG+3].copy()
            zeta0 = self._env[PTR_RINV_ZETA].copy()

            def set_rinv(z, r):
                self._env[PTR_RINV_ZETA] = z
                self._env[PTR_RINV_ORIG:PTR_RINV_ORIG+3] = r
            return self._TemporaryMoleContext(set_rinv, (zeta,rinv), (zeta0,rinv0))
    with_rinv_as_nucleus = with_rinv_at_nucleus  # For backward compatibility

    def with_integral_screen(self, threshold):
        '''Return a temporary mol context which has the rquired integral
        screen threshold
        '''
        if threshold is None:
            # This calls the default cutoff settings in cint library
            expcutoff = 0
        else:
            expcutoff = abs(numpy.log(threshold))
        expcutoff0 = self._env[PTR_EXPCUTOFF]

        def set_cutoff(cut):
            self._env[PTR_EXPCUTOFF] = cut
        return self._TemporaryMoleContext(set_cutoff, (expcutoff,), (expcutoff0,))

    def set_geom_(self, atoms_or_coords, unit=None, symmetry=None,
                  inplace=True):
        '''Update geometry
        '''
        import copy
        if inplace:
            mol = self
        else:
            mol = copy.copy(self)
            mol._env = mol._env.copy()
        if unit is None:
            unit = mol.unit
        else:
            mol.unit = unit
        if symmetry is None:
            symmetry = mol.symmetry

        if isinstance(atoms_or_coords, numpy.ndarray):
            mol.atom = list(zip([x[0] for x in mol._atom],
                                atoms_or_coords.tolist()))
        else:
            mol.atom = atoms_or_coords

        if isinstance(atoms_or_coords, numpy.ndarray) and not symmetry:
            if isinstance(unit, (str, unicode)):
                if unit.upper().startswith(('B', 'AU')):
                    unit = 1.
                else: #unit[:3].upper() == 'ANG':
                    unit = 1./param.BOHR
            else:
                unit = 1./unit

            mol._atom = list(zip([x[0] for x in mol._atom],
                                 (atoms_or_coords * unit).tolist()))
            ptr = mol._atm[:,PTR_COORD]
            mol._env[ptr+0] = unit * atoms_or_coords[:,0]
            mol._env[ptr+1] = unit * atoms_or_coords[:,1]
            mol._env[ptr+2] = unit * atoms_or_coords[:,2]
        else:
            mol.symmetry = symmetry
            mol.build(False, False)

        if mol.verbose >= logger.INFO:
            logger.info(mol, 'New geometry')
            for ia, atom in enumerate(mol._atom):
                coorda = tuple([x * param.BOHR for x in atom[1]])
                coordb = tuple([x for x in atom[1]])
                coords = coorda + coordb
                logger.info(mol, ' %3d %-4s %16.12f %16.12f %16.12f AA  '
                            '%16.12f %16.12f %16.12f Bohr\n',
                            ia+1, mol.atom_symbol(ia), *coords)
        return mol

    def update(self, chkfile):
        return self.update_from_chk(chkfile)
    def update_from_chk(self, chkfile):
        with h5py.File(chkfile, 'r') as fh5:
            mol = loads(fh5['mol'][()])
            self.__dict__.update(mol.__dict__)
        return self

    def has_ecp(self):
        '''Whether pseudo potential is used in the system.'''
        return len(self._ecpbas) > 0 or self._pseudo

    def has_ecp_soc(self):
        '''Whether spin-orbit coupling is enabled in ECP.'''
        return (len(self._ecpbas) > 0 and
                numpy.any(self._ecpbas[:,SO_TYPE_OF] == 1))


#######################################################
#NOTE: atm_id or bas_id start from 0
    def atom_symbol(self, atm_id):
        r'''For the given atom id, return the input symbol (without striping special characters)

        Args:
            atm_id : int
                0-based

        Examples:

        >>> mol.build(atom='H^2 0 0 0; H 0 0 1.1')
        >>> mol.atom_symbol(0)
        H^2
        '''
        return _symbol(self._atom[atm_id][0])

    def atom_pure_symbol(self, atm_id):
        r'''For the given atom id, return the standard symbol (striping special characters)

        Args:
            atm_id : int
                0-based

        Examples:

        >>> mol.build(atom='H^2 0 0 0; H 0 0 1.1')
        >>> mol.atom_symbol(0)
        H
        '''
        return _std_symbol(self._atom[atm_id][0])

    @property
    def elements(self):
        '''A list of elements in the molecule'''
        return [self.atom_pure_symbol(i) for i in range(self.natm)]

    def atom_charge(self, atm_id):
        r'''Nuclear effective charge of the given atom id
        Note "atom_charge /= charge(atom_symbol)" when ECP is enabled.
        Number of electrons screened by ECP can be obtained by charge(atom_symbol)-atom_charge

        Args:
            atm_id : int
                0-based

        Examples:

        >>> mol.build(atom='H 0 0 0; Cl 0 0 1.1')
        >>> mol.atom_charge(1)
        17
        '''
        if self._atm[atm_id,NUC_MOD_OF] != NUC_FRAC_CHARGE:
            # regular QM atoms
            return self._atm[atm_id,CHARGE_OF]
        else:
            # MM atoms with fractional charges
            return self._env[self._atm[atm_id,PTR_FRAC_CHARGE]]

    def atom_charges(self):
        '''np.asarray([mol.atom_charge(i) for i in range(mol.natm)])'''
        z = self._atm[:,CHARGE_OF]
        if numpy.any(self._atm[:,NUC_MOD_OF] == NUC_FRAC_CHARGE):
            # Create the integer nuclear charges first then replace the MM
            # particles with the MM charges that saved in _env[PTR_FRAC_CHARGE]
            z = numpy.array(z, dtype=numpy.double)
            idx = self._atm[:,NUC_MOD_OF] == NUC_FRAC_CHARGE
            # MM fractional charges can be positive or negative
            z[idx] = self._env[self._atm[idx,PTR_FRAC_CHARGE]]
        return z

    def atom_nelec_core(self, atm_id):
        '''Number of core electrons for pseudo potential.
        '''
        return charge(self.atom_symbol(atm_id)) - self.atom_charge(atm_id)

    def atom_coord(self, atm_id, unit='Bohr'):
        r'''Coordinates (ndarray) of the given atom id

        Args:
            atm_id : int
                0-based

        Examples:

        >>> mol.build(atom='H 0 0 0; Cl 0 0 1.1')
        >>> mol.atom_coord(1)
        [ 0.          0.          2.07869874]
        '''
        ptr = self._atm[atm_id,PTR_COORD]
        if unit[:3].upper() == 'ANG':
            return self._env[ptr:ptr+3] * param.BOHR
        else:
            return self._env[ptr:ptr+3].copy()

    def atom_coords(self, unit='Bohr'):
        '''np.asarray([mol.atom_coords(i) for i in range(mol.natm)])'''
        ptr = self._atm[:,PTR_COORD]
        c = self._env[numpy.vstack((ptr,ptr+1,ptr+2)).T].copy()
        if unit[:3].upper() == 'ANG':
            c *= param.BOHR
        return c

    atom_mass_list = atom_mass_list

    def atom_nshells(self, atm_id):
        r'''Number of basis/shells of the given atom

        Args:
            atm_id : int
                0-based

        Examples:

        >>> mol.build(atom='H 0 0 0; Cl 0 0 1.1')
        >>> mol.atom_nshells(1)
        5
        '''
        return (self._bas[:,ATOM_OF] == atm_id).sum()

    def atom_shell_ids(self, atm_id):
        r'''A list of the shell-ids of the given atom

        Args:
            atm_id : int
                0-based

        Examples:

        >>> mol.build(atom='H 0 0 0; Cl 0 0 1.1', basis='cc-pvdz')
        >>> mol.atom_shell_ids(1)
        [3, 4, 5, 6, 7]
        '''
        return numpy.where(self._bas[:,ATOM_OF] == atm_id)[0]

    def bas_coord(self, bas_id):
        r'''Coordinates (ndarray) associated with the given basis id

        Args:
            bas_id : int
                0-based

        Examples:

        >>> mol.build(atom='H 0 0 0; Cl 0 0 1.1')
        >>> mol.bas_coord(1)
        [ 0.          0.          2.07869874]
        '''
        atm_id = self.bas_atom(bas_id)
        ptr = self._atm[atm_id,PTR_COORD]
        return self._env[ptr:ptr+3].copy()

    def bas_atom(self, bas_id):
        r'''The atom (0-based id) that the given basis sits on

        Args:
            bas_id : int
                0-based

        Examples:

        >>> mol.build(atom='H 0 0 0; Cl 0 0 1.1', basis='cc-pvdz')
        >>> mol.bas_atom(7)
        1
        '''
        return self._bas[bas_id,ATOM_OF].copy()

    def bas_angular(self, bas_id):
        r'''The angular momentum associated with the given basis

        Args:
            bas_id : int
                0-based

        Examples:

        >>> mol.build(atom='H 0 0 0; Cl 0 0 1.1', basis='cc-pvdz')
        >>> mol.bas_atom(7)
        2
        '''
        return self._bas[bas_id,ANG_OF].copy()

    def bas_nctr(self, bas_id):
        r'''The number of contracted GTOs for the given shell

        Args:
            bas_id : int
                0-based

        Examples:

        >>> mol.build(atom='H 0 0 0; Cl 0 0 1.1', basis='cc-pvdz')
        >>> mol.bas_atom(3)
        3
        '''
        return self._bas[bas_id,NCTR_OF].copy()

    def bas_nprim(self, bas_id):
        r'''The number of primitive GTOs for the given shell

        Args:
            bas_id : int
                0-based

        Examples:

        >>> mol.build(atom='H 0 0 0; Cl 0 0 1.1', basis='cc-pvdz')
        >>> mol.bas_atom(3)
        11
        '''
        return self._bas[bas_id,NPRIM_OF].copy()

    def bas_kappa(self, bas_id):
        r'''Kappa (if l < j, -l-1, else l) of the given shell

        Args:
            bas_id : int
                0-based

        Examples:

        >>> mol.build(atom='H 0 0 0; Cl 0 0 1.1', basis='cc-pvdz')
        >>> mol.bas_kappa(3)
        0
        '''
        return self._bas[bas_id,KAPPA_OF].copy()

    def bas_exp(self, bas_id):
        r'''exponents (ndarray) of the given shell

        Args:
            bas_id : int
                0-based

        Examples:

        >>> mol.build(atom='H 0 0 0; Cl 0 0 1.1', basis='cc-pvdz')
        >>> mol.bas_kappa(0)
        [ 13.01     1.962    0.4446]
        '''
        nprim = self.bas_nprim(bas_id)
        ptr = self._bas[bas_id,PTR_EXP]
        return self._env[ptr:ptr+nprim].copy()

    def _libcint_ctr_coeff(self, bas_id):
        nprim = self.bas_nprim(bas_id)
        nctr = self.bas_nctr(bas_id)
        ptr = self._bas[bas_id,PTR_COEFF]
        return self._env[ptr:ptr+nprim*nctr].reshape(nctr,nprim).T

    def bas_ctr_coeff(self, bas_id):
        r'''Contract coefficients (ndarray) of the given shell

        Args:
            bas_id : int
                0-based

        Examples:

        >>> mol.M(atom='H 0 0 0; Cl 0 0 1.1', basis='cc-pvdz')
        >>> mol.bas_ctr_coeff(0)
        [[ 10.03400444]
         [  4.1188704 ]
         [  1.53971186]]
        '''
        l = self.bas_angular(bas_id)
        es = self.bas_exp(bas_id)
        cs = self._libcint_ctr_coeff(bas_id)
        cs = numpy.einsum('pi,p->pi', cs, 1/gto_norm(l, es))
        return cs

    def bas_len_spinor(self, bas_id):
        '''The number of spinor associated with given basis
        If kappa is 0, return 4l+2
        '''
        l = self.bas_angular(bas_id)
        k = self.bas_kappa(bas_id)
        return len_spinor(l, k)

    def bas_len_cart(self, bas_id):
        '''The number of Cartesian function associated with given basis
        '''
        return len_cart(self._bas[bas_id,ANG_OF])


    npgto_nr = npgto_nr

    nao_nr = nao_nr
    nao_2c = nao_2c
    nao_cart = nao_cart

    nao_nr_range = nao_nr_range
    nao_2c_range = nao_2c_range

    ao_loc_nr = ao_loc_nr
    ao_loc_2c = ao_loc_2c

    @property
    def nao(self):
        if self._nao is None:
            return self.nao_nr()
        else:
            return self._nao
    @nao.setter
    def nao(self, x):
        self._nao = x

    ao_loc = property(ao_loc_nr)

    tmap = time_reversal_map = time_reversal_map

    inertia_moment = inertia_moment

    tostring = tostring
    tofile = tofile

    def fromstring(self, string, format='xyz'):
        '''Update the Mole object based on the input geometry string'''
        atom = self.format_atom(fromstring(string, format), unit=1)
        self.set_geom_(atom, unit='Angstrom', inplace=True)
        if format == 'sdf' and 'M  CHG' in string:
            raise NotImplementedError
            #FIXME self.charge = 0
        return self

    def fromfile(self, filename, format=None):
        '''Update the Mole object based on the input geometry file'''
        atom = self.format_atom(fromfile(filename, format), unit=1)
        self.set_geom_(atom, unit='Angstrom', inplace=True)
        if format == 'sdf':
            raise NotImplementedError
        return self

    def intor(self, intor, comp=None, hermi=0, aosym='s1', out=None,
              shls_slice=None, grids=None):
        '''Integral generator.

        Args:
            intor : str
                Name of the 1e or 2e AO integrals.  Ref to :func:`getints` for the
                complete list of available 1-electron integral names

        Kwargs:
            comp : int
                Components of the integrals, e.g. int1e_ipovlp_sph has 3 components.
            hermi : int
                Symmetry of the integrals

                | 0 : no symmetry assumed (default)
                | 1 : hermitian
                | 2 : anti-hermitian

            grids : ndarray
                Coordinates of grids for the int1e_grids integrals

        Returns:
            ndarray of 1-electron integrals, can be either 2-dim or 3-dim, depending on comp

        Examples:

        >>> mol.build(atom='H 0 0 0; H 0 0 1.1', basis='sto-3g')
        >>> mol.intor('int1e_ipnuc_sph', comp=3) # <nabla i | V_nuc | j>
        [[[ 0.          0.        ]
          [ 0.          0.        ]]
         [[ 0.          0.        ]
          [ 0.          0.        ]]
         [[ 0.10289944  0.48176097]
          [-0.48176097 -0.10289944]]]
        >>> mol.intor('int1e_nuc_spinor')
        [[-1.69771092+0.j  0.00000000+0.j -0.67146312+0.j  0.00000000+0.j]
         [ 0.00000000+0.j -1.69771092+0.j  0.00000000+0.j -0.67146312+0.j]
         [-0.67146312+0.j  0.00000000+0.j -1.69771092+0.j  0.00000000+0.j]
         [ 0.00000000+0.j -0.67146312+0.j  0.00000000+0.j -1.69771092+0.j]]
        '''
        if not self._built:
            logger.warn(self, 'Warning: intor envs of %s not initialized.', self)
            # FIXME: Whether to check _built and call build?  ._bas and .basis
            # may not be consistent. calling .build() may leads to wrong intor env.
            #self.build(False, False)
        intor = self._add_suffix(intor)
        bas = self._bas
        env = self._env
        if 'ECP' in intor:
            assert(self._ecp is not None)
            bas = numpy.vstack((self._bas, self._ecpbas))
            env[AS_ECPBAS_OFFSET] = len(self._bas)
            env[AS_NECPBAS] = len(self._ecpbas)
            if shls_slice is None:
                shls_slice = (0, self.nbas, 0, self.nbas)
        elif '_grids' in intor:
            assert grids is not None
            env = numpy.append(env, grids.ravel())
            env[NGRIDS] = grids.shape[0]
            env[PTR_GRIDS] = env.size - grids.size
        return moleintor.getints(intor, self._atm, bas, env,
                                 shls_slice, comp, hermi, aosym, out=out)

    def _add_suffix(self, intor, cart=None):
        if not (intor[:4] == 'cint' or
                intor.endswith(('_sph', '_cart', '_spinor', '_ssc'))):
            if cart is None:
                cart = self.cart
            if cart:
                intor = intor + '_cart'
            else:
                intor = intor + '_sph'
        return intor

    def intor_symmetric(self, intor, comp=None, grids=None):
        '''One-electron integral generator. The integrals are assumed to be hermitian

        Args:
            intor : str
                Name of the 1-electron integral.  Ref to :func:`getints` for the
                complete list of available 1-electron integral names

        Kwargs:
            comp : int
                Components of the integrals, e.g. int1e_ipovlp_sph has 3 components.
            grids : ndarray
                Coordinates of grids for the int1e_grids integrals

        Returns:
            ndarray of 1-electron integrals, can be either 2-dim or 3-dim, depending on comp

        Examples:

        >>> mol.build(atom='H 0 0 0; H 0 0 1.1', basis='sto-3g')
        >>> mol.intor_symmetric('int1e_nuc_spinor')
        [[-1.69771092+0.j  0.00000000+0.j -0.67146312+0.j  0.00000000+0.j]
         [ 0.00000000+0.j -1.69771092+0.j  0.00000000+0.j -0.67146312+0.j]
         [-0.67146312+0.j  0.00000000+0.j -1.69771092+0.j  0.00000000+0.j]
         [ 0.00000000+0.j -0.67146312+0.j  0.00000000+0.j -1.69771092+0.j]]
        '''
        return self.intor(intor, comp, 1, aosym='s4', grids=grids)

    def intor_asymmetric(self, intor, comp=None, grids=None):
        '''One-electron integral generator. The integrals are assumed to be anti-hermitian

        Args:
            intor : str
                Name of the 1-electron integral.  Ref to :func:`getints` for the
                complete list of available 1-electron integral names

        Kwargs:
            comp : int
                Components of the integrals, e.g. int1e_ipovlp has 3 components.
            grids : ndarray
                Coordinates of grids for the int1e_grids integrals

        Returns:
            ndarray of 1-electron integrals, can be either 2-dim or 3-dim, depending on comp

        Examples:

        >>> mol.build(atom='H 0 0 0; H 0 0 1.1', basis='sto-3g')
        >>> mol.intor_asymmetric('int1e_nuc_spinor')
        [[-1.69771092+0.j  0.00000000+0.j  0.67146312+0.j  0.00000000+0.j]
         [ 0.00000000+0.j -1.69771092+0.j  0.00000000+0.j  0.67146312+0.j]
         [-0.67146312+0.j  0.00000000+0.j -1.69771092+0.j  0.00000000+0.j]
         [ 0.00000000+0.j -0.67146312+0.j  0.00000000+0.j -1.69771092+0.j]]
        '''
        return self.intor(intor, comp, 2, aosym='a4', grids=grids)

    @lib.with_doc(moleintor.getints_by_shell.__doc__)
    def intor_by_shell(self, intor, shells, comp=None, grids=None):
        intor = self._add_suffix(intor)
        if 'ECP' in intor:
            assert(self._ecp is not None)
            bas = numpy.vstack((self._bas, self._ecpbas))
            self._env[AS_ECPBAS_OFFSET] = len(self._bas)
            self._env[AS_NECPBAS] = len(self._ecpbas)
        else:
            bas = self._bas
        return moleintor.getints_by_shell(intor, shells, self._atm, bas,
                                          self._env, comp)

    eval_ao = eval_gto = eval_gto

    energy_nuc = energy_nuc
    def get_enuc(self):
        return self.energy_nuc()

    sph_labels = spheric_labels = sph_labels
    cart_labels = cart_labels
    ao_labels = ao_labels

    spinor_labels = spinor_labels

    search_ao_label = search_ao_label

    def search_shell_id(self, atm_id, l):
        return search_shell_id(self, atm_id, l)

    search_ao_nr = search_ao_nr
    search_ao_r = search_ao_r

    aoslice_by_atom = aoslice_nr_by_atom = offset_ao_by_atom = offset_nr_by_atom = aoslice_by_atom
    aoslice_2c_by_atom = offset_2c_by_atom = offset_2c_by_atom

    condense_to_shell = condense_to_shell

    to_uncontracted_cartesian_basis = to_uncontracted_cartesian_basis

    __add__ = conc_mol

    def cart2sph_coeff(self, normalized='sp'):
        '''Transformation matrix that transforms Cartesian GTOs to spherical
        GTOs for all basis functions

        Kwargs:
            normalized : string or boolean
                How the Cartesian GTOs are normalized.  Except s and p functions,
                Cartesian GTOs do not have the universal normalization coefficients
                for the different components of the same shell.  The value of this
                argument can be one of 'sp', 'all', None.  'sp' means the Cartesian s
                and p basis are normalized.  'all' means all Cartesian functions are
                normalized.  None means none of the Cartesian functions are normalized.
                The default value 'sp' is the convention used by libcint library.

        Examples:

        >>> mol = gto.M(atom='H 0 0 0; F 0 0 1', basis='ccpvtz')
        >>> c = mol.cart2sph_coeff()
        >>> s0 = mol.intor('int1e_ovlp_sph')
        >>> s1 = c.T.dot(mol.intor('int1e_ovlp_cart')).dot(c)
        >>> print(abs(s1-s0).sum())
        >>> 4.58676826646e-15
        '''
        c2s_l = [cart2sph(l, normalized=normalized) for l in range(12)]
        c2s = []
        for ib in range(self.nbas):
            l = self.bas_angular(ib)
            for n in range(self.bas_nctr(ib)):
                c2s.append(c2s_l[l])
        return scipy.linalg.block_diag(*c2s)

    def sph2spinor_coeff(self):
        '''Transformation matrix that transforms real-spherical GTOs to spinor
        GTOs for all basis functions

        Examples:

        >>> mol = gto.M(atom='H 0 0 0; F 0 0 1', basis='ccpvtz')
        >>> ca, cb = mol.sph2spinor_coeff()
        >>> 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
        '''
        from pyscf.symm import sph
        return sph.sph2spinor_coeff(self)


    def apply(self, fn, *args, **kwargs):
        if callable(fn):
            return lib.StreamObject.apply(self, fn, *args, **kwargs)
        elif isinstance(fn, (str, unicode)):
            method = getattr(self, fn.upper())
            return method(*args, **kwargs)
        else:
            raise TypeError('First argument of .apply method must be a '
                            'function/class or a name (string) of a method.')

    def ao2mo(self, mo_coeffs, erifile=None, dataname='eri_mo', intor='int2e',
              **kwargs):
        '''Integral transformation for arbitrary orbitals and arbitrary
        integrals.  See more detalied documentation in func:`ao2mo.kernel`.

        Args:
            mo_coeffs (an np array or a list of arrays) : A matrix of orbital
                coefficients if it is a numpy ndarray, or four sets of orbital
                coefficients, corresponding to the four indices of (ij|kl).

        Kwargs:
            erifile (str or h5py File or h5py Group object) : The file/object
                to store the transformed integrals.  If not given, the return
                value is an array (in memory) of the transformed integrals.
            dataname : str
                *Note* this argument is effective if erifile is given.
                The dataset name in the erifile (ref the hierarchy of HDF5 format
                http://www.hdfgroup.org/HDF5/doc1.6/UG/09_Groups.html).  By assigning
                different dataname, the existed integral file can be reused.  If
                the erifile contains the specified dataname, the old integrals
                will be replaced by the new one under the key dataname.
            intor (str) : integral name Name of the 2-electron integral.  Ref
                to :func:`getints_by_shell`
                for the complete list of available 2-electron integral names

        Returns:
            An array of transformed integrals if erifile is not given.
            Otherwise, return the file/fileobject if erifile is assigned.


        Examples:

        >>> import pyscf
        >>> mol = pyscf.M(atom='O 0 0 0; H 0 1 0; H 0 0 1', basis='sto3g')
        >>> mo1 = numpy.random.random((mol.nao_nr(), 10))
        >>> mo2 = numpy.random.random((mol.nao_nr(), 8))

        >>> eri1 = mol.ao2mo(mo1)
        >>> print(eri1.shape)
        (55, 55)

        >>> eri1 = mol.ao2mo(mo1, compact=False)
        >>> print(eri1.shape)
        (100, 100)

        >>> eri1 = mol.ao2mo(eri, (mo1,mo2,mo2,mo2))
        >>> print(eri1.shape)
        (80, 36)

        >>> eri1 = mol.ao2mo(eri, (mo1,mo2,mo2,mo2), erifile='water.h5')
        '''
        from pyscf import ao2mo
        return ao2mo.kernel(self, mo_coeffs, erifile, dataname, intor, **kwargs)

    @contextlib.contextmanager
    def _TemporaryMoleContext(self, method, args, args_bak):
        '''Almost every method depends on the Mole environment. Ensure the
        modification in temporary environment being thread safe
        '''
        haslock = self._ctx_lock
        if haslock is None:
            self._ctx_lock = threading.RLock()

        with self._ctx_lock:
            method(*args)
            try:
                yield
            finally:
                method(*args_bak)
                if haslock is None:
                    self._ctx_lock = None

def _parse_nuc_mod(str_or_int_or_fn):
    nucmod = NUC_POINT
    if callable(str_or_int_or_fn):
        nucmod = str_or_int_or_fn
    elif (isinstance(str_or_int_or_fn, (str, unicode)) and
          str_or_int_or_fn[0].upper() == 'G'): # 'gauss_nuc'
        nucmod = NUC_GAUSS
    elif str_or_int_or_fn != 0:
        nucmod = NUC_GAUSS
    return nucmod

def _update_from_cmdargs_(mol):
    try:
        # Detect whether in Ipython shell
        __IPYTHON__  # noqa:
        return
    except Exception:
        pass

    if not mol._built: # parse cmdline args only once
        opts = cmd_args.cmd_args()

        if opts.verbose:
            mol.verbose = opts.verbose
        if opts.max_memory:
            mol.max_memory = opts.max_memory

        if opts.output:
            mol.output = opts.output


def from_zmatrix(atomstr):
    '''>>> a = """H
    H 1 2.67247631453057
    H 1 4.22555607338457 2 50.7684795164077
    H 1 2.90305235726773 2 79.3904651036893 3 6.20854462618583"""
    >>> for x in zmat2cart(a): print(x)
    ['H', array([ 0.,  0.,  0.])]
    ['H', array([ 2.67247631,  0.        ,  0.        ])]
    ['H', array([ 2.67247631,  0.        ,  3.27310166])]
    ['H', array([ 0.53449526,  0.30859098,  2.83668811])]
    '''
    from pyscf.symm import rotation_mat
    atomstr = atomstr.replace(';','\n').replace(',',' ')
    symbols = []
    coord = []
    min_items_per_line = 1
    for line_id, line in enumerate(atomstr.splitlines()):
        line = line.strip()
        if line and line[0] != '#':
            rawd = line.split()
            if len(rawd) < min_items_per_line:
                raise ValueError('Zmatrix format error at L%d %s' % (line_id, line))

            symbols.append(rawd[0])
            if len(rawd) < 3:
                coord.append(numpy.zeros(3))
                min_items_per_line = 3
            elif len(rawd) == 3:
                if DISABLE_EVAL:
                    coord.append(numpy.array((float(rawd[2]), 0, 0)))
                else:
                    coord.append(numpy.array((eval(rawd[2]), 0, 0)))
                min_items_per_line = 5
            elif len(rawd) == 5:
                if DISABLE_EVAL:
                    vals = rawd[1:]
                else:
                    vals = eval(','.join(rawd[1:]))
                bonda = int(vals[0]) - 1
                bond  = float(vals[1])
                anga  = int(vals[2]) - 1
                ang   = float(vals[3])/180*numpy.pi
                assert(ang >= 0)
                v1 = coord[anga] - coord[bonda]
                if not numpy.allclose(v1[:2], 0):
                    vecn = numpy.cross(v1, numpy.array((0.,0.,1.)))
                else: # on z
                    vecn = numpy.array((0.,0.,1.))
                rmat = rotation_mat(vecn, ang)
                c = numpy.dot(rmat, v1) * (bond/numpy.linalg.norm(v1))
                coord.append(coord[bonda]+c)
                min_items_per_line = 7
            else:
                if DISABLE_EVAL:
                    vals = rawd[1:]
                else:
                    vals = eval(','.join(rawd[1:]))
                bonda = int(vals[0]) - 1
                bond  = float(vals[1])
                anga  = int(vals[2]) - 1
                ang   = float(vals[3])/180*numpy.pi
                assert(ang >= 0 and ang <= numpy.pi)
                v1 = coord[anga] - coord[bonda]
                v1 /= numpy.linalg.norm(v1)
                if ang < 1e-7:
                    c = v1 * bond
                elif numpy.pi-ang < 1e-7:
                    c = -v1 * bond
                else:
                    diha  = int(vals[4]) - 1
                    dih   = float(vals[5])/180*numpy.pi
                    v2 = coord[diha] - coord[anga]
                    vecn = numpy.cross(v2, -v1)
                    vecn_norm = numpy.linalg.norm(vecn)
                    if vecn_norm < 1e-7:
                        if not numpy.allclose(v1[:2], 0):
                            vecn = numpy.cross(v1, numpy.array((0.,0.,1.)))
                        else: # on z
                            vecn = numpy.array((0.,0.,1.))
                        rmat = rotation_mat(vecn, ang)
                        c = numpy.dot(rmat, v1) * bond
                    else:
                        rmat = rotation_mat(v1, -dih)
                        vecn = numpy.dot(rmat, vecn) / vecn_norm
                        rmat = rotation_mat(vecn, ang)
                        c = numpy.dot(rmat, v1) * bond
                coord.append(coord[bonda]+c)
    atoms = list(zip([_atom_symbol(x) for x in symbols], coord))
    return atoms
zmat2cart = zmat = from_zmatrix

def cart2zmat(coord):
    '''>>> c = numpy.array((
    (0.000000000000,  1.889726124565,  0.000000000000),
    (0.000000000000,  0.000000000000, -1.889726124565),
    (1.889726124565, -1.889726124565,  0.000000000000),
    (1.889726124565,  0.000000000000,  1.133835674739)))
    >>> print(cart2zmat(c))
    1
    1 2.67247631453057
    1 4.22555607338457 2 50.7684795164077
    1 2.90305235726773 2 79.3904651036893 3 6.20854462618583
    '''
    zstr = []
    zstr.append('1')
    if len(coord) > 1:
        r1 = coord[1] - coord[0]
        nr1 = numpy.linalg.norm(r1)
        zstr.append('1 %.15g' % nr1)
    if len(coord) > 2:
        r2 = coord[2] - coord[0]
        nr2 = numpy.linalg.norm(r2)
        a = numpy.arccos(numpy.dot(r1,r2)/(nr1*nr2))
        zstr.append('1 %.15g 2 %.15g' % (nr2, a*180/numpy.pi))
    if len(coord) > 3:
        o0, o1, o2 = coord[:3]
        p0, p1, p2 = 1, 2, 3
        for k, c in enumerate(coord[3:]):
            r0 = c - o0
            nr0 = numpy.linalg.norm(r0)
            r1 = o1 - o0
            nr1 = numpy.linalg.norm(r1)
            a1 = numpy.arccos(numpy.dot(r0,r1)/(nr0*nr1))
            b0 = numpy.cross(r0, r1)
            nb0 = numpy.linalg.norm(b0)

            if abs(nb0) < 1e-7: # o0, o1, c in line
                a2 = 0
                zstr.append('%d %.15g %d %.15g %d %.15g' %
                            (p0, nr0, p1, a1*180/numpy.pi, p2, a2))
            else:
                b1 = numpy.cross(o2-o0, r1)
                nb1 = numpy.linalg.norm(b1)

                if abs(nb1) < 1e-7:  # o0 o1 o2 in line
                    a2 = 0
                    zstr.append('%d %.15g %d %.15g %d %.15g' %
                                (p0, nr0, p1, a1*180/numpy.pi, p2, a2))
                    o2 = c
                    p2 = 4 + k
                else:
                    if numpy.dot(numpy.cross(b1, b0), r1) < 0:
                        a2 = numpy.arccos(numpy.dot(b1, b0) / (nb0*nb1))
                    else:
                        a2 =-numpy.arccos(numpy.dot(b1, b0) / (nb0*nb1))
                    zstr.append('%d %.15g %d %.15g %d %.15g' %
                                (p0, nr0, p1, a1*180/numpy.pi, p2, a2*180/numpy.pi))

    return '\n'.join(zstr)

def dyall_nuc_mod(nuc_charge, nucprop={}):
    ''' Generate the nuclear charge distribution parameter zeta
    rho(r) = nuc_charge * Norm * exp(-zeta * r^2)

    Ref. L. Visscher and K. Dyall, At. Data Nucl. Data Tables, 67, 207 (1997)
    '''
    mass = nucprop.get('mass', elements.ISOTOPE_MAIN[nuc_charge])
    r = (0.836 * mass**(1./3) + 0.570) / 52917.7249
    zeta = 1.5 / (r**2)
    return zeta

def filatov_nuc_mod(nuc_charge, nucprop={}):
    ''' Generate the nuclear charge distribution parameter zeta
    rho(r) = nuc_charge * Norm * exp(-zeta * r^2)

    Ref. M. Filatov and D. Cremer, Theor. Chem. Acc. 108, 168 (2002)
         M. Filatov and D. Cremer, Chem. Phys. Lett. 351, 259 (2002)
    '''
    c = param.LIGHT_SPEED
    nuc_charge = charge(nuc_charge)
    r = (-0.263188*nuc_charge + 106.016974 + 138.985999/nuc_charge) / c**2
    zeta = 1 / (r**2)
    return zeta

def fakemol_for_charges(coords, expnt=1e16):
    '''Construct a fake Mole object that holds the charges on the given
    coordinates (coords).  The shape of the charge can be a normal
    distribution with the Gaussian exponent (expnt).
    '''
    nbas = coords.shape[0]
    fakeatm = numpy.zeros((nbas,ATM_SLOTS), dtype=numpy.int32)
    fakebas = numpy.zeros((nbas,BAS_SLOTS), dtype=numpy.int32)
    fakeenv = [0] * PTR_ENV_START
    ptr = PTR_ENV_START
    fakeatm[:,PTR_COORD] = numpy.arange(ptr, ptr+nbas*3, 3)
    fakeenv.append(coords.ravel())
    ptr += nbas*3
    fakebas[:,ATOM_OF] = numpy.arange(nbas)
    fakebas[:,NPRIM_OF] = 1
    fakebas[:,NCTR_OF] = 1
# approximate point charge with gaussian distribution exp(-1e16*r^2)
    fakebas[:,PTR_EXP] = ptr
    fakebas[:,PTR_COEFF] = ptr+1
    fakeenv.append([expnt, 1/(2*numpy.sqrt(numpy.pi)*gaussian_int(2,expnt))])
    ptr += 2
    fakemol = Mole()
    fakemol._atm = fakeatm
    fakemol._bas = fakebas
    fakemol._env = numpy.hstack(fakeenv)
    fakemol._built = True
    return fakemol

del(BASE)


# pyscf/gto/moleintor.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>
#

'''
A low level interface to libcint library. It's recommended to use the
Mole.intor method to drive the integral evaluation funcitons.
'''

import warnings
import ctypes
import numpy
from pyscf import lib

libcgto = lib.load_library('libcgto')

ANG_OF     = 1
NPRIM_OF   = 2
NCTR_OF    = 3
KAPPA_OF   = 4
PTR_EXP    = 5
PTR_COEFF  = 6
BAS_SLOTS  = 8
NGRIDS     = 11
PTR_GRIDS  = 12

def getints(intor_name, atm, bas, env, shls_slice=None, comp=None, hermi=0,
            aosym='s1', ao_loc=None, cintopt=None, out=None):
    r'''1e and 2e integral generator.

    Args:
        intor_name : str

            ================================  =============
            Function                          Expression
            ================================  =============
            "int1e_ovlp"                      ( \| \)
            "int1e_nuc"                       ( \| nuc \| \)
            "int1e_kin"                       (.5 \| p dot p\)
            "int1e_ia01p"                     (#C(0 1) \| nabla-rinv \| cross p\)
            "int1e_giao_irjxp"                (#C(0 1) \| r cross p\)
            "int1e_cg_irxp"                   (#C(0 1) \| rc cross p\)
            "int1e_giao_a11part"              (-.5 \| nabla-rinv \| r\)
            "int1e_cg_a11part"                (-.5 \| nabla-rinv \| rc\)
            "int1e_a01gp"                     (g \| nabla-rinv cross p \|\)
            "int1e_igkin"                     (#C(0 .5) g \| p dot p\)
            "int1e_igovlp"                    (#C(0 1) g \|\)
            "int1e_ignuc"                     (#C(0 1) g \| nuc \|\)
            "int1e_z"                         ( \| zc \| \)
            "int1e_zz"                        ( \| zc zc \| \)
            "int1e_r"                         ( \| rc \| \)
            "int1e_r2"                        ( \| rc dot rc \| \)
            "int1e_rr"                        ( \| rc rc \| \)
            "int1e_rrr"                       ( \| rc rc rc \| \)
            "int1e_rrrr"                      ( \| rc rc rc rc \| \)
            "int1e_pnucp"                     (p* \| nuc dot p \| \)
            "int1e_prinvxp"                   (p* \| rinv cross p \| \)
            "int1e_ipovlp"                    (nabla \|\)
            "int1e_ipkin"                     (.5 nabla \| p dot p\)
            "int1e_ipnuc"                     (nabla \| nuc \|\)
            "int1e_iprinv"                    (nabla \| rinv \|\)
            "int1e_rinv"                      (\| rinv \|\)
            "int1e_pnucxp"                    (p* \| nuc cross p \| \)
            "int1e_irp"                       ( \| rc nabla \| \)
            "int1e_irrp"                      ( \| rc rc nabla \| \)
            "int1e_irpr"                      ( \| rc nabla rc \| \)
            "int1e_ggovlp"                    ( \| g g \| \)
            "int1e_ggkin"                     (.5 \| g g p dot p \| \)
            "int1e_ggnuc"                     ( \| g g nuc \| \)
            "int1e_grjxp"                     ( \| g r cross p \| \)
            "ECPscalar"                       AREP ECP integrals, similar to int1e_nuc
            "ECPscalar_ipnuc"                 (nabla i | ECP | ), similar to int1e_ipnuc
            "ECPscalar_iprinv"                similar to int1e_iprinv for a specific atom
            "ECPscalar_ignuc"                 similar to int1e_ignuc
            "ECPscalar_iprinvip"              similar to int1e_iprinvip
            "ECPso"                           < | Spin-orbit ECP | >
            "int1e_ovlp_spinor"               ( \| \)
            "int1e_nuc_spinor"                ( \| nuc \|\)
            "int1e_srsr_spinor"               (sigma dot r \| sigma dot r\)
            "int1e_sr_spinor"                 (sigma dot r \|\)
            "int1e_srsp_spinor"               (sigma dot r \| sigma dot p\)
            "int1e_spsp_spinor"               (sigma dot p \| sigma dot p\)
            "int1e_sp_spinor"                 (sigma dot p \|\)
            "int1e_spnucsp_spinor"            (sigma dot p \| nuc \| sigma dot p\)
            "int1e_srnucsr_spinor"            (sigma dot r \| nuc \| sigma dot r\)
            "int1e_govlp_spinor"              (g \|\)
            "int1e_gnuc_spinor"               (g \| nuc \|\)
            "int1e_cg_sa10sa01_spinor"        (.5 sigma cross rc \| sigma cross nabla-rinv \|\)
            "int1e_cg_sa10sp_spinor"          (.5 rc cross sigma \| sigma dot p\)
            "int1e_cg_sa10nucsp_spinor"       (.5 rc cross sigma \| nuc \| sigma dot p\)
            "int1e_giao_sa10sa01_spinor"      (.5 sigma cross r \| sigma cross nabla-rinv \|\)
            "int1e_giao_sa10sp_spinor"        (.5 r cross sigma \| sigma dot p\)
            "int1e_giao_sa10nucsp_spinor"     (.5 r cross sigma \| nuc \| sigma dot p\)
            "int1e_sa01sp_spinor"             (\| nabla-rinv cross sigma \| sigma dot p\)
            "int1e_spgsp_spinor"              (g sigma dot p \| sigma dot p\)
            "int1e_spgnucsp_spinor"           (g sigma dot p \| nuc \| sigma dot p\)
            "int1e_spgsa01_spinor"            (g sigma dot p \| nabla-rinv cross sigma \|\)
            "int1e_spspsp_spinor"             (sigma dot p \| sigma dot p sigma dot p\)
            "int1e_spnuc_spinor"              (sigma dot p \| nuc \|\)
            "int1e_ipovlp_spinor"             (nabla \|\)
            "int1e_ipkin_spinor"              (.5 nabla \| p dot p\)
            "int1e_ipnuc_spinor"              (nabla \| nuc \|\)
            "int1e_iprinv_spinor"             (nabla \| rinv \|\)
            "int1e_ipspnucsp_spinor"          (nabla sigma dot p \| nuc \| sigma dot p\)
            "int1e_ipsprinvsp_spinor"         (nabla sigma dot p \| rinv \| sigma dot p\)
            "int1e_grids"                     ( \| 1/r_grids \| \)
            "int1e_grids_spinor"              ( \| 1/r_grids \| \)
            "int1e_grids_ip"                  (nabla \| 1/r_grids \| \)
            "int1e_grids_ip_spinor"           (nabla \| 1/r_grids \| \)
            "int1e_grids_spvsp_spinor"        (sigma dot p \| 1/r_grids \| sigma dot p\)
            "int2e"                           ( \, \| \, \)
            "int2e_ig1"                       (#C(0 1) g \, \| \, \)
            "int2e_gg1"                       (g g \, \| \, \)
            "int2e_g1g2"                      (g \, \| g \, \)
            "int2e_p1vxp1"                    ( p* \, cross p \| \, \) ; SSO
            "int2e_spinor"                    (, \| \, \)
            "int2e_spsp1_spinor"              (sigma dot p \, sigma dot p \| \, \)
            "int2e_spsp1spsp2_spinor"         (sigma dot p \, sigma dot p \| sigma dot p \, sigma dot p \)
            "int2e_srsr1_spinor"              (sigma dot r \, sigma dot r \| \,\)
            "int2e_srsr1srsr2_spinor"         (sigma dot r \, sigma dot r \| sigma dot r \, sigma dot r\)
            "int2e_cg_sa10sp1_spinor"         (.5 rc cross sigma \, sigma dot p \| \,\)
            "int2e_cg_sa10sp1spsp2_spinor"    (.5 rc cross sigma \, sigma dot p \| sigma dot p \, sigma dot p \)
            "int2e_giao_sa10sp1_spinor"       (.5 r cross sigma \, sigma dot p \| \,\)
            "int2e_giao_sa10sp1spsp2_spinor"  (.5 r cross sigma \, sigma dot p \| sigma dot p \, sigma dot p \)
            "int2e_g1_spinor"                 (g \, \| \,\)
            "int2e_spgsp1_spinor"             (g sigma dot p \, sigma dot p \| \,\)
            "int2e_g1spsp2_spinor"            (g \, \| sigma dot p \, sigma dot p\)
            "int2e_spgsp1spsp2_spinor"        (g sigma dot p \, sigma dot p \| sigma dot p \, sigma dot p\)
            "int2e_spv1_spinor"               (sigma dot p \, \| \,\)
            "int2e_vsp1_spinor"               (\, sigma dot p \| \,\)
            "int2e_spsp2_spinor"              (\, \| sigma dot p \, sigma dot p\)
            "int2e_spv1spv2_spinor"           (sigma dot p \, \| sigma dot p \,\)
            "int2e_vsp1spv2_spinor"           (\, sigma dot p \| sigma dot p \,\)
            "int2e_spv1vsp2_spinor"           (sigma dot p \, \| \, sigma dot p\)
            "int2e_vsp1vsp2_spinor"           (\, sigma dot p \| \, sigma dot p\)
            "int2e_spv1spsp2_spinor"          (sigma dot p \, \| sigma dot p \, sigma dot p\)
            "int2e_vsp1spsp2_spinor"          (\, sigma dot p \| sigma dot p \, sigma dot p\)
            "int2e_ig1"                       (#C(0 1) g \, \| \, \)
            "int2e_ip1"                       (nabla \, \| \,\)
            "int2e_ip1_spinor"                (nabla \, \| \,\)
            "int2e_ipspsp1_spinor"            (nabla sigma dot p \, sigma dot p \| \,\)
            "int2e_ip1spsp2_spinor"           (nabla \, \| sigma dot p \, sigma dot p\)
            "int2e_ipspsp1spsp2_spinor"       (nabla sigma dot p \, sigma dot p \| sigma dot p \, sigma dot p\)
            "int2e_ipsrsr1_spinor"            (nabla sigma dot r \, sigma dot r \| \,\)
            "int2e_ip1srsr2_spinor"           (nabla \, \| sigma dot r \, sigma dot r\)
            "int2e_ipsrsr1srsr2_spinor"       (nabla sigma dot r \, sigma dot r \| sigma dot r \, sigma dot r\)
            "int2e_ip1"                       (nabla \, \| \,\)
            "int2e_ssp1ssp2_spinor"           ( \, sigma dot p \| gaunt \| \, sigma dot p\)
            "int2e_cg_ssa10ssp2_spinor"       (rc cross sigma \, \| gaunt \| \, sigma dot p\)
            "int2e_giao_ssa10ssp2_spinor"     (r cross sigma  \, \| gaunt \| \, sigma dot p\)
            "int2e_gssp1ssp2_spinor"          (g \, sigma dot p  \| gaunt \| \, sigma dot p\)
            "int2e_ipip1"                     ( nabla nabla \, \| \, \)
            "int2e_ipvip1"                    ( nabla \, nabla \| \, \)
            "int2e_ip1ip2"                    ( nabla \, \| nabla \, \)
            "int3c2e_ip1"                     (nabla \, \| \)
            "int3c2e_ip2"                     ( \, \| nabla\)
            "int2c2e_ip1"                     (nabla \| r12 \| \)
            "int3c2e_spinor"                  (nabla \, \| \)
            "int3c2e_spsp1_spinor"            (nabla \, \| \)
            "int3c2e_ip1_spinor"              (nabla \, \| \)
            "int3c2e_ip2_spinor"              ( \, \| nabla\)
            "int3c2e_ipspsp1_spinor"          (nabla sigma dot p \, sigma dot p \| \)
            "int3c2e_spsp1ip2_spinor"         (sigma dot p \, sigma dot p \| nabla \)
            "ECPscalar_spinor"                AREP ECP integrals, similar to int1e_nuc
            "ECPscalar_ipnuc_spinor"          (nabla i | ECP | ), similar to int1e_ipnuc
            "ECPscalar_iprinv_spinor"         similar to int1e_iprinv for a specific atom
            "ECPscalar_ignuc_spinor"          similar to int1e_ignuc
            "ECPscalar_iprinvip_spinor"       similar to int1e_iprinvip
            "ECPso_spinor"                    < | sigam dot Spin-orbit ECP | >
            ================================  =============

        atm : int32 ndarray
            libcint integral function argument
        bas : int32 ndarray
            libcint integral function argument
        env : float64 ndarray
            libcint integral function argument

    Kwargs:
        shls_slice : 8-element list
            (ish_start, ish_end, jsh_start, jsh_end, ksh_start, ksh_end, lsh_start, lsh_end)
        comp : int
            Components of the integrals, e.g. int1e_ipovlp has 3 components.
        hermi : int (1e integral only)
            Symmetry of the 1e integrals

            | 0 : no symmetry assumed (default)
            | 1 : hermitian
            | 2 : anti-hermitian

        aosym : str (2e integral only)
            Symmetry of the 2e integrals

            | 4 or '4' or 's4': 4-fold symmetry (default)
            | '2ij' or 's2ij' : symmetry between i, j in (ij|kl)
            | '2kl' or 's2kl' : symmetry between k, l in (ij|kl)
            | 1 or '1' or 's1': no symmetry

        out : ndarray (2e integral only)
            array to store the 2e AO integrals

    Returns:
        ndarray of 1-electron integrals, can be either 2-dim or 3-dim, depending on comp

    Examples:

    >>> mol.build(atom='H 0 0 0; H 0 0 1.1', basis='sto-3g')
    >>> gto.getints('int1e_ipnuc_sph', mol._atm, mol._bas, mol._env, comp=3) # <nabla i | V_nuc | j>
    [[[ 0.          0.        ]
      [ 0.          0.        ]]
     [[ 0.          0.        ]
      [ 0.          0.        ]]
     [[ 0.10289944  0.48176097]
      [-0.48176097 -0.10289944]]]
    '''
    intor_name, comp = _get_intor_and_comp(intor_name, comp)
    if any(bas[:,ANG_OF] > 12):
        raise NotImplementedError('cint library does not support high angular (l>12) GTOs')

    if (intor_name.startswith('int1e') or
        intor_name.startswith('ECP') or
        intor_name.startswith('int2c2e')):
        return getints2c(intor_name, atm, bas, env, shls_slice, comp,
                         hermi, ao_loc, cintopt, out)
    elif intor_name.startswith('int2e') or intor_name.startswith('int4c1e'):
        return getints4c(intor_name, atm, bas, env, shls_slice, comp,
                         aosym, ao_loc, cintopt, out)
    elif intor_name.startswith('int3c'):
        return getints3c(intor_name, atm, bas, env, shls_slice, comp,
                         aosym, ao_loc, cintopt, out)
    else:
        raise KeyError('Unknown intor %s' % intor_name)

def _get_intor_and_comp(intor_name, comp=None):
    intor_name = ascint3(intor_name)
    if comp is None:
        try:
            if '_spinor' in intor_name:
                fname = intor_name.replace('_spinor', '')
                comp = _INTOR_FUNCTIONS[fname][1]
            else:
                fname = intor_name.replace('_sph', '').replace('_cart', '')
                comp = _INTOR_FUNCTIONS[fname][0]
        except KeyError:
            warnings.warn('Function %s not found.  Set its comp to 1' % intor_name)
            comp = 1
    return intor_name, comp

_INTOR_FUNCTIONS = {
    # Functiona name            : (comp-for-scalar, comp-for-spinor)
    'int1e_ovlp'                : (1, 1),
    'int1e_nuc'                 : (1, 1),
    'int1e_kin'                 : (1, 1),
    'int1e_ia01p'               : (3, 3),
    'int1e_giao_irjxp'          : (3, 3),
    'int1e_cg_irxp'             : (3, 3),
    'int1e_giao_a11part'        : (9, 9),
    'int1e_cg_a11part'          : (9, 9),
    'int1e_a01gp'               : (9, 9),
    'int1e_igkin'               : (3, 3),
    'int1e_igovlp'              : (3, 3),
    'int1e_ignuc'               : (3, 3),
    'int1e_pnucp'               : (1, 1),
    'int1e_z'                   : (1, 1),
    'int1e_zz'                  : (1, 1),
    'int1e_r'                   : (3, 3),
    'int1e_r2'                  : (1, 1),
    'int1e_rr'                  : (9, 9),
    'int1e_rrr'                 : (27, 27),
    'int1e_rrrr'                : (81, 81),
    'int1e_z_origj'             : (1, 1),
    'int1e_zz_origj'            : (1, 1),
    'int1e_r_origj'             : (3, 3),
    'int1e_rr_origj'            : (9, 9),
    'int1e_r2_origj'            : (1, 1),
    'int1e_r4_origj'            : (1, 1),
    'int1e_p4'                  : (1, 1),
    'int1e_prinvp'              : (1, 1),
    'int1e_prinvxp'             : (3, 3),
    'int1e_pnucxp'              : (3, 3),
    'int1e_irp'                 : (9, 9),
    'int1e_irrp'                : (27, 27),
    'int1e_irpr'                : (27, 27),
    'int1e_ggovlp'              : (9, 9),
    'int1e_ggkin'               : (9, 9),
    'int1e_ggnuc'               : (9, 9),
    'int1e_grjxp'               : (9, 9),
    'int2e'                     : (1, 1),
    'int2e_ig1'                 : (3, 3),
    'int2e_gg1'                 : (9, 9),
    'int2e_g1g2'                : (9, 9),
    'int2e_ip1v_rc1'            : (9, 9),
    'int2e_ip1v_r1'             : (9, 9),
    'int2e_ipvg1_xp1'           : (9, 9),
    'int2e_ipvg2_xp1'           : (9, 9),
    'int2e_p1vxp1'              : (3, 3),
    'int1e_inuc_rcxp'           : (3, 3),
    'int1e_inuc_rxp'            : (3, 3),
    'int1e_sigma'               : (12,3),
    'int1e_spsigmasp'           : (12,3),
    'int1e_srsr'                : (4, 1),
    'int1e_sr'                  : (4, 1),
    'int1e_srsp'                : (4, 1),
    'int1e_spsp'                : (4, 1),
    'int1e_sp'                  : (4, 1),
    'int1e_spnucsp'             : (4, 1),
    'int1e_sprinvsp'            : (4, 1),
    'int1e_srnucsr'             : (4, 1),
    'int1e_sprsp'               : (12,3),
    'int1e_govlp'               : (3, 3),
    'int1e_gnuc'                : (3, 3),
    'int1e_cg_sa10sa01'         : (36,9),
    'int1e_cg_sa10sp'           : (12,3),
    'int1e_cg_sa10nucsp'        : (12,3),
    'int1e_giao_sa10sa01'       : (36,9),
    'int1e_giao_sa10sp'         : (12,3),
    'int1e_giao_sa10nucsp'      : (12,3),
    'int1e_sa01sp'              : (12,3),
    'int1e_spgsp'               : (12,3),
    'int1e_spgnucsp'            : (12,3),
    'int1e_spgsa01'             : (36,9),
    'int2e_spsp1'               : (4, 1),
    'int2e_spsp1spsp2'          : (16,1),
    'int2e_srsr1'               : (4, 1),
    'int2e_srsr1srsr2'          : (16,1),
    'int2e_cg_sa10sp1'          : (12,3),
    'int2e_cg_sa10sp1spsp2'     : (48,3),
    'int2e_giao_sa10sp1'        : (12,3),
    'int2e_giao_sa10sp1spsp2'   : (48,3),
    'int2e_g1'                  : (12,3),
    'int2e_spgsp1'              : (12,3),
    'int2e_g1spsp2'             : (12,3),
    'int2e_spgsp1spsp2'         : (48,3),
    'int2e_pp1'                 : (1, 1),
    'int2e_pp2'                 : (1, 1),
    'int2e_pp1pp2'              : (1, 1),
    'int1e_spspsp'              : (4, 1),
    'int1e_spnuc'               : (4, 1),
    'int2e_spv1'                : (4, 1),
    'int2e_vsp1'                : (4, 1),
    'int2e_spsp2'               : (4, 1),
    'int2e_spv1spv2'            : (16,1),
    'int2e_vsp1spv2'            : (16,1),
    'int2e_spv1vsp2'            : (16,1),
    'int2e_vsp1vsp2'            : (16,1),
    'int2e_spv1spsp2'           : (16,1),
    'int2e_vsp1spsp2'           : (16,1),
    'int1e_ipovlp'              : (3, 3),
    'int1e_ipkin'               : (3, 3),
    'int1e_ipnuc'               : (3, 3),
    'int1e_iprinv'              : (3, 3),
    'int1e_rinv'                : (1, 1),
    'int1e_ipspnucsp'           : (12,3),
    'int1e_ipsprinvsp'          : (12,3),
    'int1e_ippnucp'             : (3, 3),
    'int1e_ipprinvp'            : (3, 3),
    'int2e_ip1'                 : (3, 3),
    'int2e_ip2'                 : (3, 3),
    'int2e_ipspsp1'             : (12,3),
    'int2e_ip1spsp2'            : (12,3),
    'int2e_ipspsp1spsp2'        : (48,3),
    'int2e_ipsrsr1'             : (12,3),
    'int2e_ip1srsr2'            : (12,3),
    'int2e_ipsrsr1srsr2'        : (48,3),
    'int2e_ssp1ssp2'            : (16,1),
    'int2e_ssp1sps2'            : (16,1),
    'int2e_sps1ssp2'            : (16,1),
    'int2e_sps1sps2'            : (16,1),
    'int2e_cg_ssa10ssp2'        : (48,3),
    'int2e_giao_ssa10ssp2'      : (18,3),
    'int2e_gssp1ssp2'           : (18,3),
    'int2e_gauge_r1_ssp1ssp2'   : (None, 1),
    'int2e_gauge_r1_ssp1sps2'   : (None, 1),
    'int2e_gauge_r1_sps1ssp2'   : (None, 1),
    'int2e_gauge_r1_sps1sps2'   : (None, 1),
    'int2e_gauge_r2_ssp1ssp2'   : (None, 1),
    'int2e_gauge_r2_ssp1sps2'   : (None, 1),
    'int2e_gauge_r2_sps1ssp2'   : (None, 1),
    'int2e_gauge_r2_sps1sps2'   : (None, 1),
    'int1e_ipipovlp'            : (9, 9),
    'int1e_ipovlpip'            : (9, 9),
    'int1e_ipipkin'             : (9, 9),
    'int1e_ipkinip'             : (9, 9),
    'int1e_ipipnuc'             : (9, 9),
    'int1e_ipnucip'             : (9, 9),
    'int1e_ipiprinv'            : (9, 9),
    'int1e_iprinvip'            : (9, 9),
    'int2e_ipip1'               : (9, 9),
    'int2e_ipvip1'              : (9, 9),
    'int2e_ip1ip2'              : (9, 9),
    'int1e_ipippnucp'           : (9, 9),
    'int1e_ippnucpip'           : (9, 9),
    'int1e_ipipprinvp'          : (9, 9),
    'int1e_ipprinvpip'          : (9, 9),
    'int1e_ipipspnucsp'         : (36,9),
    'int1e_ipspnucspip'         : (36,9),
    'int1e_ipipsprinvsp'        : (36,9),
    'int1e_ipsprinvspip'        : (36,9),
    'int3c2e'                   : (1, 1),
    'int3c2e_ip1'               : (3, 3),
    'int3c2e_ip2'               : (3, 3),
    'int3c2e_pvp1'              : (1, 1),
    'int3c2e_pvxp1'             : (3, 3),
    'int2c2e_ip1'               : (3, 3),
    'int2c2e_ip2'               : (3, 3),
    'int3c2e_ig1'               : (3, 3),
    'int3c2e_spsp1'             : (4, 1),
    'int3c2e_ipspsp1'           : (12,3),
    'int3c2e_spsp1ip2'          : (12,3),
    'int3c2e_ipip1'             : (9, 9),
    'int3c2e_ipip2'             : (9, 9),
    'int3c2e_ipvip1'            : (9, 9),
    'int3c2e_ip1ip2'            : (9, 9),
    'int2c2e_ip1ip2'            : (9, 9),
    'int2c2e_ipip1'             : (9, 9),
    'int3c1e'                   : (1, 1),
    'int3c1e_p2'                : (1, 1),
    'int3c1e_iprinv'            : (3, 3),
    'int2c2e'                   : (1, 1),
    'int2e_yp'                  : (1, 1),
    'int2e_stg'                 : (1, 1),
    'int2e_coulerf'             : (1, 1),
    "int1e_grids"               : (1, 1),
    "int1e_grids_ip"            : (3, 3),
    "int1e_grids_spvsp"         : (4, 1),
    'ECPscalar'                 : (1, None),
    'ECPscalar_ipnuc'           : (3, None),
    'ECPscalar_iprinv'          : (3, None),
    'ECPscalar_ignuc'           : (3, None),
    'ECPscalar_iprinvip'        : (9, None),
    'ECPso'                     : (3, 1),
}

def getints2c(intor_name, atm, bas, env, shls_slice=None, comp=1, hermi=0,
              ao_loc=None, cintopt=None, out=None):
    atm = numpy.asarray(atm, dtype=numpy.int32, order='C')
    bas = numpy.asarray(bas, dtype=numpy.int32, order='C')
    env = numpy.asarray(env, dtype=numpy.double, order='C')
    natm = atm.shape[0]
    nbas = bas.shape[0]
    if shls_slice is None:
        shls_slice = (0, nbas, 0, nbas)
    else:
        assert(shls_slice[1] <= nbas and shls_slice[3] <= nbas)
    if ao_loc is None:
        ao_loc = make_loc(bas, intor_name)

    i0, i1, j0, j1 = shls_slice[:4]
    naoi = ao_loc[i1] - ao_loc[i0]
    naoj = ao_loc[j1] - ao_loc[j0]

    if 'int1e_grids' in intor_name:
        shape = (int(env[NGRIDS]), naoi, naoj, comp)
        prefix = 'GTOgrids_'
    else:
        shape = (naoi, naoj, comp)
        prefix = 'GTO'

    if intor_name.endswith('_cart') or intor_name.endswith('_sph'):
        dtype = numpy.double
        drv_name = prefix + 'int2c'
    else:
        dtype = numpy.complex128
        if '2c2e' in intor_name:
            assert(hermi != lib.HERMITIAN and
                   hermi != lib.ANTIHERMI)
        drv_name = prefix + 'int2c_spinor'

    mat = numpy.ndarray(shape, dtype, out, order='F')

    if mat.size > 0:
        if cintopt is None:
            cintopt = make_cintopt(atm, bas, env, intor_name)

        fn = getattr(libcgto, drv_name)
        fn(getattr(libcgto, intor_name), mat.ctypes.data_as(ctypes.c_void_p),
           ctypes.c_int(comp), ctypes.c_int(hermi),
           (ctypes.c_int*4)(*(shls_slice[:4])),
           ao_loc.ctypes.data_as(ctypes.c_void_p), cintopt,
           atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(natm),
           bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nbas),
           env.ctypes.data_as(ctypes.c_void_p))

    mat = numpy.rollaxis(mat, -1, 0)
    if comp == 1:
        mat = mat[0]
    return mat

def getints3c(intor_name, atm, bas, env, shls_slice=None, comp=1,
              aosym='s1', ao_loc=None, cintopt=None, out=None):
    atm = numpy.asarray(atm, dtype=numpy.int32, order='C')
    bas = numpy.asarray(bas, dtype=numpy.int32, order='C')
    env = numpy.asarray(env, dtype=numpy.double, order='C')
    natm = atm.shape[0]
    nbas = bas.shape[0]
    if shls_slice is None:
        shls_slice = (0, nbas, 0, nbas, 0, nbas)
        if 'ssc' in intor_name or 'spinor' in intor_name:
            bas = numpy.asarray(numpy.vstack((bas,bas)), dtype=numpy.int32)
            shls_slice = (0, nbas, 0, nbas, nbas, nbas*2)
            nbas = bas.shape[0]
    else:
        assert(shls_slice[1] <= nbas and
               shls_slice[3] <= nbas and
               shls_slice[5] <= nbas)

    i0, i1, j0, j1, k0, k1 = shls_slice[:6]
    if ao_loc is None:
        ao_loc = make_loc(bas, intor_name)
        if 'ssc' in intor_name:
            ao_loc[k0:] = ao_loc[k0] + make_loc(bas[k0:], 'cart')
        elif 'spinor' in intor_name:
            # The auxbasis for electron-2 is in real spherical representation
            ao_loc[k0:] = ao_loc[k0] + make_loc(bas[k0:], 'sph')

    naok = ao_loc[k1] - ao_loc[k0]

    if aosym in ('s1',):
        naoi = ao_loc[i1] - ao_loc[i0]
        naoj = ao_loc[j1] - ao_loc[j0]
        shape = (naoi, naoj, naok, comp)
    else:
        aosym = 's2ij'
        nij = ao_loc[i1]*(ao_loc[i1]+1)//2 - ao_loc[i0]*(ao_loc[i0]+1)//2
        shape = (nij, naok, comp)

    if 'spinor' in intor_name:
        mat = numpy.ndarray(shape, numpy.complex128, out, order='F')
        drv = libcgto.GTOr3c_drv
        fill = getattr(libcgto, 'GTOr3c_fill_'+aosym)
    else:
        mat = numpy.ndarray(shape, numpy.double, out, order='F')
        drv = libcgto.GTOnr3c_drv
        fill = getattr(libcgto, 'GTOnr3c_fill_'+aosym)

    if mat.size > 0:
        # Generating opt for all indices leads to large overhead and poor OMP
        # speedup for solvent model and COSX functions. In these methods,
        # the third index of the three center integrals corresponds to a
        # large number of grids. Initializing the opt for the third index is
        # not necessary.
        if cintopt is None:
            if '3c2e' in intor_name:
                # int3c2e opt without the 3rd index.
                cintopt = make_cintopt(atm, bas[:max(i1, j1)], env, intor_name)
            else:
                cintopt = make_cintopt(atm, bas, env, intor_name)

        drv(getattr(libcgto, intor_name), fill,
            mat.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(comp),
            (ctypes.c_int*6)(*(shls_slice[:6])),
            ao_loc.ctypes.data_as(ctypes.c_void_p), cintopt,
            atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(natm),
            bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nbas),
            env.ctypes.data_as(ctypes.c_void_p))

    mat = numpy.rollaxis(mat, -1, 0)
    if comp == 1:
        mat = mat[0]
    return mat

def getints4c(intor_name, atm, bas, env, shls_slice=None, comp=1,
              aosym='s1', ao_loc=None, cintopt=None, out=None):
    aosym = _stand_sym_code(aosym)
    atm = numpy.asarray(atm, dtype=numpy.int32, order='C')
    bas = numpy.asarray(bas, dtype=numpy.int32, order='C')
    env = numpy.asarray(env, dtype=numpy.double, order='C')
    c_atm = atm.ctypes.data_as(ctypes.c_void_p)
    c_bas = bas.ctypes.data_as(ctypes.c_void_p)
    c_env = env.ctypes.data_as(ctypes.c_void_p)
    natm = atm.shape[0]
    nbas = bas.shape[0]
    ao_loc = make_loc(bas, intor_name)

    if '_spinor' in intor_name:
        assert(aosym == 's1')

    if aosym == 's8':
        assert(shls_slice is None)
        from pyscf.scf import _vhf
        nao = ao_loc[-1]
        nao_pair = nao*(nao+1)//2
        out = numpy.ndarray((nao_pair*(nao_pair+1)//2), buffer=out)
        if nao_pair == 0:
            return out

        if cintopt is None:
            cintopt = make_cintopt(atm, bas, env, intor_name)
        drv = _vhf.libcvhf.GTO2e_cart_or_sph
        drv(getattr(libcgto, intor_name), cintopt,
            out.ctypes.data_as(ctypes.c_void_p),
            ao_loc.ctypes.data_as(ctypes.c_void_p),
            c_atm, ctypes.c_int(natm), c_bas, ctypes.c_int(nbas), c_env)
        return out

    else:
        if shls_slice is None:
            shls_slice = (0, nbas, 0, nbas, 0, nbas, 0, nbas)
        elif len(shls_slice) == 4:
            shls_slice = shls_slice + (0, nbas, 0, nbas)
        else:
            assert(shls_slice[1] <= nbas and shls_slice[3] <= nbas and
                   shls_slice[5] <= nbas and shls_slice[7] <= nbas)
        i0, i1, j0, j1, k0, k1, l0, l1 = shls_slice
        naoi = ao_loc[i1] - ao_loc[i0]
        naoj = ao_loc[j1] - ao_loc[j0]
        naok = ao_loc[k1] - ao_loc[k0]
        naol = ao_loc[l1] - ao_loc[l0]
        if aosym in ('s4', 's2ij'):
            nij = [naoi * (naoi + 1) // 2]
            assert(numpy.all(ao_loc[i0:i1]-ao_loc[i0] == ao_loc[j0:j1]-ao_loc[j0]))
        else:
            nij = [naoi, naoj]
        if aosym in ('s4', 's2kl'):
            nkl = [naok * (naok + 1) // 2]
            assert(numpy.all(ao_loc[k0:k1]-ao_loc[k0] == ao_loc[l0:l1]-ao_loc[l0]))
        else:
            nkl = [naok, naol]
        shape = [comp] + nij + nkl

        if '_spinor' in intor_name:
            drv = libcgto.GTOr4c_drv
            fill = libcgto.GTOr4c_fill_s1
            out = numpy.ndarray(shape[::-1], dtype=numpy.complex128, buffer=out, order='F')
            out = numpy.rollaxis(out, -1, 0)
        else:
            drv = libcgto.GTOnr2e_fill_drv
            fill = getattr(libcgto, 'GTOnr2e_fill_'+aosym)
            out = numpy.ndarray(shape, buffer=out)

        if out.size > 0:
            if cintopt is None:
                cintopt = make_cintopt(atm, bas, env, intor_name)
            prescreen = lib.c_null_ptr()
            drv(getattr(libcgto, intor_name), fill, prescreen,
                out.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(comp),
                (ctypes.c_int*8)(*shls_slice),
                ao_loc.ctypes.data_as(ctypes.c_void_p), cintopt,
                c_atm, ctypes.c_int(natm), c_bas, ctypes.c_int(nbas), c_env)

        if comp == 1:
            out = out[0]
        return out

def getints_by_shell(intor_name, shls, atm, bas, env, comp=1):
    r'''For given 2, 3 or 4 shells, interface for libcint to get 1e, 2e,
    2-center-2e or 3-center-2e integrals

    Args:
        intor_name : str
            See also :func:`getints` for the supported intor_name
        shls : list of int
            The AO shell-ids of the integrals
        atm : int32 ndarray
            libcint integral function argument
        bas : int32 ndarray
            libcint integral function argument
        env : float64 ndarray
            libcint integral function argument

    Kwargs:
        comp : int
            Components of the integrals, e.g. int1e_ipovlp has 3 components.

    Returns:
        ndarray of 2-dim to 5-dim, depending on the integral type (1e,
        2e, 3c-2e, 2c2e) and the value of comp

    Examples:
        The gradients of the spherical 2e integrals

    >>> mol.build(atom='H 0 0 0; H 0 0 1.1', basis='sto-3g')
    >>> gto.getints_by_shell('int2e_ip1_sph', (0,1,0,1), mol._atm, mol._bas, mol._env, comp=3)
    [[[[[-0.        ]]]]
      [[[[-0.        ]]]]
      [[[[-0.08760462]]]]]
    '''
    intor_name, comp = _get_intor_and_comp(intor_name, comp)

    atm = numpy.asarray(atm, dtype=numpy.int32, order='C')
    bas = numpy.asarray(bas, dtype=numpy.int32, order='C')
    env = numpy.asarray(env, dtype=numpy.double, order='C')
    natm = ctypes.c_int(atm.shape[0])
    nbas = ctypes.c_int(bas.shape[0])
    if intor_name.endswith('_cart'):
        dtype = numpy.double
        def num_cgto_of(basid):
            l = bas[basid,ANG_OF]
            return (l+1)*(l+2)//2 * bas[basid,NCTR_OF]
    elif intor_name.endswith('_sph'):
        dtype = numpy.double
        def num_cgto_of(basid):
            l = bas[basid,ANG_OF]
            return (l*2+1) * bas[basid,NCTR_OF]
    else:
        from pyscf.gto.mole import len_spinor
        dtype = numpy.complex128
        def num_cgto_of(basid):
            l = bas[basid,ANG_OF]
            k = bas[basid,KAPPA_OF]
            return len_spinor(l,k) * bas[basid,NCTR_OF]

    null = lib.c_null_ptr()
    if intor_name.startswith('int3c'):
        assert(len(shls) == 3)
        di = num_cgto_of(shls[0])
        dj = num_cgto_of(shls[1])
        l = bas[shls[2],ANG_OF]
        if intor_name.endswith('_ssc'): # mixed spherical-cartesian
            dk = (l+1)*(l+2)//2 * bas[shls[2],NCTR_OF]
        else:
            dk = (l*2+1) * bas[shls[2],NCTR_OF]
        buf = numpy.empty((di,dj,dk,comp), dtype, order='F')
        fintor = getattr(libcgto, intor_name)
        fintor(buf.ctypes.data_as(ctypes.c_void_p),
               null, (ctypes.c_int*3)(*shls),
               atm.ctypes.data_as(ctypes.c_void_p), natm,
               bas.ctypes.data_as(ctypes.c_void_p), nbas,
               env.ctypes.data_as(ctypes.c_void_p), null, null)
        if comp == 1:
            return buf.reshape(di,dj,dk)
        else:
            return buf.transpose(3,0,1,2)

    elif intor_name.startswith('int2e') or intor_name.startswith('int4c'):
        assert(len(shls) == 4)
        di, dj, dk, dl = [num_cgto_of(x) for x in shls]
        buf = numpy.empty((di,dj,dk,dl,comp), dtype, order='F')
        fintor = getattr(libcgto, intor_name)
        fintor(buf.ctypes.data_as(ctypes.c_void_p),
               null, (ctypes.c_int*4)(*shls),
               atm.ctypes.data_as(ctypes.c_void_p), natm,
               bas.ctypes.data_as(ctypes.c_void_p), nbas,
               env.ctypes.data_as(ctypes.c_void_p), null, null)
        if comp == 1:
            return buf.reshape(di,dj,dk,dl)
        else:
            return buf.transpose(4,0,1,2,3)

    elif (intor_name.startswith('int2c') or '1e' in intor_name or
          'ECP' in intor_name):
        assert(len(shls) == 2)
        di = num_cgto_of(shls[0])
        dj = num_cgto_of(shls[1])
        buf = numpy.empty((di,dj,comp), dtype, order='F')
        fintor = getattr(libcgto, intor_name)
        fintor(buf.ctypes.data_as(ctypes.c_void_p),
               null, (ctypes.c_int*2)(*shls),
               atm.ctypes.data_as(ctypes.c_void_p), natm,
               bas.ctypes.data_as(ctypes.c_void_p), nbas,
               env.ctypes.data_as(ctypes.c_void_p), null, null)
        if comp == 1:
            return buf.reshape(di,dj)
        else:
            return buf.transpose(2,0,1)

    else:
        raise RuntimeError('Unknown intor %s' % intor_name)


def make_loc(bas, key):
    if 'cart' in key:
        l = bas[:,ANG_OF]
        dims = (l+1)*(l+2)//2 * bas[:,NCTR_OF]
    elif 'sph' in key:
        dims = (bas[:,ANG_OF]*2+1) * bas[:,NCTR_OF]
    else:  # spinor
        l = bas[:,ANG_OF]
        k = bas[:,KAPPA_OF]
        dims = (l*4+2) * bas[:,NCTR_OF]
        dims[k<0] = (l[k<0] * 2 + 2) * bas[k<0,NCTR_OF]
        dims[k>0] = (l[k>0] * 2    ) * bas[k>0,NCTR_OF]

    ao_loc = numpy.empty(len(dims)+1, dtype=numpy.int32)
    ao_loc[0] = 0
    dims.cumsum(dtype=numpy.int32, out=ao_loc[1:])
    return ao_loc

def make_cintopt(atm, bas, env, intor):
    intor = intor.replace('_sph','').replace('_cart','').replace('_spinor','')
    c_atm = numpy.asarray(atm, dtype=numpy.int32, order='C')
    c_bas = numpy.asarray(bas, dtype=numpy.int32, order='C')
    c_env = numpy.asarray(env, dtype=numpy.double, order='C')
    natm = c_atm.shape[0]
    nbas = c_bas.shape[0]
    cintopt = lib.c_null_ptr()
    # TODO: call specific ECP optimizers for each intor.
    if intor[:3] == 'ECP':
        foptinit = libcgto.ECPscalar_optimizer
        foptinit(ctypes.byref(cintopt),
                 c_atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(natm),
                 c_bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nbas),
                 c_env.ctypes.data_as(ctypes.c_void_p))
        return ctypes.cast(cintopt, _ecpoptHandler)
    else:
        foptinit = getattr(libcgto, intor+'_optimizer')
        foptinit(ctypes.byref(cintopt),
                 c_atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(natm),
                 c_bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nbas),
                 c_env.ctypes.data_as(ctypes.c_void_p))
        return ctypes.cast(cintopt, _cintoptHandler)
class _cintoptHandler(ctypes.c_void_p):
    def __del__(self):
        try:
            libcgto.CINTdel_optimizer(ctypes.byref(self))
        except AttributeError:
            pass
class _ecpoptHandler(ctypes.c_void_p):
    def __del__(self):
        try:
            libcgto.ECPdel_optimizer(ctypes.byref(self))
        except AttributeError:
            pass

def _stand_sym_code(sym):
    if isinstance(sym, int):
        return 's%d' % sym
    elif sym[0] in 'sS':
        return sym.lower()
    else:
        return 's' + sym.lower()

def ascint3(intor_name):
    '''convert cint2 function name to cint3 function name'''
    if intor_name.startswith('cint'):
        intor_name = intor_name[1:]
    if not intor_name.endswith(('_sph', '_cart', '_spinor', '_ssc')):
        intor_name = intor_name + '_spinor'
    return intor_name


if __name__ == '__main__':
    from pyscf import gto
    mol = gto.Mole()
    mol.verbose = 0
    mol.output = None

    mol.atom.extend([
        ["H", (0,  0, 0  )],
        ["H", (0,  0, 1  )],
    ])
    mol.basis = {"H": 'cc-pvdz'}
    mol.build()
    mol.set_rinv_origin(mol.atom_coord(0))
    for i in range(mol.nbas):
        for j in range(mol.nbas):
            print(i, j, getints_by_shell('int1e_prinvxp_sph', (i,j),
                                         mol._atm, mol._bas, mol._env, 3))


# pyscf/gto/pp_int.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>
#

'''Analytic GTH-PP integrals for open boundary conditions.

See also pyscf/pbc/gto/pseudo/pp_int.py
'''

import numpy
from pyscf import lib

def get_gth_pp(mol):
    from pyscf.gto import ATOM_OF
    from pyscf.pbc.gto import Cell
    from pyscf.pbc.gto.pseudo import pp_int
    from pyscf.df import incore

    # Analytical integration for get_pp_loc_part1(cell).
    fakemol = pp_int.fake_cell_vloc(mol, 0)
    vpploc = 0
    if fakemol.nbas > 0:
        charges = fakemol.atom_charges()
        atmlst = fakemol._bas[:,ATOM_OF]
        v = incore.aux_e2(mol, fakemol, 'int3c2e', aosym='s2', comp=1)
        v = numpy.einsum('...i,i->...', v, -charges[atmlst])
        vpploc += lib.unpack_tril(v)

    # To compute the rest part of GTH PP, mimic the mol with a 0D cell.
    cell_0D = mol.view(Cell)
    cell_0D.dimension = 0
    cell_0D.a = numpy.eye(3)
    vpploc += pp_int.get_pp_loc_part2(cell_0D).real
    vpploc += pp_int.get_pp_nl(cell_0D).real
    return vpploc


if __name__ == '__main__':
    from pyscf import scf
    from pyscf.pbc import gto as pbcgto
    from pyscf.pbc import scf as pbcscf
    from pyscf.pbc import df
    cell = pbcgto.Cell()
    cell.atom = 'He 1. .5 .5; C .1 1.3 2.1'
    cell.basis = {'He': [(0, (2.5, 1)), (0, (1., 1))],
                  'C' :'gth-szv',}
    cell.pseudo = {'C':'gth-pade',
                   'He': pbcgto.pseudo.parse('''He
    2
     0.40000000    3    -1.98934751    -0.75604821    0.95604821
    2
     0.29482550    3     1.23870466    .855         .3
                                       .71         -1.1
                                                    .9
     0.32235865    2     2.25670239    -0.39677748
                                        0.93894690
                                                 ''')}
    cell.a = numpy.eye(3)
    cell.dimension = 0
    cell.build()
    mf = pbcscf.RHF(cell)
    mf.with_df = df.AFTDF(cell)
    mf.run()

    mol = cell.to_mol()
    scf.RHF(mol).run()


# pyscf/gto/basis/__init__.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 os
import sys
from os.path import join
if sys.version_info < (2,7):
    import imp
else:
    import importlib
from pyscf.gto.basis import parse_nwchem
from pyscf.lib.exceptions import BasisNotFoundError
from pyscf import __config__

ALIAS = {
    'ano'        : 'ano.dat'        ,
    'anorcc'     : 'ano.dat'        ,
    'anoroosdz'  : 'roos-dz.dat'    ,
    'anoroostz'  : 'roos-tz.dat'    ,
    'roosdz'     : 'roos-dz.dat'    ,
    'roostz'     : 'roos-tz.dat'    ,
    'ccpvdz'     : 'cc-pvdz.dat'    ,
    'ccpvtz'     : 'cc-pvtz.dat'    ,
    'ccpvqz'     : 'cc-pvqz.dat'    ,
    'ccpv5z'     : 'cc-pv5z.dat'    ,
    'ccpvdpdz'   : 'cc-pvdpdz.dat'  ,
    'augccpvdz'  : 'aug-cc-pvdz.dat',
    'augccpvtz'  : 'aug-cc-pvtz.dat',
    'augccpvqz'  : 'aug-cc-pvqz.dat',
    'augccpv5z'  : 'aug-cc-pv5z.dat',
    'augccpvdpdz': 'aug-cc-pvdpdz.dat',
    'ccpvdzdk'   : 'cc-pvdz-dk.dat' ,
    'ccpvtzdk'   : 'cc-pvtz-dk.dat' ,
    'ccpvqzdk'   : 'cc-pvqz-dk.dat' ,
    'ccpv5zdk'   : 'cc-pv5z-dk.dat' ,
    'ccpvdzdkh'  : 'cc-pvdz-dk.dat' ,
    'ccpvtzdkh'  : 'cc-pvtz-dk.dat' ,
    'ccpvqzdkh'  : 'cc-pvqz-dk.dat' ,
    'ccpv5zdkh'  : 'cc-pv5z-dk.dat' ,
    'augccpvdzdk' : 'aug-cc-pvdz-dk.dat',
    'augccpvtzdk' : 'aug-cc-pvtz-dk.dat',
    'augccpvqzdk' : 'aug-cc-pvqz-dk.dat',
    'augccpv5zdk' : 'aug-cc-pv5z-dk.dat',
    'augccpvdzdkh': 'aug-cc-pvdz-dk.dat',
    'augccpvtzdkh': 'aug-cc-pvtz-dk.dat',
    'augccpvqzdkh': 'aug-cc-pvqz-dk.dat',
    'augccpv5zdkh': 'aug-cc-pv5z-dk.dat',
    'ccpvdzjkfit' : 'cc-pvdz-jkfit.dat' ,
    'ccpvtzjkfit' : 'cc-pvtz-jkfit.dat' ,
    'ccpvqzjkfit' : 'cc-pvqz-jkfit.dat' ,
    'ccpv5zjkfit' : 'cc-pv5z-jkfit.dat' ,
    'ccpvdzri'    : 'cc-pvdz-ri.dat'    ,
    'ccpvtzri'    : 'cc-pvtz-ri.dat'    ,
    'ccpvqzri'    : 'cc-pvqz-ri.dat'    ,
    'ccpv5zri'    : 'cc-pv5z-ri.dat'    ,
    'augccpvdzjkfit' : 'aug-cc-pvdz-jkfit.dat' ,
    'augccpvdzpjkfit': 'aug-cc-pvdzp-jkfit.dat',
    'augccpvtzjkfit' : 'aug-cc-pvtz-jkfit.dat' ,
    'augccpvqzjkfit' : 'aug-cc-pvqz-jkfit.dat' ,
    'augccpv5zjkfit' : 'aug-cc-pv5z-jkfit.dat' ,
    'heavyaugccpvdzjkfit' : 'heavy-aug-cc-pvdz-jkfit.dat',
    'heavyaugccpvtzjkfit' : 'heavy-aug-cc-pvtz-jkfit.dat',
    'heavyaugccpvdzri' : 'heavy-aug-cc-pvdz-ri.dat',
    'heavyaugccpvtzri' : 'heavy-aug-cc-pvtz-ri.dat',
    'augccpvdzri'    : 'aug-cc-pvdz-ri.dat'    ,
    'augccpvdzpri'   : 'aug-cc-pvdzp-ri.dat'   ,
    'augccpvqzri'    : 'aug-cc-pvqz-ri.dat'    ,
    'augccpvtzri'    : 'aug-cc-pvtz-ri.dat'    ,
    'augccpv5zri'    : 'aug-cc-pv5z-ri.dat'    ,
    'ccpvtzdk3'   : 'cc-pVTZ-DK3.dat'   ,
    'ccpvqzdk3'   : 'cc-pVQZ-DK3.dat'   ,
    'augccpvtzdk3': 'aug-cc-pVTZ-DK3.dat',
    'augccpvqzdk3': 'aug-cc-pVQZ-DK3.dat',
    'dyalldz'    : 'dyall_dz'       ,
    'dyallqz'    : 'dyall_qz'       ,
    'dyalltz'    : 'dyall_tz'       ,
    'faegredz'   : 'faegre_dz'      ,
    'iglo'       : 'iglo3'          ,
    'iglo3'      : 'iglo3'          ,
    '321++g'     : join('pople-basis', '3-21++G.dat'   ),
    '321++g*'    : join('pople-basis', '3-21++Gs.dat'  ),
    '321++gs'    : join('pople-basis', '3-21++Gs.dat'  ),
    '321g'       : join('pople-basis', '3-21G.dat'     ),
    '321g*'      : join('pople-basis', '3-21Gs.dat'    ),
    '321gs'      : join('pople-basis', '3-21Gs.dat'    ),
    '431g'       : join('pople-basis', '4-31G.dat'     ),
    '631++g'     : join('pople-basis', '6-31++G.dat'   ),
    '631++g*'    : join('pople-basis', '6-31++Gs.dat'  ),
    '631++gs'    : join('pople-basis', '6-31++Gs.dat'  ),
    '631++g**'   : join('pople-basis', '6-31++Gss.dat' ),
    '631++gss'   : join('pople-basis', '6-31++Gss.dat' ),
    '631+g'      : join('pople-basis', '6-31+G.dat'    ),
    '631+g*'     : join('pople-basis', '6-31+Gs.dat'   ),
    '631+gs'     : join('pople-basis', '6-31+Gs.dat'   ),
    '631+g**'    : join('pople-basis', '6-31+Gss.dat'  ),
    '631+gss'    : join('pople-basis', '6-31+Gss.dat'  ),
    '6311++g'    : join('pople-basis', '6-311++G.dat'  ),
    '6311++g*'   : join('pople-basis', '6-311++Gs.dat' ),
    '6311++gs'   : join('pople-basis', '6-311++Gs.dat' ),
    '6311++g**'  : join('pople-basis', '6-311++Gss.dat'),
    '6311++gss'  : join('pople-basis', '6-311++Gss.dat'),
    '6311+g'     : join('pople-basis', '6-311+G.dat'   ),
    '6311+g*'    : join('pople-basis', '6-311+Gs.dat'  ),
    '6311+gs'    : join('pople-basis', '6-311+Gs.dat'  ),
    '6311+g**'   : join('pople-basis', '6-311+Gss.dat' ),
    '6311+gss'   : join('pople-basis', '6-311+Gss.dat' ),
    '6311g'      : join('pople-basis', '6-311G.dat'    ),
    '6311g*'     : join('pople-basis', '6-311Gs.dat'   ),
    '6311gs'     : join('pople-basis', '6-311Gs.dat'   ),
    '6311g**'    : join('pople-basis', '6-311Gss.dat'  ),
    '6311gss'    : join('pople-basis', '6-311Gss.dat'  ),
    '631g'       : join('pople-basis', '6-31G.dat'     ),
    '631g*'      : join('pople-basis', '6-31Gs.dat'    ),
    '631gs'      : join('pople-basis', '6-31Gs.dat'    ),
    '631g**'     : join('pople-basis', '6-31Gss.dat'   ),
    '631gss'     : join('pople-basis', '6-31Gss.dat'   ),
    'sto3g'      : 'sto-3g.dat'     ,
    'sto6g'      : 'sto-6g.dat'     ,
    'minao'      : 'minao'          ,
    'dz'         : 'dz.dat'         ,
    'dzpdunning' : 'dzp_dunning'    ,
    'dzvp'       : 'dzvp.dat'       ,
    'dzvp2'      : 'dzvp2.dat'      ,
    'dzp'        : 'dzp.dat'        ,
    'tzp'        : 'tzp.dat'        ,
    'qzp'        : 'qzp.dat'        ,
    'adzp'       : 'adzp.dat'       ,
    'atzp'       : 'atzp.dat'       ,
    'aqzp'       : 'aqzp.dat'       ,
    'dzpdk'      : 'dzp-dkh.dat'    ,
    'tzpdk'      : 'tzp-dkh.dat'    ,
    'qzpdk'      : 'qzp-dkh.dat'    ,
    'dzpdkh'     : 'dzp-dkh.dat'    ,
    'tzpdkh'     : 'tzp-dkh.dat'    ,
    'qzpdkh'     : 'qzp-dkh.dat'    ,
    'def2svp'    : 'def2-svp.dat'   ,
    'def2svpd'   : 'def2-svpd.dat'  ,
    'def2tzvpd'  : 'def2-tzvpd.dat' ,
    'def2tzvppd' : 'def2-tzvppd.dat',
    'def2tzvpp'  : 'def2-tzvpp.dat' ,
    'def2tzvp'   : 'def2-tzvp.dat'  ,
    'def2qzvpd'  : 'def2-qzvpd.dat' ,
    'def2qzvppd' : 'def2-qzvppd.dat',
    'def2qzvpp'  : 'def2-qzvpp.dat' ,
    'def2qzvp'   : 'def2-qzvp.dat'  ,
    'def2svpjfit'    : 'def2-universal-jfit.dat',
    'def2svpjkfit'   : 'def2-universal-jkfit.dat',
    'def2tzvpjfit'   : 'def2-universal-jfit.dat',
    'def2tzvpjkfit'  : 'def2-universal-jkfit.dat',
    'def2tzvppjfit'  : 'def2-universal-jfit.dat',
    'def2tzvppjkfit' : 'def2-universal-jkfit.dat',
    'def2qzvpjfit'   : 'def2-universal-jfit.dat',
    'def2qzvpjkfit'  : 'def2-universal-jkfit.dat',
    'def2qzvppjfit'  : 'def2-universal-jfit.dat',
    'def2qzvppjkfit' : 'def2-universal-jkfit.dat',
    'def2universaljfit'  : 'def2-universal-jfit.dat',
    'def2universaljkfit' : 'def2-universal-jkfit.dat',
    'def2svpri'      : 'def2-svp-ri.dat'     ,
    'def2svpdri'     : 'def2-svpd-ri.dat'    ,
    'def2tzvpri'     : 'def2-tzvp-ri.dat'    ,
    'def2tzvpdri'    : 'def2-tzvpd-ri.dat'   ,
    'def2tzvppri'    : 'def2-tzvpp-ri.dat'   ,
    'def2tzvppdri'   : 'def2-tzvppd-ri.dat'  ,
    'def2qzvpri'     : 'def2-qzvp-ri.dat'    ,
    'def2qzvppri'    : 'def2-qzvpp-ri.dat'   ,
    'def2qzvppdri'   : 'def2-qzvppd-ri.dat'  ,
    'tzv'        : 'tzv.dat'        ,
    'weigend'     : 'def2-universal-jfit.dat',
    'weigend+etb' : 'def2-universal-jfit.dat',
    'weigendcfit' : 'def2-universal-jfit.dat',
    'weigendjfit' : 'def2-universal-jfit.dat',
    'weigendjkfit': 'def2-universal-jkfit.dat',
    'demon'      : 'demon_cfit.dat' ,
    'demoncfit'  : 'demon_cfit.dat' ,
    'ahlrichs'   : 'ahlrichs_cfit.dat',
    'ahlrichscfit': 'ahlrichs_cfit.dat',
    'ccpvtzfit'  : 'cc-pvtz_fit.dat',
    'ccpvdzfit'  : 'cc-pvdz_fit.dat',
    'ccpwcvtzmp2fit': 'cc-pwCVTZ_MP2FIT.dat',
    'ccpvqzmp2fit': 'cc-pVQZ_MP2FIT.dat',
    'ccpv5zmp2fit': 'cc-pV5Z_MP2FIT.dat',
    'augccpwcvtzmp2fit': 'aug-cc-pwCVTZ_MP2FIT.dat',
    'augccpvqzmp2fit': 'aug-cc-pVQZ_MP2FIT.dat',
    'augccpv5zmp2fit': 'aug-cc-pV5Z_MP2FIT.dat',
    'ccpcvdz'    : ('cc-pvdz.dat', 'cc-pCVDZ.dat'),
    'ccpcvtz'    : ('cc-pvtz.dat', 'cc-pCVTZ.dat'),
    'ccpcvqz'    : ('cc-pvqz.dat', 'cc-pCVQZ.dat'),
    'ccpcv5z'    : 'cc-pCV5Z.dat',
    'ccpcv6z'    : 'cc-pCV6Z.dat',
    'ccpwcvdz'   : 'cc-pwCVDZ.dat',
    'ccpwcvtz'   : 'cc-pwCVTZ.dat',
    'ccpwcvqz'   : 'cc-pwCVQZ.dat',
    'ccpwcv5z'   : 'cc-pwCV5Z.dat',
    'ccpwcvdzdk' : 'cc-pwCVDZ-DK.dat',
    'ccpwcvtzdk' : 'cc-pwCVTZ-DK.dat',
    'ccpwcvqzdk' : 'cc-pwCVQZ-DK.dat',
    'ccpwcv5zdk' : 'cc-pwCV5Z-DK.dat',
    'ccpwcvtzdk3': 'cc-pwCVTZ-DK3.dat',
    'ccpwcvqzdk3': 'cc-pwCVQZ-DK3.dat',
    'augccpwcvdz': 'aug-cc-pwcvtz.dat',
    'augccpwcvtz': 'aug-cc-pwcvtz.dat',
    'augccpwcvqz': 'aug-cc-pwcvqz.dat',
    'augccpwcv5z': 'aug-cc-pwcv5z.dat',
    'augccpwcvtzdk' : 'aug-cc-pwCVTZ-DK.dat',
    'augccpwcvqzdk' : 'aug-cc-pwCVQZ-DK.dat',
    'augccpwcv5zdk' : 'aug-cc-pwcv5z-dk.dat',
    'augccpwcvtzdk3': 'aug-cc-pwCVTZ-DK3.dat',
    'augccpwcvqzdk3': 'aug-cc-pwCVQZ-DK3.dat',
    'dgaussa1cfit': 'DgaussA1_dft_cfit.dat',
    'dgaussa1xfit': 'DgaussA1_dft_xfit.dat',
    'dgaussa2cfit': 'DgaussA2_dft_cfit.dat',
    'dgaussa2xfit': 'DgaussA2_dft_xfit.dat',
    'ccpvdzpp'   : 'cc-pvdz-pp.dat' ,
    'ccpvtzpp'   : 'cc-pvtz-pp.dat' ,
    'ccpvqzpp'   : 'cc-pvqz-pp.dat' ,
    'ccpv5zpp'   : 'cc-pv5z-pp.dat' ,
    'crenbl'     : 'crenbl.dat'     ,
    'crenbs'     : 'crenbs.dat'     ,
    'lanl2dz'    : 'lanl2dz.dat'    ,
    'lanl2tz'    : 'lanl2tz.dat'    ,
    'lanl08'     : 'lanl08.dat'     ,
    'sbkjc'      : 'sbkjc.dat'      ,
    # Stuttgart ECP http://www.tc.uni-koeln.de/PP/clickpse.en.html
    'stuttgart'  : 'stuttgart_dz.dat',
    'stuttgartdz': 'stuttgart_dz.dat',
    'stuttgartrlc': 'stuttgart_dz.dat',
    'stuttgartrsc': 'stuttgart_rsc.dat',
    'stuttgartrsc_mdf': 'cc-pvdz-pp.dat',
    #'stuttgartrsc_mwb': 'stuttgart_rsc.dat',
    'ccpwcvdzpp' : 'cc-pwCVDZ-PP.dat',
    'ccpwcvtzpp' : 'cc-pwCVTZ-PP.dat',
    'ccpwcvqzpp' : 'cc-pwCVQZ-PP.dat',
    'ccpwcv5zpp' : 'cc-pwCV5Z-PP.dat',
    'ccpvdzppnr' : 'cc-pVDZ-PP-NR.dat',
    'ccpvtzppnr' : 'cc-pVTZ-PP-NR.dat',
    'augccpvdzpp': ('cc-pvdz-pp.dat', 'aug-cc-pVDZ-PP.dat'),
    'augccpvtzpp': ('cc-pvtz-pp.dat', 'aug-cc-pVTZ-PP.dat'),
    'augccpvqzpp': ('cc-pvqz-pp.dat', 'aug-cc-pVQZ-PP.dat'),
    'augccpv5zpp': ('cc-pv5z-pp.dat', 'aug-cc-pV5Z-PP.dat'),
    'pc0' : 'pc-0.dat',
    'pc1' : 'pc-1.dat',
    'pc2' : 'pc-2.dat',
    'pc3' : 'pc-3.dat',
    'pc4' : 'pc-4.dat',
    'augpc0' : 'aug-pc-0.dat',
    'augpc1' : 'aug-pc-1.dat',
    'augpc2' : 'aug-pc-2.dat',
    'augpc3' : 'aug-pc-3.dat',
    'augpc4' : 'aug-pc-4.dat',
    'pcseg0' : 'pcseg-0.dat',
    'pcseg1' : 'pcseg-1.dat',
    'pcseg2' : 'pcseg-2.dat',
    'pcseg3' : 'pcseg-3.dat',
    'pcseg4' : 'pcseg-4.dat',
    'augpcseg0' : 'aug-pcseg-0.dat',
    'augpcseg1' : 'aug-pcseg-1.dat',
    'augpcseg2' : 'aug-pcseg-2.dat',
    'augpcseg3' : 'aug-pcseg-3.dat',
    'augpcseg4' : 'aug-pcseg-4.dat',
    'sarcdkh'   : 'sarc-dkh2.dat',
# Burkatzki-Filippi-Dolg pseudo potential
    'bfdvdz'     : 'bfd_vdz.dat',
    'bfdvtz'     : 'bfd_vtz.dat',
    'bfdvqz'     : 'bfd_vqz.dat',
    'bfdv5z'     : 'bfd_v5z.dat',
    'bfd'        : 'bfd_pp.dat',
    'bfdpp'      : 'bfd_pp.dat',
#
    'ccpcvdzf12optri': os.path.join('f12-basis', 'cc-pCVDZ-F12-OptRI.dat'),
    'ccpcvtzf12optri': os.path.join('f12-basis', 'cc-pCVTZ-F12-OptRI.dat'),
    'ccpcvqzf12optri': os.path.join('f12-basis', 'cc-pCVQZ-F12-OptRI.dat'),
    'ccpvdzf12optri' : os.path.join('f12-basis', 'cc-pVDZ-F12-OptRI.dat' ),
    'ccpvtzf12optri' : os.path.join('f12-basis', 'cc-pVTZ-F12-OptRI.dat' ),
    'ccpvqzf12optri' : os.path.join('f12-basis', 'cc-pVQZ-F12-OptRI.dat' ),
    'ccpv5zf12'      : os.path.join('f12-basis', 'cc-pV5Z-F12.dat'       ),
    'ccpvdzf12rev2'  : os.path.join('f12-basis', 'cc-pVDZ-F12rev2.dat'   ),
    'ccpvtzf12rev2'  : os.path.join('f12-basis', 'cc-pVTZ-F12rev2.dat'   ),
    'ccpvqzf12rev2'  : os.path.join('f12-basis', 'cc-pVQZ-F12rev2.dat'   ),
    'ccpv5zf12rev2'  : os.path.join('f12-basis', 'cc-pV5Z-F12rev2.dat'   ),
    'ccpvdzf12nz'    : os.path.join('f12-basis', 'cc-pVDZ-F12-nZ.dat'    ),
    'ccpvtzf12nz'    : os.path.join('f12-basis', 'cc-pVTZ-F12-nZ.dat'    ),
    'ccpvqzf12nz'    : os.path.join('f12-basis', 'cc-pVQZ-F12-nZ.dat'    ),
    'augccpvdzoptri' : os.path.join('f12-basis', 'aug-cc-pVDZ-OptRI.dat' ),
    'augccpvtzoptri' : os.path.join('f12-basis', 'aug-cc-pVTZ-OptRI.dat' ),
    'augccpvqzoptri' : os.path.join('f12-basis', 'aug-cc-pVQZ-OptRI.dat' ),
    'augccpv5zoptri' : os.path.join('f12-basis', 'aug-cc-pV5Z-OptRI.dat' ),
# All-electron basis designed for periodic calculations, available in Crystal
    'pobtzvp'       :  'pob-tzvp.dat',
    'pobtzvpp'      :  'pob-tzvpp.dat',
    'crystalccpvdz' :  'crystal-cc-pvdz.dat',
# ccECP 
    'ccecp'         : join('ccecp-basis', 'ccECP', 'ccECP.dat'   ),
    'ccecpccpvdz'   : join('ccecp-basis', 'ccECP', 'ccECP_cc-pVDZ.dat'),
    'ccecpccpvtz'   : join('ccecp-basis', 'ccECP', 'ccECP_cc-pVTZ.dat'),
    'ccecpccpvqz'   : join('ccecp-basis', 'ccECP', 'ccECP_cc-pVQZ.dat'),
    'ccecpccpv5z'   : join('ccecp-basis', 'ccECP', 'ccECP_cc-pV5Z.dat'),
    'ccecpccpv6z'   : join('ccecp-basis', 'ccECP', 'ccECP_cc-pV6Z.dat'),
    'ccecpaugccpvdz': join('ccecp-basis', 'ccECP', 'ccECP_aug-cc-pVDZ.dat'),
    'ccecpaugccpvtz': join('ccecp-basis', 'ccECP', 'ccECP_aug-cc-pVTZ.dat'),
    'ccecpaugccpvqz': join('ccecp-basis', 'ccECP', 'ccECP_aug-cc-pVQZ.dat'),
    'ccecpaugccpv5z': join('ccecp-basis', 'ccECP', 'ccECP_aug-cc-pV5Z.dat'),
    'ccecpaugccpv6z': join('ccecp-basis', 'ccECP', 'ccECP_aug-cc-pV6Z.dat'),
# ccECP_He_core 
    'ccecphe'         : join('ccecp-basis', 'ccECP_He_core', 'ccECP.dat'   ),
    'ccecpheccpvdz'   : join('ccecp-basis', 'ccECP_He_core', 'ccECP_cc-pVDZ.dat'),
    'ccecpheccpvtz'   : join('ccecp-basis', 'ccECP_He_core', 'ccECP_cc-pVTZ.dat'),
    'ccecpheccpvqz'   : join('ccecp-basis', 'ccECP_He_core', 'ccECP_cc-pVQZ.dat'),
    'ccecpheccpv5z'   : join('ccecp-basis', 'ccECP_He_core', 'ccECP_cc-pV5Z.dat'),
    'ccecpheccpv6z'   : join('ccecp-basis', 'ccECP_He_core', 'ccECP_cc-pV6Z.dat'),
    'ccecpheaugccpvdz': join('ccecp-basis', 'ccECP_He_core', 'ccECP_aug-cc-pVDZ.dat'),
    'ccecpheaugccpvtz': join('ccecp-basis', 'ccECP_He_core', 'ccECP_aug-cc-pVTZ.dat'),
    'ccecpheaugccpvqz': join('ccecp-basis', 'ccECP_He_core', 'ccECP_aug-cc-pVQZ.dat'),
    'ccecpheaugccpv5z': join('ccecp-basis', 'ccECP_He_core', 'ccECP_aug-cc-pV5Z.dat'),
    'ccecpheaugccpv6z': join('ccecp-basis', 'ccECP_He_core', 'ccECP_aug-cc-pV6Z.dat'),
# ccECP_reg
    'ccecpreg'         : join('ccecp-basis', 'ccECP_reg', 'ccECP.dat'   ),
    'ccecpregccpvdz'   : join('ccecp-basis', 'ccECP_reg', 'ccECP_cc-pVDZ.dat'),
    'ccecpregccpvtz'   : join('ccecp-basis', 'ccECP_reg', 'ccECP_cc-pVTZ.dat'),
    'ccecpregccpvqz'   : join('ccecp-basis', 'ccECP_reg', 'ccECP_cc-pVQZ.dat'),
    'ccecpregccpv5z'   : join('ccecp-basis', 'ccECP_reg', 'ccECP_cc-pV5Z.dat'),
    'ccecpregaugccpvdz': join('ccecp-basis', 'ccECP_reg', 'ccECP_aug-cc-pVDZ.dat'),
    'ccecpregaugccpvtz': join('ccecp-basis', 'ccECP_reg', 'ccECP_aug-cc-pVTZ.dat'),
    'ccecpregaugccpvqz': join('ccecp-basis', 'ccECP_reg', 'ccECP_aug-cc-pVQZ.dat'),
    'ccecpregaugccpv5z': join('ccecp-basis', 'ccECP_reg', 'ccECP_aug-cc-pV5Z.dat'),
#spin-orbit ECPs
    'ecpds10mdfso' : os.path.join('soecp', 'ECPDS10MDFSO.dat'),
    'ecpds28mdfso' : os.path.join('soecp', 'ECPDS28MDFSO.dat'),
    'ecpds28mwbso' : os.path.join('soecp', 'ECPDS28MWBSO.dat'),
    'ecpds46mdfso' : os.path.join('soecp', 'ECPDS46MDFSO.dat'),
    'ecpds60mdfso' : os.path.join('soecp', 'ECPDS60MDFSO.dat'),
    'ecpds60mwbso' : os.path.join('soecp', 'ECPDS60MWBSO.dat'),
    'ecpds78mdfso' : os.path.join('soecp', 'ECPDS78MDFSO.dat'),
    'ecpds92mdfbso' : os.path.join('soecp', 'ECPDS92MDFBSO.dat'),
    'ecpds92mdfbqso' : os.path.join('soecp', 'ECPDS92MDFBQSO.dat'),
}

def _is_pople_basis(basis):
    return (basis.startswith('631') or
            basis.startswith('321') or
            basis.startswith('431'))

_BASIS_DIR = os.path.dirname(__file__)

def _parse_pople_basis(basis, symb):
    if '(' in basis:
        mbas = basis[:basis.find('(')]
        extension = basis[basis.find('(')+1:basis.find(')')]
    else:
        mbas = basis
        extension = ''

    # polarized functions are defined based on pople basis name prefix like
    # 6-31G, 6-311G etc.
    basename = mbas[0] + '-' + mbas[1:].upper()
    basename = basename.replace('+', '').replace('*', '')
    pathtmp = join('pople-basis', basename + '-polarization-%s.dat')
    def convert(s):
        if len(s) == 0:
            return []
        elif s[0].isalpha():
            return [pathtmp % s[0]] + convert(s[1:])
        else:
            return [pathtmp % s[:2]] + convert(s[2:])

    if symb in ('H', 'He'):
        if ',' in extension:
            return tuple([ALIAS[mbas]] + convert(extension.split(',')[1]))
        else:
            return ALIAS[mbas]
    else:
        return tuple([ALIAS[mbas]] + convert(extension.split(',')[0]))

OPTIMIZE_CONTRACTION = getattr(__config__, 'gto_basis_parse_optimize', False)
def parse(string, symb=None, optimize=OPTIMIZE_CONTRACTION):
    if 'ECP' in string:
        return parse_nwchem.parse_ecp(string, symb)
    else:
        return parse_nwchem.parse(string, symb, optimize)
parse.__doc__ = parse_nwchem.parse.__doc__

def parse_ecp(string, symb=None):
    # TODO: catch KeyError and provide suggestion for the possible keys
    return parse_nwchem.parse_ecp(string, symb)
parse_ecp.__doc__ = parse_nwchem.parse_ecp.__doc__

def _convert_contraction(contr_string):
    '''Parse contraction scheme string into a list

    Args:
        contr_string : str
            Desired contraction scheme in conventional 'XsYpZd...' form,
            where X, Y, Z are the total numbers of corresponding functions.
    '''
    import re
    l_fun={'s': 0, 'p': 1, 'd': 2, 'f': 3, 'g': 4, 'h': 5, 'i': 6}
    l_fun.update({chr(num): num-100 for num in range(ord('k'), ord('p'))})
    num_contr = [int(digit) for digit in re.findall(r'\d+', contr_string)]
    basis_labels = re.findall(r'[d-z]+', contr_string)
    assert len(num_contr)==len(basis_labels)
    basis_l = [l_fun[basis] for basis in basis_labels]
    assert basis_l == sorted(basis_l), 'Contraction scheme ' + contr_string +\
        ' has to be ordered by l'
    assert len(basis_l) == len(set(basis_l)), 'Some of l in ' + contr_string +\
        ' appears more than once'
    # Prepare zero list to ensure the total length is equal to the highest l+1
    contraction_list = [0] * (1+max(basis_l))
    for (l, n_contr) in zip(basis_l, num_contr):
        contraction_list[l] = n_contr
    return contraction_list

def _truncate(basis, contr_scheme, symb, split_name):
    # keep only first n_keep contractions for each l
    contr_b = []
    b_index = 0
    for l, n_keep in enumerate(contr_scheme):
        n_saved = 0
        if n_keep > 0:
            for segm in basis:
                segm_l = segm[0]
                if segm_l == l:
                    segm_len = len(segm[1][1:])
                    n_save = min(segm_len, n_keep - n_saved)
                    if n_save > 0:
                        save_segm = [line[:n_save+1] for line in
                                     segm[:][1:]]
                        contr_b.append([l] + save_segm)
                        n_saved += n_save
            assert n_saved == n_keep, ("@{} implies {} l={} function(s), but" +
                                       "only {} in {}:{}").format(split_name[1],
                                                                  contr_scheme[l],
                                                                  l, n_saved, symb,
                                                                  split_name[0])
    return contr_b

optimize_contraction = parse_nwchem.optimize_contraction
to_general_contraction = parse_nwchem.to_general_contraction


def load(filename_or_basisname, symb, optimize=OPTIMIZE_CONTRACTION):
    '''Convert the basis of the given symbol to internal format

    Args:
        filename_or_basisname : str
            Case insensitive basis set name. Special characters will be removed.
            or a string of "path/to/file" which stores the basis functions
        symb : str
            Atomic symbol, Special characters will be removed.

    Examples:
        Load STO 3G basis of carbon to oxygen atom

    >>> mol = gto.Mole()
    >>> mol.basis = {'O': load('sto-3g', 'C')}
    '''
    symb = ''.join([i for i in symb if i.isalpha()])
    if '@' in filename_or_basisname:
        split_name = filename_or_basisname.split('@')
        assert len(split_name) == 2
        filename_or_basisname = split_name[0]
        contr_scheme = _convert_contraction(split_name[1].lower())
    else:
        contr_scheme = 'Full'
    if os.path.isfile(filename_or_basisname):
        # read basis from given file
        try:
            b = parse_nwchem.load(filename_or_basisname, symb, optimize)
        except BasisNotFoundError:
            with open(filename_or_basisname, 'r') as fin:
                b =  parse_nwchem.parse(fin.read(), symb)
        if contr_scheme != 'Full':
            b = _truncate(b, contr_scheme, symb, split_name)
        return b

    name = _format_basis_name(filename_or_basisname)

    if not (name in ALIAS or _is_pople_basis(name)):
        try:
            return parse_nwchem.parse(filename_or_basisname, symb)
        except BasisNotFoundError:
            try:
                return parse_nwchem.parse(filename_or_basisname)
            except IndexError:
                raise BasisNotFoundError('Invalid basis name %s' % filename_or_basisname)
        except IndexError:
            raise BasisNotFoundError(filename_or_basisname)

    if name in ALIAS:
        basmod = ALIAS[name]
    elif _is_pople_basis(name):
        basmod = _parse_pople_basis(name, symb)
    else:
        raise BasisNotFoundError(filename_or_basisname)

    if 'dat' in basmod:
        b = parse_nwchem.load(join(_BASIS_DIR, basmod), symb, optimize)
    elif isinstance(basmod, (tuple, list)) and isinstance(basmod[0], str):
        b = []
        for f in basmod:
            b += parse_nwchem.load(join(_BASIS_DIR, f), symb, optimize)
    else:
        if sys.version_info < (2,7):
            fp, pathname, description = imp.find_module(basmod, __path__)
            mod = imp.load_module(name, fp, pathname, description)
            b = mod.__getattribute__(symb)
            fp.close()
        else:
            mod = importlib.import_module('.'+basmod, __package__)
            b = mod.__getattribute__(symb)

    if contr_scheme != 'Full':
        b = _truncate(b, contr_scheme, symb, split_name)
    return b

def load_ecp(filename_or_basisname, symb):
    '''Convert the basis of the given symbol to internal format
    '''
    symb = ''.join([i for i in symb if i.isalpha()])
    if os.path.isfile(filename_or_basisname):
        # read basis from given file
        try:
            return parse_nwchem.load_ecp(filename_or_basisname, symb)
        except BasisNotFoundError:
            with open(filename_or_basisname, 'r') as fin:
                return parse_ecp(fin.read(), symb)

    name = _format_basis_name(filename_or_basisname)

    if name in ALIAS:
        basmod = ALIAS[name]
        return parse_nwchem.load_ecp(join(_BASIS_DIR, basmod), symb)
    else:
        return parse_ecp(filename_or_basisname, symb)

def _format_basis_name(basisname):
    return basisname.lower().replace('-', '').replace('_', '').replace(' ', '')

del(OPTIMIZE_CONTRACTION)


# pyscf/gto/basis/parse_bfd_pp.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>
#

'''
Parsers for basis set associated with Burkatzi-Filippi-Dolg  pseudo potential
'''

import os
from pyscf.lib.exceptions import BasisNotFoundError

MAPSPDF = {'S': 0,
           'P': 1,
           'D': 2,
           'F': 3,
           'G': 4,
           'H': 5,
           'I': 6,
           'K': 7}

# basistype is one of 'vdz', 'vtz', 'vqz', 'v5z'
def parse(symb, basistype):
    basisfile = os.path.join(os.path.dirname(__file__), 'Burkatzi-Filippi-Dolg-PP.dat')
    raw = ''.join(search_seg(basisfile, symb))
    bas_dat = raw.split('Fit-data:')[1]
    for bas_seg in bas_dat.split('Basis-name: ')[1:]:
        if bas_seg.startswith(basistype.lower()):
            break

    bas_seg = bas_seg.split('\n')[3:]
    bas_block = []
    for bas in bas_seg:
        bas = bas.strip(' ')
        if bas:
            if bas.isdigit():
                bas_block.append([])
            else:
                bas_block[-1].append(bas)

    basis = []
    for dat in bas_block:
        l = MAPSPDF[dat[0][0].upper()]
        bas_by_l = [l]
        basis.append(bas_by_l)
        for bas in reversed(dat):
            expnt, coeff = [float(x) for x in bas.split()[1:]]
            bas_by_l.append([expnt, coeff])
    return basis

def parse_ecp(symb):
    basisfile = os.path.join(os.path.dirname(__file__), 'Burkatzi-Filippi-Dolg-PP.dat')
    raw = ''.join(search_seg(basisfile, symb))
    pp_dat = raw.split('Fit-data:')[0]
    info, loc_dat, nloc_dat = pp_dat.split('component')

    for dat in info.split('\n'):
        if 'Number of replaced protons' in dat:
            ncore = int(dat.split(' ')[-1])
        elif 'Number of projectors' in dat:
            nloc_max = int(dat.split(' ')[-1])

    loc = [[]]  # r^0 is empty
    loc_dat = loc_dat.replace('\t', ' ').split('\n')[2:5]
    zeff, rorder, expnta = [float(x) for x in loc_dat[0].split(' ')]
    assert(rorder == -1)
    loc.append([[expnta, zeff]])
    gamma, rorder, expnt = [float(x) for x in loc_dat[2].split(' ')]
    assert(rorder == 0)
    loc.append([[expnt, gamma]])
    alpha, rorder, expnt = [float(x) for x in loc_dat[1].split(' ')]
    assert(rorder == 1)
    assert(abs(zeff*expnta-alpha) < 1e-7)
    loc.append([[expnt, alpha]])

    nloc = []
    for dat in nloc_dat.replace('\t', ' ').split('\n')[2:]:
        if dat:  # remove blank lines
            coeff, rorder, expnt = [float(x) for x in dat.split(' ')[:3]]
            assert(rorder == 0)
            nloc.append([[], [],  # r^0, r^1 are empty
                         [[expnt, coeff]]])
    assert(len(nloc) == nloc_max)

    ecp = [[-1, loc]]
    for l, dat in enumerate(nloc):
        ecp.append([l, dat])

    return [ncore, ecp]

# basisfile = Burkatzi-Filippi-Dolg-PP.dat
def search_seg(basisfile, symb):
    with open(basisfile, 'r') as fin:
        # ignore head
        dat = fin.readline().lstrip(' ')
        while dat.startswith('#'):
            dat = fin.readline().lstrip(' ')

        # searching
        while dat:
            if dat[:-1].rstrip(' ') == symb:
                seg = []
                fin.readline()  # pass the "back to top" mark
                while (dat and not dat.startswith('back to top')):
                    seg.append(dat)
                    dat = fin.readline().lstrip(' ')
                return seg[:-1]
            dat = fin.readline().lstrip(' ')
    raise BasisNotFoundError('Basis not found for  %s  in  %s' % (symb, basisfile))


if __name__ == '__main__':
    #print(parse('Li', 'vtz'))
    #print(parse_ecp('Ga'))
    from pyscf.data import elements
    from pyscf.gto.basis import parse_nwchem

#    for bastype in 'vdz', 'vtz', 'vqz', 'v5z':
#        dat = []
#        for atom in elements.ELEMENTS[1:]:
#            try:
#                bas = parse(atom, bastype)
#                dat.append(parse_nwchem.convert_basis_to_nwchem(atom[0], bas))
#            except RuntimeError:
#                pass
#        with open('bfd_%s.dat'%bastype, 'w') as f:
#            f.write('#ftp://ftp.aip.org/epaps/journ_chem_phys/E-JCPSA6-126-315722/epaps_material.html\n')
#            f.write('# M. Burkatzki, C. Filippi, M. Dolg in J. Chem. Phys. 126, 234105 (2007)\n\n')
#            f.write('BASIS "ao basis" PRINT\n')
#            f.write('\n'.join(dat))
#            f.write('END\n')

    dat = []
    for atom in elements.ELEMENTS[1:]:
        try:
            ecp = parse_ecp(atom)
            dat.append(parse_nwchem.convert_ecp_to_nwchem(atom[0], ecp))
        except RuntimeError:
            pass
    with open('bfd_pp.dat', 'w') as f:
        f.write('#ftp://ftp.aip.org/epaps/journ_chem_phys/E-JCPSA6-126-315722/epaps_material.html\n')
        f.write('# M. Burkatzki, C. Filippi, M. Dolg in J. Chem. Phys. 126, 234105 (2007)\n\n')
        f.write('ECP\n')
        f.write('\n'.join(dat))
        f.write('END\n')


# pyscf/gto/basis/parse_gaussian.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>
#

'''
Parsers for basis set in Gaussian program format
'''

__all__ = ['parse', 'load']

try:
    from pyscf.gto.basis.parse_nwchem import optimize_contraction
    from pyscf.gto.basis.parse_nwchem import remove_zero
except ImportError:
    optimize_contraction = lambda basis: basis
    remove_zero = lambda basis: basis

MAXL = 12
SPDF = 'SPDFGHIJKLMN'
MAPSPDF = {key: l for l, key in enumerate(SPDF)}

DELIMETER = '****'

def parse(string, optimize=True):
    '''Parse the basis text which is in NWChem format, return an internal
    basis format which can be assigned to :attr:`Mole.basis`
    Lines started with # are ignored.
    '''
    raw_basis = []
    for dat in string.splitlines():
        x = dat.split('!', 1)[0].strip()
        if x and x != DELIMETER:
            raw_basis.append(x)
    return _parse(raw_basis, optimize)

def load(basisfile, symb, optimize=True):
    raw_basis = search_seg(basisfile, symb)
    #if not raw_basis:
    #    raise BasisNotFoundError('Basis not found for  %s  in  %s' % (symb, basisfile))
    return _parse(raw_basis, optimize)

def search_seg(basisfile, symb):
    with open(basisfile, 'r') as fin:

        def _seek(test_str):
            raw_basis = []
            dat = fin.readline()
            while dat:
                if test_str in dat:
                    return True, raw_basis
                elif dat.strip():  # Skip empty lines
                    raw_basis.append(dat)
                dat = fin.readline()
            return False, raw_basis

        has_delimeter, raw_basis = _seek(DELIMETER)
        if has_delimeter:
            dat = fin.readline()
            while dat:
                if dat.strip().split(' ', 1)[0].upper() == symb.upper():
                    raw_basis = _seek(DELIMETER)[1]
                    break
                else:
                    _seek(DELIMETER)
                dat = fin.readline()
    return raw_basis

def _parse(raw_basis, optimize=True):
    basis_add = []
    for line in raw_basis:
        dat = line.strip()
        if dat.startswith('!'):
            continue
        elif dat[0].isalpha():
            key = dat.split()
            if len(key) == 2:
                # skip the line which has only two items. It's the line for
                # element symbol
                continue
            elif key[0] == 'SP':
                basis_add.append([0])
                basis_add.append([1])
            elif len(key[0])>2 and key[0][:2] in ['l=', 'L=']:
                # Angular momentum defined explicitly
                basis_add.append([int(key[0][2:])])
            else:
                basis_add.append([MAPSPDF[key[0]]])
        else:
            line = [float(x) for x in dat.replace('D','e').split()]
            if key[0] == 'SP':
                basis_add[-2].append([line[0], line[1]])
                basis_add[-1].append([line[0], line[2]])
            else:
                basis_add[-1].append(line)
    basis_sorted = []
    for l in range(MAXL):
        basis_sorted.extend([b for b in basis_add if b[0] == l])

    if optimize:
        basis_sorted = optimize_contraction(basis_sorted)

    basis_sorted = remove_zero(basis_sorted)
    return basis_sorted

if __name__ == '__main__':
    print(load('def2-qzvp-jkfit.gbs', 'C'))


# pyscf/gto/basis/parse_molpro.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>
#

'''
Parses for basis set in the Molpro format
'''

__all__ = ['parse', 'load']

import re
import numpy

try:
    from pyscf.gto.basis.parse_nwchem import optimize_contraction
    from pyscf.gto.basis.parse_nwchem import remove_zero
except ImportError:
    optimize_contraction = lambda basis: basis
    remove_zero = lambda basis: basis

from pyscf import __config__
DISABLE_EVAL = getattr(__config__, 'DISABLE_EVAL', False)

MAXL = 12
SPDF = 'SPDFGHIKLMNO'
MAPSPDF = {key: l for l, key in enumerate(SPDF)}
COMMENT_KEYWORDS = '!*#'

# parse the basis text which is in Molpro format, return an internal basis
# format which can be assigned to gto.mole.basis
def parse(string, optimize=True):
    raw_basis = []
    for x in string.splitlines():
        x = x.strip()
        if x and x[0] not in COMMENT_KEYWORDS:
            raw_basis.append(x)
    return _parse(raw_basis, optimize)

def load(basisfile, symb, optimize=True):
    raw_basis = search_seg(basisfile, symb)
    #if not raw_basis:
    #    raise BasisNotFoundError('Basis not found for  %s  in  %s' % (symb, basisfile))
    return _parse(raw_basis, optimize)

def search_seg(basisfile, symb):
    raw_basis = []
    with open(basisfile, 'r') as fin:
        dat = fin.readline()
        while dat:
            if dat[0] in COMMENT_KEYWORDS:
                dat = fin.readline()
                continue
            elif dat[0].isalpha():
                if dat.startswith(symb+' '):
                    raw_basis.append(dat.splitlines()[0])
                elif raw_basis:
                    return raw_basis
                fin.readline()  # line for references
            elif raw_basis:
                raw_basis.append(dat.splitlines()[0])
            dat = fin.readline()
    return raw_basis


def _parse(raw_basis, optimize=True):
    # pass 1
    basis_add = []
    for dat in raw_basis:
        dat = dat.upper()
        if dat[0].isalpha():
            if ' ' not in dat:
                # Skip the line of comments
                continue
            status = dat
            val = []
            basis_add.append([status, val])
        else:
            val.append(dat)
    raw_basis = [[k, ' '.join(v)] for k,v in basis_add]

    # pass 2
    basis_add = []
    for status, valstring in raw_basis:
        tmp = status.split(':')
        key = tmp[0].split()
        l = MAPSPDF[key[1].upper()]
        #TODO if key[-1] == 'SV'
        val = tmp[1].split()
        np = int(val[0])
        nc = int(val[1])

        rawd = [float(x) for x in valstring.replace('D','e').split()]
        if nc == 0:
            for e in rawd:
                basis_add.append([l, [e, 1.]])
        else:
            exps = numpy.array(rawd[:np])
            coeff = numpy.zeros((np,nc))
            p1 = np
            for i in range(nc):
                start, end = val[2+i].split('.')
                start, end = int(start), int(end)
                nd = end - start + 1
                p0, p1 = p1, p1 + nd
                coeff[start-1:end,i] = rawd[p0:p1]

            bval = numpy.hstack((exps[:,None], coeff))
            basis_add.append([l] + bval.tolist())

    basis_sorted = []
    for l in range(MAXL):
        basis_sorted.extend([b for b in basis_add if b[0] == l])

    if optimize:
        basis_sorted = optimize_contraction(basis_sorted)

    basis_sorted = remove_zero(basis_sorted)
    return basis_sorted

def parse_ecp(string):
    ecptxt = []
    for x in string.splitlines():
        x = x.strip()
        if x and x[0] not in COMMENT_KEYWORDS:
            ecptxt.append(x)
    return _parse_ecp(ecptxt)

def _parse_ecp(raw_ecp):
    symb, nelec, nshell, nso = re.split(',|;', raw_ecp[0])[1:5]
    nelec = int(nelec)
    nshell = int(nshell)
    nso = int(nso)
    assert len(raw_ecp) == (nshell + nso + 2), "ecp info doesn't match with data"

    def parse_terms(terms):
        r_orders = [[] for i in range(7)]  # up to r^6
        for term in terms:
            line = term.split(',')
            order = int(line[0])
            try:
                coef = [float(x) for x in line[1:]]
            except ValueError:
                if DISABLE_EVAL:
                    raise ValueError('Failed to parse ecp %s' % line)
                else:
                    coef = list(eval(','.join(line[1:])))
            r_orders[order].append(coef)
        return r_orders

    ecp_add = {}
    ul = [x.strip() for x in raw_ecp[1].replace('D', 'e').split(';') if x.strip()]
    assert int(ul[0]) + 1 == len(ul), "UL doesn't match data"
    ecp_add[-1] = parse_terms(ul[1:])

    for i, sf_terms in enumerate(raw_ecp[2:2+nshell]):
        terms = [x.strip() for x in sf_terms.replace('D', 'e').split(';') if x.strip()]
        assert int(terms[0]) + 1 == len(terms), \
                "ECP %s Shell doesn't match data" % SPDF[i]
        ecp_add[i] = parse_terms(terms[1:])

    if nso > 0:
        for i, so_terms in enumerate(raw_ecp[2+nshell:]):
            terms = [x.strip() for x in so_terms.replace('D', 'e').split(';') if x.strip()]
            assert int(terms[0]) + 1 == len(terms), \
                    "ECP-SOC Shell %s doesn't match data" % SPDF[i+1]
            soc_data = parse_terms(terms[1:])
            for order, coefs in enumerate(soc_data):
                if not coefs:
                    continue
                sf_coefs = ecp_add[i+1][order]
                assert ([x[0] for x in sf_coefs] == [x[0] for x in coefs]), \
                        "In ECP Shell %s order %d, SF and SOC do not match" % (SPDF[i+1], order)
                for j, c in enumerate(coefs):
                    sf_coefs[j].append(c[1])

    bsort = []
    for l in range(-1, MAXL):
        if l in ecp_add:
            bsort.append([l, ecp_add[l]])
    return [nelec, bsort]

if __name__ == '__main__':
    #print(search_seg('minao.libmol', 'C'))
    print(load('cc_pvdz.libmol', 'C'))


# pyscf/gto/basis/parse_nwchem.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>
#

'''
Parsers for basis set in the NWChem format
'''

__all__ = ['parse', 'load', 'parse_ecp', 'load_ecp',
           'convert_basis_to_nwchem', 'convert_ecp_to_nwchem',
           'optimize_contraction', 'remove_zero', 'to_general_contraction']

import re
import numpy
import numpy as np
import scipy.linalg
from pyscf.data.elements import _std_symbol
from pyscf.lib.exceptions import BasisNotFoundError
from pyscf import __config__

DISABLE_EVAL = getattr(__config__, 'DISABLE_EVAL', False)

MAXL = 15
SPDF = 'SPDFGHIKLMNORTU'
MAPSPDF = {key: l for l, key in enumerate(SPDF)}

BASIS_SET_DELIMITER = re.compile('# *BASIS SET.*\n|END\n')
ECP_DELIMITER = re.compile('\n *ECP *\n')

def parse(string, symb=None, optimize=True):
    '''Parse the basis text which is in NWChem format. Return an internal
    basis format which can be assigned to attribute :attr:`Mole.basis`
    Empty lines, or the lines started with #, or the lines of "BASIS SET" and
    "END" will be ignored are ignored.

    Args:
        string : A string in NWChem basis format. Empty links and the lines of
        "BASIS SET" and "END" will be ignored

    Kwargs:
        optimize : Optimize basis contraction.  Convert the segment contracted
            basis to the general contracted basis.

    Examples:

    >>> mol = gto.Mole()
    >>> mol.basis = {'O': gto.basis.parse("""
    ... #BASIS SET: (6s,3p) -> [2s,1p]
    ... C    S
    ...      71.6168370              0.15432897
    ...      13.0450960              0.53532814
    ...       3.5305122              0.44463454
    ... C    SP
    ...       2.9412494             -0.09996723             0.15591627
    ...       0.6834831              0.39951283             0.60768372
    ...       0.2222899              0.70011547             0.39195739
    ... """)}

    >>> gto.basis.parse("""
    ... He    S
    ...      13.6267000              0.1752300
    ...       1.9993500              0.8934830
    ...       0.3829930              0.0000000
    ... He    S
    ...      13.6267000              0.0000000
    ...       1.9993500              0.0000000
    ...       0.3829930              1.0000000
    ... """, optimize=True)
    [[0, [13.6267, 0.17523, 0.0], [1.99935, 0.893483, 0.0], [0.382993, 0.0, 1.0]]]
    '''
    if symb is not None:
        symb = _std_symbol(symb)
        string = _search_basis_block(re.split(BASIS_SET_DELIMITER, string), symb)
        if not string:
            raise BasisNotFoundError('Basis not found for %s' % symb)

    raw_basis = []
    for dat in string.splitlines():
        dat = dat.split('#')[0].strip()  # Use # to start comments
        dat_upper = dat.upper()
        if (dat and not dat_upper.startswith('END') and not dat_upper.startswith('BASIS')):
            raw_basis.append(dat)
    return _parse(raw_basis, optimize)

def load(basisfile, symb, optimize=True):
    raw_basis = search_seg(basisfile, symb)
    return _parse(raw_basis, optimize)

def _parse(raw_basis, optimize=True):
    basis_parsed = [[] for l in range(MAXL)]
    key = None
    for line in raw_basis:
        dat = line.strip()
        if not dat or dat.startswith('#'):
            continue
        elif dat[0].isalpha():
            key = dat.split()[1].upper()
            if key == 'SP':
                basis_parsed[0].append([0])
                basis_parsed[1].append([1])
            else:
                l = MAPSPDF[key]
                current_basis = [l]
                basis_parsed[l].append(current_basis)
        else:
            dat = dat.replace('D','e').split()
            try:
                dat = [float(x) for x in dat]
            except ValueError:
                if DISABLE_EVAL:
                    raise ValueError('Failed to parse basis %s' % line)
                else:
                    dat = list(eval(','.join(dat)))
            except Exception as e:
                raise BasisNotFoundError('\n' + str(e) +
                                         '\nor the required basis file not existed.')
            if key is None:
                raise RuntimeError('Failed to parse basis')
            elif key == 'SP':
                basis_parsed[0][-1].append([dat[0], dat[1]])
                basis_parsed[1][-1].append([dat[0], dat[2]])
            else:
                current_basis.append(dat)
    basis_sorted = [b for bs in basis_parsed for b in bs]

    if optimize:
        basis_sorted = optimize_contraction(basis_sorted)

    basis_sorted = remove_zero(basis_sorted)
    return basis_sorted

def parse_ecp(string, symb=None):
    if symb is not None:
        symb = _std_symbol(symb)
        raw_data = string.splitlines()
        for i, dat in enumerate(raw_data):
            dat0 = dat.split(None, 1)
            if dat0 and dat0[0] == symb:
                break
        if i+1 == len(raw_data):
            raise BasisNotFoundError('ECP not found for  %s' % symb)
        seg = []
        for dat in raw_data[i:]:
            dat = dat.strip()
            if dat: # remove empty lines
                if ((dat[0].isalpha() and dat.split(None, 1)[0].upper() != symb.upper())):
                    break
                else:
                    seg.append(dat)
    else:
        seg = string.splitlines()

    ecptxt = []
    for dat in seg:
        dat = dat.split('#')[0].strip()
        dat_upper = dat.upper()
        if (dat and not dat_upper.startswith('END') and not dat_upper.startswith('ECP')):
            ecptxt.append(dat)
    return _parse_ecp(ecptxt)

def _parse_ecp(raw_ecp):
    ecp_add = []
    nelec = None
    for line in raw_ecp:
        dat = line.strip()
        if not dat or dat.startswith('#'): # comment line
            continue
        elif dat[0].isalpha():
            key = dat.split()[1].upper()
            if key == 'NELEC':
                nelec = int(dat.split()[2])
                continue
            elif key == 'UL':
                ecp_add.append([-1])
            else:
                ecp_add.append([MAPSPDF[key]])
            # up to r^6
            by_ang = [[] for i in range(7)]
            ecp_add[-1].append(by_ang)
        else:
            line = dat.replace('D','e').split()
            l = int(line[0])
            try:
                coef = [float(x) for x in line[1:]]
            except ValueError:
                if DISABLE_EVAL:
                    raise ValueError('Failed to parse ecp %s' % line)
                else:
                    coef = list(eval(','.join(line[1:])))
            by_ang[l].append(coef)

    if nelec is None:
        return []
    else:
        bsort = []
        for l in range(-1, MAXL):
            bsort.extend([b for b in ecp_add if b[0] == l])
        return [nelec, bsort]

def load_ecp(basisfile, symb):
    return _parse_ecp(search_ecp(basisfile, symb))

def search_seg(basisfile, symb):
    symb = _std_symbol(symb)
    with open(basisfile, 'r') as fin:
        fdata = re.split(BASIS_SET_DELIMITER, fin.read())
    raw_basis = _search_basis_block(fdata, symb)
    return [x for x in raw_basis.splitlines() if x and 'END' not in x]

def _search_basis_block(raw_data, symb):
    raw_basis = ''
    for dat in raw_data:
        dat0 = dat.split(None, 1)
        if dat0 and dat0[0] == symb:
            raw_basis = dat
            break
    return raw_basis

def search_ecp(basisfile, symb):
    symb = _std_symbol(symb)
    with open(basisfile, 'r') as fin:
        fdata = re.split(ECP_DELIMITER, fin.read())
    if len(fdata) <= 1:
        return []

    fdata = fdata[1].splitlines()
    for i, dat in enumerate(fdata):
        dat0 = dat.split(None, 1)
        if dat0 and dat0[0] == symb:
            break
    seg = []
    for dat in fdata[i:]:
        dat = dat.strip()
        if dat:  # remove empty lines
            if ((dat[0].isalpha() and dat.split(None, 1)[0].upper() != symb.upper())):
                return seg
            else:
                seg.append(dat)
    return []


def convert_basis_to_nwchem(symb, basis):
    '''Convert the internal basis format to NWChem format string'''
    res = []
    symb = _std_symbol(symb)

    # pass 1: comment line
    ls = [b[0] for b in basis]
    nprims = [len(b[1:]) for b in basis]
    nctrs = [len(b[1])-1 for b in basis]
    prim_to_ctr = {}
    for i, l in enumerate(ls):
        if l in prim_to_ctr:
            prim_to_ctr[l][0] += nprims[i]
            prim_to_ctr[l][1] += nctrs[i]
        else:
            prim_to_ctr[l] = [nprims[i], nctrs[i]]
    nprims = []
    nctrs = []
    for l in set(ls):
        nprims.append(str(prim_to_ctr[l][0])+SPDF[l].lower())
        nctrs.append(str(prim_to_ctr[l][1])+SPDF[l].lower())
    res.append('#BASIS SET: (%s) -> [%s]' % (','.join(nprims), ','.join(nctrs)))

    # pass 2: basis data
    for bas in basis:
        res.append('%-2s    %s' % (symb, SPDF[bas[0]]))
        for dat in bas[1:]:
            res.append(' '.join('%15.9f'%x for x in dat))
    return '\n'.join(res)

def convert_ecp_to_nwchem(symb, ecp):
    '''Convert the internal ecp format to NWChem format string'''
    symb = _std_symbol(symb)
    res = ['%-2s nelec %d' % (symb, ecp[0])]

    for ecp_block in ecp[1]:
        l = ecp_block[0]
        if l == -1:
            res.append('%-2s ul' % symb)
        else:
            res.append('%-2s %s' % (symb, SPDF[l].lower()))
        for r_order, dat in enumerate(ecp_block[1]):
            for e,c in dat:
                res.append('%d    %15.9f  %15.9f' % (r_order, e, c))
    return '\n'.join(res)

def optimize_contraction(basis):
    '''Search the basis segments which have the same exponents then merge them
    to the general contracted sets.

    Note the difference to the function :func:`to_general_contraction`. The
    return value of this function may still have multiple segments for each
    angular moment section.
    '''
    basdic = {}
    for b in basis:
        if isinstance(b[1], int):  # kappa = b[1]
            key = tuple(b[:2])
            ec = numpy.array(b[2:]).T
        else:
            key = tuple(b[:1])
            ec = numpy.array(b[1:]).T
        es = ec[0]
        cs = [c for c in ec[1:]]

        if key not in basdic:
            basdic[key] = []

        if basdic[key]:
            for e_cs in basdic[key]:
                if numpy.array_equal(e_cs[0], es):
                    e_cs.extend(cs)
                    break
            else:
                basdic[key].append([es] + cs)
        else:
            basdic[key].append([es] + cs)

    basis = []
    for key in sorted(basdic.keys()):
        l_kappa = list(key)
        for e_cs in basdic[key]:
            b = l_kappa + numpy.array(e_cs).T.tolist()
            basis.append(b)
    return basis

def to_general_contraction(basis):
    '''Segmented contracted basis -> general contracted basis.

    Combine multiple basis segments to one segment for each angular moment
    section.

    Examples:

    >>> gto.contract(gto.uncontract(gto.load('sto3g', 'He')))
    [[0, [6.36242139, 1.0, 0.0, 0.0], [1.158923, 0.0, 1.0, 0.0], [0.31364979, 0.0, 0.0, 1.0]]]
    '''
    basdic = {}
    for b in basis:
        if isinstance(b[1], int):  # kappa = b[1]
            key = tuple(b[:2])
            ec = numpy.array(b[2:])
        else:
            key = tuple(b[:1])
            ec = numpy.array(b[1:])
        if key in basdic:
            basdic[key].append(ec)
        else:
            basdic[key] = [ec]

    basis = []
    for key in sorted(basdic.keys()):
        l_kappa = list(key)

        es = numpy.hstack([ec[:,0] for ec in basdic[key]])
        cs = scipy.linalg.block_diag(*[ec[:,1:] for ec in basdic[key]])

        es, e_idx, rev_idx = numpy.unique(es.round(9), True, True)
        es = es[::-1]  # sort the exponents from large to small
        bcoeff = numpy.zeros((e_idx.size, cs.shape[1]))
        for i, j in enumerate(rev_idx):
            bcoeff[j] += cs[i]
        bcoeff = bcoeff[::-1]
        ec = numpy.hstack((es[:,None], bcoeff))

        basis.append(l_kappa + ec.tolist())

    return basis

def remove_zero(basis):
    '''
    Remove exponents if their contraction coefficients are all zeros.
    '''
    new_basis = []
    for b in basis:
        if isinstance(b[1], int):  # kappa = b[1]
            key = list(b[:2])
            ec = b[2:]
        else:
            key = list(b[:1])
            ec = b[1:]

        new_ec = [e_c for e_c in ec if any(c!=0 for c in e_c[1:])]
        if new_ec:
            new_basis.append(key + new_ec)
    return new_basis

if __name__ == '__main__':
    from pyscf import gto
    mol = gto.M(atom='O', basis='6-31g')
    print(load_ecp('lanl2dz.dat', 'Na'))
    b = load('ano.dat', 'Na')
    print(convert_basis_to_nwchem('Na', b))
    b = load_ecp('lanl2dz.dat', 'Na')
    print(convert_ecp_to_nwchem('Na', b))


# pyscf/gto/basis/dyall-basis/dyall_av2z.py
#  
# Elements Supported
# H  He
# B  C  N  O  F  Ne Al Si P  S  Cl Ar
# Ga Ge As Se Br Kr In Sn Sb Te I  Xe Tl Pb Bi Po At Rn 
# Nh Fl Mc Lv Ts Og
#  
#  REFERENCE
#  
# 1s:    K.G. Dyall, unpublished.
# 2p-3p: K.G. Dyall, Theor. Chem. Acc. (2016) 135:128
# 4p-6p: K.G. Dyall, Theor. Chem. Acc. (1998) 99:366;
#        revision K.G. Dyall, Theor. Chem. Acc. (2006) 115:441.
#    7p: K.G. Dyall, Theor. Chem. Acc. (2012) 131:1172.
# Available from the Dirac web site, http://dirac.chem.sdu.dk.
#    
#  
#

H  = [[0, 0, [8.29687389e+01, 1.0]],
      [0, 0, [1.24571508e+01, 1.0]],
      [0, 0, [2.83382422e+00, 1.0]],
      [0, 0, [8.00164139e-01, 1.0]],
      [0, 0, [2.58629441e-01, 1.0]],
      [0, 0, [8.99766770e-02, 1.0]],
      [0, 0, [2.35902587e-02, 1.0]],
      [1, 0, [5.02448897e-01, 1.0]],
      [1, 0, [1.14691054e-01, 1.0]]]

He = [[0, 0, [2.34522333e+02, 1.0]],
      [0, 0, [3.52200252e+01, 1.0]],
      [0, 0, [7.99840299e+00, 1.0]],
      [0, 0, [2.21385929e+00, 1.0]],
      [0, 0, [6.67371603e-01, 1.0]],
      [0, 0, [2.09005872e-01, 1.0]],
      [0, 0, [5.97344735e-02, 1.0]],
      [1, 0, [1.27228449e+00, 1.0]],
      [1, 0, [2.90416896e-01, 1.0]]]

B  = [[0, 0, [9.31890933e+03, 1.0]],
      [0, 0, [1.38381575e+03, 1.0]],
      [0, 0, [3.13231718e+02, 1.0]],
      [0, 0, [8.83973383e+01, 1.0]],
      [0, 0, [2.87378723e+01, 1.0]],
      [0, 0, [1.02925898e+01, 1.0]],
      [0, 0, [3.97389889e+00, 1.0]],
      [0, 0, [1.61011130e+00, 1.0]],
      [0, 0, [3.27059892e-01, 1.0]],
      [0, 0, [1.02704305e-01, 1.0]],
      [0, 0, [2.96312407e-02, 1.0]],
      [1, 0, [2.25301829e+01, 1.0]],
      [1, 0, [5.11445067e+00, 1.0]],
      [1, 0, [1.50095842e+00, 1.0]],
      [1, 0, [5.09574383e-01, 1.0]],
      [1, 0, [1.81570551e-01, 1.0]],
      [1, 0, [6.48912221e-02, 1.0]],
      [1, 0, [1.72291541e-02, 1.0]],
      [2, 0, [3.46511210e-01, 1.0]],
      [2, 0, [8.82389623e-02, 1.0]]]

C  = [[0, 0, [8.42995851e+03, 1.0]],
      [0, 0, [1.25053402e+03, 1.0]],
      [0, 0, [2.82743110e+02, 1.0]],
      [0, 0, [7.96080882e+01, 1.0]],
      [0, 0, [2.56825450e+01, 1.0]],
      [0, 0, [9.02650296e+00, 1.0]],
      [0, 0, [3.32774491e+00, 1.0]],
      [0, 0, [9.06582082e-01, 1.0]],
      [0, 0, [3.65254808e-01, 1.0]],
      [0, 0, [1.28879080e-01, 1.0]],
      [0, 0, [4.20847835e-02, 1.0]],
      [1, 0, [3.40303695e+01, 1.0]],
      [1, 0, [7.77836858e+00, 1.0]],
      [1, 0, [2.32153370e+00, 1.0]],
      [1, 0, [7.93710219e-01, 1.0]],
      [1, 0, [2.79198301e-01, 1.0]],
      [1, 0, [9.60203767e-02, 1.0]],
      [1, 0, [2.84450359e-02, 1.0]],
      [2, 0, [5.52361971e-01, 1.0]],
      [2, 0, [1.45443262e-01, 1.0]]]

N  = [[0, 0, [1.18811750e+04, 1.0]],
      [0, 0, [1.75538845e+03, 1.0]],
      [0, 0, [3.95938764e+02, 1.0]],
      [0, 0, [1.11307086e+02, 1.0]],
      [0, 0, [3.58959691e+01, 1.0]],
      [0, 0, [1.26367612e+01, 1.0]],
      [0, 0, [4.66951503e+00, 1.0]],
      [0, 0, [1.29793688e+00, 1.0]],
      [0, 0, [5.15279750e-01, 1.0]],
      [0, 0, [1.79915698e-01, 1.0]],
      [0, 0, [5.72423049e-02, 1.0]],
      [1, 0, [4.76179541e+01, 1.0]],
      [1, 0, [1.09319427e+01, 1.0]],
      [1, 0, [3.29673806e+00, 1.0]],
      [1, 0, [1.12932532e+00, 1.0]],
      [1, 0, [3.93173563e-01, 1.0]],
      [1, 0, [1.31750516e-01, 1.0]],
      [1, 0, [4.00347301e-02, 1.0]],
      [2, 0, [8.26922465e-01, 1.0]],
      [2, 0, [2.25482421e-01, 1.0]]]

O  = [[0, 0, [1.59925081e+04, 1.0]],
      [0, 0, [2.35199842e+03, 1.0]],
      [0, 0, [5.29060079e+02, 1.0]],
      [0, 0, [1.48462548e+02, 1.0]],
      [0, 0, [4.78417313e+01, 1.0]],
      [0, 0, [1.68559662e+01, 1.0]],
      [0, 0, [6.23711537e+00, 1.0]],
      [0, 0, [1.75808802e+00, 1.0]],
      [0, 0, [6.90688830e-01, 1.0]],
      [0, 0, [2.39078229e-01, 1.0]],
      [0, 0, [7.37104770e-02, 1.0]],
      [1, 0, [6.33409745e+01, 1.0]],
      [1, 0, [1.45830479e+01, 1.0]],
      [1, 0, [4.42797374e+00, 1.0]],
      [1, 0, [1.51782600e+00, 1.0]],
      [1, 0, [5.24175730e-01, 1.0]],
      [1, 0, [1.72133105e-01, 1.0]],
      [1, 0, [5.24505855e-02, 1.0]],
      [2, 0, [1.17766132e+00, 1.0]],
      [2, 0, [3.30124445e-01, 1.0]]]

F  = [[0, 0, [2.08122193e+04, 1.0]],
      [0, 0, [3.04524252e+03, 1.0]],
      [0, 0, [6.82912687e+02, 1.0]],
      [0, 0, [1.91245323e+02, 1.0]],
      [0, 0, [6.15619638e+01, 1.0]],
      [0, 0, [2.16954416e+01, 1.0]],
      [0, 0, [8.03372746e+00, 1.0]],
      [0, 0, [2.28806696e+00, 1.0]],
      [0, 0, [8.91891561e-01, 1.0]],
      [0, 0, [3.06490573e-01, 1.0]],
      [0, 0, [9.14278738e-02, 1.0]],
      [1, 0, [8.12382078e+01, 1.0]],
      [1, 0, [1.87382214e+01, 1.0]],
      [1, 0, [5.71665823e+00, 1.0]],
      [1, 0, [1.95959041e+00, 1.0]],
      [1, 0, [6.72331544e-01, 1.0]],
      [1, 0, [2.17227007e-01, 1.0]],
      [1, 0, [6.58062230e-02, 1.0]],
      [2, 0, [1.61811961e+00, 1.0]],
      [2, 0, [4.59945026e-01, 1.0]]]

Ne = [[0, 0, [2.63937466e+04, 1.0]],
      [0, 0, [3.84039579e+03, 1.0]],
      [0, 0, [8.58347380e+02, 1.0]],
      [0, 0, [2.39831047e+02, 1.0]],
      [0, 0, [7.70985719e+01, 1.0]],
      [0, 0, [2.71659873e+01, 1.0]],
      [0, 0, [1.00621973e+01, 1.0]],
      [0, 0, [2.88859221e+00, 1.0]],
      [0, 0, [1.11904975e+00, 1.0]],
      [0, 0, [3.82176942e-01, 1.0]],
      [0, 0, [1.13403907e-01, 1.0]],
      [1, 0, [1.01356529e+02, 1.0]],
      [1, 0, [2.34055004e+01, 1.0]],
      [1, 0, [7.16482787e+00, 1.0]],
      [1, 0, [2.45525847e+00, 1.0]],
      [1, 0, [8.37825292e-01, 1.0]],
      [1, 0, [2.67057270e-01, 1.0]],
      [1, 0, [8.25626434e-02, 1.0]],
      [2, 0, [2.16640868e+00, 1.0]],
      [2, 0, [6.40817214e-01, 1.0]]]

Al = [[0, 0, [8.22157890e+04, 1.0]],
      [0, 0, [1.15957649e+04, 1.0]],
      [0, 0, [2.54453086e+03, 1.0]],
      [0, 0, [7.02795217e+02, 1.0]],
      [0, 0, [2.24966679e+02, 1.0]],
      [0, 0, [7.99049065e+01, 1.0]],
      [0, 0, [3.07705104e+01, 1.0]],
      [0, 0, [1.24983416e+01, 1.0]],
      [0, 0, [3.36586330e+00, 1.0]],
      [0, 0, [1.18932643e+00, 1.0]],
      [0, 0, [1.75866307e-01, 1.0]],
      [0, 0, [6.51081752e-02, 1.0]],
      [0, 0, [2.17724199e-02, 1.0]],
      [1, 0, [2.70038076e+02, 1.0]],
      [1, 0, [6.26443560e+01, 1.0]],
      [1, 0, [1.95414963e+01, 1.0]],
      [1, 0, [6.98745273e+00, 1.0]],
      [1, 0, [2.60395194e+00, 1.0]],
      [1, 0, [9.65061208e-01, 1.0]],
      [1, 0, [2.11056456e-01, 1.0]],
      [1, 0, [6.00913311e-02, 1.0]],
      [1, 0, [1.41539264e-02, 1.0]],
      [2, 0, [1.93201508e-01, 1.0]],
      [2, 0, [5.27316775e-02, 1.0]]]

Si = [[0, 0, [1.04052681e+05, 1.0]],
      [0, 0, [1.45312628e+04, 1.0]],
      [0, 0, [3.16940829e+03, 1.0]],
      [0, 0, [8.71719057e+02, 1.0]],
      [0, 0, [2.78191231e+02, 1.0]],
      [0, 0, [9.85902809e+01, 1.0]],
      [0, 0, [3.78633488e+01, 1.0]],
      [0, 0, [1.52888958e+01, 1.0]],
      [0, 0, [4.12846838e+00, 1.0]],
      [0, 0, [1.50246568e+00, 1.0]],
      [0, 0, [2.53084395e-01, 1.0]],
      [0, 0, [9.29886769e-02, 1.0]],
      [0, 0, [3.13172058e-02, 1.0]],
      [1, 0, [3.30667196e+02, 1.0]],
      [1, 0, [7.65812543e+01, 1.0]],
      [1, 0, [2.39331074e+01, 1.0]],
      [1, 0, [8.59974926e+00, 1.0]],
      [1, 0, [3.24463955e+00, 1.0]],
      [1, 0, [1.23477396e+00, 1.0]],
      [1, 0, [2.95467725e-01, 1.0]],
      [1, 0, [8.50710464e-02, 1.0]],
      [1, 0, [2.24201687e-02, 1.0]],
      [2, 0, [2.73978756e-01, 1.0]],
      [2, 0, [8.00180146e-02, 1.0]]]

P  = [[0, 0, [1.27817378e+05, 1.0]],
      [0, 0, [1.76768640e+04, 1.0]],
      [0, 0, [3.83218251e+03, 1.0]],
      [0, 0, [1.04954834e+03, 1.0]],
      [0, 0, [3.33885131e+02, 1.0]],
      [0, 0, [1.18038894e+02, 1.0]],
      [0, 0, [4.52126521e+01, 1.0]],
      [0, 0, [1.81818021e+01, 1.0]],
      [0, 0, [4.94970020e+00, 1.0]],
      [0, 0, [1.84448181e+00, 1.0]],
      [0, 0, [3.39615003e-01, 1.0]],
      [0, 0, [1.23961699e-01, 1.0]],
      [0, 0, [4.09756058e-02, 1.0]],
      [1, 0, [3.85869276e+02, 1.0]],
      [1, 0, [8.91826403e+01, 1.0]],
      [1, 0, [2.78899440e+01, 1.0]],
      [1, 0, [1.00455314e+01, 1.0]],
      [1, 0, [3.82167337e+00, 1.0]],
      [1, 0, [1.48745844e+00, 1.0]],
      [1, 0, [3.87903897e-01, 1.0]],
      [1, 0, [1.12515635e-01, 1.0]],
      [1, 0, [3.05818876e-02, 1.0]],
      [2, 0, [3.69960709e-01, 1.0]],
      [2, 0, [1.13423367e-01, 1.0]]]

S  = [[0, 0, [1.54184920e+05, 1.0]],
      [0, 0, [2.11164907e+04, 1.0]],
      [0, 0, [4.54982671e+03, 1.0]],
      [0, 0, [1.24071939e+03, 1.0]],
      [0, 0, [3.93423320e+02, 1.0]],
      [0, 0, [1.38737228e+02, 1.0]],
      [0, 0, [5.30098147e+01, 1.0]],
      [0, 0, [2.12522226e+01, 1.0]],
      [0, 0, [5.83532969e+00, 1.0]],
      [0, 0, [2.21651914e+00, 1.0]],
      [0, 0, [4.35406263e-01, 1.0]],
      [0, 0, [1.58029412e-01, 1.0]],
      [0, 0, [5.10255140e-02, 1.0]],
      [1, 0, [4.28066668e+02, 1.0]],
      [1, 0, [9.87268972e+01, 1.0]],
      [1, 0, [3.08585687e+01, 1.0]],
      [1, 0, [1.11045374e+01, 1.0]],
      [1, 0, [4.23339093e+00, 1.0]],
      [1, 0, [1.66758667e+00, 1.0]],
      [1, 0, [4.77518751e-01, 1.0]],
      [1, 0, [1.40320203e-01, 1.0]],
      [1, 0, [3.88040255e-02, 1.0]],
      [2, 0, [4.79833420e-01, 1.0]],
      [2, 0, [1.52442063e-01, 1.0]]]

Cl = [[0, 0, [1.83564141e+05, 1.0]],
      [0, 0, [2.48976149e+04, 1.0]],
      [0, 0, [5.33136839e+03, 1.0]],
      [0, 0, [1.44747214e+03, 1.0]],
      [0, 0, [4.57469222e+02, 1.0]],
      [0, 0, [1.60909254e+02, 1.0]],
      [0, 0, [6.13372865e+01, 1.0]],
      [0, 0, [2.45295705e+01, 1.0]],
      [0, 0, [6.78828036e+00, 1.0]],
      [0, 0, [2.61928140e+00, 1.0]],
      [0, 0, [5.40489018e-01, 1.0]],
      [0, 0, [1.95237127e-01, 1.0]],
      [0, 0, [6.15087703e-02, 1.0]],
      [1, 0, [4.56419165e+02, 1.0]],
      [1, 0, [1.05052504e+02, 1.0]],
      [1, 0, [3.27850628e+01, 1.0]],
      [1, 0, [1.17457206e+01, 1.0]],
      [1, 0, [4.45062217e+00, 1.0]],
      [1, 0, [1.71836439e+00, 1.0]],
      [1, 0, [5.49194218e-01, 1.0]],
      [1, 0, [1.65501477e-01, 1.0]],
      [1, 0, [4.68308242e-02, 1.0]],
      [2, 0, [6.03155407e-01, 1.0]],
      [2, 0, [1.97360252e-01, 1.0]]]

Ar = [[0, 0, [2.16279421e+05, 1.0]],
      [0, 0, [2.90551539e+04, 1.0]],
      [0, 0, [6.18297423e+03, 1.0]],
      [0, 0, [1.67123592e+03, 1.0]],
      [0, 0, [5.26420305e+02, 1.0]],
      [0, 0, [1.84680306e+02, 1.0]],
      [0, 0, [7.02381594e+01, 1.0]],
      [0, 0, [2.80283299e+01, 1.0]],
      [0, 0, [7.80955209e+00, 1.0]],
      [0, 0, [3.05293044e+00, 1.0]],
      [0, 0, [6.55213131e-01, 1.0]],
      [0, 0, [2.35714250e-01, 1.0]],
      [0, 0, [7.41458248e-02, 1.0]],
      [1, 0, [4.94782136e+02, 1.0]],
      [1, 0, [1.13634523e+02, 1.0]],
      [1, 0, [3.54180102e+01, 1.0]],
      [1, 0, [1.26478698e+01, 1.0]],
      [1, 0, [4.76934758e+00, 1.0]],
      [1, 0, [1.76329840e+00, 1.0]],
      [1, 0, [6.10804371e-01, 1.0]],
      [1, 0, [1.89399130e-01, 1.0]],
      [1, 0, [5.65180048e-02, 1.0]],
      [2, 0, [7.40546739e-01, 1.0]],
      [2, 0, [2.55513919e-01, 1.0]]]

Ga = [[0, 0, [2.53393780e+06, 1.0]],
      [0, 0, [3.01895586e+05, 1.0]],
      [0, 0, [5.83285650e+04, 1.0]],
      [0, 0, [1.46132080e+04, 1.0]],
      [0, 0, [4.33631952e+03, 1.0]],
      [0, 0, [1.45157280e+03, 1.0]],
      [0, 0, [5.32161744e+02, 1.0]],
      [0, 0, [2.09781232e+02, 1.0]],
      [0, 0, [8.69253704e+01, 1.0]],
      [0, 0, [2.78236187e+01, 1.0]],
      [0, 0, [1.17396176e+01, 1.0]],
      [0, 0, [3.39595177e+00, 1.0]],
      [0, 0, [1.36900330e+00, 1.0]],
      [0, 0, [2.01619295e-01, 1.0]],
      [0, 0, [7.35445663e-02, 1.0]],
      [0, 0, [2.34775500e-02, 1.0]],
      [1, 0, [5.12265211e+03, 1.0]],
      [1, 0, [1.04125569e+03, 1.0]],
      [1, 0, [3.10993408e+02, 1.0]],
      [1, 0, [1.11520023e+02, 1.0]],
      [1, 0, [4.42424276e+01, 1.0]],
      [1, 0, [1.85498567e+01, 1.0]],
      [1, 0, [7.69458328e+00, 1.0]],
      [1, 0, [3.14658456e+00, 1.0]],
      [1, 0, [1.19940686e+00, 1.0]],
      [1, 0, [2.39820927e-01, 1.0]],
      [1, 0, [6.51562571e-02, 1.0]],
      [1, 0, [1.42967940e-02, 1.0]],
      [2, 0, [1.23474619e+02, 1.0]],
      [2, 0, [3.60838769e+01, 1.0]],
      [2, 0, [1.29999784e+01, 1.0]],
      [2, 0, [5.07899106e+00, 1.0]],
      [2, 0, [1.95065344e+00, 1.0]],
      [2, 0, [6.96231835e-01, 1.0]],
      [2, 0, [1.91389544e-01, 1.0]],
      [2, 0, [5.23958680e-02, 1.0]]]

Ge = [[0, 0, [2.81751730e+06, 1.0]],
      [0, 0, [3.34730369e+05, 1.0]],
      [0, 0, [6.44026065e+04, 1.0]],
      [0, 0, [1.60691681e+04, 1.0]],
      [0, 0, [4.75135539e+03, 1.0]],
      [0, 0, [1.58563409e+03, 1.0]],
      [0, 0, [5.79749083e+02, 1.0]],
      [0, 0, [2.27937937e+02, 1.0]],
      [0, 0, [9.41654307e+01, 1.0]],
      [0, 0, [3.01238979e+01, 1.0]],
      [0, 0, [1.27489877e+01, 1.0]],
      [0, 0, [3.75624900e+00, 1.0]],
      [0, 0, [1.54617527e+00, 1.0]],
      [0, 0, [2.54786891e-01, 1.0]],
      [0, 0, [9.46052718e-02, 1.0]],
      [0, 0, [3.24354100e-02, 1.0]],
      [1, 0, [5.83184125e+03, 1.0]],
      [1, 0, [1.17093418e+03, 1.0]],
      [1, 0, [3.47659896e+02, 1.0]],
      [1, 0, [1.24278408e+02, 1.0]],
      [1, 0, [4.92665014e+01, 1.0]],
      [1, 0, [2.06827795e+01, 1.0]],
      [1, 0, [8.66977158e+00, 1.0]],
      [1, 0, [3.55761073e+00, 1.0]],
      [1, 0, [1.37832276e+00, 1.0]],
      [1, 0, [2.93966317e-01, 1.0]],
      [1, 0, [8.41712236e-02, 1.0]],
      [1, 0, [2.17519380e-02, 1.0]],
      [2, 0, [1.42828940e+02, 1.0]],
      [2, 0, [4.18532206e+01, 1.0]],
      [2, 0, [1.51670034e+01, 1.0]],
      [2, 0, [5.98781134e+00, 1.0]],
      [2, 0, [2.34286056e+00, 1.0]],
      [2, 0, [8.62070416e-01, 1.0]],
      [2, 0, [2.43551427e-01, 1.0]],
      [2, 0, [7.49176230e-02, 1.0]]]

As = [[0, 0, [3.12433336e+06, 1.0]],
      [0, 0, [3.70089365e+05, 1.0]],
      [0, 0, [7.09073834e+04, 1.0]],
      [0, 0, [1.76195211e+04, 1.0]],
      [0, 0, [5.19096180e+03, 1.0]],
      [0, 0, [1.72695276e+03, 1.0]],
      [0, 0, [6.29694027e+02, 1.0]],
      [0, 0, [2.46916828e+02, 1.0]],
      [0, 0, [1.01707962e+02, 1.0]],
      [0, 0, [3.25284663e+01, 1.0]],
      [0, 0, [1.38072306e+01, 1.0]],
      [0, 0, [4.13706450e+00, 1.0]],
      [0, 0, [1.73598747e+00, 1.0]],
      [0, 0, [3.11494199e-01, 1.0]],
      [0, 0, [1.16970807e-01, 1.0]],
      [0, 0, [4.08761420e-02, 1.0]],
      [1, 0, [6.60866375e+03, 1.0]],
      [1, 0, [1.31073090e+03, 1.0]],
      [1, 0, [3.86811211e+02, 1.0]],
      [1, 0, [1.37819778e+02, 1.0]],
      [1, 0, [5.45796177e+01, 1.0]],
      [1, 0, [2.29356558e+01, 1.0]],
      [1, 0, [9.69590458e+00, 1.0]],
      [1, 0, [3.99224595e+00, 1.0]],
      [1, 0, [1.57217638e+00, 1.0]],
      [1, 0, [3.59097324e-01, 1.0]],
      [1, 0, [1.05285744e-01, 1.0]],
      [1, 0, [2.85666170e-02, 1.0]],
      [2, 0, [1.63399733e+02, 1.0]],
      [2, 0, [4.79699638e+01, 1.0]],
      [2, 0, [1.74693862e+01, 1.0]],
      [2, 0, [6.95636942e+00, 1.0]],
      [2, 0, [2.76398160e+00, 1.0]],
      [2, 0, [1.04220715e+00, 1.0]],
      [2, 0, [3.02713194e-01, 1.0]],
      [2, 0, [1.00852340e-01, 1.0]]]

Se = [[0, 0, [3.44182887e+06, 1.0]],
      [0, 0, [4.07068138e+05, 1.0]],
      [0, 0, [7.77244849e+04, 1.0]],
      [0, 0, [1.92426126e+04, 1.0]],
      [0, 0, [5.65024729e+03, 1.0]],
      [0, 0, [1.87425360e+03, 1.0]],
      [0, 0, [6.81631156e+02, 1.0]],
      [0, 0, [2.66608978e+02, 1.0]],
      [0, 0, [1.09522841e+02, 1.0]],
      [0, 0, [3.50346902e+01, 1.0]],
      [0, 0, [1.49130743e+01, 1.0]],
      [0, 0, [4.53761700e+00, 1.0]],
      [0, 0, [1.93762065e+00, 1.0]],
      [0, 0, [3.71820397e-01, 1.0]],
      [0, 0, [1.40721423e-01, 1.0]],
      [0, 0, [4.92783830e-02, 1.0]],
      [1, 0, [7.41469836e+03, 1.0]],
      [1, 0, [1.45355104e+03, 1.0]],
      [1, 0, [4.26408063e+02, 1.0]],
      [1, 0, [1.51422736e+02, 1.0]],
      [1, 0, [5.98886667e+01, 1.0]],
      [1, 0, [2.51773684e+01, 1.0]],
      [1, 0, [1.07012630e+01, 1.0]],
      [1, 0, [4.42319557e+00, 1.0]],
      [1, 0, [1.77020284e+00, 1.0]],
      [1, 0, [4.28396295e-01, 1.0]],
      [1, 0, [1.27540235e-01, 1.0]],
      [1, 0, [3.52657730e-02, 1.0]],
      [2, 0, [1.84809810e+02, 1.0]],
      [2, 0, [5.43191300e+01, 1.0]],
      [2, 0, [1.98619765e+01, 1.0]],
      [2, 0, [7.96442614e+00, 1.0]],
      [2, 0, [3.20429727e+00, 1.0]],
      [2, 0, [1.23191148e+00, 1.0]],
      [2, 0, [3.65720969e-01, 1.0]],
      [2, 0, [1.28304400e-01, 1.0]]]

Br = [[0, 0, [3.78958133e+06, 1.0]],
      [0, 0, [4.47177995e+05, 1.0]],
      [0, 0, [8.50617210e+04, 1.0]],
      [0, 0, [2.09773970e+04, 1.0]],
      [0, 0, [6.13808262e+03, 1.0]],
      [0, 0, [2.02984561e+03, 1.0]],
      [0, 0, [7.36219834e+02, 1.0]],
      [0, 0, [2.87212821e+02, 1.0]],
      [0, 0, [1.17667917e+02, 1.0]],
      [0, 0, [3.76490147e+01, 1.0]],
      [0, 0, [1.60686781e+01, 1.0]],
      [0, 0, [4.95782770e+00, 1.0]],
      [0, 0, [2.15077460e+00, 1.0]],
      [0, 0, [4.35777194e-01, 1.0]],
      [0, 0, [1.65871691e-01, 1.0]],
      [0, 0, [5.77845230e-02, 1.0]],
      [1, 0, [8.22636816e+03, 1.0]],
      [1, 0, [1.59529212e+03, 1.0]],
      [1, 0, [4.65309949e+02, 1.0]],
      [1, 0, [1.64692205e+02, 1.0]],
      [1, 0, [6.50352075e+01, 1.0]],
      [1, 0, [2.73373645e+01, 1.0]],
      [1, 0, [1.16494781e+01, 1.0]],
      [1, 0, [4.83745152e+00, 1.0]],
      [1, 0, [1.96737053e+00, 1.0]],
      [1, 0, [5.00650311e-01, 1.0]],
      [1, 0, [1.50759924e-01, 1.0]],
      [1, 0, [4.19952920e-02, 1.0]],
      [2, 0, [2.07075719e+02, 1.0]],
      [2, 0, [6.09053896e+01, 1.0]],
      [2, 0, [2.23455412e+01, 1.0]],
      [2, 0, [9.01205758e+00, 1.0]],
      [2, 0, [3.66355605e+00, 1.0]],
      [2, 0, [1.43085343e+00, 1.0]],
      [2, 0, [4.32158952e-01, 1.0]],
      [2, 0, [1.59020310e-01, 1.0]]]

Kr = [[0, 0, [4.14098736e+06, 1.0]],
      [0, 0, [4.88473543e+05, 1.0]],
      [0, 0, [9.26640299e+04, 1.0]],
      [0, 0, [2.27776348e+04, 1.0]],
      [0, 0, [6.64418072e+03, 1.0]],
      [0, 0, [2.19109855e+03, 1.0]],
      [0, 0, [7.92722388e+02, 1.0]],
      [0, 0, [3.08512335e+02, 1.0]],
      [0, 0, [1.26083423e+02, 1.0]],
      [0, 0, [4.03665991e+01, 1.0]],
      [0, 0, [1.72720374e+01, 1.0]],
      [0, 0, [5.39736984e+00, 1.0]],
      [0, 0, [2.37509131e+00, 1.0]],
      [0, 0, [5.03394409e-01, 1.0]],
      [0, 0, [1.92437314e-01, 1.0]],
      [0, 0, [6.77589420e-02, 1.0]],
      [1, 0, [9.01590791e+03, 1.0]],
      [1, 0, [1.73131083e+03, 1.0]],
      [1, 0, [5.02261944e+02, 1.0]],
      [1, 0, [1.77201567e+02, 1.0]],
      [1, 0, [6.98507883e+01, 1.0]],
      [1, 0, [2.93416011e+01, 1.0]],
      [1, 0, [1.25014174e+01, 1.0]],
      [1, 0, [5.21990010e+00, 1.0]],
      [1, 0, [2.15741984e+00, 1.0]],
      [1, 0, [5.74981830e-01, 1.0]],
      [1, 0, [1.74772323e-01, 1.0]],
      [1, 0, [5.00089580e-02, 1.0]],
      [2, 0, [2.25939509e+02, 1.0]],
      [2, 0, [6.64569801e+01, 1.0]],
      [2, 0, [2.44283436e+01, 1.0]],
      [2, 0, [9.88093527e+00, 1.0]],
      [2, 0, [4.03887103e+00, 1.0]],
      [2, 0, [1.58810227e+00, 1.0]],
      [2, 0, [4.65251237e-01, 1.0]],
      [2, 0, [1.97089570e-01, 1.0]]]

In = [[0, 0, [4.47385919e+07, 1.0]],
      [0, 0, [7.84746953e+06, 1.0]],
      [0, 0, [1.81845148e+06, 1.0]],
      [0, 0, [4.92160052e+05, 1.0]],
      [0, 0, [1.50183486e+05, 1.0]],
      [0, 0, [5.03830604e+04, 1.0]],
      [0, 0, [1.82783592e+04, 1.0]],
      [0, 0, [7.07947323e+03, 1.0]],
      [0, 0, [2.89756257e+03, 1.0]],
      [0, 0, [1.24246788e+03, 1.0]],
      [0, 0, [5.53880248e+02, 1.0]],
      [0, 0, [2.53316236e+02, 1.0]],
      [0, 0, [1.09014545e+02, 1.0]],
      [0, 0, [5.37405939e+01, 1.0]],
      [0, 0, [2.64745385e+01, 1.0]],
      [0, 0, [1.13185663e+01, 1.0]],
      [0, 0, [5.60695039e+00, 1.0]],
      [0, 0, [2.03797049e+00, 1.0]],
      [0, 0, [8.97637569e-01, 1.0]],
      [0, 0, [1.70387417e-01, 1.0]],
      [0, 0, [6.30363231e-02, 1.0]],
      [0, 0, [1.99791250e-02, 1.0]],
      [1, 0, [1.52811632e+05, 1.0]],
      [1, 0, [2.31102862e+04, 1.0]],
      [1, 0, [5.48315297e+03, 1.0]],
      [1, 0, [1.69249982e+03, 1.0]],
      [1, 0, [6.16342890e+02, 1.0]],
      [1, 0, [2.50226953e+02, 1.0]],
      [1, 0, [1.08979898e+02, 1.0]],
      [1, 0, [4.96529293e+01, 1.0]],
      [1, 0, [2.24311921e+01, 1.0]],
      [1, 0, [1.04552603e+01, 1.0]],
      [1, 0, [4.93647170e+00, 1.0]],
      [1, 0, [2.08406173e+00, 1.0]],
      [1, 0, [8.38431026e-01, 1.0]],
      [1, 0, [1.88292155e-01, 1.0]],
      [1, 0, [5.50115276e-02, 1.0]],
      [1, 0, [1.28549670e-02, 1.0]],
      [2, 0, [1.21604575e+03, 1.0]],
      [2, 0, [3.49566731e+02, 1.0]],
      [2, 0, [1.31845522e+02, 1.0]],
      [2, 0, [5.62844092e+01, 1.0]],
      [2, 0, [2.57312214e+01, 1.0]],
      [2, 0, [1.20908676e+01, 1.0]],
      [2, 0, [5.75825571e+00, 1.0]],
      [2, 0, [2.61040238e+00, 1.0]],
      [2, 0, [1.12466116e+00, 1.0]],
      [2, 0, [4.51079303e-01, 1.0]],
      [2, 0, [1.50882920e-01, 1.0]],
      [2, 0, [4.37084690e-02, 1.0]]]

Sn = [[0, 0, [4.52454474e+07, 1.0]],
      [0, 0, [8.06913583e+06, 1.0]],
      [0, 0, [1.89030230e+06, 1.0]],
      [0, 0, [5.15164594e+05, 1.0]],
      [0, 0, [1.57917873e+05, 1.0]],
      [0, 0, [5.31299541e+04, 1.0]],
      [0, 0, [1.93090892e+04, 1.0]],
      [0, 0, [7.48643551e+03, 1.0]],
      [0, 0, [3.06581053e+03, 1.0]],
      [0, 0, [1.31489101e+03, 1.0]],
      [0, 0, [5.86133042e+02, 1.0]],
      [0, 0, [2.68053609e+02, 1.0]],
      [0, 0, [1.15960596e+02, 1.0]],
      [0, 0, [5.72055564e+01, 1.0]],
      [0, 0, [2.81847965e+01, 1.0]],
      [0, 0, [1.20444518e+01, 1.0]],
      [0, 0, [5.99464606e+00, 1.0]],
      [0, 0, [2.21813817e+00, 1.0]],
      [0, 0, [9.96628933e-01, 1.0]],
      [0, 0, [2.06852767e-01, 1.0]],
      [0, 0, [7.81075870e-02, 1.0]],
      [0, 0, [2.66623380e-02, 1.0]],
      [1, 0, [1.75024680e+05, 1.0]],
      [1, 0, [2.62827500e+04, 1.0]],
      [1, 0, [6.16829182e+03, 1.0]],
      [1, 0, [1.88575085e+03, 1.0]],
      [1, 0, [6.81964315e+02, 1.0]],
      [1, 0, [2.75536358e+02, 1.0]],
      [1, 0, [1.19641083e+02, 1.0]],
      [1, 0, [5.44395323e+01, 1.0]],
      [1, 0, [2.47637780e+01, 1.0]],
      [1, 0, [1.15612755e+01, 1.0]],
      [1, 0, [5.49201008e+00, 1.0]],
      [1, 0, [2.30932764e+00, 1.0]],
      [1, 0, [9.44728464e-01, 1.0]],
      [1, 0, [2.21380094e-01, 1.0]],
      [1, 0, [6.81775450e-02, 1.0]],
      [1, 0, [1.87074110e-02, 1.0]],
      [2, 0, [1.35467891e+03, 1.0]],
      [2, 0, [3.88072853e+02, 1.0]],
      [2, 0, [1.46260893e+02, 1.0]],
      [2, 0, [6.24965159e+01, 1.0]],
      [2, 0, [2.86563738e+01, 1.0]],
      [2, 0, [1.35268000e+01, 1.0]],
      [2, 0, [6.48705608e+00, 1.0]],
      [2, 0, [2.97708822e+00, 1.0]],
      [2, 0, [1.30747793e+00, 1.0]],
      [2, 0, [5.38655350e-01, 1.0]],
      [2, 0, [1.84430760e-01, 1.0]],
      [2, 0, [6.09867920e-02, 1.0]]]

Sb = [[0, 0, [4.63134065e+07, 1.0]],
      [0, 0, [8.36845994e+06, 1.0]],
      [0, 0, [1.97735147e+06, 1.0]],
      [0, 0, [5.41778775e+05, 1.0]],
      [0, 0, [1.66635891e+05, 1.0]],
      [0, 0, [5.61734474e+04, 1.0]],
      [0, 0, [2.04366457e+04, 1.0]],
      [0, 0, [7.92710090e+03, 1.0]],
      [0, 0, [3.24642723e+03, 1.0]],
      [0, 0, [1.39205923e+03, 1.0]],
      [0, 0, [6.20277009e+02, 1.0]],
      [0, 0, [2.83575206e+02, 1.0]],
      [0, 0, [1.23269536e+02, 1.0]],
      [0, 0, [6.08236360e+01, 1.0]],
      [0, 0, [2.99592243e+01, 1.0]],
      [0, 0, [1.27965036e+01, 1.0]],
      [0, 0, [6.39705264e+00, 1.0]],
      [0, 0, [2.40357154e+00, 1.0]],
      [0, 0, [1.09988836e+00, 1.0]],
      [0, 0, [2.44704564e-01, 1.0]],
      [0, 0, [9.36768848e-02, 1.0]],
      [0, 0, [3.27250760e-02, 1.0]],
      [1, 0, [2.00137952e+05, 1.0]],
      [1, 0, [2.98583849e+04, 1.0]],
      [1, 0, [6.93385476e+03, 1.0]],
      [1, 0, [2.09933878e+03, 1.0]],
      [1, 0, [7.53837270e+02, 1.0]],
      [1, 0, [3.03065872e+02, 1.0]],
      [1, 0, [1.31189362e+02, 1.0]],
      [1, 0, [5.96180158e+01, 1.0]],
      [1, 0, [2.72946125e+01, 1.0]],
      [1, 0, [1.27429080e+01, 1.0]],
      [1, 0, [6.07654083e+00, 1.0]],
      [1, 0, [2.54816358e+00, 1.0]],
      [1, 0, [1.06056905e+00, 1.0]],
      [1, 0, [2.62801937e-01, 1.0]],
      [1, 0, [8.28808970e-02, 1.0]],
      [1, 0, [2.39100260e-02, 1.0]],
      [2, 0, [1.49993943e+03, 1.0]],
      [2, 0, [4.28185737e+02, 1.0]],
      [2, 0, [1.61218619e+02, 1.0]],
      [2, 0, [6.89311912e+01, 1.0]],
      [2, 0, [3.16808458e+01, 1.0]],
      [2, 0, [1.50116370e+01, 1.0]],
      [2, 0, [7.24046742e+00, 1.0]],
      [2, 0, [3.35932073e+00, 1.0]],
      [2, 0, [1.49989506e+00, 1.0]],
      [2, 0, [6.31672285e-01, 1.0]],
      [2, 0, [2.21234800e-01, 1.0]],
      [2, 0, [7.98754890e-02, 1.0]]]

Te = [[0, 0, [4.60876438e+07, 1.0]],
      [0, 0, [8.48185531e+06, 1.0]],
      [0, 0, [2.02935319e+06, 1.0]],
      [0, 0, [5.60649984e+05, 1.0]],
      [0, 0, [1.73423633e+05, 1.0]],
      [0, 0, [5.86865858e+04, 1.0]],
      [0, 0, [2.14064360e+04, 1.0]],
      [0, 0, [8.31765178e+03, 1.0]],
      [0, 0, [3.41026858e+03, 1.0]],
      [0, 0, [1.46339902e+03, 1.0]],
      [0, 0, [6.52363087e+02, 1.0]],
      [0, 0, [2.98384036e+02, 1.0]],
      [0, 0, [1.30289566e+02, 1.0]],
      [0, 0, [6.43051359e+01, 1.0]],
      [0, 0, [3.16829645e+01, 1.0]],
      [0, 0, [1.35535279e+01, 1.0]],
      [0, 0, [6.80430789e+00, 1.0]],
      [0, 0, [2.59330173e+00, 1.0]],
      [0, 0, [1.20660620e+00, 1.0]],
      [0, 0, [2.84056038e-01, 1.0]],
      [0, 0, [1.09830964e-01, 1.0]],
      [0, 0, [3.86945480e-02, 1.0]],
      [1, 0, [2.26468078e+05, 1.0]],
      [1, 0, [3.35991490e+04, 1.0]],
      [1, 0, [7.72803459e+03, 1.0]],
      [1, 0, [2.31840208e+03, 1.0]],
      [1, 0, [8.26824561e+02, 1.0]],
      [1, 0, [3.30800651e+02, 1.0]],
      [1, 0, [1.42756809e+02, 1.0]],
      [1, 0, [6.47839168e+01, 1.0]],
      [1, 0, [2.97970683e+01, 1.0]],
      [1, 0, [1.38969575e+01, 1.0]],
      [1, 0, [6.64134641e+00, 1.0]],
      [1, 0, [2.78449372e+00, 1.0]],
      [1, 0, [1.17847600e+00, 1.0]],
      [1, 0, [3.06494951e-01, 1.0]],
      [1, 0, [9.81746816e-02, 1.0]],
      [1, 0, [2.81622530e-02, 1.0]],
      [2, 0, [1.64872625e+03, 1.0]],
      [2, 0, [4.69035378e+02, 1.0]],
      [2, 0, [1.76389881e+02, 1.0]],
      [2, 0, [7.54443007e+01, 1.0]],
      [2, 0, [3.47355997e+01, 1.0]],
      [2, 0, [1.65110620e+01, 1.0]],
      [2, 0, [8.00141748e+00, 1.0]],
      [2, 0, [3.74770432e+00, 1.0]],
      [2, 0, [1.69671413e+00, 1.0]],
      [2, 0, [7.27350979e-01, 1.0]],
      [2, 0, [2.59354130e-01, 1.0]],
      [2, 0, [1.00387500e-01, 1.0]]]

I  = [[0, 0, [4.76716569e+07, 1.0]],
      [0, 0, [8.85780213e+06, 1.0]],
      [0, 0, [2.13271716e+06, 1.0]],
      [0, 0, [5.91444786e+05, 1.0]],
      [0, 0, [1.83355348e+05, 1.0]],
      [0, 0, [6.21146979e+04, 1.0]],
      [0, 0, [2.26645336e+04, 1.0]],
      [0, 0, [8.80519941e+03, 1.0]],
      [0, 0, [3.60858100e+03, 1.0]],
      [0, 0, [1.54755208e+03, 1.0]],
      [0, 0, [6.89378471e+02, 1.0]],
      [0, 0, [3.15140312e+02, 1.0]],
      [0, 0, [1.38182652e+02, 1.0]],
      [0, 0, [6.81673701e+01, 1.0]],
      [0, 0, [3.35643795e+01, 1.0]],
      [0, 0, [1.43531864e+01, 1.0]],
      [0, 0, [7.23366317e+00, 1.0]],
      [0, 0, [2.78850745e+00, 1.0]],
      [0, 0, [1.31752658e+00, 1.0]],
      [0, 0, [3.24888716e-01, 1.0]],
      [0, 0, [1.26580436e-01, 1.0]],
      [0, 0, [4.56490070e-02, 1.0]],
      [1, 0, [2.53817535e+05, 1.0]],
      [1, 0, [3.74699676e+04, 1.0]],
      [1, 0, [8.54270467e+03, 1.0]],
      [1, 0, [2.54055332e+03, 1.0]],
      [1, 0, [9.00086660e+02, 1.0]],
      [1, 0, [3.58408502e+02, 1.0]],
      [1, 0, [1.54196337e+02, 1.0]],
      [1, 0, [6.98648702e+01, 1.0]],
      [1, 0, [3.22303175e+01, 1.0]],
      [1, 0, [1.50109163e+01, 1.0]],
      [1, 0, [7.18477925e+00, 1.0]],
      [1, 0, [3.01682750e+00, 1.0]],
      [1, 0, [1.29734703e+00, 1.0]],
      [1, 0, [3.51707219e-01, 1.0]],
      [1, 0, [1.13965453e-01, 1.0]],
      [1, 0, [3.38228330e-02, 1.0]],
      [2, 0, [1.80142819e+03, 1.0]],
      [2, 0, [5.10725112e+02, 1.0]],
      [2, 0, [1.91813939e+02, 1.0]],
      [2, 0, [8.20523022e+01, 1.0]],
      [2, 0, [3.78288870e+01, 1.0]],
      [2, 0, [1.80291432e+01, 1.0]],
      [2, 0, [8.77223438e+00, 1.0]],
      [2, 0, [4.14272362e+00, 1.0]],
      [2, 0, [1.89787537e+00, 1.0]],
      [2, 0, [8.25505789e-01, 1.0]],
      [2, 0, [2.98523850e-01, 1.0]],
      [2, 0, [1.23237630e-01, 1.0]]]

Xe = [[0, 0, [4.79578326e+07, 1.0]],
      [0, 0, [9.04292743e+06, 1.0]],
      [0, 0, [2.19959406e+06, 1.0]],
      [0, 0, [6.14117878e+05, 1.0]],
      [0, 0, [1.91258945e+05, 1.0]],
      [0, 0, [6.49876195e+04, 1.0]],
      [0, 0, [2.37589185e+04, 1.0]],
      [0, 0, [9.24138595e+03, 1.0]],
      [0, 0, [3.78988485e+03, 1.0]],
      [0, 0, [1.62580877e+03, 1.0]],
      [0, 0, [7.24292309e+02, 1.0]],
      [0, 0, [3.31147981e+02, 1.0]],
      [0, 0, [1.45745038e+02, 1.0]],
      [0, 0, [7.18803328e+01, 1.0]],
      [0, 0, [3.53927479e+01, 1.0]],
      [0, 0, [1.51571648e+01, 1.0]],
      [0, 0, [7.66704355e+00, 1.0]],
      [0, 0, [2.98773483e+00, 1.0]],
      [0, 0, [1.43153305e+00, 1.0]],
      [0, 0, [3.67400732e-01, 1.0]],
      [0, 0, [1.44011855e-01, 1.0]],
      [0, 0, [5.38533710e-02, 1.0]],
      [1, 0, [2.81773747e+05, 1.0]],
      [1, 0, [4.14212882e+04, 1.0]],
      [1, 0, [9.36856349e+03, 1.0]],
      [1, 0, [2.76346063e+03, 1.0]],
      [1, 0, [9.72894041e+02, 1.0]],
      [1, 0, [3.85623880e+02, 1.0]],
      [1, 0, [1.65397231e+02, 1.0]],
      [1, 0, [7.48094176e+01, 1.0]],
      [1, 0, [3.45694513e+01, 1.0]],
      [1, 0, [1.60788107e+01, 1.0]],
      [1, 0, [7.70648406e+00, 1.0]],
      [1, 0, [3.24418792e+00, 1.0]],
      [1, 0, [1.41638808e+00, 1.0]],
      [1, 0, [3.98205340e-01, 1.0]],
      [1, 0, [1.30222232e-01, 1.0]],
      [1, 0, [4.06211830e-02, 1.0]],
      [2, 0, [1.95959651e+03, 1.0]],
      [2, 0, [5.53672740e+02, 1.0]],
      [2, 0, [2.07645956e+02, 1.0]],
      [2, 0, [8.88217231e+01, 1.0]],
      [2, 0, [4.09925340e+01, 1.0]],
      [2, 0, [1.95816201e+01, 1.0]],
      [2, 0, [9.56107055e+00, 1.0]],
      [2, 0, [4.54825152e+00, 1.0]],
      [2, 0, [2.10524778e+00, 1.0]],
      [2, 0, [9.27044060e-01, 1.0]],
      [2, 0, [3.39242330e-01, 1.0]],
      [2, 0, [1.51288890e-01, 1.0]]]

Tl = [[0, 0, [5.80338285e+07, 1.0]],
      [0, 0, [1.47025619e+07, 1.0]],
      [0, 0, [4.63555754e+06, 1.0]],
      [0, 0, [1.58527398e+06, 1.0]],
      [0, 0, [5.79246257e+05, 1.0]],
      [0, 0, [2.21879438e+05, 1.0]],
      [0, 0, [8.87307796e+04, 1.0]],
      [0, 0, [3.68480778e+04, 1.0]],
      [0, 0, [1.58432979e+04, 1.0]],
      [0, 0, [7.03105139e+03, 1.0]],
      [0, 0, [3.21218877e+03, 1.0]],
      [0, 0, [1.50707828e+03, 1.0]],
      [0, 0, [7.23010621e+02, 1.0]],
      [0, 0, [3.39805146e+02, 1.0]],
      [0, 0, [1.70939014e+02, 1.0]],
      [0, 0, [8.80173226e+01, 1.0]],
      [0, 0, [4.30162250e+01, 1.0]],
      [0, 0, [2.26098628e+01, 1.0]],
      [0, 0, [1.02665686e+01, 1.0]],
      [0, 0, [5.26482246e+00, 1.0]],
      [0, 0, [1.89889444e+00, 1.0]],
      [0, 0, [8.84613719e-01, 1.0]],
      [0, 0, [2.01065745e-01, 1.0]],
      [0, 0, [7.32370907e-02, 1.0]],
      [0, 0, [1.91277593e-02, 1.0]],
      [1, 0, [1.37282646e+07, 1.0]],
      [1, 0, [2.35721283e+06, 1.0]],
      [1, 0, [5.23566782e+05, 1.0]],
      [1, 0, [1.36290763e+05, 1.0]],
      [1, 0, [4.01480894e+04, 1.0]],
      [1, 0, [1.31820210e+04, 1.0]],
      [1, 0, [4.78798622e+03, 1.0]],
      [1, 0, [1.90284415e+03, 1.0]],
      [1, 0, [8.15022877e+02, 1.0]],
      [1, 0, [3.69676016e+02, 1.0]],
      [1, 0, [1.75043773e+02, 1.0]],
      [1, 0, [8.44562670e+01, 1.0]],
      [1, 0, [4.20432285e+01, 1.0]],
      [1, 0, [2.11720749e+01, 1.0]],
      [1, 0, [1.01075987e+01, 1.0]],
      [1, 0, [4.90314928e+00, 1.0]],
      [1, 0, [1.96328877e+00, 1.0]],
      [1, 0, [8.05496991e-01, 1.0]],
      [1, 0, [2.00496286e-01, 1.0]],
      [1, 0, [5.59278837e-02, 1.0]],
      [1, 0, [1.14312811e-02, 1.0]],
      [2, 0, [9.80668020e+03, 1.0]],
      [2, 0, [2.53461202e+03, 1.0]],
      [2, 0, [9.01015229e+02, 1.0]],
      [2, 0, [3.75502849e+02, 1.0]],
      [2, 0, [1.71843268e+02, 1.0]],
      [2, 0, [8.32163373e+01, 1.0]],
      [2, 0, [4.17508976e+01, 1.0]],
      [2, 0, [2.11097058e+01, 1.0]],
      [2, 0, [1.05041178e+01, 1.0]],
      [2, 0, [5.14706485e+00, 1.0]],
      [2, 0, [2.34214699e+00, 1.0]],
      [2, 0, [1.03451933e+00, 1.0]],
      [2, 0, [4.21552202e-01, 1.0]],
      [2, 0, [1.50267783e-01, 1.0]],
      [2, 0, [4.10356483e-02, 1.0]],
      [3, 0, [4.89995595e+02, 1.0]],
      [3, 0, [1.66331967e+02, 1.0]],
      [3, 0, [6.99800850e+01, 1.0]],
      [3, 0, [3.21750898e+01, 1.0]],
      [3, 0, [1.54175772e+01, 1.0]],
      [3, 0, [7.36995778e+00, 1.0]],
      [3, 0, [3.40866133e+00, 1.0]],
      [3, 0, [1.41384687e+00, 1.0]]]

Pb = [[0, 0, [5.78790468e+07, 1.0]],
      [0, 0, [1.47277856e+07, 1.0]],
      [0, 0, [4.66980788e+06, 1.0]],
      [0, 0, [1.60593717e+06, 1.0]],
      [0, 0, [5.89847781e+05, 1.0]],
      [0, 0, [2.26912673e+05, 1.0]],
      [0, 0, [9.10631045e+04, 1.0]],
      [0, 0, [3.79208693e+04, 1.0]],
      [0, 0, [1.63391338e+04, 1.0]],
      [0, 0, [7.26252891e+03, 1.0]],
      [0, 0, [3.32169850e+03, 1.0]],
      [0, 0, [1.55964749e+03, 1.0]],
      [0, 0, [7.48637891e+02, 1.0]],
      [0, 0, [3.52613276e+02, 1.0]],
      [0, 0, [1.77504674e+02, 1.0]],
      [0, 0, [9.14632745e+01, 1.0]],
      [0, 0, [4.47101722e+01, 1.0]],
      [0, 0, [2.35272879e+01, 1.0]],
      [0, 0, [1.07214573e+01, 1.0]],
      [0, 0, [5.52715061e+00, 1.0]],
      [0, 0, [2.02746272e+00, 1.0]],
      [0, 0, [9.62040984e-01, 1.0]],
      [0, 0, [2.32944034e-01, 1.0]],
      [0, 0, [8.68598864e-02, 1.0]],
      [0, 0, [2.50419695e-02, 1.0]],
      [1, 0, [1.44315847e+07, 1.0]],
      [1, 0, [2.50743314e+06, 1.0]],
      [1, 0, [5.60444246e+05, 1.0]],
      [1, 0, [1.46372787e+05, 1.0]],
      [1, 0, [4.31579089e+04, 1.0]],
      [1, 0, [1.41524402e+04, 1.0]],
      [1, 0, [5.12587435e+03, 1.0]],
      [1, 0, [2.03021930e+03, 1.0]],
      [1, 0, [8.66820906e+02, 1.0]],
      [1, 0, [3.92148422e+02, 1.0]],
      [1, 0, [1.85356353e+02, 1.0]],
      [1, 0, [8.95192825e+01, 1.0]],
      [1, 0, [4.44857645e+01, 1.0]],
      [1, 0, [2.24228444e+01, 1.0]],
      [1, 0, [1.06751765e+01, 1.0]],
      [1, 0, [5.19177307e+00, 1.0]],
      [1, 0, [2.09083067e+00, 1.0]],
      [1, 0, [8.71717990e-01, 1.0]],
      [1, 0, [2.16745290e-01, 1.0]],
      [1, 0, [6.58254189e-02, 1.0]],
      [1, 0, [1.67280721e-02, 1.0]],
      [2, 0, [1.06108755e+04, 1.0]],
      [2, 0, [2.73011313e+03, 1.0]],
      [2, 0, [9.67314733e+02, 1.0]],
      [2, 0, [4.02389437e+02, 1.0]],
      [2, 0, [1.83970261e+02, 1.0]],
      [2, 0, [8.91221384e+01, 1.0]],
      [2, 0, [4.47577918e+01, 1.0]],
      [2, 0, [2.27108628e+01, 1.0]],
      [2, 0, [1.13491680e+01, 1.0]],
      [2, 0, [5.59531052e+00, 1.0]],
      [2, 0, [2.58639830e+00, 1.0]],
      [2, 0, [1.16289182e+00, 1.0]],
      [2, 0, [4.85219166e-01, 1.0]],
      [2, 0, [1.75459932e-01, 1.0]],
      [2, 0, [5.57726242e-02, 1.0]],
      [3, 0, [5.20138448e+02, 1.0]],
      [3, 0, [1.76510085e+02, 1.0]],
      [3, 0, [7.43430142e+01, 1.0]],
      [3, 0, [3.42407420e+01, 1.0]],
      [3, 0, [1.64610261e+01, 1.0]],
      [3, 0, [7.90770330e+00, 1.0]],
      [3, 0, [3.68433522e+00, 1.0]],
      [3, 0, [1.54783108e+00, 1.0]]]

Bi = [[0, 0, [5.80226841e+07, 1.0]],
      [0, 0, [1.48197988e+07, 1.0]],
      [0, 0, [4.72238743e+06, 1.0]],
      [0, 0, [1.63218333e+06, 1.0]],
      [0, 0, [6.02327402e+05, 1.0]],
      [0, 0, [2.32625999e+05, 1.0]],
      [0, 0, [9.36570190e+04, 1.0]],
      [0, 0, [3.90989196e+04, 1.0]],
      [0, 0, [1.68788402e+04, 1.0]],
      [0, 0, [7.51276737e+03, 1.0]],
      [0, 0, [3.43938033e+03, 1.0]],
      [0, 0, [1.61581494e+03, 1.0]],
      [0, 0, [7.75851541e+02, 1.0]],
      [0, 0, [3.66121900e+02, 1.0]],
      [0, 0, [1.84400480e+02, 1.0]],
      [0, 0, [9.50651171e+01, 1.0]],
      [0, 0, [4.64662683e+01, 1.0]],
      [0, 0, [2.44772899e+01, 1.0]],
      [0, 0, [1.11901610e+01, 1.0]],
      [0, 0, [5.79784882e+00, 1.0]],
      [0, 0, [2.15870772e+00, 1.0]],
      [0, 0, [1.04226135e+00, 1.0]],
      [0, 0, [2.65908662e-01, 1.0]],
      [0, 0, [1.00734400e-01, 1.0]],
      [0, 0, [3.00281802e-02, 1.0]],
      [1, 0, [1.52678508e+07, 1.0]],
      [1, 0, [2.68636941e+06, 1.0]],
      [1, 0, [6.04560600e+05, 1.0]],
      [1, 0, [1.58483354e+05, 1.0]],
      [1, 0, [4.67858698e+04, 1.0]],
      [1, 0, [1.53250769e+04, 1.0]],
      [1, 0, [5.53455235e+03, 1.0]],
      [1, 0, [2.18424309e+03, 1.0]],
      [1, 0, [9.29423926e+02, 1.0]],
      [1, 0, [4.19320068e+02, 1.0]],
      [1, 0, [1.97851207e+02, 1.0]],
      [1, 0, [9.56967864e+01, 1.0]],
      [1, 0, [4.74356303e+01, 1.0]],
      [1, 0, [2.39296309e+01, 1.0]],
      [1, 0, [1.13516041e+01, 1.0]],
      [1, 0, [5.53461227e+00, 1.0]],
      [1, 0, [2.24108719e+00, 1.0]],
      [1, 0, [9.53174998e-01, 1.0]],
      [1, 0, [2.48293103e-01, 1.0]],
      [1, 0, [7.79158225e-02, 1.0]],
      [1, 0, [2.13172095e-02, 1.0]],
      [2, 0, [1.14563420e+04, 1.0]],
      [2, 0, [2.93479123e+03, 1.0]],
      [2, 0, [1.03643045e+03, 1.0]],
      [2, 0, [4.30323545e+02, 1.0]],
      [2, 0, [1.96540411e+02, 1.0]],
      [2, 0, [9.52357533e+01, 1.0]],
      [2, 0, [4.78674481e+01, 1.0]],
      [2, 0, [2.43647577e+01, 1.0]],
      [2, 0, [1.22239217e+01, 1.0]],
      [2, 0, [6.06071971e+00, 1.0]],
      [2, 0, [2.84128694e+00, 1.0]],
      [2, 0, [1.29733557e+00, 1.0]],
      [2, 0, [5.52008084e-01, 1.0]],
      [2, 0, [2.02232067e-01, 1.0]],
      [2, 0, [6.96834540e-02, 1.0]],
      [3, 0, [5.51023047e+02, 1.0]],
      [3, 0, [1.86924359e+02, 1.0]],
      [3, 0, [7.88015411e+01, 1.0]],
      [3, 0, [3.63519861e+01, 1.0]],
      [3, 0, [1.75278866e+01, 1.0]],
      [3, 0, [8.45836180e+00, 1.0]],
      [3, 0, [3.96752275e+00, 1.0]],
      [3, 0, [1.68654417e+00, 1.0]]]

Po = [[0, 0, [5.83099270e+07, 1.0]],
      [0, 0, [1.49430679e+07, 1.0]],
      [0, 0, [4.78316183e+06, 1.0]],
      [0, 0, [1.66076999e+06, 1.0]],
      [0, 0, [6.15531036e+05, 1.0]],
      [0, 0, [2.38575062e+05, 1.0]],
      [0, 0, [9.63301741e+04, 1.0]],
      [0, 0, [4.03038171e+04, 1.0]],
      [0, 0, [1.74276022e+04, 1.0]],
      [0, 0, [7.76599981e+03, 1.0]],
      [0, 0, [3.55802853e+03, 1.0]],
      [0, 0, [1.67229646e+03, 1.0]],
      [0, 0, [8.03189111e+02, 1.0]],
      [0, 0, [3.79694900e+02, 1.0]],
      [0, 0, [1.91317460e+02, 1.0]],
      [0, 0, [9.86768775e+01, 1.0]],
      [0, 0, [4.82455671e+01, 1.0]],
      [0, 0, [2.54409112e+01, 1.0]],
      [0, 0, [1.16672459e+01, 1.0]],
      [0, 0, [6.07428396e+00, 1.0]],
      [0, 0, [2.29230811e+00, 1.0]],
      [0, 0, [1.12470444e+00, 1.0]],
      [0, 0, [2.99794160e-01, 1.0]],
      [0, 0, [1.14926500e-01, 1.0]],
      [0, 0, [3.47359516e-02, 1.0]],
      [1, 0, [1.60947307e+07, 1.0]],
      [1, 0, [2.86516751e+06, 1.0]],
      [1, 0, [6.48913936e+05, 1.0]],
      [1, 0, [1.70700193e+05, 1.0]],
      [1, 0, [5.04504867e+04, 1.0]],
      [1, 0, [1.65086424e+04, 1.0]],
      [1, 0, [5.94571483e+03, 1.0]],
      [1, 0, [2.33841649e+03, 1.0]],
      [1, 0, [9.91719770e+02, 1.0]],
      [1, 0, [4.46207589e+02, 1.0]],
      [1, 0, [2.10147178e+02, 1.0]],
      [1, 0, [1.01715388e+02, 1.0]],
      [1, 0, [5.02831330e+01, 1.0]],
      [1, 0, [2.53679474e+01, 1.0]],
      [1, 0, [1.20014535e+01, 1.0]],
      [1, 0, [5.86779855e+00, 1.0]],
      [1, 0, [2.38995036e+00, 1.0]],
      [1, 0, [1.03602518e+00, 1.0]],
      [1, 0, [2.82313749e-01, 1.0]],
      [1, 0, [9.04113582e-02, 1.0]],
      [1, 0, [2.56072955e-02, 1.0]],
      [2, 0, [1.23256609e+04, 1.0]],
      [2, 0, [3.14433523e+03, 1.0]],
      [2, 0, [1.10687289e+03, 1.0]],
      [2, 0, [4.58688520e+02, 1.0]],
      [2, 0, [2.09270647e+02, 1.0]],
      [2, 0, [1.01414646e+02, 1.0]],
      [2, 0, [5.10053207e+01, 1.0]],
      [2, 0, [2.60297016e+01, 1.0]],
      [2, 0, [1.31054355e+01, 1.0]],
      [2, 0, [6.53090018e+00, 1.0]],
      [2, 0, [3.09916893e+00, 1.0]],
      [2, 0, [1.43365559e+00, 1.0]],
      [2, 0, [6.19815647e-01, 1.0]],
      [2, 0, [2.29483247e-01, 1.0]],
      [2, 0, [9.04070040e-02, 1.0]],
      [3, 0, [5.82612869e+02, 1.0]],
      [3, 0, [1.97561084e+02, 1.0]],
      [3, 0, [8.33493584e+01, 1.0]],
      [3, 0, [3.85056868e+01, 1.0]],
      [3, 0, [1.86165214e+01, 1.0]],
      [3, 0, [9.02103234e+00, 1.0]],
      [3, 0, [4.25765078e+00, 1.0]],
      [3, 0, [1.82952431e+00, 1.0]]]

At = [[0, 0, [5.84244092e+07, 1.0]],
      [0, 0, [1.50209341e+07, 1.0]],
      [0, 0, [4.82965478e+06, 1.0]],
      [0, 0, [1.68476525e+06, 1.0]],
      [0, 0, [6.27250219e+05, 1.0]],
      [0, 0, [2.44042340e+05, 1.0]],
      [0, 0, [9.88469174e+04, 1.0]],
      [0, 0, [4.14580280e+04, 1.0]],
      [0, 0, [1.79601077e+04, 1.0]],
      [0, 0, [8.01410800e+03, 1.0]],
      [0, 0, [3.67511528e+03, 1.0]],
      [0, 0, [1.72833709e+03, 1.0]],
      [0, 0, [8.30437479e+02, 1.0]],
      [0, 0, [3.93276071e+02, 1.0]],
      [0, 0, [1.98241093e+02, 1.0]],
      [0, 0, [1.02297017e+02, 1.0]],
      [0, 0, [5.00479865e+01, 1.0]],
      [0, 0, [2.64182357e+01, 1.0]],
      [0, 0, [1.21526000e+01, 1.0]],
      [0, 0, [6.35636699e+00, 1.0]],
      [0, 0, [2.42798951e+00, 1.0]],
      [0, 0, [1.20909391e+00, 1.0]],
      [0, 0, [3.34586720e-01, 1.0]],
      [0, 0, [1.29471622e-01, 1.0]],
      [0, 0, [3.92682625e-02, 1.0]],
      [1, 0, [1.68580977e+07, 1.0]],
      [1, 0, [3.03531193e+06, 1.0]],
      [1, 0, [6.91745658e+05, 1.0]],
      [1, 0, [1.82596030e+05, 1.0]],
      [1, 0, [5.40346837e+04, 1.0]],
      [1, 0, [1.76678260e+04, 1.0]],
      [1, 0, [6.34774612e+03, 1.0]],
      [1, 0, [2.48856435e+03, 1.0]],
      [1, 0, [1.05207443e+03, 1.0]],
      [1, 0, [4.72117514e+02, 1.0]],
      [1, 0, [2.21927629e+02, 1.0]],
      [1, 0, [1.07421196e+02, 1.0]],
      [1, 0, [5.29690220e+01, 1.0]],
      [1, 0, [2.67152734e+01, 1.0]],
      [1, 0, [1.26155116e+01, 1.0]],
      [1, 0, [6.18686221e+00, 1.0]],
      [1, 0, [2.53471710e+00, 1.0]],
      [1, 0, [1.11839341e+00, 1.0]],
      [1, 0, [3.17263922e-01, 1.0]],
      [1, 0, [1.03102418e-01, 1.0]],
      [1, 0, [2.97478162e-02, 1.0]],
      [2, 0, [1.32162902e+04, 1.0]],
      [2, 0, [3.35817578e+03, 1.0]],
      [2, 0, [1.17845424e+03, 1.0]],
      [2, 0, [4.87406344e+02, 1.0]],
      [2, 0, [2.22124347e+02, 1.0]],
      [2, 0, [1.07639309e+02, 1.0]],
      [2, 0, [5.41610092e+01, 1.0]],
      [2, 0, [2.77000451e+01, 1.0]],
      [2, 0, [1.39904848e+01, 1.0]],
      [2, 0, [7.00395005e+00, 1.0]],
      [2, 0, [3.35866386e+00, 1.0]],
      [2, 0, [1.57110943e+00, 1.0]],
      [2, 0, [6.88260345e-01, 1.0]],
      [2, 0, [2.56992818e-01, 1.0]],
      [2, 0, [1.09389179e-01, 1.0]],
      [3, 0, [6.14902927e+02, 1.0]],
      [3, 0, [2.08417199e+02, 1.0]],
      [3, 0, [8.79848050e+01, 1.0]],
      [3, 0, [4.07008509e+01, 1.0]],
      [3, 0, [1.97263216e+01, 1.0]],
      [3, 0, [9.59533206e+00, 1.0]],
      [3, 0, [4.55444773e+00, 1.0]],
      [3, 0, [1.97652644e+00, 1.0]]]

Rn = [[0, 0, [5.68476598e+07, 1.0]],
      [0, 0, [1.46794059e+07, 1.0]],
      [0, 0, [4.74912556e+06, 1.0]],
      [0, 0, [1.66787655e+06, 1.0]],
      [0, 0, [6.25170530e+05, 1.0]],
      [0, 0, [2.44681432e+05, 1.0]],
      [0, 0, [9.96160398e+04, 1.0]],
      [0, 0, [4.19595350e+04, 1.0]],
      [0, 0, [1.82417354e+04, 1.0]],
      [0, 0, [8.16322270e+03, 1.0]],
      [0, 0, [3.75224966e+03, 1.0]],
      [0, 0, [1.76799035e+03, 1.0]],
      [0, 0, [8.50979574e+02, 1.0]],
      [0, 0, [4.04118111e+02, 1.0]],
      [0, 0, [2.03907122e+02, 1.0]],
      [0, 0, [1.05354324e+02, 1.0]],
      [0, 0, [5.17227450e+01, 1.0]],
      [0, 0, [2.73375318e+01, 1.0]],
      [0, 0, [1.26282083e+01, 1.0]],
      [0, 0, [6.63534680e+00, 1.0]],
      [0, 0, [2.56553897e+00, 1.0]],
      [0, 0, [1.29475083e+00, 1.0]],
      [0, 0, [3.70279586e-01, 1.0]],
      [0, 0, [1.44373172e-01, 1.0]],
      [0, 0, [4.43919447e-02, 1.0]],
      [1, 0, [1.72824122e+07, 1.0]],
      [1, 0, [3.15641728e+06, 1.0]],
      [1, 0, [7.25282723e+05, 1.0]],
      [1, 0, [1.92412503e+05, 1.0]],
      [1, 0, [5.70898829e+04, 1.0]],
      [1, 0, [1.86765023e+04, 1.0]],
      [1, 0, [6.70190435e+03, 1.0]],
      [1, 0, [2.62169163e+03, 1.0]],
      [1, 0, [1.10574454e+03, 1.0]],
      [1, 0, [4.95179408e+02, 1.0]],
      [1, 0, [2.32404364e+02, 1.0]],
      [1, 0, [1.12459428e+02, 1.0]],
      [1, 0, [5.53557248e+01, 1.0]],
      [1, 0, [2.79136374e+01, 1.0]],
      [1, 0, [1.31728352e+01, 1.0]],
      [1, 0, [6.48261416e+00, 1.0]],
      [1, 0, [2.67263353e+00, 1.0]],
      [1, 0, [1.19863600e+00, 1.0]],
      [1, 0, [3.52585561e-01, 1.0]],
      [1, 0, [1.15893147e-01, 1.0]],
      [1, 0, [3.45578302e-02, 1.0]],
      [2, 0, [1.41613025e+04, 1.0]],
      [2, 0, [3.58429411e+03, 1.0]],
      [2, 0, [1.25385658e+03, 1.0]],
      [2, 0, [5.17561771e+02, 1.0]],
      [2, 0, [2.35591634e+02, 1.0]],
      [2, 0, [1.14150621e+02, 1.0]],
      [2, 0, [5.74581199e+01, 1.0]],
      [2, 0, [2.94435853e+01, 1.0]],
      [2, 0, [1.49162359e+01, 1.0]],
      [2, 0, [7.49983429e+00, 1.0]],
      [2, 0, [3.63156854e+00, 1.0]],
      [2, 0, [1.71598621e+00, 1.0]],
      [2, 0, [7.60663337e-01, 1.0]],
      [2, 0, [2.86675318e-01, 1.0]],
      [2, 0, [1.32356919e-01, 1.0]],
      [3, 0, [6.47900273e+02, 1.0]],
      [3, 0, [2.19493865e+02, 1.0]],
      [3, 0, [9.27081990e+01, 1.0]],
      [3, 0, [4.29375191e+01, 1.0]],
      [3, 0, [2.08572780e+01, 1.0]],
      [3, 0, [1.01812229e+01, 1.0]],
      [3, 0, [4.85784113e+00, 1.0]],
      [3, 0, [2.12743536e+00, 1.0]]]

Nh = [[0, 0, [5.26564207e+07, 1.0]],
      [0, 0, [1.39866817e+07, 1.0]],
      [0, 0, [4.74422515e+06, 1.0]],
      [0, 0, [1.77365178e+06, 1.0]],
      [0, 0, [7.21158475e+05, 1.0]],
      [0, 0, [3.07295461e+05, 1.0]],
      [0, 0, [1.36171416e+05, 1.0]],
      [0, 0, [6.18619002e+04, 1.0]],
      [0, 0, [2.87076767e+04, 1.0]],
      [0, 0, [1.35459317e+04, 1.0]],
      [0, 0, [6.49347521e+03, 1.0]],
      [0, 0, [3.16091577e+03, 1.0]],
      [0, 0, [1.56640294e+03, 1.0]],
      [0, 0, [7.82892664e+02, 1.0]],
      [0, 0, [4.03205829e+02, 1.0]],
      [0, 0, [2.12975482e+02, 1.0]],
      [0, 0, [1.13031738e+02, 1.0]],
      [0, 0, [6.15748622e+01, 1.0]],
      [0, 0, [3.08168254e+01, 1.0]],
      [0, 0, [1.74596207e+01, 1.0]],
      [0, 0, [8.40649632e+00, 1.0]],
      [0, 0, [4.56996553e+00, 1.0]],
      [0, 0, [1.83861196e+00, 1.0]],
      [0, 0, [8.93962268e-01, 1.0]],
      [0, 0, [2.57722155e-01, 1.0]],
      [0, 0, [9.51258475e-02, 1.0]],
      [0, 0, [1.83512150e-02, 1.0]],
      [1, 0, [4.22950772e+07, 1.0]],
      [1, 0, [1.14195506e+07, 1.0]],
      [1, 0, [3.38548142e+06, 1.0]],
      [1, 0, [1.07780748e+06, 1.0]],
      [1, 0, [3.62617690e+05, 1.0]],
      [1, 0, [1.27813994e+05, 1.0]],
      [1, 0, [4.70083150e+04, 1.0]],
      [1, 0, [1.80335129e+04, 1.0]],
      [1, 0, [7.23642488e+03, 1.0]],
      [1, 0, [3.04722041e+03, 1.0]],
      [1, 0, [1.34609780e+03, 1.0]],
      [1, 0, [6.21238577e+02, 1.0]],
      [1, 0, [2.97454758e+02, 1.0]],
      [1, 0, [1.44471995e+02, 1.0]],
      [1, 0, [7.28158737e+01, 1.0]],
      [1, 0, [3.55970861e+01, 1.0]],
      [1, 0, [1.84700824e+01, 1.0]],
      [1, 0, [8.78221533e+00, 1.0]],
      [1, 0, [4.36436593e+00, 1.0]],
      [1, 0, [1.80752239e+00, 1.0]],
      [1, 0, [7.63569979e-01, 1.0]],
      [1, 0, [2.23571621e-01, 1.0]],
      [1, 0, [5.49490045e-02, 1.0]],
      [1, 0, [8.20016250e-03, 1.0]],
      [2, 0, [7.11290116e+04, 1.0]],
      [2, 0, [1.66859908e+04, 1.0]],
      [2, 0, [5.39220281e+03, 1.0]],
      [2, 0, [2.08768842e+03, 1.0]],
      [2, 0, [9.08703746e+02, 1.0]],
      [2, 0, [4.28499285e+02, 1.0]],
      [2, 0, [2.13339353e+02, 1.0]],
      [2, 0, [1.10374338e+02, 1.0]],
      [2, 0, [5.81724949e+01, 1.0]],
      [2, 0, [3.06216577e+01, 1.0]],
      [2, 0, [1.63711945e+01, 1.0]],
      [2, 0, [8.52644077e+00, 1.0]],
      [2, 0, [4.33515906e+00, 1.0]],
      [2, 0, [2.10433759e+00, 1.0]],
      [2, 0, [9.58005214e-01, 1.0]],
      [2, 0, [3.99541705e-01, 1.0]],
      [2, 0, [1.49209297e-01, 1.0]],
      [2, 0, [5.23860200e-02, 1.0]],
      [3, 0, [1.29101515e+03, 1.0]],
      [3, 0, [4.29840881e+02, 1.0]],
      [3, 0, [1.79894990e+02, 1.0]],
      [3, 0, [8.33836896e+01, 1.0]],
      [3, 0, [4.05824687e+01, 1.0]],
      [3, 0, [2.02169215e+01, 1.0]],
      [3, 0, [9.96237638e+00, 1.0]],
      [3, 0, [4.75970910e+00, 1.0]],
      [3, 0, [2.14831388e+00, 1.0]],
      [3, 0, [8.42948503e-01, 1.0]]]

Fl = [[0, 0, [5.25502639e+07, 1.0]],
      [0, 0, [1.39595194e+07, 1.0]],
      [0, 0, [4.73540204e+06, 1.0]],
      [0, 0, [1.77067315e+06, 1.0]],
      [0, 0, [7.20542815e+05, 1.0]],
      [0, 0, [3.07470753e+05, 1.0]],
      [0, 0, [1.36535866e+05, 1.0]],
      [0, 0, [6.21737444e+04, 1.0]],
      [0, 0, [2.89220455e+04, 1.0]],
      [0, 0, [1.36768679e+04, 1.0]],
      [0, 0, [6.56867354e+03, 1.0]],
      [0, 0, [3.20255515e+03, 1.0]],
      [0, 0, [1.58931895e+03, 1.0]],
      [0, 0, [7.96091240e+02, 1.0]],
      [0, 0, [4.10483800e+02, 1.0]],
      [0, 0, [2.16908478e+02, 1.0]],
      [0, 0, [1.15644477e+02, 1.0]],
      [0, 0, [6.31209060e+01, 1.0]],
      [0, 0, [3.17161143e+01, 1.0]],
      [0, 0, [1.80024256e+01, 1.0]],
      [0, 0, [8.73072229e+00, 1.0]],
      [0, 0, [4.76363329e+00, 1.0]],
      [0, 0, [1.93814058e+00, 1.0]],
      [0, 0, [9.54669754e-01, 1.0]],
      [0, 0, [2.86374137e-01, 1.0]],
      [0, 0, [1.07521373e-01, 1.0]],
      [0, 0, [2.21291140e-02, 1.0]],
      [1, 0, [4.28248145e+07, 1.0]],
      [1, 0, [1.16539717e+07, 1.0]],
      [1, 0, [3.47376142e+06, 1.0]],
      [1, 0, [1.11041385e+06, 1.0]],
      [1, 0, [3.74684888e+05, 1.0]],
      [1, 0, [1.32320817e+05, 1.0]],
      [1, 0, [4.87108090e+04, 1.0]],
      [1, 0, [1.86846501e+04, 1.0]],
      [1, 0, [7.48927959e+03, 1.0]],
      [1, 0, [3.14764942e+03, 1.0]],
      [1, 0, [1.38726175e+03, 1.0]],
      [1, 0, [6.38739584e+02, 1.0]],
      [1, 0, [3.05193815e+02, 1.0]],
      [1, 0, [1.48004065e+02, 1.0]],
      [1, 0, [7.45062744e+01, 1.0]],
      [1, 0, [3.63892502e+01, 1.0]],
      [1, 0, [1.88902011e+01, 1.0]],
      [1, 0, [8.97228497e+00, 1.0]],
      [1, 0, [4.46826548e+00, 1.0]],
      [1, 0, [1.85130238e+00, 1.0]],
      [1, 0, [7.85413909e-01, 1.0]],
      [1, 0, [2.18861843e-01, 1.0]],
      [1, 0, [6.09423019e-02, 1.0]],
      [1, 0, [1.29943290e-02, 1.0]],
      [2, 0, [7.46077358e+04, 1.0]],
      [2, 0, [1.74711158e+04, 1.0]],
      [2, 0, [5.63264536e+03, 1.0]],
      [2, 0, [2.17573481e+03, 1.0]],
      [2, 0, [9.45254444e+02, 1.0]],
      [2, 0, [4.45092382e+02, 1.0]],
      [2, 0, [2.21401105e+02, 1.0]],
      [2, 0, [1.14479054e+02, 1.0]],
      [2, 0, [6.03368457e+01, 1.0]],
      [2, 0, [3.17718763e+01, 1.0]],
      [2, 0, [1.70024067e+01, 1.0]],
      [2, 0, [8.87166791e+00, 1.0]],
      [2, 0, [4.52530825e+00, 1.0]],
      [2, 0, [2.20707128e+00, 1.0]],
      [2, 0, [1.01493811e+00, 1.0]],
      [2, 0, [4.29553469e-01, 1.0]],
      [2, 0, [1.61789962e-01, 1.0]],
      [2, 0, [4.51817250e-02, 1.0]],
      [3, 0, [1.35864531e+03, 1.0]],
      [3, 0, [4.51803345e+02, 1.0]],
      [3, 0, [1.89101027e+02, 1.0]],
      [3, 0, [8.77131238e+01, 1.0]],
      [3, 0, [4.27786264e+01, 1.0]],
      [3, 0, [2.13716351e+01, 1.0]],
      [3, 0, [1.05894354e+01, 1.0]],
      [3, 0, [5.09611076e+00, 1.0]],
      [3, 0, [2.32511881e+00, 1.0]],
      [3, 0, [9.31317463e-01, 1.0]]]

Mc = [[0, 0, [5.27716144e+07, 1.0]],
      [0, 0, [1.40183307e+07, 1.0]],
      [0, 0, [4.75524325e+06, 1.0]],
      [0, 0, [1.77809769e+06, 1.0]],
      [0, 0, [7.23967370e+05, 1.0]],
      [0, 0, [3.09280864e+05, 1.0]],
      [0, 0, [1.37600124e+05, 1.0]],
      [0, 0, [6.28046699e+04, 1.0]],
      [0, 0, [2.92903043e+04, 1.0]],
      [0, 0, [1.38838077e+04, 1.0]],
      [0, 0, [6.68134579e+03, 1.0]],
      [0, 0, [3.26243412e+03, 1.0]],
      [0, 0, [1.62105900e+03, 1.0]],
      [0, 0, [8.13470672e+02, 1.0]],
      [0, 0, [4.19940640e+02, 1.0]],
      [0, 0, [2.22127358e+02, 1.0]],
      [0, 0, [1.18708746e+02, 1.0]],
      [0, 0, [6.48872165e+01, 1.0]],
      [0, 0, [3.27179656e+01, 1.0]],
      [0, 0, [1.85990233e+01, 1.0]],
      [0, 0, [9.06637472e+00, 1.0]],
      [0, 0, [4.96566370e+00, 1.0]],
      [0, 0, [2.04003435e+00, 1.0]],
      [0, 0, [1.01772027e+00, 1.0]],
      [0, 0, [3.16405429e-01, 1.0]],
      [0, 0, [1.20026752e-01, 1.0]],
      [0, 0, [2.57013710e-02, 1.0]],
      [1, 0, [4.34918164e+07, 1.0]],
      [1, 0, [1.19496194e+07, 1.0]],
      [1, 0, [3.58591482e+06, 1.0]],
      [1, 0, [1.15223503e+06, 1.0]],
      [1, 0, [3.90329429e+05, 1.0]],
      [1, 0, [1.38232687e+05, 1.0]],
      [1, 0, [5.09730225e+04, 1.0]],
      [1, 0, [1.95623370e+04, 1.0]],
      [1, 0, [7.83558761e+03, 1.0]],
      [1, 0, [3.28762007e+03, 1.0]],
      [1, 0, [1.44574490e+03, 1.0]],
      [1, 0, [6.64103144e+02, 1.0]],
      [1, 0, [3.16618543e+02, 1.0]],
      [1, 0, [1.53305357e+02, 1.0]],
      [1, 0, [7.70867381e+01, 1.0]],
      [1, 0, [3.76643092e+01, 1.0]],
      [1, 0, [1.95553550e+01, 1.0]],
      [1, 0, [9.30698171e+00, 1.0]],
      [1, 0, [4.65188068e+00, 1.0]],
      [1, 0, [1.93992868e+00, 1.0]],
      [1, 0, [8.36628356e-01, 1.0]],
      [1, 0, [2.41153965e-01, 1.0]],
      [1, 0, [7.02069403e-02, 1.0]],
      [1, 0, [1.68164390e-02, 1.0]],
      [2, 0, [7.89160022e+04, 1.0]],
      [2, 0, [1.84439609e+04, 1.0]],
      [2, 0, [5.93073054e+03, 1.0]],
      [2, 0, [2.28492725e+03, 1.0]],
      [2, 0, [9.90628169e+02, 1.0]],
      [2, 0, [4.65712715e+02, 1.0]],
      [2, 0, [2.31443508e+02, 1.0]],
      [2, 0, [1.19601763e+02, 1.0]],
      [2, 0, [6.30468393e+01, 1.0]],
      [2, 0, [3.32158051e+01, 1.0]],
      [2, 0, [1.77974545e+01, 1.0]],
      [2, 0, [9.31155755e+00, 1.0]],
      [2, 0, [4.76967207e+00, 1.0]],
      [2, 0, [2.34357067e+00, 1.0]],
      [2, 0, [1.08804086e+00, 1.0]],
      [2, 0, [4.66643181e-01, 1.0]],
      [2, 0, [1.76404306e-01, 1.0]],
      [2, 0, [3.98283780e-02, 1.0]],
      [3, 0, [1.42553854e+03, 1.0]],
      [3, 0, [4.73475164e+02, 1.0]],
      [3, 0, [1.98163904e+02, 1.0]],
      [3, 0, [9.19700000e+01, 1.0]],
      [3, 0, [4.49363423e+01, 1.0]],
      [3, 0, [2.25058281e+01, 1.0]],
      [3, 0, [1.12049794e+01, 1.0]],
      [3, 0, [5.42672660e+00, 1.0]],
      [3, 0, [2.49940949e+00, 1.0]],
      [3, 0, [1.01858987e+00, 1.0]]]

Lv = [[0, 0, [5.23395635e+07, 1.0]],
      [0, 0, [1.39042308e+07, 1.0]],
      [0, 0, [4.71663685e+06, 1.0]],
      [0, 0, [1.76377617e+06, 1.0]],
      [0, 0, [7.18589017e+05, 1.0]],
      [0, 0, [3.07335338e+05, 1.0]],
      [0, 0, [1.36970937e+05, 1.0]],
      [0, 0, [6.26396175e+04, 1.0]],
      [0, 0, [2.92750025e+04, 1.0]],
      [0, 0, [1.39042228e+04, 1.0]],
      [0, 0, [6.70288368e+03, 1.0]],
      [0, 0, [3.27766374e+03, 1.0]],
      [0, 0, [1.63087763e+03, 1.0]],
      [0, 0, [8.20370801e+02, 1.0]],
      [0, 0, [4.24028714e+02, 1.0]],
      [0, 0, [2.24257275e+02, 1.0]],
      [0, 0, [1.20648388e+02, 1.0]],
      [0, 0, [6.61447606e+01, 1.0]],
      [0, 0, [3.35432929e+01, 1.0]],
      [0, 0, [1.91090728e+01, 1.0]],
      [0, 0, [9.37631555e+00, 1.0]],
      [0, 0, [5.15488955e+00, 1.0]],
      [0, 0, [2.14108026e+00, 1.0]],
      [0, 0, [1.08035187e+00, 1.0]],
      [0, 0, [3.47368104e-01, 1.0]],
      [0, 0, [1.32716064e-01, 1.0]],
      [0, 0, [2.90857190e-02, 1.0]],
      [1, 0, [4.41478101e+07, 1.0]],
      [1, 0, [1.22444768e+07, 1.0]],
      [1, 0, [3.69876288e+06, 1.0]],
      [1, 0, [1.19459836e+06, 1.0]],
      [1, 0, [4.06258118e+05, 1.0]],
      [1, 0, [1.44274530e+05, 1.0]],
      [1, 0, [5.32907112e+04, 1.0]],
      [1, 0, [2.04626398e+04, 1.0]],
      [1, 0, [8.19071685e+03, 1.0]],
      [1, 0, [3.43086834e+03, 1.0]],
      [1, 0, [1.50538611e+03, 1.0]],
      [1, 0, [6.89856467e+02, 1.0]],
      [1, 0, [3.28166812e+02, 1.0]],
      [1, 0, [1.58651064e+02, 1.0]],
      [1, 0, [7.96804421e+01, 1.0]],
      [1, 0, [3.89427462e+01, 1.0]],
      [1, 0, [2.02214028e+01, 1.0]],
      [1, 0, [9.64080707e+00, 1.0]],
      [1, 0, [4.83575422e+00, 1.0]],
      [1, 0, [2.02868007e+00, 1.0]],
      [1, 0, [8.89231478e-01, 1.0]],
      [1, 0, [2.66801723e-01, 1.0]],
      [1, 0, [7.97707589e-02, 1.0]],
      [1, 0, [2.02401690e-02, 1.0]],
      [2, 0, [8.35091124e+04, 1.0]],
      [2, 0, [1.94803968e+04, 1.0]],
      [2, 0, [6.24777184e+03, 1.0]],
      [2, 0, [2.40079648e+03, 1.0]],
      [2, 0, [1.03866895e+03, 1.0]],
      [2, 0, [4.87503404e+02, 1.0]],
      [2, 0, [2.42041207e+02, 1.0]],
      [2, 0, [1.24999999e+02, 1.0]],
      [2, 0, [6.58998208e+01, 1.0]],
      [2, 0, [3.47378902e+01, 1.0]],
      [2, 0, [1.86358336e+01, 1.0]],
      [2, 0, [9.77729485e+00, 1.0]],
      [2, 0, [5.02893584e+00, 1.0]],
      [2, 0, [2.48933959e+00, 1.0]],
      [2, 0, [1.16591219e+00, 1.0]],
      [2, 0, [5.05800523e-01, 1.0]],
      [2, 0, [1.91490581e-01, 1.0]],
      [2, 0, [6.21776100e-02, 1.0]],
      [3, 0, [1.49214833e+03, 1.0]],
      [3, 0, [4.95004732e+02, 1.0]],
      [3, 0, [2.07147264e+02, 1.0]],
      [3, 0, [9.61846084e+01, 1.0]],
      [3, 0, [4.70708711e+01, 1.0]],
      [3, 0, [2.36274420e+01, 1.0]],
      [3, 0, [1.18132991e+01, 1.0]],
      [3, 0, [5.75382195e+00, 1.0]],
      [3, 0, [2.67230702e+00, 1.0]],
      [3, 0, [1.10531308e+00, 1.0]]]

Ts = [[0, 0, [5.25589660e+07, 1.0]],
      [0, 0, [1.39622807e+07, 1.0]],
      [0, 0, [4.73580777e+06, 1.0]],
      [0, 0, [1.77055155e+06, 1.0]],
      [0, 0, [7.21438834e+05, 1.0]],
      [0, 0, [3.08728868e+05, 1.0]],
      [0, 0, [1.37769819e+05, 1.0]],
      [0, 0, [6.31102557e+04, 1.0]],
      [0, 0, [2.95481902e+04, 1.0]],
      [0, 0, [1.40571584e+04, 1.0]],
      [0, 0, [6.78665037e+03, 1.0]],
      [0, 0, [3.32284017e+03, 1.0]],
      [0, 0, [1.65532803e+03, 1.0]],
      [0, 0, [8.34251093e+02, 1.0]],
      [0, 0, [4.31714084e+02, 1.0]],
      [0, 0, [2.28450220e+02, 1.0]],
      [0, 0, [1.23329160e+02, 1.0]],
      [0, 0, [6.77383376e+01, 1.0]],
      [0, 0, [3.44769397e+01, 1.0]],
      [0, 0, [1.96687234e+01, 1.0]],
      [0, 0, [9.71913239e+00, 1.0]],
      [0, 0, [5.35865868e+00, 1.0]],
      [0, 0, [2.24658072e+00, 1.0]],
      [0, 0, [1.14580781e+00, 1.0]],
      [0, 0, [3.79419796e-01, 1.0]],
      [0, 0, [1.45755955e-01, 1.0]],
      [0, 0, [3.23551340e-02, 1.0]],
      [1, 0, [4.50079474e+07, 1.0]],
      [1, 0, [1.25887001e+07, 1.0]],
      [1, 0, [3.82527515e+06, 1.0]],
      [1, 0, [1.24107592e+06, 1.0]],
      [1, 0, [4.23497767e+05, 1.0]],
      [1, 0, [1.50750108e+05, 1.0]],
      [1, 0, [5.57553472e+04, 1.0]],
      [1, 0, [2.14132811e+04, 1.0]],
      [1, 0, [8.56301977e+03, 1.0]],
      [1, 0, [3.57984955e+03, 1.0]],
      [1, 0, [1.56686094e+03, 1.0]],
      [1, 0, [7.16151398e+02, 1.0]],
      [1, 0, [3.39848795e+02, 1.0]],
      [1, 0, [1.64014961e+02, 1.0]],
      [1, 0, [8.22626489e+01, 1.0]],
      [1, 0, [4.02034342e+01, 1.0]],
      [1, 0, [2.08756913e+01, 1.0]],
      [1, 0, [9.96460376e+00, 1.0]],
      [1, 0, [5.01432190e+00, 1.0]],
      [1, 0, [2.11448402e+00, 1.0]],
      [1, 0, [9.40561230e-01, 1.0]],
      [1, 0, [2.93186959e-01, 1.0]],
      [1, 0, [8.92968449e-02, 1.0]],
      [1, 0, [2.34310560e-02, 1.0]],
      [2, 0, [8.83574648e+04, 1.0]],
      [2, 0, [2.05729720e+04, 1.0]],
      [2, 0, [6.58130697e+03, 1.0]],
      [2, 0, [2.52237884e+03, 1.0]],
      [2, 0, [1.08894908e+03, 1.0]],
      [2, 0, [5.10258847e+02, 1.0]],
      [2, 0, [2.53088086e+02, 1.0]],
      [2, 0, [1.30617020e+02, 1.0]],
      [2, 0, [6.88645118e+01, 1.0]],
      [2, 0, [3.63215132e+01, 1.0]],
      [2, 0, [1.95080977e+01, 1.0]],
      [2, 0, [1.02633964e+01, 1.0]],
      [2, 0, [5.29989718e+00, 1.0]],
      [2, 0, [2.64229268e+00, 1.0]],
      [2, 0, [1.24755275e+00, 1.0]],
      [2, 0, [5.46674430e-01, 1.0]],
      [2, 0, [2.07169726e-01, 1.0]],
      [2, 0, [8.19822830e-02, 1.0]],
      [3, 0, [1.55874574e+03, 1.0]],
      [3, 0, [5.16482517e+02, 1.0]],
      [3, 0, [2.16090441e+02, 1.0]],
      [3, 0, [1.00375884e+02, 1.0]],
      [3, 0, [4.91918626e+01, 1.0]],
      [3, 0, [2.47415547e+01, 1.0]],
      [3, 0, [1.24172009e+01, 1.0]],
      [3, 0, [6.07887380e+00, 1.0]],
      [3, 0, [2.84454763e+00, 1.0]],
      [3, 0, [1.19184727e+00, 1.0]]]

Og = [[0, 0, [5.24543434e+07, 1.0]],
      [0, 0, [1.39342753e+07, 1.0]],
      [0, 0, [4.72586931e+06, 1.0]],
      [0, 0, [1.76655977e+06, 1.0]],
      [0, 0, [7.20004310e+05, 1.0]],
      [0, 0, [3.08342403e+05, 1.0]],
      [0, 0, [1.37804335e+05, 1.0]],
      [0, 0, [6.32550050e+04, 1.0]],
      [0, 0, [2.96875536e+04, 1.0]],
      [0, 0, [1.41569741e+04, 1.0]],
      [0, 0, [6.84986815e+03, 1.0]],
      [0, 0, [3.36021593e+03, 1.0]],
      [0, 0, [1.67684236e+03, 1.0]],
      [0, 0, [8.47152669e+02, 1.0]],
      [0, 0, [4.39087321e+02, 1.0]],
      [0, 0, [2.32521218e+02, 1.0]],
      [0, 0, [1.25971789e+02, 1.0]],
      [0, 0, [6.93405284e+01, 1.0]],
      [0, 0, [3.54206791e+01, 1.0]],
      [0, 0, [2.02418111e+01, 1.0]],
      [0, 0, [1.00527536e+01, 1.0]],
      [0, 0, [5.56202638e+00, 1.0]],
      [0, 0, [2.35158724e+00, 1.0]],
      [0, 0, [1.21100065e+00, 1.0]],
      [0, 0, [4.12227893e-01, 1.0]],
      [0, 0, [1.59031719e-01, 1.0]],
      [0, 0, [4.35402970e-02, 1.0]],
      [1, 0, [4.56226785e+07, 1.0]],
      [1, 0, [1.28729373e+07, 1.0]],
      [1, 0, [3.93590515e+06, 1.0]],
      [1, 0, [1.28310998e+06, 1.0]],
      [1, 0, [4.39439154e+05, 1.0]],
      [1, 0, [1.56830707e+05, 1.0]],
      [1, 0, [5.80943914e+04, 1.0]],
      [1, 0, [2.23217925e+04, 1.0]],
      [1, 0, [8.92013696e+03, 1.0]],
      [1, 0, [3.72283531e+03, 1.0]],
      [1, 0, [1.62573603e+03, 1.0]],
      [1, 0, [7.41242740e+02, 1.0]],
      [1, 0, [3.50951511e+02, 1.0]],
      [1, 0, [1.69104303e+02, 1.0]],
      [1, 0, [8.47054959e+01, 1.0]],
      [1, 0, [4.13905180e+01, 1.0]],
      [1, 0, [2.14933937e+01, 1.0]],
      [1, 0, [1.02679675e+01, 1.0]],
      [1, 0, [5.18229819e+00, 1.0]],
      [1, 0, [2.19602479e+00, 1.0]],
      [1, 0, [9.89384744e-01, 1.0]],
      [1, 0, [3.19468837e-01, 1.0]],
      [1, 0, [9.86706817e-02, 1.0]],
      [1, 0, [2.70144270e-02, 1.0]],
      [2, 0, [9.34050542e+04, 1.0]],
      [2, 0, [2.17092956e+04, 1.0]],
      [2, 0, [6.92742965e+03, 1.0]],
      [2, 0, [2.64817876e+03, 1.0]],
      [2, 0, [1.14081718e+03, 1.0]],
      [2, 0, [5.33672197e+02, 1.0]],
      [2, 0, [2.64430103e+02, 1.0]],
      [2, 0, [1.36373690e+02, 1.0]],
      [2, 0, [7.18993191e+01, 1.0]],
      [2, 0, [3.79447993e+01, 1.0]],
      [2, 0, [2.04023819e+01, 1.0]],
      [2, 0, [1.07633686e+01, 1.0]],
      [2, 0, [5.57883329e+00, 1.0]],
      [2, 0, [2.80009817e+00, 1.0]],
      [2, 0, [1.33179571e+00, 1.0]],
      [2, 0, [5.88745938e-01, 1.0]],
      [2, 0, [2.23304938e-01, 1.0]],
      [2, 0, [8.78787900e-02, 1.0]],
      [3, 0, [1.62557486e+03, 1.0]],
      [3, 0, [5.37987566e+02, 1.0]],
      [3, 0, [2.25026762e+02, 1.0]],
      [3, 0, [1.04559255e+02, 1.0]],
      [3, 0, [5.13069042e+01, 1.0]],
      [3, 0, [2.58520536e+01, 1.0]],
      [3, 0, [1.30188065e+01, 1.0]],
      [3, 0, [6.40298978e+00, 1.0]],
      [3, 0, [3.01668444e+00, 1.0]],
      [3, 0, [1.27846072e+00, 1.0]]]
