text: #!/usr/bin/env python
  #
  # Author: Qiming Sun <osirpt.sun@gmail.com>
  #
  
  '''
  Non-relativistic restricted Hartree-Fock with point group symmetry.
  
  The symmetry are not handled in a separate data structure.  Note that during
  the SCF iteration,  the orbitals are grouped in terms of symmetry irreps.
  But the orbitals in the result are sorted based on the orbital energies.
  Function symm.label_orb_symm can be used to detect the symmetry of the
  molecular orbitals.
  '''
  
  import time
  from functools import reduce
  import numpy
  import scipy.linalg
  from pyscf import lib
  from pyscf import symm
  from pyscf.lib import logger
  from pyscf.scf import hf
  from pyscf.scf import rohf
  from pyscf.scf import chkfile
  
  # mo_energy, mo_coeff, mo_occ are all in nosymm representation
  
  def analyze(mf, verbose=logger.DEBUG, **kwargs):
      '''Analyze the given SCF object:  print orbital energies, occupancies;
      print orbital coefficients; Occupancy for each irreps; Mulliken population analysis
      '''
      from pyscf.lo import orth
      from pyscf.tools import dump_mat
      mol = mf.mol
      if not mol.symmetry:
          return hf.analyze(mf, verbose, **kwargs)
  
      mo_energy = mf.mo_energy
      mo_occ = mf.mo_occ
      mo_coeff = mf.mo_coeff
      ovlp_ao = mf.get_ovlp()
      log = logger.new_logger(mf, verbose)
      if log.verbose >= logger.NOTE:
          nirrep = len(mol.irrep_id)
          orbsym = get_orbsym(mf.mol, mo_coeff, ovlp_ao, False)
          wfnsym = 0
          noccs = [sum(orbsym[mo_occ>0]==ir) for ir in mol.irrep_id]
          log.note('total symmetry = %s', symm.irrep_id2name(mol.groupname, wfnsym))
          log.note('occupancy for each irrep:  ' + (' %4s'*nirrep), *mol.irrep_name)
          log.note('                           ' + (' %4d'*nirrep), *noccs)
          log.note('**** MO energy ****')
          irname_full = {}
          for k,ir in enumerate(mol.irrep_id):
              irname_full[ir] = mol.irrep_name[k]
          irorbcnt = {}
          for k, j in enumerate(orbsym):
              if j in irorbcnt:
                  irorbcnt[j] += 1
              else:
                  irorbcnt[j] = 1
              log.note('MO #%d (%s #%d), energy= %.15g occ= %g',
                       k+1, irname_full[j], irorbcnt[j], mo_energy[k], mo_occ[k])
  
      if log.verbose >= logger.DEBUG:
          label = mol.ao_labels()
          molabel = []
          irorbcnt = {}
          for k, j in enumerate(orbsym):
              if j in irorbcnt:
                  irorbcnt[j] += 1
              else:
                  irorbcnt[j] = 1
              molabel.append('#%-d(%s #%d)' % (k+1, irname_full[j], irorbcnt[j]))
          log.debug(' ** MO coefficients (expansion on meta-Lowdin AOs) **')
          orth_coeff = orth.orth_ao(mol, 'meta_lowdin', s=ovlp_ao)
          c = reduce(numpy.dot, (orth_coeff.T, ovlp_ao, mo_coeff))
          dump_mat.dump_rec(mf.stdout, c, label, molabel, start=1, **kwargs)
  
      dm = mf.make_rdm1(mo_coeff, mo_occ)
      return mf.mulliken_meta(mol, dm, s=ovlp_ao, verbose=log)
  
  def get_irrep_nelec(mol, mo_coeff, mo_occ, s=None):
      '''Electron numbers for each irreducible representation.
  
      Args:
          mol : an instance of :class:`Mole`
              To provide irrep_id, and spin-adapted basis
          mo_coeff : 2D ndarray
              Regular orbital coefficients, without grouping for irreps
          mo_occ : 1D ndarray
              Regular occupancy, without grouping for irreps
  
      Returns:
          irrep_nelec : dict
              The number of electrons for each irrep {'ir_name':int,...}.
  
      Examples:
  
      >>> mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='ccpvdz', symmetry=True, verbose=0)
      >>> mf = scf.RHF(mol)
      >>> mf.scf()
      -76.016789472074251
      >>> scf.hf_symm.get_irrep_nelec(mol, mf.mo_coeff, mf.mo_occ)
      {'A1': 6, 'A2': 0, 'B1': 2, 'B2': 2}
      '''
      orbsym = get_orbsym(mol, mo_coeff, s, False)
      irrep_nelec = dict([(mol.irrep_name[k], int(sum(mo_occ[orbsym==ir])))
                          for k, ir in enumerate(mol.irrep_id)])
      return irrep_nelec
  
  def canonicalize(mf, mo_coeff, mo_occ, fock=None):
      '''Canonicalization diagonalizes the Fock matrix in occupied, open,
      virtual subspaces separatedly (without change occupancy).
      '''
      mol = mf.mol
      if not mol.symmetry:
          return hf.canonicalize(mf, mo_coeff, mo_occ, fock)
  
      if fock is None:
          dm = mf.make_rdm1(mo_coeff, mo_occ)
          fock = mf.get_hcore() + mf.get_veff(mf.mol, dm)
      coreidx = mo_occ == 2
      viridx = mo_occ == 0
      openidx = ~(coreidx | viridx)
      mo = numpy.empty_like(mo_coeff)
      mo_e = numpy.empty(mo_occ.size)
  
      if hasattr(mo_coeff, 'orbsym'):
          orbsym = mo_coeff.orbsym
          irreps = set(orbsym)
          for ir in irreps:
              idx0 = orbsym == ir
              for idx1 in (coreidx, openidx, viridx):
                  idx = idx0 & idx1
                  if numpy.count_nonzero(idx) > 0:
                      orb = mo_coeff[:,idx]
                      f1 = reduce(numpy.dot, (orb.T.conj(), fock, orb))
                      e, c = scipy.linalg.eigh(f1)
                      mo[:,idx] = numpy.dot(mo_coeff[:,idx], c)
                      mo_e[idx] = e
      else:
          s = mf.get_ovlp()
          for idx in (coreidx, openidx, viridx):
              if numpy.count_nonzero(idx) > 0:
                  orb = mo_coeff[:,idx]
                  f1 = reduce(numpy.dot, (orb.T.conj(), fock, orb))
                  e, c = scipy.linalg.eigh(f1)
                  c = numpy.dot(mo_coeff[:,idx], c)
                  mo[:,idx] = _symmetrize_canonicalization_(mf, e, c, s)
                  mo_e[idx] = e
          orbsym = get_orbsym(mol, mo, s, False)
  
      mo = lib.tag_array(mo, orbsym=orbsym)
      return mo_e, mo
  
  def _symmetrize_canonicalization_(mf, mo_energy, mo_coeff, s):
      '''Restore symmetry for canonicalized orbitals
      '''
      def search_for_degeneracy(mo_energy):
          idx = numpy.where(abs(mo_energy[1:] - mo_energy[:-1]) < 1e-6)[0]
          return numpy.unique(numpy.hstack((idx, idx+1)))
  
      mol = mf.mol
      degidx = search_for_degeneracy(mo_energy)
      logger.debug1(mf, 'degidx %s', degidx)
      if degidx.size > 0:
          esub = mo_energy[degidx]
          csub = mo_coeff[:,degidx]
          scsub = numpy.dot(s, csub)
          emin = abs(esub).min() * .5
          es = []
          cs = []
          for i,ir in enumerate(mol.irrep_id):
              so = mol.symm_orb[i]
              sosc = numpy.dot(so.T, scsub)
              s_ir = reduce(numpy.dot, (so.T, s, so))
              fock_ir = numpy.dot(sosc*esub, sosc.T)
              mo_energy, u = mf._eigh(fock_ir, s_ir)
              idx = abs(mo_energy) > emin
              es.append(mo_energy[idx])
              cs.append(numpy.dot(mol.symm_orb[i], u[:,idx]))
          es = numpy.hstack(es).round(7)
          idx = numpy.argsort(es)
          assert(numpy.allclose(es[idx], esub.round(7)))
          mo_coeff[:,degidx] = numpy.hstack(cs)[:,idx]
      return mo_coeff
  
  def so2ao_mo_coeff(so, irrep_mo_coeff):
      '''Transfer the basis of MO coefficients, from spin-adapted basis to AO basis
      '''
      return numpy.hstack([numpy.dot(so[ir],irrep_mo_coeff[ir]) \
                           for ir in range(so.__len__())])
  
  def check_irrep_nelec(mol, irrep_nelec, nelec):
      for irname in irrep_nelec:
          if irname not in mol.irrep_name:
              logger.warn(mol, 'Molecule does not have irrep %s', irname)
  
      float_irname = []
      fix_na = 0
      fix_nb = 0
      for i, irname in enumerate(mol.irrep_name):
          if irname in irrep_nelec:
              if isinstance(irrep_nelec[irname], (int, numpy.integer)):
                  nelecb = irrep_nelec[irname] // 2
                  neleca = irrep_nelec[irname] - nelecb
              else:
                  neleca, nelecb = irrep_nelec[irname]
              norb = mol.symm_orb[i].shape[1]
              if neleca > norb or nelecb > norb:
                  msg =('More electrons than orbitals for irrep %s '
                        'nelec = %d + %d, norb = %d' %
                        (irname, neleca,nelecb, norb))
                  raise ValueError(msg)
              fix_na += neleca
              fix_nb += nelecb
          else:
              float_irname.append(irname)
      if fix_na < fix_nb:
          raise ValueError('alpha electrons %d < beta electrons %d\n'
                           'irrep_nelec %s' % (fix_na, fix_nb, irrep_nelec))
      if mol.spin < fix_na-fix_nb:
          raise ValueError('alpha electrons %d - beta electrons %d > mol.spin %d\n'
                           'irrep_nelec %s' % (fix_na, fix_nb, mol.spin, irrep_nelec))
  
      if isinstance(nelec, (int, numpy.integer)):
          nelecb = nelec // 2
          neleca = nelec - nelecb
      else:
          neleca, nelecb = nelec
      fix_ne = fix_na + fix_nb
      if ((fix_na > neleca) or (fix_nb > nelecb) or
          (fix_na+nelecb > mol.nelectron) or
          (fix_nb+neleca > mol.nelectron)):
          msg =('More electrons defined by irrep_nelec than total num electrons. '
                'mol.nelectron = %d  irrep_nelec = %s' %
                (mol.nelectron, irrep_nelec))
          raise ValueError(msg)
      else:
          logger.info(mol, 'Freeze %d electrons in irreps %s',
                      fix_ne, irrep_nelec.keys())
  
      if len(set(float_irname)) == 0 and fix_ne != mol.nelectron:
          msg =('Num electrons defined by irrep_nelec != total num electrons. '
                'mol.nelectron = %d  irrep_nelec = %s' %
                (mol.nelectron, irrep_nelec))
          raise ValueError(msg)
      else:
          logger.info(mol, '    %d free electrons in irreps %s',
                      mol.nelectron-fix_ne, ' '.join(float_irname))
      return fix_na, fix_nb, float_irname
  
  #TODO: force E1gx/E1gy ... use the same coefficients
  def eig(mf, h, s):
      '''Solve generalized eigenvalue problem, for each irrep.  The
      eigenvalues and eigenvectors are not sorted to ascending order.
      Instead, they are grouped based on irreps.
      '''
      mol = mf.mol
      if not mol.symmetry:
          return mf._eigh(h, s)
  
      nirrep = mol.symm_orb.__len__()
      h = symm.symmetrize_matrix(h, mol.symm_orb)
      s = symm.symmetrize_matrix(s, mol.symm_orb)
      cs = []
      es = []
      orbsym = []
      for ir in range(nirrep):
          e, c = mf._eigh(h[ir], s[ir])
          cs.append(c)
          es.append(e)
          orbsym.append([mol.irrep_id[ir]] * e.size)
      e = numpy.hstack(es)
      c = so2ao_mo_coeff(mol.symm_orb, cs)
      c = lib.tag_array(c, orbsym=numpy.hstack(orbsym))
      return e, c
  
  
  class RHF(hf.RHF):
      __doc__ = hf.SCF.__doc__ + '''
      Attributes for symmetry allowed RHF:
          irrep_nelec : dict
              Specify the number of electrons for particular irrep {'ir_name':int,...}.
              For the irreps not listed in this dict, the program will choose the
              occupancy based on the orbital energies.
  
      Examples:
  
      >>> mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='ccpvdz', symmetry=True, verbose=0)
      >>> mf = scf.RHF(mol)
      >>> mf.scf()
      -76.016789472074251
      >>> mf.get_irrep_nelec()
      {'A1': 6, 'A2': 0, 'B1': 2, 'B2': 2}
      >>> mf.irrep_nelec = {'A2': 2}
      >>> mf.scf()
      -72.768201804695622
      >>> mf.get_irrep_nelec()
      {'A1': 6, 'A2': 2, 'B1': 2, 'B2': 0}
      '''
      def __init__(self, mol):
          hf.RHF.__init__(self, mol)
          # number of electrons for each irreps
          self.irrep_nelec = {} # {'ir_name':int,...}
          self._keys = self._keys.union(['irrep_nelec'])
  
      def build(self, mol=None):
          for irname in self.irrep_nelec:
              if irname not in self.mol.irrep_name:
                  logger.warn(self, 'No irrep %s', irname)
          if mol.symmetry:
              check_irrep_nelec(self.mol, self.irrep_nelec, self.mol.nelectron)
          return hf.RHF.build(self, mol)
  
      eig = eig
  
      def get_grad(self, mo_coeff, mo_occ, fock=None):
          g = hf.RHF.get_grad(self, mo_coeff, mo_occ, fock)
          if self.mol.symmetry:
              occidx = mo_occ > 0
              viridx = ~occidx
              orbsym = get_orbsym(self.mol, mo_coeff)
              sym_forbid = orbsym[viridx].reshape(-1,1) != orbsym[occidx]
              g[sym_forbid.ravel()] = 0
          return g
  
      def get_occ(self, mo_energy=None, mo_coeff=None):
          ''' We assumed mo_energy are grouped by symmetry irreps, (see function
          self.eig). The orbitals are sorted after SCF.
          '''
          if mo_energy is None: mo_energy = self.mo_energy
          mol = self.mol
          if not mol.symmetry:
              return hf.RHF.get_occ(self, mo_energy, mo_coeff)
  
          orbsym = get_orbsym(self.mol, mo_coeff)
          mo_occ = numpy.zeros_like(mo_energy)
          rest_idx = numpy.ones(mo_occ.size, dtype=bool)
          nelec_fix = 0
          for i, ir in enumerate(mol.irrep_id):
              irname = mol.irrep_name[i]
              if irname in self.irrep_nelec:
                  ir_idx = numpy.where(orbsym == ir)[0]
                  n = self.irrep_nelec[irname]
                  occ_sort = numpy.argsort(mo_energy[ir_idx].round(9))
                  occ_idx  = ir_idx[occ_sort[:n//2]]
                  mo_occ[occ_idx] = 2
                  nelec_fix += n
                  rest_idx[ir_idx] = False
          nelec_float = mol.nelectron - nelec_fix
          assert(nelec_float >= 0)
          if nelec_float > 0:
              rest_idx = numpy.where(rest_idx)[0]
              occ_sort = numpy.argsort(mo_energy[rest_idx].round(9))
              occ_idx  = rest_idx[occ_sort[:nelec_float//2]]
              mo_occ[occ_idx] = 2
  
          vir_idx = (mo_occ==0)
          if self.verbose >= logger.INFO and numpy.count_nonzero(vir_idx) > 0:
              ehomo = max(mo_energy[~vir_idx])
              elumo = min(mo_energy[ vir_idx])
              noccs = []
              for i, ir in enumerate(mol.irrep_id):
                  irname = mol.irrep_name[i]
                  ir_idx = (orbsym == ir)
  
                  noccs.append(int(mo_occ[ir_idx].sum()))
                  if ehomo in mo_energy[ir_idx]:
                      irhomo = irname
                  if elumo in mo_energy[ir_idx]:
                      irlumo = irname
              logger.info(self, 'HOMO (%s) = %.15g  LUMO (%s) = %.15g',
                          irhomo, ehomo, irlumo, elumo)
  
              logger.debug(self, 'irrep_nelec = %s', noccs)
              _dump_mo_energy(mol, mo_energy, mo_occ, ehomo, elumo, orbsym,
                              verbose=self.verbose)
          return mo_occ
  
      def _finalize(self):
          hf.RHF._finalize(self)
  
          # sort MOs wrt orbital energies, it should be done last.
          o_sort = numpy.argsort(self.mo_energy[self.mo_occ> 0].round(9))
          v_sort = numpy.argsort(self.mo_energy[self.mo_occ==0].round(9))
          idx = numpy.arange(self.mo_energy.size)
          idx = numpy.hstack((idx[self.mo_occ> 0][o_sort],
                              idx[self.mo_occ==0][v_sort]))
          self.mo_energy = self.mo_energy[idx]
          orbsym = get_orbsym(self.mol, self.mo_coeff)
          self.mo_coeff = lib.tag_array(self.mo_coeff[:,idx], orbsym=orbsym[idx])
          self.mo_occ = self.mo_occ[idx]
          if self.chkfile:
              chkfile.dump_scf(self.mol, self.chkfile, self.e_tot, self.mo_energy,
                               self.mo_coeff, self.mo_occ, overwrite_mol=False)
          return self
  
      def analyze(self, verbose=None, **kwargs):
          if verbose is None: verbose = self.verbose
          return analyze(self, verbose, **kwargs)
  
      @lib.with_doc(get_irrep_nelec.__doc__)
      def get_irrep_nelec(self, mol=None, mo_coeff=None, mo_occ=None, s=None):
          if mol is None: mol = self.mol
          if mo_occ is None: mo_occ = self.mo_occ
          if mo_coeff is None: mo_coeff = self.mo_coeff
          if s is None: s = self.get_ovlp()
          return get_irrep_nelec(mol, mo_coeff, mo_occ, s)
  
      canonicalize = canonicalize
  
  
  class ROHF(rohf.ROHF):
      __doc__ = hf.SCF.__doc__ + '''
      Attributes for symmetry allowed ROHF:
          irrep_nelec : dict
              Specify the number of alpha/beta electrons for particular irrep
              {'ir_name':(int,int), ...}.
              For the irreps not listed in these dicts, the program will choose the
              occupancy based on the orbital energies.
  
      Examples:
  
      >>> mol = gto.M(atom='O 0 0 0; H 0 0 1; H 0 1 0', basis='ccpvdz', symmetry=True, charge=1, spin=1, verbose=0)
      >>> mf = scf.RHF(mol)
      >>> mf.scf()
      -75.619358861084052
      >>> mf.get_irrep_nelec()
      {'A1': (3, 3), 'A2': (0, 0), 'B1': (1, 1), 'B2': (1, 0)}
      >>> mf.irrep_nelec = {'B1': (1, 0)}
      >>> mf.scf()
      -75.425669486776457
      >>> mf.get_irrep_nelec()
      {'A1': (3, 3), 'A2': (0, 0), 'B1': (1, 0), 'B2': (1, 1)}
      '''
      def __init__(self, mol):
          rohf.ROHF.__init__(self, mol)
          self.irrep_nelec = {}
  # use _irrep_doccs and _irrep_soccs help self.eig to compute orbital energy,
  # do not overwrite them
          self._irrep_doccs = []
          self._irrep_soccs = []
          self._keys = self._keys.union(['irrep_nelec'])
  
      def dump_flags(self):
          rohf.ROHF.dump_flags(self)
          if self.irrep_nelec:
              logger.info(self, 'irrep_nelec %s', self.irrep_nelec)
          return self
  
      def build(self, mol=None):
          if mol is None: mol = self.mol
          if mol.symmetry:
              if self.nelec is None:
                  nelec = self.mol.nelec
              else:
                  nelec = self.nelec
              check_irrep_nelec(mol, self.irrep_nelec, nelec)
          return hf.RHF.build(self, mol)
  
      eig = eig
  
      def get_grad(self, mo_coeff, mo_occ, fock=None):
          g = rohf.ROHF.get_grad(self, mo_coeff, mo_occ, fock)
          if self.mol.symmetry:
              occidxa = mo_occ > 0
              occidxb = mo_occ == 2
              viridxa = ~occidxa
              viridxb = ~occidxb
              uniq_var_a = viridxa.reshape(-1,1) & occidxa
              uniq_var_b = viridxb.reshape(-1,1) & occidxb
  
              orbsym = get_orbsym(self.mol, mo_coeff)
              sym_forbid = orbsym.reshape(-1,1) != orbsym
              sym_forbid = sym_forbid[uniq_var_a | uniq_var_b]
              g[sym_forbid.ravel()] = 0
          return g
  
      def get_occ(self, mo_energy=None, mo_coeff=None):
          if mo_energy is None: mo_energy = self.mo_energy
          mol = self.mol
          if not self.mol.symmetry:
              return rohf.ROHF.get_occ(self, mo_energy, mo_coeff)
  
          if hasattr(mo_energy, 'mo_ea'):
              mo_ea = mo_energy.mo_ea
              mo_eb = mo_energy.mo_eb
          else:
              mo_ea = mo_eb = mo_energy
          nmo = mo_ea.size
          mo_occ = numpy.zeros(nmo)
  
          orbsym = get_orbsym(self.mol, mo_coeff)
  
          rest_idx = numpy.ones(mo_occ.size, dtype=bool)
          neleca_fix = 0
          nelecb_fix = 0
          for i, ir in enumerate(mol.irrep_id):
              irname = mol.irrep_name[i]
              if irname in self.irrep_nelec:
                  ir_idx = numpy.where(orbsym == ir)[0]
                  if isinstance(self.irrep_nelec[irname], (int, numpy.integer)):
                      nelecb = self.irrep_nelec[irname] // 2
                      neleca = self.irrep_nelec[irname] - nelecb
                  else:
                      neleca, nelecb = self.irrep_nelec[irname]
                  mo_occ[ir_idx] = rohf._fill_rohf_occ(mo_energy[ir_idx],
                                                       mo_ea[ir_idx], mo_eb[ir_idx],
                                                       nelecb, neleca-nelecb)
                  neleca_fix += neleca
                  nelecb_fix += nelecb
                  rest_idx[ir_idx] = False
  
          nelec_float = mol.nelectron - neleca_fix - nelecb_fix
          assert(nelec_float >= 0)
          if len(rest_idx) > 0:
              rest_idx = numpy.where(rest_idx)[0]
              nopen = mol.spin - (neleca_fix - nelecb_fix)
              ncore = (nelec_float - nopen)//2
              mo_occ[rest_idx] = rohf._fill_rohf_occ(mo_energy[rest_idx],
                                                     mo_ea[rest_idx], mo_eb[rest_idx],
                                                     ncore, nopen)
  
          if self.nelec is None:
              nelec = self.mol.nelec
          else:
              nelec = self.nelec
          ncore = nelec[1]
          nocc  = nelec[0]
          nopen = nocc - ncore
          vir_idx = (mo_occ==0)
          if self.verbose >= logger.INFO and nocc < nmo and ncore > 0:
              ehomo = max(mo_energy[~vir_idx])
              elumo = min(mo_energy[ vir_idx])
              ndoccs = []
              nsoccs = []
              for i, ir in enumerate(mol.irrep_id):
                  irname = mol.irrep_name[i]
                  ir_idx = (orbsym == ir)
  
                  ndoccs.append(numpy.count_nonzero(mo_occ[ir_idx]==2))
                  nsoccs.append(numpy.count_nonzero(mo_occ[ir_idx]==1))
                  if ehomo in mo_energy[ir_idx]:
                      irhomo = irname
                  if elumo in mo_energy[ir_idx]:
                      irlumo = irname
  
              # to help self.eigh compute orbital energy
              self._irrep_doccs = ndoccs
              self._irrep_soccs = nsoccs
  
              logger.info(self, 'HOMO (%s) = %.15g  LUMO (%s) = %.15g',
                          irhomo, ehomo, irlumo, elumo)
  
              logger.debug(self, 'double occ irrep_nelec = %s', ndoccs)
              logger.debug(self, 'single occ irrep_nelec = %s', nsoccs)
              #_dump_mo_energy(mol, mo_energy, mo_occ, ehomo, elumo, orbsym,
              #                verbose=self.verbose)
              if nopen > 0:
                  core_idx = mo_occ == 2
                  open_idx = mo_occ == 1
                  vir_idx = mo_occ == 0
                  logger.debug(self, '                  Roothaan           | alpha              | beta')
                  logger.debug(self, '  Highest 2-occ = %18.15g | %18.15g | %18.15g',
                               max(mo_energy[core_idx]),
                               max(mo_ea[core_idx]), max(mo_eb[core_idx]))
                  logger.debug(self, '  Lowest 0-occ =  %18.15g | %18.15g | %18.15g',
                               min(mo_energy[vir_idx]),
                               min(mo_ea[vir_idx]), min(mo_eb[vir_idx]))
                  for i in numpy.where(open_idx)[0]:
                      logger.debug(self, '  1-occ =         %18.15g | %18.15g | %18.15g',
                                   mo_energy[i], mo_ea[i], mo_eb[i])
  
              numpy.set_printoptions(threshold=nmo)
              logger.debug(self, '  Roothaan mo_energy =\n%s', mo_energy)
              logger.debug1(self, '  alpha mo_energy =\n%s', mo_ea)
              logger.debug1(self, '  beta  mo_energy =\n%s', mo_eb)
              numpy.set_printoptions(threshold=1000)
          return mo_occ
  
      def make_rdm1(self, mo_coeff=None, mo_occ=None):
          if mo_coeff is None:
              mo_coeff = self.mo_coeff
          if mo_occ is None:
              mo_occ = self.mo_occ
          mo_a = mo_coeff[:,mo_occ>0]
          mo_b = mo_coeff[:,mo_occ==2]
          dm_a = numpy.dot(mo_a, mo_a.T)
          dm_b = numpy.dot(mo_b, mo_b.T)
          return numpy.array((dm_a, dm_b))
  
      def _finalize(self):
          rohf.ROHF._finalize(self)
  
          # sort MOs wrt orbital energies, it should be done last.
          c_sort = numpy.argsort(self.mo_energy[self.mo_occ==2].round(9))
          o_sort = numpy.argsort(self.mo_energy[self.mo_occ==1].round(9))
          v_sort = numpy.argsort(self.mo_energy[self.mo_occ==0].round(9))
          idx = numpy.arange(self.mo_energy.size)
          idx = numpy.hstack((idx[self.mo_occ==2][c_sort],
                              idx[self.mo_occ==1][o_sort],
                              idx[self.mo_occ==0][v_sort]))
          if hasattr(self.mo_energy, 'mo_ea'):
              mo_ea = self.mo_energy.mo_ea[idx]
              mo_eb = self.mo_energy.mo_eb[idx]
              self.mo_energy = lib.tag_array(self.mo_energy[idx],
                                             mo_ea=mo_ea, mo_eb=mo_eb)
          else:
              self.mo_energy = self.mo_energy[idx]
          orbsym = get_orbsym(self.mol, self.mo_coeff)
          self.mo_coeff = lib.tag_array(self.mo_coeff[:,idx], orbsym=orbsym[idx])
          self.mo_occ = self.mo_occ[idx]
          if self.chkfile:
              chkfile.dump_scf(self.mol, self.chkfile, self.e_tot, self.mo_energy,
                               self.mo_coeff, self.mo_occ, overwrite_mol=False)
          return self
  
      def analyze(self, verbose=None, **kwargs):
          if verbose is None: verbose = self.verbose
          from pyscf.lo import orth
          from pyscf.tools import dump_mat
          if not self.mol.symmetry:
              return rohf.ROHF.analyze(self, verbose, **kwargs)
  
          mol = self.mol
          mo_energy = self.mo_energy
          mo_occ = self.mo_occ
          mo_coeff = self.mo_coeff
          ovlp_ao = self.get_ovlp()
          log = logger.new_logger(self, verbose)
          if log.verbose >= logger.NOTE:
              nirrep = len(mol.irrep_id)
              orbsym = get_orbsym(self.mol, mo_coeff)
              wfnsym = 0
              ndoccs = []
              nsoccs = []
              for k,ir in enumerate(mol.irrep_id):
                  ndoccs.append(sum(orbsym[mo_occ==2] == ir))
                  nsoccs.append(sum(orbsym[mo_occ==1] == ir))
                  if nsoccs[k] % 2:
                      wfnsym ^= ir
              if mol.groupname in ('Dooh', 'Coov'):
                  log.info('TODO: total symmetry for %s', mol.groupname)
              else:
                  log.info('total symmetry = %s',
                           symm.irrep_id2name(mol.groupname, wfnsym))
              log.info('occupancy for each irrep:  ' + (' %4s'*nirrep),
                       *mol.irrep_name)
              log.info('double occ                 ' + (' %4d'*nirrep), *ndoccs)
              log.info('single occ                 ' + (' %4d'*nirrep), *nsoccs)
              log.info('**** MO energy ****')
              irname_full = {}
              for k,ir in enumerate(mol.irrep_id):
                  irname_full[ir] = mol.irrep_name[k]
              irorbcnt = {}
              if hasattr(mo_energy, 'mo_ea'):
                  mo_ea = mo_energy.mo_ea
                  mo_eb = mo_energy.mo_eb
                  log.note('                          Roothaan           | alpha              | beta')
                  for k, j in enumerate(orbsym):
                      if j in irorbcnt:
                          irorbcnt[j] += 1
                      else:
                          irorbcnt[j] = 1
                      log.note('MO #%-4d(%-3s #%-2d) energy= %-18.15g | %-18.15g | %-18.15g occ= %g',
                               k+1, irname_full[j], irorbcnt[j],
                               mo_energy[k], mo_ea[k], mo_eb[k], mo_occ[k])
              else:
                  for k, j in enumerate(orbsym):
                      if j in irorbcnt:
                          irorbcnt[j] += 1
                      else:
                          irorbcnt[j] = 1
                      log.note('MO #%-3d (%s #%-2d), energy= %-18.15g occ= %g',
                               k+1, irname_full[j], irorbcnt[j],
                               mo_energy[k], mo_occ[k])
  
          if log.verbose >= logger.DEBUG:
              label = mol.ao_labels()
              molabel = []
              irorbcnt = {}
              for k, j in enumerate(orbsym):
                  if j in irorbcnt:
                      irorbcnt[j] += 1
                  else:
                      irorbcnt[j] = 1
                  molabel.append('#%-d(%s #%d)' % (k+1, irname_full[j], irorbcnt[j]))
              log.debug(' ** MO coefficients (expansion on meta-Lowdin AOs) **')
              orth_coeff = orth.orth_ao(mol, 'meta_lowdin', s=ovlp_ao)
              c = reduce(numpy.dot, (orth_coeff.T, ovlp_ao, mo_coeff))
              dump_mat.dump_rec(self.stdout, c, label, molabel, start=1, **kwargs)
  
          dm = self.make_rdm1(mo_coeff, mo_occ)
          return self.mulliken_meta(mol, dm, s=ovlp_ao, verbose=verbose)
  
      def get_irrep_nelec(self, mol=None, mo_coeff=None, mo_occ=None):
          from pyscf.scf import uhf_symm
          if mol is None: mol = self.mol
          if mo_coeff is None: mo_coeff = (self.mo_coeff,self.mo_coeff)
          if mo_occ is None: mo_occ = ((self.mo_occ>0), (self.mo_occ==2))
          return uhf_symm.get_irrep_nelec(mol, mo_coeff, mo_occ)
  
      @lib.with_doc(canonicalize.__doc__)
      def canonicalize(self, mo_coeff, mo_occ, fock=None):
          if not hasattr(fock, 'focka'):
              fock = mf.get_fock(dm=dm)
          mo_e, mo_coeff = canonicalize(self, mo_coeff, mo_occ, fock)
          mo_ea = numpy.einsum('pi,pi->i', mo_coeff, fock.focka.dot(mo_coeff))
          mo_eb = numpy.einsum('pi,pi->i', mo_coeff, fock.fockb.dot(mo_coeff))
          mo_e = lib.tag_array(mo_e, mo_ea=mo_ea, mo_eb=mo_eb)
          return mo_e, mo_coeff
  
  
  def _dump_mo_energy(mol, mo_energy, mo_occ, ehomo, elumo, orbsym, title='',
                      verbose=logger.DEBUG):
      if isinstance(verbose, logger.Logger):
          log = verbose
      else:
          log = logger.Logger(mol.stdout, verbose)
      nirrep = mol.symm_orb.__len__()
      for i, ir in enumerate(mol.irrep_id):
          irname = mol.irrep_name[i]
          ir_idx = (orbsym == ir)
          nso = numpy.count_nonzero(ir_idx)
          nocc = numpy.count_nonzero(mo_occ[ir_idx])
          e_ir = mo_energy[ir_idx]
          if nocc == 0:
              log.debug('%s%s nocc = 0', title, irname)
          elif nocc == nso:
              log.debug('%s%s nocc = %d  HOMO = %.15g',
                        title, irname, nocc, e_ir[nocc-1])
          else:
              log.debug('%s%s nocc = %d  HOMO = %.15g  LUMO = %.15g',
                        title, irname, nocc, e_ir[nocc-1], e_ir[nocc])
              if e_ir[nocc-1]+1e-3 > elumo:
                  log.warn('%s%s HOMO %.15g > system LUMO %.15g',
                           title, irname, e_ir[nocc-1], elumo)
              if e_ir[nocc] < ehomo+1e-3:
                  log.warn('%s%s LUMO %.15g < system HOMO %.15g',
                           title, irname, e_ir[nocc], ehomo)
          log.debug('   mo_energy = %s', e_ir)
  
  
  class HF1e(ROHF):
      def scf(self, *args):
          logger.info(self, '\n')
          logger.info(self, '******** 1 electron system ********')
          self.converged = True
          h1e = self.get_hcore(self.mol)
          s1e = self.get_ovlp(self.mol)
          self.mo_energy, self.mo_coeff = self.eig(h1e, s1e)
          self.mo_occ = self.get_occ(self.mo_energy, self.mo_coeff)
          self.e_tot = self.mo_energy[self.mo_occ>0][0] + self.mol.energy_nuc()
          self._finalize()
          return self.e_tot
  
  def get_orbsym(mol, mo_coeff, s=None, check=False):
      if mo_coeff is None:
          orbsym = numpy.hstack([[ir] * mol.symm_orb[i].shape[1]
                                 for i, ir in enumerate(mol.irrep_id)])
      elif hasattr(mo_coeff, 'orbsym'):
          orbsym = mo_coeff.orbsym
      else:
          orbsym = symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb,
                                       mo_coeff, s, check)
      return numpy.asarray(orbsym)
  
  
  if __name__ == '__main__':
      from pyscf import gto
      mol = gto.Mole()
      mol.build(
          verbose = 1,
          output = None,
          atom = [['H', (0.,0.,0.)],
                  ['H', (0.,0.,1.)], ],
          basis = {'H': 'ccpvdz'},
          symmetry = True
      )
  
      method = RHF(mol)
      method.verbose = 5
      method.irrep_nelec['A1u'] = 2
      energy = method.kernel()
      print(energy)
      method.analyze()
  