text: #!/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.
  
  '''
  MO integrals for UCASSCF methods
  '''
  
  import sys
  import ctypes
  import time
  import tempfile
  import numpy
  import h5py
  from pyscf import lib
  from pyscf.lib import logger
  from pyscf import ao2mo
  from pyscf.ao2mo import _ao2mo
  
  libmcscf = lib.load_library('libmcscf')
  
  # least memory requirements:
  #       ncore**2*(nmo-ncore)**2 + ncas**2*nmo**2*2 + nmo**3   words
  # nmo  ncore  ncas  outcore  incore
  # 200  40     16    2.4GB    5.3 GB (_eri 1.6GB )
  # 250  50     16    4.9GB   12.0 GB (_eri 3.9GB )
  # 300  60     16    9.0GB   23.7 GB (_eri 8.1GB )
  # 400  80     16   24.6GB
  # 500  100    16   54.8GB
  # 600  120    16   107 GB
  
  
  def trans_e1_incore(eri_ao, mo, ncore, ncas):
      nmo = mo[0].shape[1]
      nocc = (ncore[0] + ncas, ncore[1] + ncas)
  
      erib = ao2mo.incore.half_e1(eri_ao, (mo[1][:,:nocc[1]],mo[1]),
                                  compact=False)
      load_buf = lambda bufid: erib[bufid*nmo:bufid*nmo+nmo].copy()
      AAPP, AApp, APPA, tmp, IAPCV, APcv = \
              _trans_aapp_((mo[1],mo[0]), (ncore[1],ncore[0]), ncas, load_buf)
      jC_PP, jC_pp, kC_PP, ICVCV = \
              _trans_cvcv_((mo[1],mo[0]), (ncore[1],ncore[0]), ncas, load_buf)[:4]
      erib = tmp = None
  
      eria = ao2mo.incore.half_e1(eri_ao, (mo[0][:,:nocc[0]],mo[0]),
                                  compact=False)
      load_buf = lambda bufid: eria[bufid*nmo:bufid*nmo+nmo].copy()
      aapp, aaPP, appa, apPA, Iapcv, apCV = \
              _trans_aapp_(mo, ncore, ncas, load_buf)
      jc_pp, jc_PP, kc_pp, Icvcv, cvCV = \
              _trans_cvcv_(mo, ncore, ncas, load_buf)
  
      jkcpp = jc_pp - kc_pp
      jkcPP = jC_PP - kC_PP
      return jkcpp, jkcPP, jC_pp, jc_PP, \
              aapp, aaPP, AApp, AAPP, appa, apPA, APPA, \
              Iapcv, IAPCV, apCV, APcv, Icvcv, ICVCV, cvCV
  
  
  def trans_e1_outcore(mol, mo, ncore, ncas,
                       max_memory=None, ioblk_size=512, verbose=logger.WARN):
      time0 = (time.clock(), time.time())
      if isinstance(verbose, logger.Logger):
          log = verbose
      else:
          log = logger.Logger(mol.stdout, verbose)
      nao, nmo = mo[0].shape
      nao_pair = nao*(nao+1)//2
      nocc = (ncore[0] + ncas, ncore[1] + ncas)
  
      swapfile = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
      ao2mo.outcore.half_e1(mol, (mo[1][:,:nocc[1]],mo[1]), swapfile.name,
                            verbose=log, compact=False)
  
      fswap = h5py.File(swapfile.name, 'r')
      klaoblks = len(fswap['0'])
      def load_bufa(bfn_id):
          if log.verbose >= logger.DEBUG1:
              time1[:] = log.timer('between load_buf', *tuple(time1))
          buf = numpy.empty((nmo,nao_pair))
          col0 = 0
          for ic in range(klaoblks):
              dat = fswap['0/%d'%ic]
              col1 = col0 + dat.shape[1]
              buf[:nmo,col0:col1] = dat[bfn_id*nmo:(bfn_id+1)*nmo]
              col0 = col1
          if log.verbose >= logger.DEBUG1:
              time1[:] = log.timer('load_buf', *tuple(time1))
          return buf
      time0 = log.timer('halfe1-beta', *time0)
      time1 = [time.clock(), time.time()]
      ao_loc = numpy.array(mol.ao_loc_nr(), dtype=numpy.int32)
      AAPP, AApp, APPA, tmp, IAPCV, APcv = \
              _trans_aapp_((mo[1],mo[0]), (ncore[1],ncore[0]), ncas, load_bufa,
                           ao_loc)
      time0 = log.timer('trans_AAPP', *time0)
      jC_PP, jC_pp, kC_PP, ICVCV = \
              _trans_cvcv_((mo[1],mo[0]), (ncore[1],ncore[0]), ncas, load_bufa,
                           ao_loc)[:4]
      time0 = log.timer('trans_CVCV', *time0)
      tmp = None
      fswap.close()
  
      ###########################
  
      swapfile = tempfile.NamedTemporaryFile(dir=lib.param.TMPDIR)
      ao2mo.outcore.half_e1(mol, (mo[0][:,:nocc[0]],mo[0]), swapfile.name,
                            verbose=log, compact=False)
  
      fswap = h5py.File(swapfile.name, 'r')
      klaoblks = len(fswap['0'])
      def load_bufb(bfn_id):
          if log.verbose >= logger.DEBUG1:
              time1[:] = log.timer('between load_buf', *tuple(time1))
          buf = numpy.empty((nmo,nao_pair))
          col0 = 0
          for ic in range(klaoblks):
              dat = fswap['0/%d'%ic]
              col1 = col0 + dat.shape[1]
              buf[:nmo,col0:col1] = dat[bfn_id*nmo:(bfn_id+1)*nmo]
              col0 = col1
          if log.verbose >= logger.DEBUG1:
              time1[:] = log.timer('load_buf', *tuple(time1))
          return buf
      time0 = log.timer('halfe1-alpha', *time0)
      time1 = [time.clock(), time.time()]
      aapp, aaPP, appa, apPA, Iapcv, apCV = \
              _trans_aapp_(mo, ncore, ncas, load_bufb, ao_loc)
      time0 = log.timer('trans_aapp', *time0)
      jc_pp, jc_PP, kc_pp, Icvcv, cvCV = \
              _trans_cvcv_(mo, ncore, ncas, load_bufb, ao_loc)
      time0 = log.timer('trans_cvcv', *time0)
      fswap.close()
  
      jkcpp = jc_pp - kc_pp
      jkcPP = jC_PP - kC_PP
      return jkcpp, jkcPP, jC_pp, jc_PP, \
              aapp, aaPP, AApp, AAPP, appa, apPA, APPA, \
              Iapcv, IAPCV, apCV, APcv, Icvcv, ICVCV, cvCV
  
  
  def _trans_aapp_(mo, ncore, ncas, fload, ao_loc=None):
      nmo = mo[0].shape[1]
      nocc = (ncore[0] + ncas, ncore[1] + ncas)
      c_nmo = ctypes.c_int(nmo)
  
      klshape = (0, nmo, 0, nmo)
  
      japcv = numpy.empty((ncas,nmo,ncore[0],nmo-ncore[0]))
      aapp = numpy.empty((ncas,ncas,nmo,nmo))
      aaPP = numpy.empty((ncas,ncas,nmo,nmo))
      appa = numpy.empty((ncas,nmo,nmo,ncas))
      apPA = numpy.empty((ncas,nmo,nmo,ncas))
      apCV = numpy.empty((ncas,nmo,ncore[1],nmo-ncore[1]))
      ppp = numpy.empty((nmo,nmo,nmo))
      for i in range(ncas):
          buf = _ao2mo.nr_e2(fload(ncore[0]+i), mo[0], klshape,
                              aosym='s4', mosym='s2', ao_loc=ao_loc)
          lib.unpack_tril(buf, out=ppp)
          aapp[i] = ppp[ncore[0]:nocc[0]]
          appa[i] = ppp[:,:,ncore[0]:nocc[0]]
          #japcp = avcp * 2 - acpv.transpose(0,2,1,3) - avcp.transpose(0,3,2,1)
          japcv[i] = ppp[:,:ncore[0],ncore[0]:] * 2 \
                   - ppp[:ncore[0],:,ncore[0]:].transpose(1,0,2) \
                   - ppp[ncore[0]:,:ncore[0],:].transpose(2,1,0)
  
          buf = _ao2mo.nr_e2(fload(ncore[0]+i), mo[1], klshape,
                              aosym='s4', mosym='s2', ao_loc=ao_loc)
          lib.unpack_tril(buf, out=ppp)
          aaPP[i] = ppp[ncore[0]:nocc[0]]
          apPA[i] = ppp[:,:,ncore[1]:nocc[1]]
          apCV[i] = ppp[:,:ncore[1],ncore[1]:]
  
      return aapp, aaPP, appa, apPA, japcv, apCV
  
  def _trans_cvcv_(mo, ncore, ncas, fload, ao_loc=None):
      nmo = mo[0].shape[1]
      c_nmo = ctypes.c_int(nmo)
  
      jc_pp = numpy.empty((ncore[0],nmo,nmo))
      jc_PP = numpy.zeros((nmo,nmo))
      kc_pp = numpy.empty((ncore[0],nmo,nmo))
      jcvcv = numpy.zeros((ncore[0],nmo-ncore[0],ncore[0],nmo-ncore[0]))
      cvCV = numpy.empty((ncore[0],nmo-ncore[0],ncore[1],nmo-ncore[1]))
      vcp = numpy.empty((nmo-ncore[0],ncore[0],nmo))
      cpp = numpy.empty((ncore[0],nmo,nmo))
      for i in range(ncore[0]):
          buf = fload(i)
          klshape = (0, ncore[1], ncore[1], nmo)
          _ao2mo.nr_e2(buf[ncore[0]:nmo], mo[1], klshape,
                        aosym='s4', mosym='s1', out=cvCV[i], ao_loc=ao_loc)
  
          klshape = (0, nmo, 0, nmo)
          tmp = _ao2mo.nr_e2(buf[i:i+1], mo[1], klshape, aosym='s4',
                              mosym='s1', ao_loc=ao_loc)
          jc_PP += tmp.reshape(nmo,nmo)
  
          klshape = (0, ncore[0], 0, nmo)
          _ao2mo.nr_e2(buf[ncore[0]:nmo], mo[0], klshape,
                        aosym='s4', mosym='s1', out=vcp, ao_loc=ao_loc)
          kc_pp[i,ncore[0]:] = vcp[:,i]
  
          klshape = (0, nmo, 0, nmo)
          _ao2mo.nr_e2(buf[:ncore[0]], mo[0], klshape,
                        aosym='s4', mosym='s2', out=buf[:ncore[0]],
                        ao_loc=ao_loc)
          lib.unpack_tril(buf[:ncore[0]], out=cpp)
          jc_pp[i] = cpp[i]
          kc_pp[i,:ncore[0]] = cpp[:,i]
  
          #jcvcv = cvcv * 2 - cvcv.transpose(2,1,0,3) - ccvv.transpose(0,2,1,3)
          jcvcv[i] = vcp[:,:,ncore[0]:] * 2 \
                   - vcp[:,:,ncore[0]:].transpose(2,1,0) \
                   - cpp[:,ncore[0]:,ncore[0]:].transpose(1,0,2)
  
      return jc_pp, jc_PP, kc_pp, jcvcv, cvCV
  
  
  
  class _ERIS(object):
      def __init__(self, casscf, mo, method='incore'):
          mol = casscf.mol
          self.ncore = casscf.ncore
          self.ncas = casscf.ncas
          nmo = mo[0].shape[1]
          ncore = self.ncore
          ncas = self.ncas
          mem_incore, mem_outcore, mem_basic = _mem_usage(ncore, ncas, nmo)
          mem_now = lib.current_memory()[0]
  
          eri = casscf._scf._eri
          if (method == 'incore' and eri is not None and
              ((mem_incore+mem_now) < casscf.max_memory*.9) or
              mol.incore_anyway):
              if eri is None:
                  eri = mol.intor('int2e', aosym='s8')
              self.jkcpp, self.jkcPP, self.jC_pp, self.jc_PP, \
              self.aapp, self.aaPP, self.AApp, self.AAPP, \
              self.appa, self.apPA, self.APPA, \
              self.Iapcv, self.IAPCV, self.apCV, self.APcv, \
              self.Icvcv, self.ICVCV, self.cvCV = \
                      trans_e1_incore(eri, mo, casscf.ncore, casscf.ncas)
              self.vhf_c = (numpy.einsum('ipq->pq', self.jkcpp) + self.jC_pp,
                            numpy.einsum('ipq->pq', self.jkcPP) + self.jc_PP)
          else:
              import gc
              gc.collect()
              log = logger.Logger(casscf.stdout, casscf.verbose)
              max_memory = max(2000, casscf.max_memory*.9-mem_now)
              if ((mem_outcore+mem_now) < casscf.max_memory*.9):
                  if max_memory < mem_basic:
                      log.warn('Calculation needs %d MB memory, over CASSCF.max_memory (%d MB) limit',
                               (mem_outcore+mem_now)/.9, casscf.max_memory)
                  self.jkcpp, self.jkcPP, self.jC_pp, self.jc_PP, \
                  self.aapp, self.aaPP, self.AApp, self.AAPP, \
                  self.appa, self.apPA, self.APPA, \
                  self.Iapcv, self.IAPCV, self.apCV, self.APcv, \
                  self.Icvcv, self.ICVCV, self.cvCV = \
                          trans_e1_outcore(mol, mo, casscf.ncore, casscf.ncas,
                                           max_memory=max_memory, verbose=log)
                  self.vhf_c = (numpy.einsum('ipq->pq', self.jkcpp) + self.jC_pp,
                                numpy.einsum('ipq->pq', self.jkcPP) + self.jc_PP)
              else:
                  raise RuntimeError('.max_memory not enough')
  
  def _mem_usage(ncore, ncas, nmo):
      ncore = (ncore[0] + ncore[1]) * .5
      nvir = nmo - ncore
      basic = (ncas**2*nmo**2*7 + nmo**3*2) * 8/1e6
      outcore = basic + (ncore**2*nvir**2*3 + ncas*nmo*ncore*nvir*4 + ncore*nmo**2*3) * 8/1e6
      incore = outcore + nmo**4/1e6 + ncore*nmo**3*4/1e6
      return incore, outcore, basic
  
  if __name__ == '__main__':
      from pyscf import scf
      from pyscf import gto
      from pyscf.mcscf import umc1step
  
      mol = gto.Mole()
      mol.verbose = 0
      mol.output = None#"out_h2o"
      mol.atom = [
          ['O', ( 0., 0.    , 0.   )],
          ['H', ( 0., -0.757, 0.587)],
          ['H', ( 0., 0.757 , 0.587)],]
      mol.basis = {'H': 'sto-3g',
                   'O': 'sto-3g',}
      mol.charge = 1
      mol.spin = 1
      mol.build()
  
      m = scf.UHF(mol)
      ehf = m.scf()
  
      mc = umc1step.UCASSCF(m, 4, 4)
      mc.verbose = 4
      mo = m.mo_coeff
  
      eris0 = _ERIS(mc, mo, 'incore')
      eris1 = _ERIS(mc, mo, 'outcore')
      print('jkcpp', numpy.allclose(eris1.jkcpp, eris0.jkcpp))
      print('jkcPP', numpy.allclose(eris1.jkcPP, eris0.jkcPP))
      print('jC_pp', numpy.allclose(eris1.jC_pp, eris0.jC_pp))
      print('jc_PP', numpy.allclose(eris1.jc_PP, eris0.jc_PP))
      print('aapp ', numpy.allclose(eris1.aapp , eris0.aapp ))
      print('aaPP ', numpy.allclose(eris1.aaPP , eris0.aaPP ))
      print('AApp ', numpy.allclose(eris1.AApp , eris0.AApp ))
      print('AAPP ', numpy.allclose(eris1.AAPP , eris0.AAPP ))
      print('appa ', numpy.allclose(eris1.appa , eris0.appa ))
      print('apPA ', numpy.allclose(eris1.apPA , eris0.apPA ))
      print('APPA ', numpy.allclose(eris1.APPA , eris0.APPA ))
      print('cvCV ', numpy.allclose(eris1.cvCV , eris0.cvCV ))
      print('Icvcv', numpy.allclose(eris1.Icvcv, eris0.Icvcv))
      print('ICVCV', numpy.allclose(eris1.ICVCV, eris0.ICVCV))
      print('Iapcv', numpy.allclose(eris1.Iapcv, eris0.Iapcv))
      print('IAPCV', numpy.allclose(eris1.IAPCV, eris0.IAPCV))
      print('apCV ', numpy.allclose(eris1.apCV , eris0.apCV ))
      print('APcv ', numpy.allclose(eris1.APcv , eris0.APcv ))
  
  
      nmo = mo[0].shape[1]
      ncore = mc.ncore
      ncas = mc.ncas
      nocc = (ncas + ncore[0], ncas + ncore[1])
      eriaa = ao2mo.incore.full(mc._scf._eri, mo[0])
      eriab = ao2mo.incore.general(mc._scf._eri, (mo[0],mo[0],mo[1],mo[1]))
      eribb = ao2mo.incore.full(mc._scf._eri, mo[1])
      eriaa = ao2mo.restore(1, eriaa, nmo)
      eriab = ao2mo.restore(1, eriab, nmo)
      eribb = ao2mo.restore(1, eribb, nmo)
      jkcpp = numpy.einsum('iipq->ipq', eriaa[:ncore[0],:ncore[0],:,:]) \
            - numpy.einsum('ipqi->ipq', eriaa[:ncore[0],:,:,:ncore[0]])
      jkcPP = numpy.einsum('iipq->ipq', eribb[:ncore[1],:ncore[1],:,:]) \
            - numpy.einsum('ipqi->ipq', eribb[:ncore[1],:,:,:ncore[1]])
      jC_pp = numpy.einsum('pqii->pq', eriab[:,:,:ncore[1],:ncore[1]])
      jc_PP = numpy.einsum('iipq->pq', eriab[:ncore[0],:ncore[0],:,:])
      aapp = numpy.copy(eriaa[ncore[0]:nocc[0],ncore[0]:nocc[0],:,:])
      aaPP = numpy.copy(eriab[ncore[0]:nocc[0],ncore[0]:nocc[0],:,:])
      AApp = numpy.copy(eriab[:,:,ncore[1]:nocc[1],ncore[1]:nocc[1]].transpose(2,3,0,1))
      AAPP = numpy.copy(eribb[ncore[1]:nocc[1],ncore[1]:nocc[1],:,:])
      appa = numpy.copy(eriaa[ncore[0]:nocc[0],:,:,ncore[0]:nocc[0]])
      apPA = numpy.copy(eriab[ncore[0]:nocc[0],:,:,ncore[1]:nocc[1]])
      APPA = numpy.copy(eribb[ncore[1]:nocc[1],:,:,ncore[1]:nocc[1]])
  
      cvCV = numpy.copy(eriab[:ncore[0],ncore[0]:,:ncore[1],ncore[1]:])
      Icvcv = eriaa[:ncore[0],ncore[0]:,:ncore[0],ncore[0]:] * 2\
            - eriaa[:ncore[0],:ncore[0],ncore[0]:,ncore[0]:].transpose(0,3,1,2) \
            - eriaa[:ncore[0],ncore[0]:,:ncore[0],ncore[0]:].transpose(0,3,2,1)
      ICVCV = eribb[:ncore[1],ncore[1]:,:ncore[1],ncore[1]:] * 2\
            - eribb[:ncore[1],:ncore[1],ncore[1]:,ncore[1]:].transpose(0,3,1,2) \
            - eribb[:ncore[1],ncore[1]:,:ncore[1],ncore[1]:].transpose(0,3,2,1)
  
      Iapcv = eriaa[ncore[0]:nocc[0],:,:ncore[0],ncore[0]:] * 2 \
            - eriaa[:,ncore[0]:,:ncore[0],ncore[0]:nocc[0]].transpose(3,0,2,1) \
            - eriaa[:,:ncore[0],ncore[0]:,ncore[0]:nocc[0]].transpose(3,0,1,2)
      IAPCV = eribb[ncore[1]:nocc[1],:,:ncore[1],ncore[1]:] * 2 \
            - eribb[:,ncore[1]:,:ncore[1],ncore[1]:nocc[1]].transpose(3,0,2,1) \
            - eribb[:,:ncore[1],ncore[1]:,ncore[1]:nocc[1]].transpose(3,0,1,2)
      apCV = numpy.copy(eriab[ncore[0]:nocc[0],:,:ncore[1],ncore[1]:])
      APcv = numpy.copy(eriab[:ncore[0],ncore[0]:,ncore[1]:nocc[1],:].transpose(2,3,0,1))
  
      print('jkcpp', numpy.allclose(jkcpp, eris0.jkcpp))
      print('jkcPP', numpy.allclose(jkcPP, eris0.jkcPP))
      print('jC_pp', numpy.allclose(jC_pp, eris0.jC_pp))
      print('jc_PP', numpy.allclose(jc_PP, eris0.jc_PP))
      print('aapp ', numpy.allclose(aapp , eris0.aapp ))
      print('aaPP ', numpy.allclose(aaPP , eris0.aaPP ))
      print('AApp ', numpy.allclose(AApp , eris0.AApp ))
      print('AAPP ', numpy.allclose(AAPP , eris0.AAPP ))
      print('appa ', numpy.allclose(appa , eris0.appa ))
      print('apPA ', numpy.allclose(apPA , eris0.apPA ))
      print('APPA ', numpy.allclose(APPA , eris0.APPA ))
      print('cvCV ', numpy.allclose(cvCV , eris0.cvCV ))
      print('Icvcv', numpy.allclose(Icvcv, eris0.Icvcv))
      print('ICVCV', numpy.allclose(ICVCV, eris0.ICVCV))
      print('Iapcv', numpy.allclose(Iapcv, eris0.Iapcv))
      print('IAPCV', numpy.allclose(IAPCV, eris0.IAPCV))
      print('apCV ', numpy.allclose(apCV , eris0.apCV ))
      print('APcv ', numpy.allclose(APcv , eris0.APcv ))
  
  