How to use the pyscf.ao2mo function in pyscf

def solve_ERI( OEI, TEI, DMguess, numPairs ):

    mol = gto.Mole()
    mol.build(verbose=3)
    mol.atom.append(('C', (0, 0, 0)))
    mol.nelectron = 2 * numPairs

    L = OEI.shape[0]
    mf = scf.RHF( mol )
    mf.get_hcore = lambda *args: OEI
    mf.get_ovlp = lambda *args: np.eye( L )
    mf._eri = ao2mo.restore(8, TEI, L)
    
    mf.scf( DMguess )
    DMloc = np.dot(np.dot( mf.mo_coeff, np.diag( mf.mo_occ )), mf.mo_coeff.T )
    if ( mf.converged == False ):
        mf = rhf_newtonraphson.solve( mf, dm_guess=DMloc )
        DMloc = np.dot(np.dot( mf.mo_coeff, np.diag( mf.mo_occ )), mf.mo_coeff.T )
    return DMloc

from functools import reduce
import numpy
from pyscf import gto, lo, fci, ao2mo, scf, lib

mol = gto.M(atom=[('H', 0, 0, i*1.8) for i in range(10)],
            basis = 'sto6g', unit='B')
s = mol.intor('cint1e_ovlp_sph')
orb = lo.lowdin(s)
#mf = scf.RHF(mol).run()
#orb = mf.mo_coeff

h1 = mol.intor('cint1e_nuc_sph')
h1+= mol.intor('cint1e_kin_sph')
h1 = reduce(numpy.dot, (orb.T, h1, orb))
h2 = ao2mo.kernel(mol, orb)

e, ci = fci.direct_spin0.kernel(h1, h2, 10, 10, ecore=mol.energy_nuc(),
                                max_cycle=500, max_space=100, verbose=5)
print(e)

e, ci = fci.direct_spin0.kernel(h1, h2, 10, 10, ecore=mol.energy_nuc(), ci0=ci,
                                max_cycle=500, max_space=100, verbose=5)
print(e)

e, ci = fci.direct_spin0.kernel(h1, h2, 10, 10, ecore=mol.energy_nuc(), ci0=ci,
                                max_cycle=500, max_space=100, verbose=5)
print(e)


#
# Reducing OMP threads can improve the numerical stability

tau = numpy.ndarray((nocc*(nocc+1)//2,nvir,nvir), buffer=outbuf)
            p0 = 0
            for i in range(nocc):
                tau[p0:p0+i+1] = t2[i,:i+1]
                p0 += i + 1
            tau = _ao2mo.nr_e2(tau.reshape(-1,nvir**2), aos, (0,nao,0,nao), 's1', 's1')
            tau = tau.reshape(-1,nao,nao)
            time0 = logger.timer_debug1(self, 'vvvv-tau', *time0)

            ao2mopt = _ao2mo.AO2MOpt(mol, 'cint2e_sph', 'CVHFnr_schwarz_cond',
                                     'CVHFsetnr_direct_scf')
            outbuf[:] = 0
            ao_loc = mol.ao_loc_nr()
            max_memory = max(0, self.max_memory - lib.current_memory()[0])
            dmax = max(4, int(numpy.sqrt(max_memory*.95e6/8/nao**2/2)))
            sh_ranges = ao2mo.outcore.balance_partition(ao_loc, dmax)
            dmax = max(x[2] for x in sh_ranges)
            eribuf = numpy.empty((dmax,dmax,nao,nao))
            loadbuf = numpy.empty((dmax,dmax,nao,nao))
            fint = gto.moleintor.getints2e

            for ip, (ish0, ish1, ni) in enumerate(sh_ranges):
                for jsh0, jsh1, nj in sh_ranges[:ip]:
                    eri = fint('cint2e_sph', mol._atm, mol._bas, mol._env,
                               shls_slice=(ish0,ish1,jsh0,jsh1), aosym='s2kl',
                               ao_loc=ao_loc, cintopt=ao2mopt._cintopt, out=eribuf)
                    i0, i1 = ao_loc[ish0], ao_loc[ish1]
                    j0, j1 = ao_loc[jsh0], ao_loc[jsh1]
                    tmp = numpy.ndarray((i1-i0,nao,j1-j0,nao), buffer=loadbuf)
                    _ccsd.libcc.CCload_eri(tmp.ctypes.data_as(ctypes.c_void_p),
                                           eri.ctypes.data_as(ctypes.c_void_p),
                                           (ctypes.c_int*4)(i0, i1, j0, j1),

def ao2mo(self, mo_coeff=None):
        if mo_coeff is None:
            mo_coeff = (self.mo_coeff[0][:,self.ncore[0]:self.ncore[0]+self.ncas],
                        self.mo_coeff[1][:,self.ncore[1]:self.ncore[1]+self.ncas])
        nao, nmo = mo_coeff[0].shape
        if self._scf._eri is not None and \
           (nao*nao*nmo*nmo*12+self._scf._eri.size)*8/1e6 < self.max_memory*.95:
            moab = numpy.hstack((mo_coeff[0], mo_coeff[1]))
            na = mo_coeff[0].shape[1]
            nab = moab.shape[1]
            eri = pyscf.ao2mo.incore.full(self._scf._eri, moab)
            eri = pyscf.ao2mo.restore(1, eri, nab)
            eri_aa = eri[:na,:na,:na,:na].copy()
            eri_ab = eri[:na,:na,na:,na:].copy()
            eri_bb = eri[na:,na:,na:,na:].copy()
        else:
            moab = numpy.hstack((mo_coeff[0], mo_coeff[1]))
            eri = pyscf.ao2mo.full(self.mol, mo_coeff, verbose=self.verbose)
            na = mo_coeff[0].shape[1]
            nab = moab.shape[1]
            eri = pyscf.ao2mo.restore(1, eri, nab)
            eri_aa = eri[:na,:na,:na,:na].copy()
            eri_ab = eri[:na,:na,na:,na:].copy()
            eri_bb = eri[na:,na:,na:,na:].copy()

        return (eri_aa, eri_ab, eri_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:])

atom = 'O 0 0 0; O 0 0 1.2',
    basis = 'ccpvdz',
    spin = 2)
myhf = scf.RHF(mol)
myhf.kernel()
mycas = mcscf.CASSCF(myhf, 6, 8) # 6 orbital, 8 electron
mycas.kernel()

# CAS space orbitals
cas_orb = mycas.mo_coeff[:,mycas.ncore:mycas.ncore+mycas.ncas]

# 2-electron part spin-same-orbit coupling
#       [ijkl] = 
# JCP, 122, 034107 Eq (3) = h2 * (-i)
# For simplicty, we didn't consider the permutation symmetry k >= l, therefore aosym='s1'
h2 = ao2mo.kernel(mol, cas_orb, intor='int2e_p1vxp1_sph', comp=3, aosym='s1')
print('SSO 2e integrals shape %s' % str(h2.shape))

# 1-electron part for atom A
#       
# JCP, 122, 034107 Eq (2) = h1 * (iZ_A)
mol.set_rinv_origin(mol.atom_coord(1))  # set the gauge origin on second atom
h1 = numpy.einsum('xpq,pi,qj->xij', mol.intor('int1e_prinvxp_sph'), cas_orb, cas_orb)
print('1e integral shape %s' % str(h1.shape))

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

['H', ( 1.,-0.5   ,-1.   )],
        ['H', ( 0.,-0.    ,-1.   )],
        ['H', ( 1.,-0.5   , 0.   )],
        ['H', ( 0., 1.    , 1.   )],
        ['H', ( 1., 2.    , 3.   )],
        ['H', ( 1., 2.    , 4.   )],
    ]
    mol.basis = 'sto-3g'
    mol.build()

    m = scf.RHF(mol)
    m.kernel()
    norb = m.mo_coeff.shape[1]
    nelec = mol.nelectron
    h1e = reduce(numpy.dot, (m.mo_coeff.T, m.get_hcore(), m.mo_coeff))
    eri = ao2mo.kernel(m._eri, m.mo_coeff, compact=False)
    eri = eri.reshape(norb,norb,norb,norb)

    e1, c1 = kernel(h1e, eri, norb, nelec)
    e2, c2 = direct_spin0.kernel(h1e, eri, norb, nelec)
    print(e1, e1 - -11.894559902235565, 'diff to FCI', e1-e2)

def _get_vvvv(eris):
    if eris.vvvv is None and getattr(eris, 'vvL', None) is not None:  # DF eris
        vvL = np.asarray(eris.vvL)
        nvir = int(np.sqrt(eris.vvL.shape[0]*2))
        return ao2mo.restore(1, lib.dot(vvL, vvL.T), nvir)
    elif len(eris.vvvv.shape) == 2:  # DO not use .ndim here for h5py library
                                     # backward compatbility
        nvir = int(np.sqrt(eris.vvvv.shape[0]*2))
        return ao2mo.restore(1, np.asarray(eris.vvvv), nvir)
    else:
        return eris.vvvv

def make_hdiag(h1e, eri, strs, norb, nelec):
    eri = ao2mo.restore(1, eri, norb)
    diagj = numpy.einsum('iijj->ij',eri)
    diagk = numpy.einsum('ijji->ij',eri)

    ndet = len(strs)
    hdiag = numpy.zeros(ndet)
    for idet, (stra, strb) in enumerate(strs.reshape(ndet,2,-1)):
        aocc = str2orblst(stra, norb)[0]
        bocc = str2orblst(strb, norb)[0]
        e1 = h1e[aocc,aocc].sum() + h1e[bocc,bocc].sum()
        e2 = diagj[aocc][:,aocc].sum() + diagj[aocc][:,bocc].sum() \
           + diagj[bocc][:,aocc].sum() + diagj[bocc][:,bocc].sum() \
           - diagk[aocc][:,aocc].sum() - diagk[bocc][:,bocc].sum()
        hdiag[idet] = e1 + e2*.5
    return hdiag