code: #!/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>
  #
  
  '''
  Selected CI
  
  Simple usage::
  
      >>> from pyscf import gto, scf, ao2mo, fci
      >>> mol = gto.M(atom='C 0 0 0; C 0 0 1')
      >>> mf = scf.RHF(mol).run()
      >>> h1 = mf.mo_coeff.T.dot(mf.get_hcore()).dot(mf.mo_coeff)
      >>> h2 = ao2mo.kernel(mol, mf.mo_coeff)
      >>> e = fci.selected_ci.kernel(h1, h2, mf.mo_coeff.shape[1], mol.nelectron)[0]
  '''
  
  import ctypes
  import numpy
  from pyscf import lib
  from pyscf.lib import logger
  from pyscf import ao2mo
  from pyscf.fci import cistring
  from pyscf.fci import direct_spin1
  from pyscf.fci import rdm
  from pyscf import __config__
  
  libfci = lib.load_library('libfci')
  
  def contract_2e(eri, civec_strs, norb, nelec, link_index=None):
      ci_coeff, nelec, ci_strs = _unpack(civec_strs, nelec)
      if link_index is None:
          link_index = _all_linkstr_index(ci_strs, norb, nelec)
      cd_indexa, dd_indexa, cd_indexb, dd_indexb = link_index
      na, nlinka = cd_indexa.shape[:2]
      nb, nlinkb = cd_indexb.shape[:2]
  
      eri = ao2mo.restore(1, eri, norb)
      eri1 = eri.transpose(0,2,1,3) - eri.transpose(0,2,3,1)
      idx,idy = numpy.tril_indices(norb, -1)
      idx = idx * norb + idy
      eri1 = lib.take_2d(eri1.reshape(norb**2,-1), idx, idx) * 2
      fcivec = ci_coeff.reshape(na,nb)
      # (bb|bb)
      if nelec[1] > 1:
          mb, mlinkb = dd_indexb.shape[:2]
          fcivecT = lib.transpose(fcivec)
          ci1T = numpy.zeros((nb,na))
          libfci.SCIcontract_2e_aaaa(eri1.ctypes.data_as(ctypes.c_void_p),
                                     fcivecT.ctypes.data_as(ctypes.c_void_p),
                                     ci1T.ctypes.data_as(ctypes.c_void_p),
                                     ctypes.c_int(norb),
                                     ctypes.c_int(nb), ctypes.c_int(na),
                                     ctypes.c_int(mb), ctypes.c_int(mlinkb),
                                     dd_indexb.ctypes.data_as(ctypes.c_void_p))
          ci1 = lib.transpose(ci1T, out=fcivecT)
      else:
          ci1 = numpy.zeros_like(fcivec)
      # (aa|aa)
      if nelec[0] > 1:
          ma, mlinka = dd_indexa.shape[:2]
          libfci.SCIcontract_2e_aaaa(eri1.ctypes.data_as(ctypes.c_void_p),
                                     fcivec.ctypes.data_as(ctypes.c_void_p),
                                     ci1.ctypes.data_as(ctypes.c_void_p),
                                     ctypes.c_int(norb),
                                     ctypes.c_int(na), ctypes.c_int(nb),
                                     ctypes.c_int(ma), ctypes.c_int(mlinka),
                                     dd_indexa.ctypes.data_as(ctypes.c_void_p))
  
      h_ps = numpy.einsum('pqqs->ps', eri)
      eri1 = eri * 2
      for k in range(norb):
          eri1[:,:,k,k] += h_ps/nelec[0]
          eri1[k,k,:,:] += h_ps/nelec[1]
      eri1 = ao2mo.restore(4, eri1, norb)
      # (bb|aa)
      libfci.SCIcontract_2e_bbaa(eri1.ctypes.data_as(ctypes.c_void_p),
                                 fcivec.ctypes.data_as(ctypes.c_void_p),
                                 ci1.ctypes.data_as(ctypes.c_void_p),
                                 ctypes.c_int(norb),
                                 ctypes.c_int(na), ctypes.c_int(nb),
                                 ctypes.c_int(nlinka), ctypes.c_int(nlinkb),
                                 cd_indexa.ctypes.data_as(ctypes.c_void_p),
                                 cd_indexb.ctypes.data_as(ctypes.c_void_p))
  
      return _as_SCIvector(ci1.reshape(ci_coeff.shape), ci_strs)
  
  def select_strs(myci, eri, eri_pq_max, civec_max, strs, norb, nelec):
      strs = numpy.asarray(strs, dtype=numpy.int64)
      nstrs = len(strs)
      nvir = norb - nelec
      strs_add = numpy.empty((nstrs*(nelec*nvir)**2//4), dtype=numpy.int64)
      libfci.SCIselect_strs.restype = ctypes.c_int
      nadd = libfci.SCIselect_strs(strs_add.ctypes.data_as(ctypes.c_void_p),
                                   strs.ctypes.data_as(ctypes.c_void_p),
                                   eri.ctypes.data_as(ctypes.c_void_p),
                                   eri_pq_max.ctypes.data_as(ctypes.c_void_p),
                                   civec_max.ctypes.data_as(ctypes.c_void_p),
                                   ctypes.c_double(myci.select_cutoff),
                                   ctypes.c_int(norb), ctypes.c_int(nelec),
                                   ctypes.c_int(nstrs))
      strs_add = sorted(set(strs_add[:nadd]) - set(strs))
      return numpy.asarray(strs_add, dtype=numpy.int64)
  
  def enlarge_space(myci, civec_strs, eri, norb, nelec):
      if isinstance(civec_strs, (tuple, list)):
          nelec, (strsa, strsb) = _unpack(civec_strs[0], nelec)[1:]
          ci_coeff = lib.asarray(civec_strs)
      else:
          ci_coeff, nelec, (strsa, strsb) = _unpack(civec_strs, nelec)
      na = len(strsa)
      nb = len(strsb)
      ci0 = ci_coeff.reshape(-1,na,nb)
      civec_a_max = lib.norm(ci0, axis=2).max(axis=0)
      civec_b_max = lib.norm(ci0, axis=1).max(axis=0)
      ci_aidx = numpy.where(civec_a_max > myci.ci_coeff_cutoff)[0]
      ci_bidx = numpy.where(civec_b_max > myci.ci_coeff_cutoff)[0]
      civec_a_max = civec_a_max[ci_aidx]
      civec_b_max = civec_b_max[ci_bidx]
      strsa = strsa[ci_aidx]
      strsb = strsb[ci_bidx]
  
      eri = ao2mo.restore(1, eri, norb)
      eri_pq_max = abs(eri.reshape(norb**2,-1)).max(axis=1).reshape(norb,norb)
  
      strsa_add = select_strs(myci, eri, eri_pq_max, civec_a_max, strsa, norb, nelec[0])
      strsb_add = select_strs(myci, eri, eri_pq_max, civec_b_max, strsb, norb, nelec[1])
      strsa = numpy.append(strsa, strsa_add)
      strsb = numpy.append(strsb, strsb_add)
      aidx = numpy.argsort(strsa)
      bidx = numpy.argsort(strsb)
      ci_strs = (strsa[aidx], strsb[bidx])
      aidx = numpy.where(aidx < len(ci_aidx))[0]
      bidx = numpy.where(bidx < len(ci_bidx))[0]
      ma = len(strsa)
      mb = len(strsb)
  
      cs = []
      for i in range(ci0.shape[0]):
          ci1 = numpy.zeros((ma,mb))
          tmp = lib.take_2d(ci0[i], ci_aidx, ci_bidx)
          lib.takebak_2d(ci1, tmp, aidx, bidx)
          cs.append(_as_SCIvector(ci1, ci_strs))
  
      if not isinstance(civec_strs, (tuple, list)) and civec_strs.ndim < 3:
          cs = cs[0]
      return cs
  
  def cre_des_linkstr(strs, norb, nelec, tril=False):
      '''Given intermediates, the link table to generate input strs
      '''
      strs = numpy.asarray(strs, dtype=numpy.int64)
      nvir = norb - nelec
      nstrs = len(strs)
      link_index = numpy.zeros((nstrs,nelec+nelec*nvir,4), dtype=numpy.int32)
      libfci.SCIcre_des_linkstr(link_index.ctypes.data_as(ctypes.c_void_p),
                                ctypes.c_int(norb), ctypes.c_int(nstrs),
                                ctypes.c_int(nelec),
                                strs.ctypes.data_as(ctypes.c_void_p),
                                ctypes.c_int(tril))
      return link_index
  
  def cre_des_linkstr_tril(strs, norb, nelec):
      '''Given intermediates, the link table to generate input strs
      '''
      return cre_des_linkstr(strs, norb, nelec, True)
  
  def des_des_linkstr(strs, norb, nelec, tril=False):
      '''Given intermediates, the link table to generate input strs
      '''
      if nelec < 2:
          return None
  
      strs = numpy.asarray(strs, dtype=numpy.int64)
      nvir = norb - nelec
      nstrs = len(strs)
      inter1 = numpy.empty((nstrs*nelec), dtype=numpy.int64)
      libfci.SCIdes_uniq_strs.restype = ctypes.c_int
      ninter = libfci.SCIdes_uniq_strs(inter1.ctypes.data_as(ctypes.c_void_p),
                                       strs.ctypes.data_as(ctypes.c_void_p),
                                       ctypes.c_int(norb), ctypes.c_int(nelec),
                                       ctypes.c_int(nstrs))
      inter1 = numpy.asarray(sorted(set(inter1[:ninter])), dtype=numpy.int64)
      ninter = len(inter1)
  
      inter = numpy.empty((ninter*nelec), dtype=numpy.int64)
      ninter = libfci.SCIdes_uniq_strs(inter.ctypes.data_as(ctypes.c_void_p),
                                       inter1.ctypes.data_as(ctypes.c_void_p),
                                       ctypes.c_int(norb), ctypes.c_int(nelec-1),
                                       ctypes.c_int(ninter))
      inter = numpy.asarray(sorted(set(inter[:ninter])), dtype=numpy.int64)
      ninter = len(inter)
  
      nvir += 2
      link_index = numpy.zeros((ninter,nvir*nvir,4), dtype=numpy.int32)
      libfci.SCIdes_des_linkstr(link_index.ctypes.data_as(ctypes.c_void_p),
                                ctypes.c_int(norb), ctypes.c_int(nelec),
                                ctypes.c_int(nstrs), ctypes.c_int(ninter),
                                strs.ctypes.data_as(ctypes.c_void_p),
                                inter.ctypes.data_as(ctypes.c_void_p),
                                ctypes.c_int(tril))
      return link_index
  
  def des_des_linkstr_tril(strs, norb, nelec):
      '''Given intermediates, the link table to generate input strs
      '''
      return des_des_linkstr(strs, norb, nelec, True)
  
  def gen_des_linkstr(strs, norb, nelec):
      '''Given intermediates, the link table to generate input strs
      '''
      if nelec < 1:
          return None
  
      strs = numpy.asarray(strs, dtype=numpy.int64)
      nvir = norb - nelec
      nstrs = len(strs)
      inter = numpy.empty((nstrs*nelec), dtype=numpy.int64)
      libfci.SCIdes_uniq_strs.restype = ctypes.c_int
      ninter = libfci.SCIdes_uniq_strs(inter.ctypes.data_as(ctypes.c_void_p),
                                       strs.ctypes.data_as(ctypes.c_void_p),
                                       ctypes.c_int(norb), ctypes.c_int(nelec),
                                       ctypes.c_int(nstrs))
      inter = numpy.asarray(sorted(set(inter[:ninter])), dtype=numpy.int64)
      ninter = len(inter)
  
      nvir += 1
      link_index = numpy.zeros((ninter,nvir,4), dtype=numpy.int32)
      libfci.SCIdes_linkstr(link_index.ctypes.data_as(ctypes.c_void_p),
                            ctypes.c_int(norb), ctypes.c_int(nelec),
                            ctypes.c_int(nstrs), ctypes.c_int(ninter),
                            strs.ctypes.data_as(ctypes.c_void_p),
                            inter.ctypes.data_as(ctypes.c_void_p))
      return link_index
  
  def gen_cre_linkstr(strs, norb, nelec):
      '''Given intermediates, the link table to generate input strs
      '''
      if nelec == norb:
          return None
  
      strs = numpy.asarray(strs, dtype=numpy.int64)
      nvir = norb - nelec
      nstrs = len(strs)
      inter = numpy.empty((nstrs*nvir), dtype=numpy.int64)
      libfci.SCIcre_uniq_strs.restype = ctypes.c_int
      ninter = libfci.SCIcre_uniq_strs(inter.ctypes.data_as(ctypes.c_void_p),
                                       strs.ctypes.data_as(ctypes.c_void_p),
                                       ctypes.c_int(norb), ctypes.c_int(nelec),
                                       ctypes.c_int(nstrs))
      inter = numpy.asarray(sorted(set(inter[:ninter])), dtype=numpy.int64)
      ninter = len(inter)
  
      link_index = numpy.zeros((ninter,nelec+1,4), dtype=numpy.int32)
      libfci.SCIcre_linkstr(link_index.ctypes.data_as(ctypes.c_void_p),
                            ctypes.c_int(norb), ctypes.c_int(nelec),
                            ctypes.c_int(nstrs), ctypes.c_int(ninter),
                            strs.ctypes.data_as(ctypes.c_void_p),
                            inter.ctypes.data_as(ctypes.c_void_p))
      return link_index
  
  
  def make_hdiag(h1e, eri, ci_strs, norb, nelec):
      ci_coeff, nelec, ci_strs = _unpack(None, nelec, ci_strs)
      na = len(ci_strs[0])
      nb = len(ci_strs[1])
      hdiag = numpy.empty(na*nb)
  
      h1e = numpy.asarray(h1e, order='C')
      eri = ao2mo.restore(1, eri, norb)
      jdiag = numpy.asarray(numpy.einsum('iijj->ij',eri), order='C')
      kdiag = numpy.asarray(numpy.einsum('ijji->ij',eri), order='C')
      c_h1e = h1e.ctypes.data_as(ctypes.c_void_p)
      c_jdiag = jdiag.ctypes.data_as(ctypes.c_void_p)
      c_kdiag = kdiag.ctypes.data_as(ctypes.c_void_p)
      occslsta = cistring._strs2occslst(ci_strs[0], norb)
      occslstb = cistring._strs2occslst(ci_strs[1], norb)
      libfci.FCImake_hdiag_uhf(hdiag.ctypes.data_as(ctypes.c_void_p),
                               c_h1e, c_h1e, c_jdiag, c_jdiag, c_jdiag, c_kdiag, c_kdiag,
                               ctypes.c_int(norb),
                               ctypes.c_int(na), ctypes.c_int(nb),
                               ctypes.c_int(nelec[0]), ctypes.c_int(nelec[1]),
                               occslsta.ctypes.data_as(ctypes.c_void_p),
                               occslstb.ctypes.data_as(ctypes.c_void_p))
      return hdiag
  
  def kernel_fixed_space(myci, h1e, eri, norb, nelec, ci_strs, ci0=None,
                         tol=None, lindep=None, max_cycle=None, max_space=None,
                         nroots=None, davidson_only=None,
                         max_memory=None, verbose=None, ecore=0, **kwargs):
      log = logger.new_logger(myci, verbose)
      if tol is None: tol = myci.conv_tol
      if lindep is None: lindep = myci.lindep
      if max_cycle is None: max_cycle = myci.max_cycle
      if max_space is None: max_space = myci.max_space
      if max_memory is None: max_memory = myci.max_memory
      if nroots is None: nroots = myci.nroots
      if myci.verbose >= logger.WARN:
          myci.check_sanity()
  
      nelec = direct_spin1._unpack_nelec(nelec, myci.spin)
      ci0, nelec, ci_strs = _unpack(ci0, nelec, ci_strs)
      na = len(ci_strs[0])
      nb = len(ci_strs[1])
      h2e = direct_spin1.absorb_h1e(h1e, eri, norb, nelec, .5)
      h2e = ao2mo.restore(1, h2e, norb)
  
      link_index = _all_linkstr_index(ci_strs, norb, nelec)
      hdiag = myci.make_hdiag(h1e, eri, ci_strs, norb, nelec)
  
      if isinstance(ci0, _SCIvector):
          if ci0.size == na*nb:
              ci0 = [ci0.ravel()]
          else:
              ci0 = [x.ravel() for x in ci0]
      else:
          ci0 = myci.get_init_guess(ci_strs, norb, nelec, nroots, hdiag)
  
      def hop(c):
          hc = myci.contract_2e(h2e, _as_SCIvector(c, ci_strs), norb, nelec, link_index)
          return hc.reshape(-1)
      precond = lambda x, e, *args: x/(hdiag-e+1e-4)
  
      #e, c = lib.davidson(hop, ci0, precond, tol=myci.conv_tol)
      e, c = myci.eig(hop, ci0, precond, tol=tol, lindep=lindep,
                      max_cycle=max_cycle, max_space=max_space, nroots=nroots,
                      max_memory=max_memory, verbose=log, **kwargs)
      if nroots > 1:
          return e+ecore, [_as_SCIvector(ci.reshape(na,nb),ci_strs) for ci in c]
      else:
          return e+ecore, _as_SCIvector(c.reshape(na,nb), ci_strs)
  
  
  def kernel_float_space(myci, h1e, eri, norb, nelec, ci0=None,
                         tol=None, lindep=None, max_cycle=None, max_space=None,
                         nroots=None, davidson_only=None,
                         max_memory=None, verbose=None, ecore=0, **kwargs):
      log = logger.new_logger(myci, verbose)
      if tol is None: tol = myci.conv_tol
      if lindep is None: lindep = myci.lindep
      if max_cycle is None: max_cycle = myci.max_cycle
      if max_space is None: max_space = myci.max_space
      if max_memory is None: max_memory = myci.max_memory
      if nroots is None: nroots = myci.nroots
      if myci.verbose >= logger.WARN:
          myci.check_sanity()
  
      nelec = direct_spin1._unpack_nelec(nelec, myci.spin)
      h2e = direct_spin1.absorb_h1e(h1e, eri, norb, nelec, .5)
      h2e = ao2mo.restore(1, h2e, norb)
  
  # TODO: initial guess from CISD
      if isinstance(ci0, _SCIvector):
          if ci0.size == len(ci0._strs[0])*len(ci0._strs[1]):
              ci0 = [ci0.ravel()]
          else:
              ci0 = [x.ravel() for x in ci0]
      else:
          ci_strs = (numpy.asarray([int('1'*nelec[0], 2)]),
                     numpy.asarray([int('1'*nelec[1], 2)]))
          ci0 = _as_SCIvector(numpy.ones((1,1)), ci_strs)
          ci0 = myci.enlarge_space(ci0, h2e, norb, nelec)
          if ci0.size < nroots:
              log.warn('''
    Selected-CI space generated from HF ground state (by double exciting) is not enough for excited states.
    HOMO->LUMO excitations are included in the initial guess.
    NOTE: This may introduce excited states of different symmetry.\n''')
              corea = '1' * (nelec[0]-1)
              coreb = '1' * (nelec[1]-1)
              ci_strs = (numpy.asarray([int('1'+corea, 2), int('10'+corea, 2)]),
                         numpy.asarray([int('1'+coreb, 2), int('10'+coreb, 2)]))
              ci0 = _as_SCIvector(numpy.ones((2,2)), ci_strs)
              ci0 = myci.enlarge_space(ci0, h2e, norb, nelec)
          if ci0.size < nroots:
              raise RuntimeError('Not enough selected-CI space for %d states' % nroots)
          ci_strs = ci0._strs
          hdiag = myci.make_hdiag(h1e, eri, ci_strs, norb, nelec)
          ci0 = myci.get_init_guess(ci_strs, norb, nelec, nroots, hdiag)
  
      def hop(c):
          hc = myci.contract_2e(h2e, _as_SCIvector(c, ci_strs), norb, nelec, link_index)
          return hc.ravel()
      precond = lambda x, e, *args: x/(hdiag-e+myci.level_shift)
  
      namax = cistring.num_strings(norb, nelec[0])
      nbmax = cistring.num_strings(norb, nelec[1])
      e_last = 0
      float_tol = myci.start_tol
      tol_decay_rate = myci.tol_decay_rate
      conv = False
      for icycle in range(norb):
          ci_strs = ci0[0]._strs
          float_tol = max(float_tol*tol_decay_rate, tol*1e2)
          log.debug('cycle %d  ci.shape %s  float_tol %g',
                    icycle, (len(ci_strs[0]), len(ci_strs[1])), float_tol)
  
          ci0 = [c.ravel() for c in ci0]
          link_index = _all_linkstr_index(ci_strs, norb, nelec)
          hdiag = myci.make_hdiag(h1e, eri, ci_strs, norb, nelec)
          #e, ci0 = lib.davidson(hop, ci0.reshape(-1), precond, tol=float_tol)
          e, ci0 = myci.eig(hop, ci0, precond, tol=float_tol, lindep=lindep,
                            max_cycle=max_cycle, max_space=max_space, nroots=nroots,
                            max_memory=max_memory, verbose=log, **kwargs)
          if nroots > 1:
              ci0 = [_as_SCIvector(c, ci_strs) for c in ci0]
              de, e_last = min(e)-e_last, min(e)
              log.info('cycle %d  E = %s  dE = %.8g', icycle, e+ecore, de)
          else:
              ci0 = [_as_SCIvector(ci0, ci_strs)]
              de, e_last = e-e_last, e
              log.info('cycle %d  E = %.15g  dE = %.8g', icycle, e+ecore, de)
  
          if ci0[0].shape == (namax,nbmax) or abs(de) < tol*1e3:
              conv = True
              break
  
          last_ci0_size = float(len(ci_strs[0])), float(len(ci_strs[1]))
          ci0 = myci.enlarge_space(ci0, h2e, norb, nelec)
          na = len(ci0[0]._strs[0])
          nb = len(ci0[0]._strs[1])
          if ((.99 < na/last_ci0_size[0] < 1.01) and
              (.99 < nb/last_ci0_size[1] < 1.01)):
              conv = True
              break
  
      ci_strs = ci0[0]._strs
      log.debug('Extra CI in selected space %s', (len(ci_strs[0]), len(ci_strs[1])))
      ci0 = [c.ravel() for c in ci0]
      link_index = _all_linkstr_index(ci_strs, norb, nelec)
      hdiag = myci.make_hdiag(h1e, eri, ci_strs, norb, nelec)
      e, c = myci.eig(hop, ci0, precond, tol=tol, lindep=lindep,
                      max_cycle=max_cycle, max_space=max_space, nroots=nroots,
                      max_memory=max_memory, verbose=log, **kwargs)
  
      na = len(ci_strs[0])
      nb = len(ci_strs[1])
      if nroots > 1:
          for i, ei in enumerate(e+ecore):
              log.info('Selected CI state %d  E = %.15g', i, ei)
          return e+ecore, [_as_SCIvector(ci.reshape(na,nb),ci_strs) for ci in c]
      else:
          log.info('Selected CI  E = %.15g', e+ecore)
          return e+ecore, _as_SCIvector(c.reshape(na,nb), ci_strs)
  
  def kernel(h1e, eri, norb, nelec, ci0=None, level_shift=1e-3, tol=1e-10,
             lindep=1e-14, max_cycle=50, max_space=12, nroots=1,
             davidson_only=False, pspace_size=400, orbsym=None, wfnsym=None,
             select_cutoff=1e-3, ci_coeff_cutoff=1e-3, ecore=0, **kwargs):
      return direct_spin1._kfactory(SelectedCI, h1e, eri, norb, nelec, ci0,
                                    level_shift, tol, lindep, max_cycle,
                                    max_space, nroots, davidson_only,
                                    pspace_size, select_cutoff=select_cutoff,
                                    ci_coeff_cutoff=ci_coeff_cutoff, ecore=ecore,
                                    **kwargs)
  
  def make_rdm1s(civec_strs, norb, nelec, link_index=None):
      '''Spin separated 1-particle density matrices.
      The return values include two density matrices: (alpha,alpha), (beta,beta)
  
      dm1[p,q] = <q^\dagger p>
  
      The convention is based on McWeeney's book, Eq (5.4.20).
      The contraction between 1-particle Hamiltonian and rdm1 is
      E = einsum('pq,qp', h1, rdm1)
      '''
      ci_coeff, nelec, ci_strs = _unpack(civec_strs, nelec)
      if link_index is None:
          cd_indexa = cre_des_linkstr(ci_strs[0], norb, nelec[0])
          cd_indexb = cre_des_linkstr(ci_strs[1], norb, nelec[1])
      else:
          cd_indexa, dd_indexa, cd_indexb, dd_indexb = link_index
      rdm1a = rdm.make_rdm1_spin1('FCImake_rdm1a', ci_coeff, ci_coeff,
                                  norb, nelec, (cd_indexa,cd_indexb))
      rdm1b = rdm.make_rdm1_spin1('FCImake_rdm1b', ci_coeff, ci_coeff,
                                  norb, nelec, (cd_indexa,cd_indexb))
      return rdm1a, rdm1b
  
  def make_rdm1(civec_strs, norb, nelec, link_index=None):
      r'''Spin-traced 1-particle density matrix.
  
      dm1[p,q] = <q_alpha^\dagger p_alpha> + <q_beta^\dagger p_beta>
  
      The convention is based on McWeeney's book, Eq (5.4.20)
      The contraction between 1-particle Hamiltonian and rdm1 is
      E = einsum('pq,qp', h1, rdm1)
      '''
      rdm1a, rdm1b = make_rdm1s(civec_strs, norb, nelec, link_index)
      return rdm1a + rdm1b
  
  # dm[p,q,r,s] = <|p^+ q r^+ s|>
  def make_rdm2s(civec_strs, norb, nelec, link_index=None, **kwargs):
      r'''Spin separated 2-particle density matrices.
      The return values include three density matrices:
      (alpha,alpha,alpha,alpha), (alpha,alpha,beta,beta), (beta,beta,beta,beta)
  
      2pdm[p,q,r,s] = :math:`\langle p^\dagger r^\dagger s q\rangle`
      '''
      ci_coeff, nelec, ci_strs = _unpack(civec_strs, nelec)
      if link_index is None:
          cd_indexa = cre_des_linkstr(ci_strs[0], norb, nelec[0])
          dd_indexa = des_des_linkstr(ci_strs[0], norb, nelec[0])
          cd_indexb = cre_des_linkstr(ci_strs[1], norb, nelec[1])
          dd_indexb = des_des_linkstr(ci_strs[1], norb, nelec[1])
      else:
          cd_indexa, dd_indexa, cd_indexb, dd_indexb = link_index
      na, nlinka = cd_indexa.shape[:2]
      nb, nlinkb = cd_indexb.shape[:2]
  
      fcivec = ci_coeff.reshape(na,nb)
      # (bb|aa) and (aa|bb)
      dm2ab = rdm.make_rdm12_spin1('FCItdm12kern_ab', fcivec, fcivec,
                                   norb, nelec, (cd_indexa,cd_indexb), 0)[1]
      # (aa|aa)
      dm2aa = numpy.zeros([norb]*4)
      if nelec[0] > 1:
          ma, mlinka = dd_indexa.shape[:2]
          libfci.SCIrdm2_aaaa(libfci.SCIrdm2kern_aaaa,
                              dm2aa.ctypes.data_as(ctypes.c_void_p),
                              fcivec.ctypes.data_as(ctypes.c_void_p),
                              fcivec.ctypes.data_as(ctypes.c_void_p),
                              ctypes.c_int(norb),
                              ctypes.c_int(na), ctypes.c_int(nb),
                              ctypes.c_int(ma), ctypes.c_int(mlinka),
                              dd_indexa.ctypes.data_as(ctypes.c_void_p))
      # (bb|bb)
      dm2bb = numpy.zeros([norb]*4)
      if nelec[1] > 1:
          mb, mlinkb = dd_indexb.shape[:2]
          fcivecT = lib.transpose(fcivec)
          libfci.SCIrdm2_aaaa(libfci.SCIrdm2kern_aaaa,
                              dm2bb.ctypes.data_as(ctypes.c_void_p),
                              fcivecT.ctypes.data_as(ctypes.c_void_p),
                              fcivecT.ctypes.data_as(ctypes.c_void_p),
                              ctypes.c_int(norb),
                              ctypes.c_int(nb), ctypes.c_int(na),
                              ctypes.c_int(mb), ctypes.c_int(mlinkb),
                              dd_indexb.ctypes.data_as(ctypes.c_void_p))
      return dm2aa, dm2ab, dm2bb
  
  def make_rdm2(civec_strs, norb, nelec, link_index=None, **kwargs):
      r'''Spin-traced two-particle density matrix.
  
      2pdm[p,q,r,s] = :math:`\langle p_\alpha^\dagger r_\alpha^\dagger s_\alpha q_\alpha\rangle +
                             \langle p_\beta^\dagger  r_\alpha^\dagger s_\alpha q_\beta\rangle +
                             \langle p_\alpha^\dagger r_\beta^\dagger  s_\beta  q_\alpha\rangle +
                             \langle p_\beta^\dagger  r_\beta^\dagger  s_\beta  q_\beta\rangle`.
      '''
      dm2aa, dm2ab, dm2bb = make_rdm2s(civec_strs, norb, nelec, link_index)
      dm2aa += dm2bb
      dm2aa += dm2ab
      dm2aa += dm2ab.transpose(2,3,0,1)
      return dm2aa
  
  def trans_rdm1s(cibra_strs, ciket_strs, norb, nelec, link_index=None):
      r'''Spin separated transition 1-particle density matrices.
      See also function :func:`make_rdm1s`
  
      1pdm[p,q] = :math:`\langle q^\dagger p \rangle`
      '''
      cibra, nelec, ci_strs = _unpack(cibra_strs, nelec)
      ciket, nelec1, ci_strs1 = _unpack(ciket_strs, nelec)
      assert(all(ci_strs[0] == ci_strs1[0]) and
             all(ci_strs[1] == ci_strs1[1]))
      if link_index is None:
          cd_indexa = cre_des_linkstr(ci_strs[0], norb, nelec[0])
          cd_indexb = cre_des_linkstr(ci_strs[1], norb, nelec[1])
      else:
          cd_indexa, dd_indexa, cd_indexb, dd_indexb = link_index
      rdm1a = rdm.make_rdm1_spin1('FCItrans_rdm1a', cibra, ciket,
                                  norb, nelec, (cd_indexa,cd_indexb))
      rdm1b = rdm.make_rdm1_spin1('FCItrans_rdm1b', cibra, ciket,
                                  norb, nelec, (cd_indexa,cd_indexb))
      return rdm1a, rdm1b
  
  def trans_rdm1(cibra_strs, ciket_strs, norb, nelec, link_index=None):
      r'''Spin traced transition 1-particle density matrices.
      See also function :func:`make_rdm1`
  
      1pdm[p,q] = :math:`\langle q_\alpha^\dagger p_\alpha \rangle
                         + \langle q_\beta^\dagger p_\beta \rangle`
      '''
      rdm1a, rdm1b = trans_rdm1s(cibra_strs, ciket_strs, norb, nelec, link_index)
      return rdm1a + rdm1b
  
  def spin_square(civec_strs, norb, nelec):
      '''Spin square for RHF-FCI CI wfn only (obtained from spin-degenerated
      Hamiltonian)'''
      ci1 = contract_ss(civec_strs, norb, nelec)
  
      ss = numpy.einsum('ij,ij->', civec_strs.reshape(ci1.shape), ci1)
      s = numpy.sqrt(ss+.25) - .5
      multip = s*2+1
      return ss, multip
  
  def contract_ss(civec_strs, norb, nelec):
      r''' S^2 |\Psi\rangle
      '''
      ci_coeff, nelec, ci_strs = _unpack(civec_strs, nelec)
      strsa, strsb = ci_strs
      neleca, nelecb = nelec
      ci_coeff = ci_coeff.reshape(len(strsa),len(strsb))
  
      def gen_map(fstr_index, strs, nelec, des=True):
          a_index = fstr_index(strs, norb, nelec)
          amap = numpy.zeros((a_index.shape[0],norb,2), dtype=numpy.int32)
          if des:
              for k, tab in enumerate(a_index):
                  sign = tab[:,3]
                  tab = tab[sign!=0]
                  amap[k,tab[:,1]] = tab[:,2:]
          else:
              for k, tab in enumerate(a_index):
                  sign = tab[:,3]
                  tab = tab[sign!=0]
                  amap[k,tab[:,0]] = tab[:,2:]
          return amap
  
      if neleca > 0:
          ades = gen_map(gen_des_linkstr, strsa, neleca)
      else:
          ades = None
  
      if nelecb > 0:
          bdes = gen_map(gen_des_linkstr, strsb, nelecb)
      else:
          bdes = None
  
      if neleca < norb:
          acre = gen_map(gen_cre_linkstr, strsa, neleca, False)
      else:
          acre = None
  
      if nelecb < norb:
          bcre = gen_map(gen_cre_linkstr, strsb, nelecb, False)
      else:
          bcre = None
  
      def trans(ci1, aindex, bindex):
          if aindex is None or bindex is None:
              return None
  
          ma = len(aindex)
          mb = len(bindex)
          t1 = numpy.zeros((ma,mb))
          for i in range(norb):
              signa = aindex[:,i,1]
              signb = bindex[:,i,1]
              maska = numpy.where(signa!=0)[0]
              maskb = numpy.where(signb!=0)[0]
              addra = aindex[maska,i,0]
              addrb = bindex[maskb,i,0]
              citmp = lib.take_2d(ci_coeff, addra, addrb)
              citmp *= signa[maska].reshape(-1,1)
              citmp *= signb[maskb]
              #: t1[addra.reshape(-1,1),addrb] += citmp
              lib.takebak_2d(t1, citmp, maska, maskb)
          for i in range(norb):
              signa = aindex[:,i,1]
              signb = bindex[:,i,1]
              maska = numpy.where(signa!=0)[0]
              maskb = numpy.where(signb!=0)[0]
              addra = aindex[maska,i,0]
              addrb = bindex[maskb,i,0]
              citmp = lib.take_2d(t1, maska, maskb)
              citmp *= signa[maska].reshape(-1,1)
              citmp *= signb[maskb]
              #: ci1[maska.reshape(-1,1), maskb] += citmp
              lib.takebak_2d(ci1, citmp, addra, addrb)
  
      ci1 = numpy.zeros_like(ci_coeff)
      trans(ci1, ades, bcre) # S+*S-
      trans(ci1, acre, bdes) # S-*S+
      ci1 *= .5
      ci1 += (neleca-nelecb)**2*.25*ci_coeff
      return _as_SCIvector(ci1, ci_strs)
  
  def to_fci(civec_strs, norb, nelec):
      ci_coeff, nelec, ci_strs = _unpack(civec_strs, nelec)
      addrsa = [cistring.str2addr(norb, nelec[0], x) for x in ci_strs[0]]
      addrsb = [cistring.str2addr(norb, nelec[1], x) for x in ci_strs[1]]
      na = cistring.num_strings(norb, nelec[0])
      nb = cistring.num_strings(norb, nelec[1])
      ci0 = numpy.zeros((na,nb))
      lib.takebak_2d(ci0, ci_coeff, addrsa, addrsb)
      return ci0
  
  def from_fci(fcivec, ci_strs, norb, nelec):
      fcivec, nelec, ci_strs = _unpack(fcivec, nelec, ci_strs)
      addrsa = [cistring.str2addr(norb, nelec[0], x) for x in ci_strs[0]]
      addrsb = [cistring.str2addr(norb, nelec[1], x) for x in ci_strs[1]]
      na = cistring.num_strings(norb, nelec[0])
      nb = cistring.num_strings(norb, nelec[1])
      fcivec = fcivec.reshape(na,nb)
      civec = lib.take_2d(fcivec, addrsa, addrsb)
      return _as_SCIvector(civec, ci_strs)
  
  
  class SelectedCI(direct_spin1.FCISolver):
  
      ci_coeff_cutoff = getattr(__config__, 'fci_selected_ci_SCI_ci_coeff_cutoff', .5e-3)
      select_cutoff = getattr(__config__, 'fci_selected_ci_SCI_select_cutoff', .5e-3)
      conv_tol = getattr(__config__, 'fci_selected_ci_SCI_conv_tol', 1e-9)
      start_tol = getattr(__config__, 'fci_selected_ci_SCI_start_tol', 3e-4)
      tol_decay_rate = getattr(__config__, 'fci_selected_ci_SCI_tol_decay_rate', 0.3)
  
      def __init__(self, mol=None):
          direct_spin1.FCISolver.__init__(self, mol)
  
  ##################################################
  # don't modify the following attributes, they are not input options
          #self.converged = False
          #self.ci = None
          self._strs = None
          keys = set(('ci_coeff_cutoff', 'select_cutoff', 'conv_tol',
                      'start_tol', 'tol_decay_rate'))
          self._keys = self._keys.union(keys)
  
      def dump_flags(self, verbose=None):
          direct_spin1.FCISolver.dump_flags(self, verbose)
          logger.info(self, 'ci_coeff_cutoff %g', self.ci_coeff_cutoff)
          logger.info(self, 'select_cutoff   %g', self.select_cutoff)
  
      def contract_2e(self, eri, civec_strs, norb, nelec, link_index=None, **kwargs):
  # The argument civec_strs is a CI vector in function FCISolver.contract_2e.
  # Save and patch self._strs to make this contract_2e function compatible to
  # FCISolver.contract_2e.
          if getattr(civec_strs, '_strs', None) is not None:
              self._strs = civec_strs._strs
          else:
              assert(civec_strs.size == len(self._strs[0])*len(self._strs[1]))
              civec_strs = _as_SCIvector(civec_strs, self._strs)
          return contract_2e(eri, civec_strs, norb, nelec, link_index)
  
      def get_init_guess(self, ci_strs, norb, nelec, nroots, hdiag):
          '''Initial guess is the single Slater determinant
          '''
          na = len(ci_strs[0])
          nb = len(ci_strs[1])
          ci0 = direct_spin1._get_init_guess(na, nb, nroots, hdiag)
          return [_as_SCIvector(x, ci_strs) for x in ci0]
  
      def make_hdiag(self, h1e, eri, ci_strs, norb, nelec):
          return make_hdiag(h1e, eri, ci_strs, norb, nelec)
  
      enlarge_space = enlarge_space
      kernel = kernel_float_space
      kernel_fixed_space = kernel_fixed_space
  
  #    def approx_kernel(self, h1e, eri, norb, nelec, ci0=None, link_index=None,
  #                      tol=None, lindep=None, max_cycle=None,
  #                      max_memory=None, verbose=None, **kwargs):
  #        ci_strs = getattr(ci0, '_strs', self._strs)
  #        return self.kernel_fixed_space(h1e, eri, norb, nelec, ci_strs,
  #                                       ci0, link_index, tol, lindep, 6,
  #                                       max_memory, verbose, **kwargs)
  
      @lib.with_doc(spin_square.__doc__)
      def spin_square(self, civec_strs, norb, nelec):
          nelec = direct_spin1._unpack_nelec(nelec, self.spin)
          return spin_square(_as_SCIvector_if_not(civec_strs, self._strs), norb, nelec)
  
      def large_ci(self, civec_strs, norb, nelec, tol=.1, return_strs=True):
          nelec = direct_spin1._unpack_nelec(nelec, self.spin)
          ci, _, (strsa, strsb) = _unpack(civec_strs, nelec, self._strs)
          addra, addrb = numpy.where(abs(ci) > tol)
          if return_strs:
              strsa = [bin(x) for x in strsa[addra]]
              strsb = [bin(x) for x in strsb[addrb]]
              return list(zip(ci[addra,addrb], strsa, strsb))
          else:
              occslsta = cistring._strs2occslst(strsa[addra], norb)
              occslstb = cistring._strs2occslst(strsb[addrb], norb)
              return list(zip(ci[addra,addrb], occslsta, occslstb))
  
      def contract_ss(self, fcivec, norb, nelec):
          return contract_ss(fcivec, norb, nelec)
  
      @lib.with_doc(make_rdm1s.__doc__)
      def make_rdm1s(self, civec_strs, norb, nelec, link_index=None):
          nelec = direct_spin1._unpack_nelec(nelec, self.spin)
          civec_strs = _as_SCIvector_if_not(civec_strs, self._strs)
          return make_rdm1s(civec_strs, norb, nelec, link_index)
  
      @lib.with_doc(make_rdm1.__doc__)
      def make_rdm1(self, civec_strs, norb, nelec, link_index=None):
          nelec = direct_spin1._unpack_nelec(nelec, self.spin)
          rdm1a, rdm1b = self.make_rdm1s(civec_strs, norb, nelec, link_index)
          return rdm1a + rdm1b
  
      @lib.with_doc(make_rdm2s.__doc__)
      def make_rdm2s(self, civec_strs, norb, nelec, link_index=None, **kwargs):
          nelec = direct_spin1._unpack_nelec(nelec, self.spin)
          civec_strs = _as_SCIvector_if_not(civec_strs, self._strs)
          return make_rdm2s(civec_strs, norb, nelec, link_index)
  
      @lib.with_doc(make_rdm2.__doc__)
      def make_rdm2(self, civec_strs, norb, nelec, link_index=None, **kwargs):
          nelec = direct_spin1._unpack_nelec(nelec, self.spin)
          civec_strs = _as_SCIvector_if_not(civec_strs, self._strs)
          return make_rdm2(civec_strs, norb, nelec, link_index)
  
      def make_rdm12s(self, civec_strs, norb, nelec, link_index=None, **kwargs):
          neleca, nelecb = nelec = direct_spin1._unpack_nelec(nelec, self.spin)
          civec_strs = _as_SCIvector_if_not(civec_strs, self._strs)
          dm2aa, dm2ab, dm2bb = make_rdm2s(civec_strs, norb, nelec, link_index)
          if neleca > 1 and nelecb > 1:
              dm1a = numpy.einsum('iikl->kl', dm2aa) / (neleca-1)
              dm1b = numpy.einsum('iikl->kl', dm2bb) / (nelecb-1)
          else:
              dm1a, dm1b = make_rdm1s(civec_strs, norb, nelec, link_index)
          return (dm1a, dm1b), (dm2aa, dm2ab, dm2bb)
  
      def make_rdm12(self, civec_strs, norb, nelec, link_index=None, **kwargs):
          nelec = direct_spin1._unpack_nelec(nelec, self.spin)
          nelec_tot = sum(nelec)
          civec_strs = _as_SCIvector_if_not(civec_strs, self._strs)
          dm2 = make_rdm2(civec_strs, norb, nelec, link_index)
          if nelec_tot > 1:
              dm1 = numpy.einsum('iikl->kl', dm2) / (nelec_tot-1)
          else:
              dm1 = make_rdm1(civec_strs, norb, nelec, link_index)
          return dm1, dm2
  
      @lib.with_doc(trans_rdm1s.__doc__)
      def trans_rdm1s(self, cibra, ciket, norb, nelec, link_index=None):
          nelec = direct_spin1._unpack_nelec(nelec, self.spin)
          cibra = _as_SCIvector_if_not(cibra, self._strs)
          ciket = _as_SCIvector_if_not(ciket, self._strs)
          return trans_rdm1s(cibra, ciket, norb, nelec, link_index)
  
      @lib.with_doc(trans_rdm1.__doc__)
      def trans_rdm1(self, cibra, ciket, norb, nelec, link_index=None):
          nelec = direct_spin1._unpack_nelec(nelec, self.spin)
          cibra = _as_SCIvector_if_not(cibra, self._strs)
          ciket = _as_SCIvector_if_not(ciket, self._strs)
          return trans_rdm1(cibra, ciket, norb, nelec, link_index)
  
      def gen_linkstr(self, norb, nelec, tril=True, spin=None, ci_strs=None):
          if spin is None:
              spin = self.spin
          if ci_strs is None:
              ci_strs = self._strs
          neleca, nelecb = direct_spin1._unpack_nelec(nelec, spin)
          if tril:
              cd_indexa = cre_des_linkstr_tril(ci_strs[0], norb, neleca)
              dd_indexa = des_des_linkstr_tril(ci_strs[0], norb, neleca)
              cd_indexb = cre_des_linkstr_tril(ci_strs[1], norb, nelecb)
              dd_indexb = des_des_linkstr_tril(ci_strs[1], norb, nelecb)
          else:
              cd_indexa = cre_des_linkstr(ci_strs[0], norb, neleca)
              dd_indexa = des_des_linkstr(ci_strs[0], norb, neleca)
              cd_indexb = cre_des_linkstr(ci_strs[1], norb, nelecb)
              dd_indexb = des_des_linkstr(ci_strs[1], norb, nelecb)
          return cd_indexa, dd_indexa, cd_indexb, dd_indexb
  
  SCI = SelectedCI
  
  
  def _unpack(civec_strs, nelec, ci_strs=None, spin=None):
      neleca, nelecb = direct_spin1._unpack_nelec(nelec, spin)
      ci_strs = getattr(civec_strs, '_strs', ci_strs)
      if ci_strs is not None:
          strsa, strsb = ci_strs
          strsa = numpy.asarray(strsa)
          strsb = numpy.asarray(strsb)
          ci_strs = (strsa, strsb)
      return civec_strs, (neleca, nelecb), ci_strs
  
  def _all_linkstr_index(ci_strs, norb, nelec):
      cd_indexa = cre_des_linkstr_tril(ci_strs[0], norb, nelec[0])
      dd_indexa = des_des_linkstr_tril(ci_strs[0], norb, nelec[0])
      cd_indexb = cre_des_linkstr_tril(ci_strs[1], norb, nelec[1])
      dd_indexb = des_des_linkstr_tril(ci_strs[1], norb, nelec[1])
      return cd_indexa, dd_indexa, cd_indexb, dd_indexb
  
  # numpy.ndarray does not allow to attach attribtues.  Overwrite the
  # numpy.ndarray class to tag the ._strs attribute
  class _SCIvector(numpy.ndarray):
      def __array_finalize__(self, obj):
          self._strs = getattr(obj, '_strs', None)
  
      # Whenever the contents of the array was modified (through ufunc), the tag
      # should be expired. Overwrite the output of ufunc to restore ndarray type.
      def __array_wrap__(self, out, context=None):
          return numpy.ndarray.__array_wrap__(self, out, context).view(numpy.ndarray)
  
  def _as_SCIvector(civec, ci_strs):
      civec = civec.view(_SCIvector)
      civec._strs = ci_strs
      return civec
  
  def _as_SCIvector_if_not(civec, ci_strs):
      if getattr(civec, '_strs', None) is None:
          civec = _as_SCIvector(civec, ci_strs)
      return civec
  
  if __name__ == '__main__':
      from functools import reduce
      from pyscf import gto
      from pyscf import scf
      from pyscf import ao2mo
      from pyscf.fci import spin_op
      from pyscf.fci import addons
  
      mol = gto.Mole()
      mol.verbose = 0
      mol.output = None
      mol.atom = [
          ['H', ( 1.,-1.    , 0.   )],
          ['H', ( 0.,-1.    ,-1.   )],
          ['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_spin1.kernel(h1e, eri, norb, nelec)
      print(e1, e1 - -11.894559902235565, 'diff to FCI', e1-e2)
  
      print(c1.shape, c2.shape)
      dm1_1 = make_rdm1(c1, norb, nelec)
      dm1_2 = direct_spin1.make_rdm1(c2, norb, nelec)
      print(abs(dm1_1 - dm1_2).sum())
      dm2_1 = make_rdm2(c1, norb, nelec)
      dm2_2 = direct_spin1.make_rdm12(c2, norb, nelec)[1]
      print(abs(dm2_1 - dm2_2).sum())
  
      myci = SelectedCI()
      e, c = kernel_fixed_space(myci, h1e, eri, norb, nelec, c1._strs)
      print(e - -11.894559902235565)
  
      print(myci.large_ci(c1, norb, nelec))
      print(myci.spin_square(c1, norb, nelec)[0] -
            spin_op.spin_square0(to_fci(c1, norb, nelec), norb, nelec)[0])
  
      myci = SelectedCI()
      myci = addons.fix_spin_(myci)
      e1, c1 = myci.kernel(h1e, eri, norb, nelec)
      print(e1, e1 - -11.89467612053687)
      print(myci.spin_square(c1, norb, nelec))
   #!/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: Timothy Berkelbach <tim.berkelbach@gmail.com>
  #         Qiming Sun <osirpt.sun@gmail.com>
  #
  
  import sys
  import ctypes
  import numpy
  from pyscf import lib
  from pyscf.dft import numint
  from pyscf.dft.numint import eval_mat, _dot_ao_ao, _dot_ao_dm
  from pyscf.dft.numint import _scale_ao, _contract_rho
  from pyscf.dft.numint import _rks_gga_wv0, _rks_gga_wv1
  from pyscf.dft.numint import _uks_gga_wv0, _uks_gga_wv1
  from pyscf.dft.numint import OCCDROP
  from pyscf.pbc.dft.gen_grid import libpbc, make_mask, BLKSIZE
  from pyscf.pbc.lib.kpts_helper import is_zero, gamma_point, member
  
  #try:
  ### Moderate speedup by caching eval_ao
  #    from pyscf import pbc
  #    from joblib import Memory
  #    memory = Memory(cachedir='./tmp/', mmap_mode='r', verbose=0)
  #    def memory_cache(f):
  #        g = memory.cache(f)
  #        def maybe_cache(*args, **kwargs):
  #            if pbc.DEBUG:
  #                return g(*args, **kwargs)
  #            else:
  #                return f(*args, **kwargs)
  #        return maybe_cache
  #except:
  #    memory_cache = lambda f: f
  
  def eval_ao(cell, coords, kpt=numpy.zeros(3), deriv=0, relativity=0, shls_slice=None,
              non0tab=None, out=None, verbose=None):
      '''Collocate AO crystal orbitals (opt. gradients) on the real-space grid.
  
      Args:
          cell : instance of :class:`Cell`
  
          coords : (nx*ny*nz, 3) ndarray
              The real-space grid point coordinates.
  
      Kwargs:
          kpt : (3,) ndarray
              The k-point corresponding to the crystal AO.
          deriv : int
              AO derivative order.  It affects the shape of the return array.
              If deriv=0, the returned AO values are stored in a (N,nao) array.
              Otherwise the AO values are stored in an array of shape (M,N,nao).
              Here N is the number of grids, nao is the number of AO functions,
              M is the size associated to the derivative deriv.
  
      Returns:
          aoR : ([4,] nx*ny*nz, nao=cell.nao_nr()) ndarray
              The value of the AO crystal orbitals on the real-space grid by default.
              If deriv=1, also contains the value of the orbitals gradient in the
              x, y, and z directions.  It can be either complex or float array,
              depending on the kpt argument.  If kpt is not given (gamma point),
              aoR is a float array.
  
      See Also:
          pyscf.dft.numint.eval_ao
  
      '''
      ao_kpts = eval_ao_kpts(cell, coords, numpy.reshape(kpt, (-1,3)), deriv,
                             relativity, shls_slice, non0tab, out, verbose)
      return ao_kpts[0]
  
  
  #@memory_cache
  def eval_ao_kpts(cell, coords, kpts=None, deriv=0, relativity=0,
                   shls_slice=None, non0tab=None, out=None, verbose=None, **kwargs):
      '''
      Returns:
          ao_kpts: (nkpts, [comp], ngrids, nao) ndarray
              AO values at each k-point
      '''
      if kpts is None:
          if 'kpt' in kwargs:
              sys.stderr.write('WARN: KNumInt.eval_ao function finds keyword '
                               'argument "kpt" and converts it to "kpts"\n')
              kpts = kwargs['kpt']
          else:
              kpts = numpy.zeros((1,3))
      kpts = numpy.reshape(kpts, (-1,3))
  
      comp = (deriv+1)*(deriv+2)*(deriv+3)//6
      if cell.cart:
          feval = 'GTOval_cart_deriv%d' % deriv
      else:
          feval = 'GTOval_sph_deriv%d' % deriv
      return cell.pbc_eval_gto(feval, coords, comp, kpts,
                               shls_slice=shls_slice, non0tab=non0tab, out=out)
  
  
  def eval_rho(cell, ao, dm, non0tab=None, xctype='LDA', hermi=0, verbose=None):
      '''Collocate the *real* density (opt. gradients) on the real-space grid.
  
      Args:
          cell : instance of :class:`Mole` or :class:`Cell`
  
          ao : ([4,] nx*ny*nz, nao=cell.nao_nr()) ndarray
              The value of the AO crystal orbitals on the real-space grid by default.
              If xctype='GGA', also contains the value of the gradient in the x, y,
              and z directions.
  
      Returns:
          rho : ([4,] nx*ny*nz) ndarray
              The value of the density on the real-space grid. If xctype='GGA',
              also contains the value of the gradient in the x, y, and z
              directions.
  
      See Also:
          pyscf.dft.numint.eval_rho
  
      '''
  
      if xctype == 'LDA' or xctype == 'HF':
          ngrids, nao = ao.shape
      else:
          ngrids, nao = ao[0].shape
  
      if non0tab is None:
          non0tab = numpy.empty(((ngrids+BLKSIZE-1)//BLKSIZE, cell.nbas),
                                dtype=numpy.uint8)
          non0tab[:] = 0xff
  
      # complex orbitals or density matrix
      if numpy.iscomplexobj(ao) or numpy.iscomplexobj(dm):
          shls_slice = (0, cell.nbas)
          ao_loc = cell.ao_loc_nr()
          dm = dm.astype(numpy.complex128)
  # For GGA, function eval_rho returns   real(|\nabla i> D_ij <j| + |i> D_ij <\nabla j|)
  #       = real(|\nabla i> D_ij <j| + |i> D_ij <\nabla j|)
  #       = real(|\nabla i> D_ij <j| + conj(|\nabla j> conj(D_ij) < i|))
  #       = real(|\nabla i> D_ij <j|) + real(|\nabla j> conj(D_ij) < i|)
  #       = real(|\nabla i> [D_ij + (D^\dagger)_ij] <j|)
  # symmetrization dm (D + D.conj().T) then /2 because the code below computes
  #       2*real(|\nabla i> D_ij <j|)
          if not hermi:
              dm = (dm + dm.conj().T) * .5
  
          def dot_bra(bra, aodm):
              #:rho  = numpy.einsum('pi,pi->p', bra.real, aodm.real)
              #:rho += numpy.einsum('pi,pi->p', bra.imag, aodm.imag)
              #:return rho
              return _contract_rho(bra, aodm)
  
          if xctype == 'LDA' or xctype == 'HF':
              c0 = _dot_ao_dm(cell, ao, dm, non0tab, shls_slice, ao_loc)
              rho = dot_bra(ao, c0)
  
          elif xctype == 'GGA':
              rho = numpy.empty((4,ngrids))
              c0 = _dot_ao_dm(cell, ao[0], dm, non0tab, shls_slice, ao_loc)
              rho[0] = dot_bra(ao[0], c0)
              for i in range(1, 4):
                  rho[i] = dot_bra(ao[i], c0) * 2
  
          else:
              # rho[4] = \nabla^2 rho, rho[5] = 1/2 |nabla f|^2
              rho = numpy.empty((6,ngrids))
              c0 = _dot_ao_dm(cell, ao[0], dm, non0tab, shls_slice, ao_loc)
              rho[0] = dot_bra(ao[0], c0)
              rho[5] = 0
              for i in range(1, 4):
                  rho[i] = dot_bra(ao[i], c0) * 2  # *2 for +c.c.
                  c1 = _dot_ao_dm(cell, ao[i], dm, non0tab, shls_slice, ao_loc)
                  rho[5] += dot_bra(ao[i], c1)
              XX, YY, ZZ = 4, 7, 9
              ao2 = ao[XX] + ao[YY] + ao[ZZ]
              rho[4] = dot_bra(ao2, c0)
              rho[4] += rho[5]
              rho[4] *= 2 # *2 for +c.c.
              rho[5] *= .5
      else:
          # real orbitals and real DM
          rho = numint.eval_rho(cell, ao, dm, non0tab, xctype, hermi, verbose)
      return rho
  
  def eval_rho2(cell, ao, mo_coeff, mo_occ, non0tab=None, xctype='LDA',
                verbose=None):
      '''Refer to `pyscf.dft.numint.eval_rho2` for full documentation.
      '''
      xctype = xctype.upper()
      if xctype == 'LDA' or xctype == 'HF':
          ngrids, nao = ao.shape
      else:
          ngrids, nao = ao[0].shape
  
      if non0tab is None:
          non0tab = numpy.empty(((ngrids+BLKSIZE-1)//BLKSIZE,cell.nbas),
                               dtype=numpy.uint8)
          non0tab[:] = 0xff
  
      # complex orbitals or density matrix
      if numpy.iscomplexobj(ao) or numpy.iscomplexobj(mo_coeff):
          def dot(bra, ket):
              #:rho  = numpy.einsum('pi,pi->p', bra.real, ket.real)
              #:rho += numpy.einsum('pi,pi->p', bra.imag, ket.imag)
              #:return rho
              return _contract_rho(bra, ket)
  
          shls_slice = (0, cell.nbas)
          ao_loc = cell.ao_loc_nr()
          pos = mo_occ > OCCDROP
          cpos = numpy.einsum('ij,j->ij', mo_coeff[:,pos], numpy.sqrt(mo_occ[pos]))
  
          if pos.sum() > 0:
              if xctype == 'LDA' or xctype == 'HF':
                  c0 = _dot_ao_dm(cell, ao, cpos, non0tab, shls_slice, ao_loc)
                  rho = dot(c0, c0)
              elif xctype == 'GGA':
                  rho = numpy.empty((4,ngrids))
                  c0 = _dot_ao_dm(cell, ao[0], cpos, non0tab, shls_slice, ao_loc)
                  rho[0] = dot(c0, c0)
                  for i in range(1, 4):
                      c1 = _dot_ao_dm(cell, ao[i], cpos, non0tab, shls_slice, ao_loc)
                      rho[i] = dot(c0, c1) * 2  # *2 for +c.c.
              else: # meta-GGA
                  # rho[4] = \nabla^2 rho, rho[5] = 1/2 |nabla f|^2
                  rho = numpy.empty((6,ngrids))
                  c0 = _dot_ao_dm(cell, ao[0], cpos, non0tab, shls_slice, ao_loc)
                  rho[0] = dot(c0, c0)
                  rho[5] = 0
                  for i in range(1, 4):
                      c1 = _dot_ao_dm(cell, ao[i], cpos, non0tab, shls_slice, ao_loc)
                      rho[i] = dot(c0, c1) * 2  # *2 for +c.c.
                      rho[5]+= dot(c1, c1)
                  XX, YY, ZZ = 4, 7, 9
                  ao2 = ao[XX] + ao[YY] + ao[ZZ]
                  c1 = _dot_ao_dm(cell, ao2, cpos, non0tab, shls_slice, ao_loc)
                  rho[4] = dot(c0, c1)
                  rho[4]+= rho[5]
                  rho[4]*= 2
                  rho[5]*= .5
          else:
              if xctype == 'LDA' or xctype == 'HF':
                  rho = numpy.zeros(ngrids)
              elif xctype == 'GGA':
                  rho = numpy.zeros((4,ngrids))
              else:
                  rho = numpy.zeros((6,ngrids))
  
          neg = mo_occ < -OCCDROP
          if neg.sum() > 0:
              cneg = numpy.einsum('ij,j->ij', mo_coeff[:,neg], numpy.sqrt(-mo_occ[neg]))
              if xctype == 'LDA' or xctype == 'HF':
                  c0 = _dot_ao_dm(cell, ao, cneg, non0tab, shls_slice, ao_loc)
                  rho -= dot(c0, c0)
              elif xctype == 'GGA':
                  c0 = _dot_ao_dm(cell, ao[0], cneg, non0tab, shls_slice, ao_loc)
                  rho[0] -= dot(c0, c0)
                  for i in range(1, 4):
                      c1 = _dot_ao_dm(cell, ao[i], cneg, non0tab, shls_slice, ao_loc)
                      rho[i] -= dot(c0, c1) * 2  # *2 for +c.c.
              else:
                  c0 = _dot_ao_dm(cell, ao[0], cneg, non0tab, shls_slice, ao_loc)
                  rho[0] -= dot(c0, c0)
                  rho5 = 0
                  for i in range(1, 4):
                      c1 = _dot_ao_dm(cell, ao[i], cneg, non0tab, shls_slice, ao_loc)
                      rho[i] -= dot(c0, c1) * 2  # *2 for +c.c.
                      rho5 -= dot(c1, c1)
                  XX, YY, ZZ = 4, 7, 9
                  ao2 = ao[XX] + ao[YY] + ao[ZZ]
                  c1 = _dot_ao_dm(cell, ao2, cneg, non0tab, shls_slice, ao_loc)
                  rho[4] -= dot(c0, c1) * 2
                  rho[4] -= rho5 * 2
                  rho[5] -= rho5 * .5
      else:
          rho = numint.eval_rho2(cell, ao, mo_coeff, mo_occ, non0tab, xctype, verbose)
      return rho
  
  
  def nr_rks(ni, cell, grids, xc_code, dms, spin=0, relativity=0, hermi=0,
             kpts=None, kpts_band=None, max_memory=2000, verbose=None):
      '''Calculate RKS XC functional and potential matrix for given meshgrids and density matrix
  
      Note: This is a replica of pyscf.dft.numint.nr_rks_vxc with kpts added.
      This implemented uses slow function in numint, which only calls eval_rho, eval_mat.
      Faster function uses eval_rho2 which is not yet implemented.
  
      Args:
          ni : an instance of :class:`NumInt` or :class:`KNumInt`
  
          cell : instance of :class:`Mole` or :class:`Cell`
  
          grids : an instance of :class:`Grids`
              grids.coords and grids.weights are needed for coordinates and weights of meshgrids.
          xc_code : str
              XC functional description.
              See :func:`parse_xc` of pyscf/dft/libxc.py for more details.
          dms : 2D/3D array or a list of 2D/3D arrays
              Density matrices (2D) / density matrices for k-points (3D)
  
      Kwargs:
          spin : int
              spin polarized if spin = 1
          relativity : int
              No effects.
          hermi : int
              No effects
          max_memory : int or float
              The maximum size of cache to use (in MB).
          verbose : int or object of :class:`Logger`
              No effects.
          kpts : (3,) ndarray or (nkpts,3) ndarray
              Single or multiple k-points sampled for the DM.  Default is gamma point.
          kpts_band : (3,) ndarray or (*,3) ndarray
              A list of arbitrary "band" k-points at which to evaluate the XC matrix.
  
      Returns:
          nelec, excsum, vmat.
          nelec is the number of electrons generated by numerical integration.
          excsum is the XC functional value.  vmat is the XC potential matrix in
          2D array of shape (nao,nao) where nao is the number of AO functions.
      '''
      if kpts is None:
          kpts = numpy.zeros((1,3))
  
      xctype = ni._xc_type(xc_code)
      make_rho, nset, nao = ni._gen_rho_evaluator(cell, dms, hermi)
  
      nelec = numpy.zeros(nset)
      excsum = numpy.zeros(nset)
      vmat = [0]*nset
      if xctype == 'LDA':
          ao_deriv = 0
          for ao_k1, ao_k2, mask, weight, coords \
                  in ni.block_loop(cell, grids, nao, ao_deriv, kpts, kpts_band,
                                   max_memory):
              for i in range(nset):
                  rho = make_rho(i, ao_k2, mask, xctype)
                  exc, vxc = ni.eval_xc(xc_code, rho, 0, relativity, 1)[:2]
                  vrho = vxc[0]
                  den = rho*weight
                  nelec[i] += den.sum()
                  excsum[i] += (den*exc).sum()
                  vmat[i] += ni.eval_mat(cell, ao_k1, weight, rho, vxc,
                                         mask, xctype, 0, verbose)
      elif xctype == 'GGA':
          ao_deriv = 1
          for ao_k1, ao_k2, mask, weight, coords \
                  in ni.block_loop(cell, grids, nao, ao_deriv, kpts, kpts_band,
                                   max_memory):
              for i in range(nset):
                  rho = make_rho(i, ao_k2, mask, xctype)
                  exc, vxc = ni.eval_xc(xc_code, rho, 0, relativity, 1)[:2]
                  den = rho[0]*weight
                  nelec[i] += den.sum()
                  excsum[i] += (den*exc).sum()
                  vmat[i] += ni.eval_mat(cell, ao_k1, weight, rho, vxc,
                                         mask, xctype, 0, verbose)
      elif xctype == 'MGGA':
          if (any(x in xc_code.upper() for x in ('CC06', 'CS', 'BR89', 'MK00'))):
              raise NotImplementedError('laplacian in meta-GGA method')
          ao_deriv = 2
          for ao_k1, ao_k2, mask, weight, coords \
                  in ni.block_loop(cell, grids, nao, ao_deriv, kpts, kpts_band,
                                   max_memory):
              for i in range(nset):
                  rho = make_rho(i, ao_k2, mask, xctype)
                  exc, vxc = ni.eval_xc(xc_code, rho, 0, relativity, 1)[:2]
                  den = rho[0]*weight
                  nelec[i] += den.sum()
                  excsum[i] += (den*exc).sum()
                  vmat[i] += ni.eval_mat(cell, ao_k1, weight, rho, vxc,
                                         mask, xctype, 0, verbose)
      if nset == 1:
          nelec = nelec[0]
          excsum = excsum[0]
          vmat = vmat[0]
      return nelec, excsum, numpy.asarray(vmat)
  
  def nr_uks(ni, cell, grids, xc_code, dms, spin=1, relativity=0, hermi=0,
             kpts=None, kpts_band=None, max_memory=2000, verbose=None):
      '''Calculate UKS XC functional and potential matrix for given meshgrids and density matrix
  
      Note: This is a replica of pyscf.dft.numint.nr_rks_vxc with kpts added.
      This implemented uses slow function in numint, which only calls eval_rho, eval_mat.
      Faster function uses eval_rho2 which is not yet implemented.
  
      Args:
          ni : an instance of :class:`NumInt` or :class:`KNumInt`
  
          cell : instance of :class:`Mole` or :class:`Cell`
  
          grids : an instance of :class:`Grids`
              grids.coords and grids.weights are needed for coordinates and weights of meshgrids.
          xc_code : str
              XC functional description.
              See :func:`parse_xc` of pyscf/dft/libxc.py for more details.
          dms :
              Density matrices
  
      Kwargs:
          spin : int
              spin polarized if spin = 1
          relativity : int
              No effects.
          hermi : int
              Input density matrices symmetric or not
          max_memory : int or float
              The maximum size of cache to use (in MB).
          verbose : int or object of :class:`Logger`
              No effects.
          kpts : (3,) ndarray or (nkpts,3) ndarray
              Single or multiple k-points sampled for the DM.  Default is gamma point.
              kpts_band : (3,) ndarray or (*,3) ndarray
              A list of arbitrary "band" k-points at which to evaluate the XC matrix.
  
      Returns:
          nelec, excsum, vmat.
          nelec is the number of electrons generated by numerical integration.
          excsum is the XC functional value.  vmat is the XC potential matrix in
          2D array of shape (nao,nao) where nao is the number of AO functions.
      '''
      if kpts is None:
          kpts = numpy.zeros((1,3))
  
      xctype = ni._xc_type(xc_code)
      dma, dmb = _format_uks_dm(dms)
      nao = dma.shape[-1]
      make_rhoa, nset = ni._gen_rho_evaluator(cell, dma, hermi)[:2]
      make_rhob       = ni._gen_rho_evaluator(cell, dmb, hermi)[0]
  
      nelec = numpy.zeros((2,nset))
      excsum = numpy.zeros(nset)
      vmata = [0]*nset
      vmatb = [0]*nset
      if xctype == 'LDA':
          ao_deriv = 0
          for ao_k1, ao_k2, mask, weight, coords \
                  in ni.block_loop(cell, grids, nao, ao_deriv, kpts, kpts_band,
                                   max_memory):
              for i in range(nset):
                  rho_a = make_rhoa(i, ao_k2, mask, xctype)
                  rho_b = make_rhob(i, ao_k2, mask, xctype)
                  exc, vxc = ni.eval_xc(xc_code, (rho_a, rho_b),
                                        1, relativity, 1, verbose)[:2]
                  vrho = vxc[0]
                  den = rho_a * weight
                  nelec[0,i] += den.sum()
                  excsum[i] += (den*exc).sum()
                  den = rho_b * weight
                  nelec[1,i] += den.sum()
                  excsum[i] += (den*exc).sum()
  
                  vmata[i] += ni.eval_mat(cell, ao_k1, weight, rho_a, vrho[:,0],
                                          mask, xctype, 1, verbose)
                  vmatb[i] += ni.eval_mat(cell, ao_k1, weight, rho_b, vrho[:,1],
                                          mask, xctype, 1, verbose)
      elif xctype == 'GGA':
          ao_deriv = 1
          for ao_k1, ao_k2, mask, weight, coords \
                  in ni.block_loop(cell, grids, nao, ao_deriv, kpts,
                                   kpts_band, max_memory):
              for i in range(nset):
                  rho_a = make_rhoa(i, ao_k2, mask, xctype)
                  rho_b = make_rhob(i, ao_k2, mask, xctype)
                  exc, vxc = ni.eval_xc(xc_code, (rho_a, rho_b),
                                        1, relativity, 1, verbose)[:2]
                  vrho, vsigma = vxc[:2]
                  den = rho_a[0]*weight
                  nelec[0,i] += den.sum()
                  excsum[i] += (den*exc).sum()
                  den = rho_b[0]*weight
                  nelec[1,i] += den.sum()
                  excsum[i] += (den*exc).sum()
  
                  vmata[i] += ni.eval_mat(cell, ao_k1, weight, (rho_a,rho_b),
                                          (vrho[:,0], (vsigma[:,0],vsigma[:,1])),
                                          mask, xctype, 1, verbose)
                  vmatb[i] += ni.eval_mat(cell, ao_k1, weight, (rho_b,rho_a),
                                          (vrho[:,1], (vsigma[:,2],vsigma[:,1])),
                                          mask, xctype, 1, verbose)
      elif xctype == 'MGGA':
          assert(all(x not in xc_code.upper() for x in ('CC06', 'CS', 'BR89', 'MK00')))
          ao_deriv = 2
          for ao_k1, ao_k2, mask, weight, coords \
                  in ni.block_loop(cell, grids, nao, ao_deriv, kpts, kpts_band,
                                   max_memory):
              for i in range(nset):
                  rho_a = make_rhoa(i, ao_k2, mask, xctype)
                  rho_b = make_rhob(i, ao_k2, mask, xctype)
                  exc, vxc = ni.eval_xc(xc_code, (rho_a, rho_b),
                                        1, relativity, 1, verbose)[:2]
                  vrho, vsigma, vlapl, vtau = vxc
                  den = rho_a[0]*weight
                  nelec[0,i] += den.sum()
                  excsum[i] += (den*exc).sum()
                  den = rho_b[0]*weight
                  nelec[1,i] += den.sum()
                  excsum[i] += (den*exc).sum()
  
                  v = (vrho[:,0], (vsigma[:,0],vsigma[:,1]), None, vtau[:,0])
                  vmata[i] += ni.eval_mat(cell, ao_k1, weight, (rho_a,rho_b), v,
                                          mask, xctype, 1, verbose)
                  v = (vrho[:,1], (vsigma[:,2],vsigma[:,1]), None, vtau[:,1])
                  vmatb[i] += ni.eval_mat(cell, ao_k1, weight, (rho_b,rho_a), v,
                                          mask, xctype, 1, verbose)
                  v = None
  
      if dma.ndim == vmata[0].ndim:  # One set of DMs in the input
          nelec = nelec[:,0]
          excsum = excsum[0]
          vmata = vmata[0]
          vmatb = vmatb[0]
      return nelec, excsum, numpy.asarray((vmata,vmatb))
  
  def _format_uks_dm(dms):
      dma, dmb = dms
      if getattr(dms, 'mo_coeff', None) is not None:
  #TODO: test whether dm.mo_coeff matching dm
          mo_coeff = dms.mo_coeff
          mo_occ = dms.mo_occ
          if (isinstance(mo_coeff[0], numpy.ndarray) and
              mo_coeff[0].ndim < dma.ndim): # handle ROKS
              mo_occa = [numpy.array(occ> 0, dtype=numpy.double) for occ in mo_occ]
              mo_occb = [numpy.array(occ==2, dtype=numpy.double) for occ in mo_occ]
              dma = lib.tag_array(dma, mo_coeff=mo_coeff, mo_occ=mo_occa)
              dmb = lib.tag_array(dmb, mo_coeff=mo_coeff, mo_occ=mo_occb)
          else:
              dma = lib.tag_array(dma, mo_coeff=mo_coeff[0], mo_occ=mo_occ[0])
              dmb = lib.tag_array(dmb, mo_coeff=mo_coeff[1], mo_occ=mo_occ[1])
      return dma, dmb
  
  nr_rks_vxc = nr_rks
  nr_uks_vxc = nr_uks
  
  def nr_rks_fxc(ni, cell, grids, xc_code, dm0, dms, relativity=0, hermi=0,
                 rho0=None, vxc=None, fxc=None, kpts=None, max_memory=2000,
                 verbose=None):
      '''Contract RKS XC kernel matrix with given density matrices
  
      Args:
          ni : an instance of :class:`NumInt` or :class:`KNumInt`
  
          cell : instance of :class:`Mole` or :class:`Cell`
  
          grids : an instance of :class:`Grids`
              grids.coords and grids.weights are needed for coordinates and weights of meshgrids.
          xc_code : str
              XC functional description.
              See :func:`parse_xc` of pyscf/dft/libxc.py for more details.
          dms : 2D/3D array or a list of 2D/3D arrays
              Density matrices (2D) / density matrices for k-points (3D)
  
      Kwargs:
          hermi : int
              Input density matrices symmetric or not
          max_memory : int or float
              The maximum size of cache to use (in MB).
          rho0 : float array
              Zero-order density (and density derivative for GGA).  Giving kwargs rho0,
              vxc and fxc to improve better performance.
          vxc : float array
              First order XC derivatives
          fxc : float array
              Second order XC derivatives
  
      Examples:
  
      '''
      if kpts is None:
          kpts = numpy.zeros((1,3))
      xctype = ni._xc_type(xc_code)
  
      make_rho, nset, nao = ni._gen_rho_evaluator(cell, dms, hermi)
      if ((xctype == 'LDA' and fxc is None) or
          (xctype == 'GGA' and rho0 is None)):
          make_rho0 = ni._gen_rho_evaluator(cell, dm0, 1)[0]
  
      ao_loc = cell.ao_loc_nr()
      vmat = [0] * nset
      if xctype == 'LDA':
          ao_deriv = 0
          ip = 0
          for ao_k1, ao_k2, mask, weight, coords \
                  in ni.block_loop(cell, grids, nao, ao_deriv, kpts, None, max_memory):
              ngrid = weight.size
              if fxc is None:
                  rho = make_rho0(0, ao_k1, mask, xctype)
                  fxc0 = ni.eval_xc(xc_code, rho, 0, relativity, 2, verbose)[2]
                  frr = fxc0[0]
              else:
                  frr = fxc[0][ip:ip+ngrid]
                  ip += ngrid
  
              for i in range(nset):
                  rho1 = make_rho(i, ao_k1, mask, xctype)
                  wv = weight * frr * rho1
                  vmat[i] += ni._fxc_mat(cell, ao_k1, wv, mask, xctype, ao_loc)
  
      elif xctype == 'GGA':
          ao_deriv = 1
          ip = 0
          for ao_k1, ao_k2, mask, weight, coords \
                  in ni.block_loop(cell, grids, nao, ao_deriv, kpts, None, max_memory):
              ngrid = weight.size
              if rho0 is None:
                  rho = make_rho0(0, ao_k1, mask, xctype)
              else:
                  rho = numpy.asarray(rho0[:,ip:ip+ngrid], order='C')
  
              if vxc is None or fxc is None:
                  vxc0, fxc0 = ni.eval_xc(xc_code, rho, 0, relativity, 2, verbose)[1:3]
              else:
                  vxc0 = (None, vxc[1][ip:ip+ngrid])
                  fxc0 = (fxc[0][ip:ip+ngrid], fxc[1][ip:ip+ngrid], fxc[2][ip:ip+ngrid])
                  ip += ngrid
  
              for i in range(nset):
                  rho1 = make_rho(i, ao_k1, mask, xctype)
                  wv = _rks_gga_wv1(rho, rho1, vxc0, fxc0, weight)
                  vmat[i] += ni._fxc_mat(cell, ao_k1, wv, mask, xctype, ao_loc)
  
          # call swapaxes method to swap last two indices because vmat may be a 3D
          # array (nset,nao,nao) in single k-point mode or a 4D array
          # (nset,nkpts,nao,nao) in k-points mode
          for i in range(nset):  # for (\nabla\mu) \nu + \mu (\nabla\nu)
              vmat[i] = vmat[i] + vmat[i].swapaxes(-2,-1).conj()
  
      elif xctype == 'MGGA':
          raise NotImplementedError('meta-GGA')
  
      if isinstance(dms, numpy.ndarray) and dms.ndim == vmat[0].ndim:
          # One set of DMs in the input
          vmat = vmat[0]
      return numpy.asarray(vmat)
  
  def nr_rks_fxc_st(ni, cell, grids, xc_code, dm0, dms_alpha, relativity=0, singlet=True,
                    rho0=None, vxc=None, fxc=None, kpts=None, max_memory=2000,
                    verbose=None):
      '''Associated to singlet or triplet Hessian
      Note the difference to nr_rks_fxc, dms_alpha is the response density
      matrices of alpha spin, alpha+/-beta DM is applied due to singlet/triplet
      coupling
  
      Ref. CPL, 256, 454
      '''
      if kpts is None:
          kpts = numpy.zeros((1,3))
      xctype = ni._xc_type(xc_code)
  
      make_rho, nset, nao = ni._gen_rho_evaluator(cell, dms_alpha)
      if ((xctype == 'LDA' and fxc is None) or
          (xctype == 'GGA' and rho0 is None)):
          make_rho0 = ni._gen_rho_evaluator(cell, dm0, 1)[0]
  
      ao_loc = cell.ao_loc_nr()
      vmat = [0] * nset
      if xctype == 'LDA':
          ao_deriv = 0
          ip = 0
          for ao_k1, ao_k2, mask, weight, coords \
                  in ni.block_loop(cell, grids, nao, ao_deriv, kpts, None, max_memory):
              ngrid = weight.size
              if fxc is None:
                  rho = make_rho0(0, ao_k1, mask, xctype)
                  rho *= .5  # alpha density
                  fxc0 = ni.eval_xc(xc_code, (rho,rho), 1, deriv=2)[2]
                  u_u, u_d, d_d = fxc0[0].T
              else:
                  u_u, u_d, d_d = fxc[0][ip:ip+ngrid].T
                  ip += ngrid
              if singlet:
                  frho = u_u + u_d
              else:
                  frho = u_u - u_d
  
              for i in range(nset):
                  rho1 = make_rho(i, ao_k1, mask, xctype)
                  wv = weight * frho * rho1
                  vmat[i] += ni._fxc_mat(cell, ao_k1, wv, mask, xctype, ao_loc)
  
      elif xctype == 'GGA':
          ao_deriv = 1
          ip = 0
          for ao_k1, ao_k2, mask, weight, coords \
                  in ni.block_loop(cell, grids, nao, ao_deriv, kpts, None, max_memory):
              ngrid = weight.size
              if vxc is None or fxc is None:
                  rho = make_rho0(0, ao_k1, mask, xctype)
                  rho *= .5  # alpha density
                  vxc0, fxc0 = ni.eval_xc(xc_code, (rho,rho), 1, deriv=2)[1:3]
  
                  vsigma = vxc0[1].T
                  u_u, u_d, d_d = fxc0[0].T  # v2rho2
                  u_uu, u_ud, u_dd, d_uu, d_ud, d_dd = fxc0[1].T  # v2rhosigma
                  uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd = fxc0[2].T  # v2sigma2
              else:
                  rho = rho0[0][:,ip:ip+ngrid]
                  vsigma = vxc[1][ip:ip+ngrid].T
                  u_u, u_d, d_d = fxc[0][ip:ip+ngrid].T  # v2rho2
                  u_uu, u_ud, u_dd, d_uu, d_ud, d_dd = fxc[1][ip:ip+ngrid].T  # v2rhosigma
                  uu_uu, uu_ud, uu_dd, ud_ud, ud_dd, dd_dd = fxc[2][ip:ip+ngrid].T  # v2sigma2
  
              if singlet:
                  fgamma = vsigma[0] + vsigma[1] * .5
                  frho = u_u + u_d
                  fgg = uu_uu + .5*ud_ud + 2*uu_ud + uu_dd
                  frhogamma = u_uu + u_dd + u_ud
              else:
                  fgamma = vsigma[0] - vsigma[1] * .5
                  frho = u_u - u_d
                  fgg = uu_uu - uu_dd
                  frhogamma = u_uu - u_dd
  
              for i in range(nset):
                  # rho1[0 ] = |b><j| z_{bj}
                  # rho1[1:] = \nabla(|b><j|) z_{bj}
                  rho1 = make_rho(i, ao_k1, mask, xctype)
                  wv = _rks_gga_wv1(rho, rho1, (None,fgamma),
                                    (frho,frhogamma,fgg), weight)
                  vmat[i] += ni._fxc_mat(cell, ao_k1, wv, mask, xctype, ao_loc)
  
          for i in range(nset):  # for (\nabla\mu) \nu + \mu (\nabla\nu)
              vmat[i] = vmat[i] + vmat[i].swapaxes(-2,-1).conj()
  
      elif xctype == 'MGGA':
          raise NotImplementedError('meta-GGA')
  
      if isinstance(dms_alpha, numpy.ndarray) and dms_alpha.ndim == vmat[0].ndim:
          vmat = vmat[0]
      return numpy.asarray(vmat)
  
  
  def nr_uks_fxc(ni, cell, grids, xc_code, dm0, dms, relativity=0, hermi=0,
                 rho0=None, vxc=None, fxc=None, kpts=None, max_memory=2000,
                 verbose=None):
      '''Contract UKS XC kernel matrix with given density matrices
  
      Args:
          ni : an instance of :class:`NumInt` or :class:`KNumInt`
  
          cell : instance of :class:`Mole` or :class:`Cell`
  
          grids : an instance of :class:`Grids`
              grids.coords and grids.weights are needed for coordinates and weights of meshgrids.
          xc_code : str
              XC functional description.
              See :func:`parse_xc` of pyscf/dft/libxc.py for more details.
          dms : 2D array a list of 2D arrays
              Density matrix or multiple density matrices
  
      Kwargs:
          hermi : int
              Input density matrices symmetric or not
          max_memory : int or float
              The maximum size of cache to use (in MB).
          rho0 : float array
              Zero-order density (and density derivative for GGA).  Giving kwargs rho0,
              vxc and fxc to improve better performance.
          vxc : float array
              First order XC derivatives
          fxc : float array
              Second order XC derivatives
  
      Returns:
          nelec, excsum, vmat.
          nelec is the number of electrons generated by numerical integration.
          excsum is the XC functional value.  vmat is the XC potential matrix in
          2D array of shape (nao,nao) where nao is the number of AO functions.
  
      Examples:
  
      '''
      if kpts is None:
          kpts = numpy.zeros((1,3))
      xctype = ni._xc_type(xc_code)
  
      dma, dmb = _format_uks_dm(dms)
      nao = dma.shape[-1]
      make_rhoa, nset = ni._gen_rho_evaluator(cell, dma, hermi)[:2]
      make_rhob       = ni._gen_rho_evaluator(cell, dmb, hermi)[0]
  
      if ((xctype == 'LDA' and fxc is None) or
          (xctype == 'GGA' and rho0 is None)):
          dm0a, dm0b = _format_uks_dm(dm0)
          make_rho0a = ni._gen_rho_evaluator(cell, dm0a, 1)[0]
          make_rho0b = ni._gen_rho_evaluator(cell, dm0b, 1)[0]
  
      shls_slice = (0, cell.nbas)
      ao_loc = cell.ao_loc_nr()
  
      vmata = [0] * nset
      vmatb = [0] * nset
      if xctype == 'LDA':
          ao_deriv = 0
          ip = 0
          for ao_k1, ao_k2, mask, weight, coords \
                  in ni.block_loop(cell, grids, nao, ao_deriv, kpts, None, max_memory):
              ngrid = weight.size
              if fxc is None:
                  rho0a = make_rho0a(0, ao_k1, mask, xctype)
                  rho0b = make_rho0b(0, ao_k1, mask, xctype)
                  fxc0 = ni.eval_xc(xc_code, (rho0a,rho0b), 1, relativity, 2, verbose)[2]
                  u_u, u_d, d_d = fxc0[0].T
              else:
                  u_u, u_d, d_d = fxc[0][ip:ip+ngrid].T
                  ip += ngrid
  
              for i in range(nset):
                  rho1a = make_rhoa(i, ao_k1, mask, xctype)
                  rho1b = make_rhob(i, ao_k1, mask, xctype)
                  wv = u_u * rho1a + u_d * rho1b
                  wv *= weight
                  vmata[i] += ni._fxc_mat(cell, ao_k1, wv, mask, xctype, ao_loc)
                  wv = u_d * rho1a + d_d * rho1b
                  wv *= weight
                  vmatb[i] += ni._fxc_mat(cell, ao_k1, wv, mask, xctype, ao_loc)
  
      elif xctype == 'GGA':
          ao_deriv = 1
          ip = 0
          for ao_k1, ao_k2, mask, weight, coords \
                  in ni.block_loop(cell, grids, nao, ao_deriv, kpts, None, max_memory):
              ngrid = weight.size
              if rho0 is None:
                  rho0a = make_rho0a(0, ao_k1, mask, xctype)
                  rho0b = make_rho0b(0, ao_k1, mask, xctype)
              else:
                  rho0a = rho0[0][:,ip:ip+ngrid]
                  rho0b = rho0[1][:,ip:ip+ngrid]
              if vxc is None or fxc is None:
                  vxc0, fxc0 = ni.eval_xc(xc_code, (rho0a,rho0b), 1, relativity, 2, verbose)[1:3]
              else:
                  vxc0 = (None, vxc[1][ip:ip+ngrid])
                  fxc0 = (fxc[0][ip:ip+ngrid], fxc[1][ip:ip+ngrid], fxc[2][ip:ip+ngrid])
                  ip += ngrid
  
              for i in range(nset):
                  rho1a = make_rhoa(i, ao_k1, mask, xctype)
                  rho1b = make_rhob(i, ao_k1, mask, xctype)
                  wva, wvb = _uks_gga_wv1((rho0a,rho0b), (rho1a,rho1b),
                                          vxc0, fxc0, weight)
                  vmata[i] += ni._fxc_mat(cell, ao_k1, wva, mask, xctype, ao_loc)
                  vmatb[i] += ni._fxc_mat(cell, ao_k1, wvb, mask, xctype, ao_loc)
  
          for i in range(nset):  # for (\nabla\mu) \nu + \mu (\nabla\nu)
              vmata[i] = vmata[i] + vmata[i].swapaxes(-1,-2).conj()
              vmatb[i] = vmatb[i] + vmatb[i].swapaxes(-1,-2).conj()
      elif xctype == 'MGGA':
          raise NotImplementedError('meta-GGA')
  
      if dma.ndim == vmata[0].ndim:  # One set of DMs in the input
          vmata = vmata[0]
          vmatb = vmatb[0]
      return numpy.asarray((vmata,vmatb))
  
  def _fxc_mat(cell, ao, wv, non0tab, xctype, ao_loc):
      shls_slice = (0, cell.nbas)
  
      if xctype == 'LDA' or xctype == 'HF':
          #:aow = numpy.einsum('pi,p->pi', ao, wv)
          aow = _scale_ao(ao, wv)
          mat = _dot_ao_ao(cell, ao, aow, non0tab, shls_slice, ao_loc)
      else:
          #:aow = numpy.einsum('npi,np->pi', ao, wv)
          aow = _scale_ao(ao, wv)
          mat = _dot_ao_ao(cell, ao[0], aow, non0tab, shls_slice, ao_loc)
      return mat
  
  def cache_xc_kernel(ni, cell, grids, xc_code, mo_coeff, mo_occ, spin=0,
                      kpts=None, max_memory=2000):
      '''Compute the 0th order density, Vxc and fxc.  They can be used in TDDFT,
      DFT hessian module etc.
      '''
      if kpts is None:
          kpts = numpy.zeros((1,3))
      xctype = ni._xc_type(xc_code)
      ao_deriv = 0
      if xctype == 'GGA':
          ao_deriv = 1
      elif xctype == 'MGGA':
          raise NotImplementedError('meta-GGA')
  
      nao = cell.nao_nr()
      if spin == 0:
          rho = []
          for ao_k1, ao_k2, mask, weight, coords \
                  in ni.block_loop(cell, grids, nao, ao_deriv, kpts, None, max_memory):
              rho.append(ni.eval_rho2(cell, ao_k1, mo_coeff, mo_occ, mask, xctype))
          rho = numpy.hstack(rho)
      else:
          rhoa = []
          rhob = []
          for ao_k1, ao_k2, mask, weight, coords \
                  in ni.block_loop(cell, grids, nao, ao_deriv, kpts, None, max_memory):
              rhoa.append(ni.eval_rho2(cell, ao_k1, mo_coeff[0], mo_occ[0], mask, xctype))
              rhob.append(ni.eval_rho2(cell, ao_k1, mo_coeff[1], mo_occ[1], mask, xctype))
          rho = (numpy.hstack(rhoa), numpy.hstack(rhob))
      vxc, fxc = ni.eval_xc(xc_code, rho, spin, 0, 2, 0)[1:3]
      return rho, vxc, fxc
  
  
  def get_rho(ni, cell, dm, grids, kpts=numpy.zeros((1,3)), max_memory=2000):
      '''Density in real space
      '''
      make_rho, nset, nao = ni._gen_rho_evaluator(cell, dm)
      assert(nset == 1)
      rho = numpy.empty(grids.weights.size)
      p1 = 0
      for ao_k1, ao_k2, mask, weight, coords \
              in ni.block_loop(cell, grids, nao, 0, kpts, None, max_memory):
          p0, p1 = p1, p1 + weight.size
          rho[p0:p1] = make_rho(0, ao_k1, mask, 'LDA')
      return rho
  
  
  class NumInt(numint.NumInt):
      '''Generalization of pyscf's NumInt class for a single k-point shift and
      periodic images.
      '''
      def eval_ao(self, cell, coords, kpt=numpy.zeros(3), deriv=0, relativity=0,
                  shls_slice=None, non0tab=None, out=None, verbose=None):
          return eval_ao(cell, coords, kpt, deriv, relativity, shls_slice,
                         non0tab, out, verbose)
  
      @lib.with_doc(make_mask.__doc__)
      def make_mask(self, cell, coords, relativity=0, shls_slice=None,
                    verbose=None):
          return make_mask(cell, coords, relativity, shls_slice, verbose)
  
      @lib.with_doc(eval_rho.__doc__)
      def eval_rho(self, cell, ao, dm, non0tab=None, xctype='LDA', hermi=0, verbose=None):
          return eval_rho(cell, ao, dm, non0tab, xctype, hermi, verbose)
  
      def eval_rho2(self, cell, ao, mo_coeff, mo_occ, non0tab=None, xctype='LDA',
                    verbose=None):
          return eval_rho2(cell, ao, mo_coeff, mo_occ, non0tab, xctype, verbose)
  
      def nr_vxc(self, cell, grids, xc_code, dms, spin=0, relativity=0, hermi=0,
                 kpt=None, kpts_band=None, max_memory=2000, verbose=None):
          '''Evaluate RKS/UKS XC functional and potential matrix.
          See :func:`nr_rks` and :func:`nr_uks` for more details.
          '''
          if spin == 0:
              return self.nr_rks(cell, grids, xc_code, dms, hermi,
                                 kpt, kpts_band, max_memory, verbose)
          else:
              return self.nr_uks(cell, grids, xc_code, dms, hermi,
                                 kpt, kpts_band, max_memory, verbose)
  
      @lib.with_doc(nr_rks.__doc__)
      def nr_rks(self, cell, grids, xc_code, dms, hermi=0,
                 kpt=numpy.zeros(3), kpts_band=None, max_memory=2000, verbose=None):
          if kpts_band is not None:
  # To compute Vxc on kpts_band, convert the NumInt object to KNumInt object.
              ni = KNumInt()
              ni.__dict__.update(self.__dict__)
              nao = dms.shape[-1]
              return ni.nr_rks(cell, grids, xc_code, dms.reshape(-1,1,nao,nao),
                               hermi, kpt.reshape(1,3), kpts_band, max_memory,
                               verbose)
          return nr_rks(self, cell, grids, xc_code, dms,
                        0, 0, hermi, kpt, kpts_band, max_memory, verbose)
  
      @lib.with_doc(nr_uks.__doc__)
      def nr_uks(self, cell, grids, xc_code, dms, hermi=0,
                 kpt=numpy.zeros(3), kpts_band=None, max_memory=2000, verbose=None):
          if kpts_band is not None:
  # To compute Vxc on kpts_band, convert the NumInt object to KNumInt object.
              ni = KNumInt()
              ni.__dict__.update(self.__dict__)
              nao = dms[0].shape[-1]
              return ni.nr_uks(cell, grids, xc_code, dms.reshape(-1,1,nao,nao),
                               hermi, kpt.reshape(1,3), kpts_band, max_memory,
                               verbose)
          return nr_uks(self, cell, grids, xc_code, dms,
                        1, 0, hermi, kpt, kpts_band, max_memory, verbose)
  
      def eval_mat(self, cell, ao, weight, rho, vxc,
                   non0tab=None, xctype='LDA', spin=0, verbose=None):
  # Guess whether ao is evaluated for kpts_band.  When xctype is LDA, ao on grids
  # should be a 2D array.  For other xc functional, ao should be a 3D array.
          if ao.ndim == 2 or (xctype != 'LDA' and ao.ndim == 3):
              mat = eval_mat(cell, ao, weight, rho, vxc, non0tab, xctype, spin, verbose)
          else:
              nkpts = len(ao)
              nao = ao[0].shape[-1]
              mat = numpy.empty((nkpts,nao,nao), dtype=numpy.complex128)
              for k in range(nkpts):
                  mat[k] = eval_mat(cell, ao[k], weight, rho, vxc,
                                    non0tab, xctype, spin, verbose)
          return mat
  
      def _fxc_mat(self, cell, ao, wv, non0tab, xctype, ao_loc):
          return _fxc_mat(cell, ao, wv, non0tab, xctype, ao_loc)
  
      def block_loop(self, cell, grids, nao, deriv=0, kpt=numpy.zeros(3),
                     kpts_band=None, max_memory=2000, non0tab=None, blksize=None):
          '''Define this macro to loop over grids by blocks.
          '''
  # For UniformGrids, grids.coords does not indicate whehter grids are initialized
          if grids.non0tab is None:
              grids.build(with_non0tab=True)
          grids_coords = grids.coords
          grids_weights = grids.weights
          ngrids = grids_coords.shape[0]
          comp = (deriv+1)*(deriv+2)*(deriv+3)//6
  # NOTE to index grids.non0tab, the blksize needs to be the integer multiplier of BLKSIZE
          if blksize is None:
              blksize = int(max_memory*1e6/(comp*2*nao*16*BLKSIZE))*BLKSIZE
              blksize = max(BLKSIZE, min(blksize, ngrids, BLKSIZE*1200))
          if non0tab is None:
              non0tab = grids.non0tab
          if non0tab is None:
              non0tab = numpy.empty(((ngrids+BLKSIZE-1)//BLKSIZE,cell.nbas),
                                    dtype=numpy.uint8)
              non0tab[:] = 0xff
          kpt = numpy.reshape(kpt, 3)
          if kpts_band is None:
              kpt1 = kpt2 = kpt
          else:
              kpt1 = kpts_band
              kpt2 = kpt
  
          for ip0 in range(0, ngrids, blksize):
              ip1 = min(ngrids, ip0+blksize)
              coords = grids_coords[ip0:ip1]
              weight = grids_weights[ip0:ip1]
              non0 = non0tab[ip0//BLKSIZE:]
              ao_k2 = self.eval_ao(cell, coords, kpt2, deriv=deriv, non0tab=non0)
              if abs(kpt1-kpt2).sum() < 1e-9:
                  ao_k1 = ao_k2
              else:
                  ao_k1 = self.eval_ao(cell, coords, kpt1, deriv=deriv)
              yield ao_k1, ao_k2, non0, weight, coords
              ao_k1 = ao_k2 = None
  
      def _gen_rho_evaluator(self, cell, dms, hermi=0):
          return numint.NumInt._gen_rho_evaluator(self, cell, dms, hermi)
  
      nr_rks_fxc = nr_rks_fxc
      nr_uks_fxc = nr_uks_fxc
      cache_xc_kernel  = cache_xc_kernel
      get_rho = get_rho
  
      def rsh_and_hybrid_coeff(self, xc_code, spin=0):
          omega, alpha, hyb = numint.NumInt.rsh_and_hybrid_coeff(self, xc_code, spin)
          if abs(omega) > 1e-10:
              raise NotImplementedError
          return omega, alpha, hyb
  _NumInt = NumInt
  
  
  class KNumInt(numint.NumInt):
      '''Generalization of pyscf's NumInt class for k-point sampling and
      periodic images.
      '''
      def __init__(self, kpts=numpy.zeros((1,3))):
          numint.NumInt.__init__(self)
          self.kpts = numpy.reshape(kpts, (-1,3))
  
      def eval_ao(self, cell, coords, kpts=numpy.zeros((1,3)), deriv=0, relativity=0,
                  shls_slice=None, non0tab=None, out=None, verbose=None, **kwargs):
          return eval_ao_kpts(cell, coords, kpts, deriv,
                              relativity, shls_slice, non0tab, out, verbose)
  
      @lib.with_doc(make_mask.__doc__)
      def make_mask(self, cell, coords, relativity=0, shls_slice=None,
                    verbose=None):
          return make_mask(cell, coords, relativity, shls_slice, verbose)
  
      def eval_rho(self, cell, ao_kpts, dm_kpts, non0tab=None, xctype='LDA',
                   hermi=0, verbose=None):
          '''Collocate the *real* density (opt. gradients) on the real-space grid.
  
          Args:
              cell : Mole or Cell object
              ao_kpts : (nkpts, ngrids, nao) ndarray
                  AO values at each k-point
              dm_kpts: (nkpts, nao, nao) ndarray
                  Density matrix at each k-point
  
          Returns:
             rhoR : (ngrids,) ndarray
          '''
          nkpts = len(ao_kpts)
          rhoR = 0
          for k in range(nkpts):
              rhoR += eval_rho(cell, ao_kpts[k], dm_kpts[k], non0tab, xctype,
                               hermi, verbose)
          rhoR *= 1./nkpts
          return rhoR
  
      def eval_rho2(self, cell, ao_kpts, mo_coeff_kpts, mo_occ_kpts,
                    non0tab=None, xctype='LDA', verbose=None):
          nkpts = len(ao_kpts)
          rhoR = 0
          for k in range(nkpts):
              rhoR += eval_rho2(cell, ao_kpts[k], mo_coeff_kpts[k],
                                mo_occ_kpts[k], non0tab, xctype, verbose)
          rhoR *= 1./nkpts
          return rhoR
  
      def nr_vxc(self, cell, grids, xc_code, dms, spin=0, relativity=0, hermi=0,
                 kpts=None, kpts_band=None, max_memory=2000, verbose=None):
          '''Evaluate RKS/UKS XC functional and potential matrix.
          See :func:`nr_rks` and :func:`nr_uks` for more details.
          '''
          if spin == 0:
              return self.nr_rks(cell, grids, xc_code, dms, hermi,
                                 kpts, kpts_band, max_memory, verbose)
          else:
              return self.nr_uks(cell, grids, xc_code, dms, hermi,
                                 kpts, kpts_band, max_memory, verbose)
  
      @lib.with_doc(nr_rks.__doc__)
      def nr_rks(self, cell, grids, xc_code, dms, hermi=0, kpts=None, kpts_band=None,
                 max_memory=2000, verbose=None, **kwargs):
          if kpts is None:
              if 'kpt' in kwargs:
                  sys.stderr.write('WARN: KNumInt.nr_rks function finds keyword '
                                   'argument "kpt" and converts it to "kpts"\n')
                  kpts = kwargs['kpt']
              else:
                  kpts = self.kpts
          kpts = kpts.reshape(-1,3)
  
          return nr_rks(self, cell, grids, xc_code, dms, 0, 0,
                        hermi, kpts, kpts_band, max_memory, verbose)
  
      @lib.with_doc(nr_uks.__doc__)
      def nr_uks(self, cell, grids, xc_code, dms, hermi=0, kpts=None, kpts_band=None,
                 max_memory=2000, verbose=None, **kwargs):
          if kpts is None:
              if 'kpt' in kwargs:
                  sys.stderr.write('WARN: KNumInt.nr_uks function finds keyword '
                                   'argument "kpt" and converts it to "kpts"\n')
                  kpts = kwargs['kpt']
              else:
                  kpts = self.kpts
          kpts = kpts.reshape(-1,3)
  
          return nr_uks(self, cell, grids, xc_code, dms, 1, 0,
                        hermi, kpts, kpts_band, max_memory, verbose)
  
      def eval_mat(self, cell, ao_kpts, weight, rho, vxc,
                   non0tab=None, xctype='LDA', spin=0, verbose=None):
          nkpts = len(ao_kpts)
          nao = ao_kpts[0].shape[-1]
          dtype = numpy.result_type(*ao_kpts)
          mat = numpy.empty((nkpts,nao,nao), dtype=dtype)
          for k in range(nkpts):
              mat[k] = eval_mat(cell, ao_kpts[k], weight, rho, vxc,
                                non0tab, xctype, spin, verbose)
          return mat
  
      def _fxc_mat(self, cell, ao_kpts, wv, non0tab, xctype, ao_loc):
          nkpts = len(ao_kpts)
          nao = ao_kpts[0].shape[-1]
          dtype = numpy.result_type(*ao_kpts)
          mat = numpy.empty((nkpts,nao,nao), dtype=dtype)
          for k in range(nkpts):
              mat[k] = _fxc_mat(cell, ao_kpts[k], wv, non0tab, xctype, ao_loc)
          return mat
  
      def block_loop(self, cell, grids, nao, deriv=0, kpts=numpy.zeros((1,3)),
                     kpts_band=None, max_memory=2000, non0tab=None, blksize=None):
          '''Define this macro to loop over grids by blocks.
          '''
          if grids.coords is None:
              grids.build(with_non0tab=True)
          grids_coords = grids.coords
          grids_weights = grids.weights
          ngrids = grids_coords.shape[0]
          nkpts = len(kpts)
          comp = (deriv+1)*(deriv+2)*(deriv+3)//6
  # NOTE to index grids.non0tab, the blksize needs to be the integer multiplier of BLKSIZE
          if blksize is None:
              blksize = int(max_memory*1e6/(comp*2*nkpts*nao*16*BLKSIZE))*BLKSIZE
              blksize = max(BLKSIZE, min(blksize, ngrids, BLKSIZE*1200))
          if non0tab is None:
              non0tab = grids.non0tab
          if non0tab is None:
              non0tab = numpy.empty(((ngrids+BLKSIZE-1)//BLKSIZE,cell.nbas),
                                    dtype=numpy.uint8)
              non0tab[:] = 0xff
          if kpts_band is not None:
              kpts_band = numpy.reshape(kpts_band, (-1,3))
              where = [member(k, kpts) for k in kpts_band]
              where = [k_id[0] if len(k_id)>0 else None for k_id in where]
  
          for ip0 in range(0, ngrids, blksize):
              ip1 = min(ngrids, ip0+blksize)
              coords = grids_coords[ip0:ip1]
              weight = grids_weights[ip0:ip1]
              non0 = non0tab[ip0//BLKSIZE:]
              ao_k1 = ao_k2 = self.eval_ao(cell, coords, kpts, deriv=deriv, non0tab=non0)
              if kpts_band is not None:
                  ao_k1 = self.eval_ao(cell, coords, kpts_band, deriv=deriv, non0tab=non0)
              yield ao_k1, ao_k2, non0, weight, coords
              ao_k1 = ao_k2 = None
  
      def _gen_rho_evaluator(self, cell, dms, hermi=0):
          if getattr(dms, 'mo_coeff', None) is not None:
              mo_coeff = dms.mo_coeff
              mo_occ = dms.mo_occ
              if isinstance(dms[0], numpy.ndarray) and dms[0].ndim == 2:
                  mo_coeff = [mo_coeff]
                  mo_occ = [mo_occ]
              nao = cell.nao_nr()
              ndms = len(mo_occ)
              def make_rho(idm, ao, non0tab, xctype):
                  return self.eval_rho2(cell, ao, mo_coeff[idm], mo_occ[idm],
                                        non0tab, xctype)
          else:
              if isinstance(dms[0], numpy.ndarray) and dms[0].ndim == 2:
                  dms = [numpy.stack(dms)]
              #if not hermi:
              # Density (or response of density) is always real for DFT.
              # Symmetrizing DM for gamma point should not change the value of
              # density. However, when k-point is considered, unless dm and
              # dm.conj().transpose produce the same real part of density, the
              # symmetrization code below may be incorrect (proof is needed).
              #    # dm.shape = (nkpts, nao, nao)
              #    dms = [(dm+dm.conj().transpose(0,2,1))*.5 for dm in dms]
              nao = dms[0].shape[-1]
              ndms = len(dms)
              def make_rho(idm, ao_kpts, non0tab, xctype):
                  return self.eval_rho(cell, ao_kpts, dms[idm], non0tab, xctype,
                                       hermi=hermi)
          return make_rho, ndms, nao
  
      nr_rks_fxc = nr_rks_fxc
      nr_uks_fxc = nr_uks_fxc
      cache_xc_kernel  = cache_xc_kernel
      get_rho = get_rho
  
      def rsh_and_hybrid_coeff(self, xc_code, spin=0):
          omega, alpha, hyb = numint.NumInt.rsh_and_hybrid_coeff(self, xc_code, spin)
          if abs(omega) > 1e-10:
              raise NotImplementedError
          return omega, alpha, hyb
  _KNumInt = KNumInt
   #!/usr/bin/env python
  # Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  import os
  import unittest
  import numpy
  from pyscf import gto, scf
  from pyscf import dft
  from pyscf import lib
  
  mol = gto.Mole()
  mol.verbose = 0
  mol.output = None
  mol.atom = 'h 0 0 0; h 1 .5 0; h 0 4 1; h 1 0 .2'
  mol.basis = 'aug-ccpvdz'
  mol.build()
  #dm = scf.RHF(mol).run(conv_tol=1e-14).make_rdm1()
  dm = numpy.load(os.path.realpath(os.path.join(__file__, '..', 'dm_h4.npy')))
  mf = dft.RKS(mol)
  mf.grids.atom_grid = {"H": (50, 110)}
  mf.prune = None
  mf.grids.build(with_non0tab=False)
  nao = mol.nao_nr()
  ao = dft.numint.eval_ao(mol, mf.grids.coords, deriv=1)
  rho = dft.numint.eval_rho(mol, ao, dm, xctype='GGA')
  
  def tearDownModule():
      global mol, mf, ao, rho
      del mol, mf, ao, rho
  
  def finger(a):
      w = numpy.cos(numpy.arange(a.size))
      return numpy.dot(w, a.ravel())
  
  class KnownValues(unittest.TestCase):
      def test_parse_xc(self):
          hyb, fn_facs = dft.xcfun.parse_xc('.5*HF+.5*B3LYP,VWN*.5')
          self.assertAlmostEqual(hyb[0], .6, 12)
          self.assertEqual([x[0] for x in fn_facs], [0,6,16,3])
          self.assertTrue(numpy.allclose([x[1] for x in fn_facs],
                                         (0.04, 0.36, 0.405, 0.595)))
          hyb, fn_facs = dft.xcfun.parse_xc('HF,')
          self.assertEqual(hyb[0], 1)
          self.assertEqual(fn_facs, [])
  
          hyb, fn_facs = dft.libxc.parse_xc('B88 - SLATER')
          self.assertEqual(fn_facs, [(106, 1), (1, -1)])
          hyb, fn_facs = dft.libxc.parse_xc('B88 -SLATER*.5')
          self.assertEqual(fn_facs, [(106, 1), (1, -0.5)])
  
          hyb, fn_facs = dft.xcfun.parse_xc('0.5*B3LYP+0.25*B3LYP')
          self.assertTrue(numpy.allclose(hyb, [.15, 0, 0]))
          hyb = dft.libxc.hybrid_coeff('0.5*B3LYP+0.25*B3LYP')
          self.assertAlmostEqual(hyb, .15, 12)
  
          hyb, fn_facs = dft.xcfun.parse_xc('CAM_B3LYP')
          self.assertTrue(numpy.allclose(hyb, [0.19, 0.65, 0.33]))
  
          hyb, fn_facs = dft.xcfun.parse_xc('0.6*CAM_B3LYP+0.4*B3P86')
          self.assertTrue(numpy.allclose(hyb, [.08+0.19*.6, 0.65*.6, 0.33]))
          self.assertTrue(numpy.allclose(fn_facs,
                                         [(9, 0.6), (3, 0.19), (16, 0.486), (0, 0.032), (6, 0.288), (46, 0.324)]))
          rsh = dft.xcfun.rsh_coeff('0.6*CAM_B3LYP+0.4*B3P86')
          self.assertTrue(numpy.allclose(rsh, (0.33, 0.39, -0.196)))
  
          hyb, fn_facs = dft.xcfun.parse_xc('0.4*B3P86+0.6*CAM_B3LYP')
          self.assertTrue(numpy.allclose(hyb, [.08+0.19*.6, 0.65*.6, 0.33]))
          self.assertTrue(numpy.allclose(fn_facs,
                                         [(0, 0.032), (6, 0.288), (46, 0.324), (3, 0.19), (9, 0.6), (16, 0.486)]))
          rsh = dft.xcfun.rsh_coeff('0.4*B3P86+0.6*CAM_B3LYP')
          self.assertTrue(numpy.allclose(rsh, (0.33, 0.39, -0.196)))
  
          hyb, fn_facs = dft.xcfun.parse_xc('0.5*SR-HF(0.3) + .8*HF + .22*LR_HF')
          self.assertEqual(hyb, [1.3, 1.02, 0.3])
  
          hyb, fn_facs = dft.xcfun.parse_xc('0.5*SR-HF + .22*LR_HF(0.3) + .8*HF')
          self.assertEqual(hyb, [1.3, 1.02, 0.3])
  
          hyb, fn_facs = dft.xcfun.parse_xc('0.5*SR-HF + .8*HF + .22*LR_HF(0.3)')
          self.assertEqual(hyb, [1.3, 1.02, 0.3])
  
          hyb, fn_facs = dft.xcfun.parse_xc('0.5*RSH(2.04;0.56;0.3) + 0.5*BP86')
          self.assertEqual(hyb, [1.3, 1.02, 0.3])
          self.assertEqual(fn_facs, [(6, 0.5), (46, 0.5)])
  
          self.assertRaises(ValueError, dft.xcfun.parse_xc, 'SR_HF(0.3) + LR_HF(.5)')
          self.assertRaises(ValueError, dft.xcfun.parse_xc, 'LR-HF(0.3) + SR-HF(.5)')
  
          hyb = dft.xcfun.hybrid_coeff('M05')
          self.assertAlmostEqual(hyb, 0.28, 9)
  
          hyb, fn_facs = dft.xcfun.parse_xc('APBE,')
          self.assertEqual(fn_facs[0][0], 58)
  
          hyb, fn_facs = dft.xcfun.parse_xc('VWN,')
          self.assertEqual(fn_facs, [(3, 1)])
  
          hyb, fn_facs = dft.xcfun.parse_xc('TF,')
          self.assertEqual(fn_facs, [(24, 1)])
  
          ref = [(0, 1), (3, 1)]
          self.assertEqual(dft.xcfun.parse_xc_name('LDA,VWN'), (0,3))
          self.assertEqual(dft.xcfun.parse_xc(('LDA','VWN'))[1], ref)
          self.assertEqual(dft.xcfun.parse_xc((0, 3))[1], ref)
          self.assertEqual(dft.xcfun.parse_xc('0, 3')[1], ref)
          self.assertEqual(dft.xcfun.parse_xc(3)[1], [(3,1)])
  
          #self.assertEqual(dft.xcfun.parse_xc('M11-L')[1], [(226,1),(75,1)])
          #self.assertEqual(dft.xcfun.parse_xc('M11L' )[1], [(226,1),(75,1)])
          #self.assertEqual(dft.xcfun.parse_xc('M11-L,M11L' )[1], [(226,1),(75,1)])
          #self.assertEqual(dft.xcfun.parse_xc('M11_L,M11-L')[1], [(226,1),(75,1)])
          #self.assertEqual(dft.xcfun.parse_xc('M11L,M11_L' )[1], [(226,1),(75,1)])
  
          #self.assertEqual(dft.xcfun.parse_xc('Xpbe,')[1], [(123,1)])
          #self.assertEqual(dft.xcfun.parse_xc('pbe,' )[1], [(101,1)])
          hyb, fn_facs = dft.xcfun.parse_xc('PBE*.4+LDA')
          self.assertEqual(fn_facs, [(5, 0.4), (4, 0.4), (0, 1)])
          hyb, fn_facs = dft.xcfun.parse_xc('PBE*.4+VWN')
          self.assertEqual(fn_facs, [(5, 0.4), (4, 0.4), (3, 1)])
  
          self.assertTrue (dft.xcfun.is_meta_gga('m05'))
          self.assertFalse(dft.xcfun.is_meta_gga('pbe0'))
          self.assertFalse(dft.xcfun.is_meta_gga('tf,'))
          self.assertFalse(dft.xcfun.is_meta_gga('vv10'))
          self.assertTrue (dft.xcfun.is_gga('PBE0'))
          self.assertFalse(dft.xcfun.is_gga('m05'))
          self.assertFalse(dft.xcfun.is_gga('tf,'))
          self.assertTrue (dft.xcfun.is_lda('tf,'))
          self.assertFalse(dft.xcfun.is_lda('vv10'))
          self.assertTrue (dft.xcfun.is_hybrid_xc('m05'))
          self.assertTrue (dft.xcfun.is_hybrid_xc('pbe0,'))
          self.assertFalse(dft.xcfun.is_hybrid_xc('m05,'))
          self.assertFalse(dft.xcfun.is_hybrid_xc('vv10'))
          self.assertTrue (dft.xcfun.is_hybrid_xc(('b3lyp',4,'vv10')))
  
      def test_nlc_coeff(self):
          self.assertEqual(dft.xcfun.nlc_coeff('vv10'), [5.9, 0.0093])
  
      def test_lda(self):
          e,v,f,k = dft.xcfun.eval_xc('lda,', rho[0][:3], deriv=3)
          self.assertAlmostEqual(lib.finger(e)   , -0.4720562542635522, 8)
          self.assertAlmostEqual(lib.finger(v[0]), -0.6294083390180697, 8)
          self.assertAlmostEqual(lib.finger(f[0]), -1.1414693830969338, 8)
          self.assertAlmostEqual(lib.finger(k[0]),  4.1402447248393921, 8)
  
          e,v,f,k = dft.xcfun.eval_xc('lda,', [rho[0][:3]*.5]*2, spin=1, deriv=3)
          self.assertAlmostEqual(lib.finger(e)   , -0.4720562542635522, 8)
          self.assertAlmostEqual(lib.finger(v[0].T[0]), -0.6294083390180697, 8)
          self.assertAlmostEqual(lib.finger(v[0].T[1]), -0.6294083390180697, 8)
          self.assertAlmostEqual(lib.finger(f[0].T[0]), -1.1414693830969338*2, 8)
          self.assertAlmostEqual(lib.finger(f[0].T[2]), -1.1414693830969338*2, 8)
          self.assertAlmostEqual(lib.finger(k[0].T[0]),  4.1402447248393921*4, 7)
          self.assertAlmostEqual(lib.finger(k[0].T[3]),  4.1402447248393921*4, 7)
  
      def test_lyp(self):
          e,v,f = dft.xcfun.eval_xc(',LYP', rho, deriv=3)[:3]
          self.assertAlmostEqual(numpy.dot(rho[0],e), -62.114576182676615, 8)
          self.assertAlmostEqual(numpy.dot(rho[0],v[0]),-81.771670866308455, 8)
          self.assertAlmostEqual(numpy.dot(rho[0],v[1]), 27.485383255125743, 8)
          self.assertAlmostEqual(numpy.dot(rho[0],f[0]), 186.823806251777, 7)
          self.assertAlmostEqual(numpy.dot(rho[0],f[1]), -3391.2428894571085, 6)
          self.assertAlmostEqual(numpy.dot(rho[0],f[2]), 0, 9)
  
      def test_beckex(self):
          rho =(numpy.array([1.    , 1., 0., 0.]).reshape(-1,1),
                numpy.array([    .8, 1., 0., 0.]).reshape(-1,1))
          e,v,f = dft.xcfun.eval_xc('b88,', rho, spin=1, deriv=3)[:3]
          self.assertAlmostEqual(lib.finger(e)   ,-0.9061911523772116   , 9)
          self.assertAlmostEqual(lib.finger(v[0]),-1.8531364353196298   , 9)
          self.assertAlmostEqual(lib.finger(v[1]),-0.0018308066137967724, 9)
          self.assertAlmostEqual(lib.finger(f[0]),-0.21602284426026866  , 9)
          self.assertAlmostEqual(lib.finger(f[1]), 0.0072053520662545617, 9)
          self.assertAlmostEqual(lib.finger(f[2]), 0.0002275350850255538, 9)
  
      def test_m05x(self):
          rho =(numpy.array([1., 1., 0., 0., 0., 0.165 ]).reshape(-1,1),
                numpy.array([.8, 1., 0., 0., 0., 0.1050]).reshape(-1,1))
          test_ref = numpy.array([-1.57876583, -2.12127045,-2.11264351,-0.00315462,
                                   0.00000000, -0.00444560, 3.45640232, 4.4349756])
          exc, vxc, fxc, kxc = dft.xcfun.eval_xc('m05,', rho, 1, deriv=3)
          self.assertAlmostEqual(float(exc)*1.8, test_ref[0], 5)
          self.assertAlmostEqual(abs(vxc[0]-test_ref[1:3]).max(), 0, 6)
          self.assertAlmostEqual(abs(vxc[1]-test_ref[3:6]).max(), 0, 6)
          self.assertAlmostEqual(abs(vxc[3]-test_ref[6:8]).max(), 0, 5)
  
          exc, vxc, fxc, kxc = dft.xcfun.eval_xc('m05,', rho[0], 0, deriv=3)
          self.assertAlmostEqual(float(exc), -0.5746231988116002, 5)
          self.assertAlmostEqual(float(vxc[0]), -0.8806121005703862, 6)
          self.assertAlmostEqual(float(vxc[1]), -0.0032300155406846756, 7)
          self.assertAlmostEqual(float(vxc[3]), 0.4474953100487698, 5)
  
      def test_camb3lyp(self):
          rho = numpy.array([1., 1., 0.1, 0.1]).reshape(-1,1)
          exc, vxc, fxc, kxc = dft.xcfun.eval_xc('camb3lyp', rho, 0, deriv=1)
          # FIXME, xcfun and libxc do not agree on camb3lyp
          # self.assertAlmostEqual(float(exc), -0.5752559666317147, 5)
          # self.assertAlmostEqual(float(vxc[0]), -0.7709812578936763, 5)
          # self.assertAlmostEqual(float(vxc[1]), -0.0029862221286189846, 7)
  
          self.assertEqual(dft.xcfun.rsh_coeff('camb3lyp'), (0.33, 0.65, -0.46))
  
      def test_define_xc(self):
          def eval_xc(xc_code, rho, spin=0, relativity=0, deriv=1, verbose=None):
              exc = vxc = fxc = kxc = None
              return exc, vxc, fxc, kxc
  
          mf = dft.RKS(mol)
          ni = dft.xcfun.define_xc(mf._numint, eval_xc, 'GGA', hyb=0.2)
          ni = dft.xcfun.define_xc(mf._numint, 'b3lyp+vwn', 'GGA', hyb=0.2)
          self.assertRaises(ValueError, dft.xcfun.define_xc, mf._numint, 0.1)
  
      def test_vs_libxc_rks(self):
          ao = dft.numint.eval_ao(mol, mf.grids.coords[:200], deriv=2)
          rho = dft.numint.eval_rho(mol, ao, dm, xctype='MGGA')
          rhoa = rho[:,:200]
          def check(xc_code, deriv=3, e_place=9, v_place=9, f_place=9, k_place=9):
              exc0, vxc0, fxc0, kxc0 = dft.libxc.eval_xc(xc_code, rhoa, 0, deriv=deriv)
              exc1, vxc1, fxc1, kxc1 = dft.xcfun.eval_xc(xc_code, rhoa, 0, deriv=deriv)
              self.assertAlmostEqual(abs(exc0-exc1).max(), 0, e_place)
              if deriv > 0:
                  for v0, v1 in zip(vxc0, vxc1):
                      if v0 is not None and v1 is not None:
                          self.assertAlmostEqual(abs(v0-v1).max(), 0, v_place)
              if deriv > 1:
                  for f0, f1 in zip(fxc0, fxc1):
                      if f0 is not None and f1 is not None:
                          self.assertAlmostEqual(abs(f0-f1).max(), 0, f_place)
              if deriv > 2:
                  for k0, k1 in zip(kxc0, kxc1):
                      if k0 is not None and k1 is not None:
                          self.assertAlmostEqual(abs(k0-k1).max(), 0, k_place)
  
          check('lda,')
  
          check('pw86,')
          check('pbe,', e_place=6, v_place=6, f_place=5, k_place=4)
          #?check('becke,')
          #?check('br,')
          #?check('LDAERF,')
          check('optx,')
          check('OPTXCORR,')
          check('RPBE,')
          check('TF,'  )
          check('PW91,' , e_place=6, v_place=4, f_place=2, k_place=-1)
          check('m05,'  , deriv=1, e_place=6, v_place=6)
          check('m052x,', deriv=1, e_place=6, v_place=6)
          check('m06,'  , deriv=1, e_place=6, v_place=6)
          check('m062x,', deriv=1, e_place=6, v_place=6)
          check('m06l,' , deriv=1, e_place=6, v_place=6)
          check('TPSS,' ,                                  k_place=-4)
          #?check('REVTPSS,', deriv=1)  # xcfun crash
          check('APBE,')
          check('BLOC,' ,                                  k_place=-5)
          check('PBEINT,', e_place=7, v_place=6, f_place=5, k_place=4)
  
          check(',vwn3')
          check(',vwn5')
          check(',pbe'  , deriv=2)
          #?check(',br')
          #?check(',LDAERF')
          check(',lyp'  , deriv=2)
          check(',SPBE' , deriv=2, e_place=1, v_place=1, f_place=0)
          check(',PW91' , deriv=2,                       f_place=3)
          check(',m052x', deriv=1)
          check(',m05'  , deriv=1)
          check(',m06'  , deriv=1)
          check(',m062x', deriv=1)
          check(',m06l' , deriv=1)
          check(',TPSS' , deriv=1,              v_place=1)
          check(',REVTPSS', deriv=1, e_place=2, v_place=1)
          check(',p86'    , deriv=2, e_place=5, v_place=5, f_place=3)
          check(',APBE'   , deriv=2)
          check(',PBEINT' , deriv=1)
          check(',TPSSLOC', deriv=1, e_place=1, v_place=0)
  
          #?check('br')
          check('revpbe', deriv=2, e_place=6, v_place=6, f_place=5)
          check('b97'   , deriv=2, e_place=6, v_place=5, f_place=3)
          #?check('b97_1')
          #?check('b97_2')
          check('SVWN')
          check('BLYP'  , deriv=2)
          check('BP86'  , deriv=2, e_place=5, v_place=5, f_place=3)
          check('OLYP'  , deriv=2)
          check('KT1'   , deriv=1)
          check('KT2'   , deriv=1)
          #?check('KT3')
          check('PBE0'   , deriv=2, e_place=6, v_place=6, f_place=5)
          check('B3P86'  , deriv=2, e_place=5, v_place=5, f_place=3)
          check('B3P86G' , deriv=2, e_place=5, v_place=5, f_place=3)
          check('B3PW91' , deriv=2,                       f_place=4)
          check('B3PW91G', deriv=2, e_place=2, v_place=2, f_place=2)
          check('B3LYP'  , deriv=2)
          check('B3LYP5' , deriv=2)
          check('B3LYPG' , deriv=2)
          check('O3LYP'  , deriv=2)
          check('X3LYP'  , deriv=2, e_place=7, v_place=5, f_place=2)
          check('CAMB3LYP', deriv=1)
          check('B97_1'   , deriv=2, e_place=6, v_place=5, f_place=3)
          check('B97_2'   , deriv=2, e_place=6, v_place=5, f_place=3)
          check('TPSSH'   , deriv=1,            v_place=1)
  
      def test_vs_libxc_uks(self):
          ao = dft.numint.eval_ao(mol, mf.grids.coords[:400], deriv=2)
          rho = dft.numint.eval_rho(mol, ao, dm, xctype='MGGA')
          rhoa = rho[:,:200]
          rhob = rhoa + rho[:,200:400]
          def check(xc_code, deriv=3, e_place=9, v_place=9, f_place=9, k_place=9):
              exc0, vxc0, fxc0, kxc0 = dft.libxc.eval_xc(xc_code, (rhoa, rhob), 1, deriv=deriv)
              exc1, vxc1, fxc1, kxc1 = dft.xcfun.eval_xc(xc_code, (rhoa, rhob), 1, deriv=deriv)
              self.assertAlmostEqual(abs(exc0-exc1).max(), 0, e_place)
              if deriv > 0:
                  for v0, v1 in zip(vxc0, vxc1):
                      if v0 is not None and v1 is not None:
                          self.assertAlmostEqual(abs(v0-v1).max(), 0, v_place)
              if deriv > 1:
                  for f0, f1 in zip(fxc0, fxc1):
                      if f0 is not None and f1 is not None:
                          self.assertAlmostEqual(abs(f0-f1).max(), 0, f_place)
              if deriv > 2 and kxc0 is not None:
                  for k0, k1 in zip(kxc0, kxc1):
                      if k0 is not None and k1 is not None:
                          self.assertAlmostEqual(abs(k0-k1).max(), 0, k_place)
  
          check('lda,')
  
          check('pw86,')
          check('pbe,', e_place=6, v_place=6, f_place=5, k_place=4)
          #?check('becke,')
          #?check('br,')
          #?check('LDAERF,')
          check('optx,')
          check('OPTXCORR,')
          check('RPBE,')
          check('TF,'  , e_place=0, v_place=-1, f_place=-2, k_place=-2)
          check('PW91,' , e_place=6, v_place=4, f_place=2, k_place=-1)
          check('m05,'  , deriv=1, e_place=6, v_place=6)
          check('m052x,', deriv=1, e_place=6, v_place=6)
          check('m06,'  , deriv=1, e_place=6, v_place=6)
          check('m062x,', deriv=1, e_place=6, v_place=6)
          check('m06l,' , deriv=1, e_place=6, v_place=6)
          check('TPSS,' ,                                  k_place=-4)
          #?check('REVTPSS,', deriv=1)  # libxc crash
          check('APBE,')
          check('BLOC,' ,                                  k_place=-5)
          check('PBEINT,', e_place=7, v_place=6, f_place=5, k_place=4)
  
          check(',vwn3', e_place=2, v_place=1, f_place=1, k_place=0)
          check(',vwn5')
          check(',pbe'  , deriv=2)
          #?check(',br')
          #?check(',LDAERF')
          check(',lyp'  , deriv=2)
          check(',SPBE' , deriv=2, e_place=1, v_place=1, f_place=0)
          check(',PW91' , deriv=2,                       f_place=3)
          check(',m052x', deriv=1)
          check(',m05'  , deriv=1)
          check(',m06'  , deriv=1)
          check(',m062x', deriv=1)
          check(',m06l' , deriv=1)
          check(',TPSS' , deriv=1,              v_place=1)
          check(',REVTPSS', deriv=1, e_place=2, v_place=1)
          check(',p86'    , deriv=2, e_place=5, v_place=5, f_place=3)
          check(',APBE'   , deriv=2)
          check(',PBEINT' , deriv=1)
          check(',TPSSLOC', deriv=1, e_place=1, v_place=0)
  
          #?check('br')
          check('revpbe', deriv=2, e_place=6, v_place=6, f_place=5)
          check('b97'   , deriv=2, e_place=6, v_place=5, f_place=3)
          #?check('b97_1')
          #?check('b97_2')
          check('SVWN')
          check('BLYP'  , deriv=2)
          check('BP86'  , deriv=2, e_place=5, v_place=5, f_place=3)
          check('OLYP'  , deriv=2)
          check('KT1'   , deriv=1)
          check('KT2'   , deriv=1)
          #?check('KT3')
          check('PBE0'   , deriv=2, e_place=6, v_place=6, f_place=5)
          check('B3P86'  , deriv=2, e_place=5, v_place=5, f_place=3)
          check('B3P86G' , deriv=2, e_place=3, v_place=2, f_place=2)
          check('B3PW91' , deriv=2,                       f_place=4)
          check('B3PW91G', deriv=2, e_place=2, v_place=2, f_place=2)
          check('B3LYP'  , deriv=2)
          check('B3LYP5' , deriv=2)
          check('B3LYPG' , deriv=2, e_place=3, v_place=2, f_place=2)
          check('O3LYP'  , deriv=2, e_place=3, v_place=2, f_place=1)
          check('X3LYP'  , deriv=2, e_place=7, v_place=5, f_place=2)
          check('CAMB3LYP', deriv=1)
          check('B97_1'   , deriv=2, e_place=6, v_place=5, f_place=3)
          check('B97_2'   , deriv=2, e_place=6, v_place=5, f_place=3)
          check('TPSSH'   , deriv=1,            v_place=1)
  
  
  if __name__ == "__main__":
      print("Test xcfun")
      unittest.main()
  
  
apis: numpy.sqrt numpy.zeros_like numpy.ones numpy.empty numpy.append numpy.zeros numpy.tril_indices numpy.asarray numpy.argsort numpy.ndarray.__array_wrap__ numpy.einsum numpy.where
apis: numpy.empty numpy.zeros numpy.result_type numpy.stack numpy.reshape numpy.asarray numpy.hstack numpy.sqrt numpy.iscomplexobj numpy.array
apis: numpy.arange numpy.allclose numpy.dot numpy.array