src/IboView/IvIao.cpp

/* Copyright (c) 2015  Gerald Knizia
 *
 * This file is part of the IboView program (see: http://www.iboview.org)
 *
 * IboView is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, version 3.
 *
 * IboView is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for details.
 *
 * You should have received a copy of the GNU General Public License
 * along with bfint (LICENSE). If not, see http://www.gnu.org/licenses/
 *
 * Please see IboView documentation in README.txt for:
 * -- A list of included external software and their licenses. The included
 *    external software's copyright is not touched by this agreement.
 * -- Notes on re-distribution and contributions to/further development of
 *    the IboView software
 */

#include "Iv.h"

#include <iostream>
#include <boost/format.hpp>
#include <algorithm>
#include <set>

#include "CxAlgebra.h"
#include "CtMatrix.h"
#include "CtTiming.h"
#include "IvIao.h"
#include "IvDocument.h"

using namespace ct;

enum FIaoBasisFlags {
   IAO_OrthType = 0x03,
   IAO_OrthNone = 0x00,
   IAO_OrthSym = 0x01,
   IAO_OrthZbd = 0x02,
   IAO_NormalizeInput = 0x04
};


// // symmetrically orthogonalize a set of vectors with respect to overlap matrix S.
// void SymOrth(FMatrixView Orbs, FMatrixView const S, FMemoryStack &Mem)
// {
//    FStackMatrix
//       SmhOrb(Orbs.nCols, Orbs.nCols, &Mem),
//       T1(Orbs.nRows, Orbs.nCols, &Mem);
//    ChainMxm(SmhOrb, Transpose(Orbs), S, Orbs, Mem);
//    CalcSmhMatrix(SmhOrb, Mem, FSmhOptions(1e-15,1e-15,0,0));
//    Mxm(T1, Orbs, SmhOrb);
//    Move(Orbs, T1);
// }


// note: CIb must be provided from outside.
void MakeIaoBasisNew(FMatrixView CIb, FBasisSet *pMinBasis,
   FAtomSet const &/*Atoms*/, FBasisSet const *pOrbBasis,
   FMatrixView const &COcc_, FMatrixView const &S1, FMatrixView const &S1cd, uint Flags,
   FMemoryStack &Mem)
{
   uint
      nBf = pOrbBasis->nFn(), // number of 'real basis' functions
      nOcc = COcc_.nCols,
      nIb = pMinBasis->nFn(); // actual AOs in minimal basis
//    CIb = MakeStackMatrix(nBf, nIb, Mem);
   assert(CIb.nRows == nBf && CIb.nCols == nIb);

   FMatrixView
      COcc;
   if (0 == (Flags&IAO_NormalizeInput)) {
      COcc = COcc_;
   } else {
      // input has absorbed occupation numbers. Remove this.
      // (todo: find a better way...)
      assert(COcc_.nRows == S1.nRows);
      COcc = MakeStackMatrix(COcc_.nRows, nOcc, Mem);
      Move(COcc, COcc_);
      SymOrth(COcc, S1, Mem);
//       Mem.Alloc(pOccScales, nOcc);
//       FStackMatrix
//          SOcc(nOcc, nOcc, &Mem);
//       ChainMxm(SOcc, Transpose(COcc_), S1, COcc_, Mem);
//       SOcc.Print(xout, "Occ self overlap");
//       for (uint iOcc = 0; iOcc < nOcc; ++ iOcc) {
//          pOccScales[iOcc] = std::sqrt(SOcc(iOcc,iOcc));
//          double f = 1./pOccScales[iOcc];
//          for (uint i = 0; i < COcc_.nRows; ++ i)
//             COcc(i,iOcc) = f * COcc_(i,iOcc);
//       }
   }

   {
      FStackMatrix
         S2cd(nIb, nIb, &Mem),
         P12(nBf, nIb, &Mem),
         COcc2(nIb, nOcc, &Mem),
         CTil(nIb, nOcc, &Mem),
         STilCd(nIb, nIb, &Mem),
         CTil2BarT(nOcc, nIb, &Mem),
         T4(nBf, nOcc, &Mem);
      MakeIntMatrix(P12, *pOrbBasis, *pMinBasis, FKrn2i_Direct(&ir::g_IrOverlapKernel), Mem);
      MakeIntMatrix(S2cd, *pMinBasis, *pMinBasis, FKrn2i_Direct(&ir::g_IrOverlapKernel), Mem);

      // COcc2 = mdot(S21, COcc)                # O(N m^2)
      Mxm(COcc2, Transpose(P12), COcc);

      // P12 = la.solve(S1, S12)   # P12 = S1^{-1} S12
      CholeskySolve(P12, S1cd);

      // CTil = la.solve(S2, COcc2)             # O(m^3)
      CalcCholeskyFactors(S2cd);
      Move(CTil, COcc2);
      CholeskySolve(CTil, S2cd);

      // STil = mdot(COcc2.T, CTil)             # O(m^3)
      Mxm(STilCd, Transpose(COcc2), CTil);

      // CTil2Bar = la.solve(STil, CTil.T).T    # O(m^3)
      CalcCholeskyFactors(STilCd);
      Move(CTil2BarT, Transpose(CTil));
      CholeskySolve(CTil2BarT, STilCd);

      // T4 = COcc - mdot(P12, CTil2Bar)        # O(N m^2)
      Mxm(T4, P12, Transpose(CTil2BarT), -1.0);
      Add(T4, COcc);

      // CIb = P12 + mdot(T4, COcc2.T)          # O(N m^2)
      Mxm(CIb, T4, Transpose(COcc2));
      Add(CIb, P12);

      if ((Flags & IAO_OrthType) == IAO_OrthSym)
         SymOrth(CIb, S1, Mem);
   }

   if (COcc.pData != COcc_.pData)
      Mem.Free(COcc.pData);

   // still on stack: CIb.
}


namespace ct {

struct FLocalizeOptions
{
   uint
      nLocExp; //  2 or 4, exponent p in localization functional: L = sum[iA] n(i,A)^p
   uint
      nMaxIt;
   double
      ThrLoc;
   std::ostream
      *pxout; // if non-0, write output here.
   uint
      Verbosity;

   FLocalizeOptions()
      : nLocExp(4), nMaxIt(2048), ThrLoc(1e-8), pxout(0), Verbosity(1)
   {}
};

inline double pow2(double x) { return x*x; }
inline double pow3(double x) { return x*x*x; }
inline double pow4(double x) { double x2 = x*x; return x2*x2; }

// execute a 2x2 rotation of the vectors pA and pB (both with length nSize)
// of angle phi. In-place.
void Rot2x2(double *RESTRICT pA, double *RESTRICT pB, size_t nSize, double phi)
{
   double
      cs = std::cos(phi),
      ss = std::sin(phi);
   for (std::size_t i = 0; i < nSize; ++ i) {
      double
         tempA =  cs * pA[i] + ss * pB[i],
         tempB = -ss * pA[i] + cs * pB[i];
      pA[i] = tempA;
      pB[i] = tempB;
   }
}

// localize a set of vectors onto
//  L = \sum[i,A] \sum[mu \in A] CVec(mu,i)^p -> min.
// in a PM-like fashion.
// Parameters:
//    - pGroupOffsets: Array of length nGroups+1.
//      [pGroupOffsets[A], pGroupOffsets[A+1]) gives the association of
//      vector components to groups.
// Note: This function can be extended to localize on other criteria
// than charge, provided a replacement of the calculation of the Cii/Cjj/Cij
// matrix elements is provided. Most useful in the local exchange case is likely
// the optimization of <i|(r-A)^4|i>. It can also be used as a replacement for
// ZBD orthogonalization then (which is equivalent to optimizing on (r-A)^2).
void LocalizeVectors(FMatrixView CVec, size_t const *pGroupOffsets, size_t nGroups, FLocalizeOptions const &Opt, FMemoryStack const &Mem)
{
   IR_SUPPRESS_UNUSED_WARNING(Mem);

   size_t
      iIt;
   double
      L = -1., fVar2 = -1.;
   if (Opt.pxout && Opt.Verbosity >= 2) {
      *Opt.pxout << "   ITER.       LOC(Orb)      GRADIENT" << std::endl;
   }
   for (iIt = 0; iIt < Opt.nMaxIt; ++ iIt) {
      // calculate the value of the functional.
      L = 0.;
      // loop over groups (normally a ``group'' is all functions on an atom)
      for (size_t iGrp = 0; iGrp < nGroups; ++ iGrp)
         for (size_t iVec = 0; iVec < CVec.nCols; ++ iVec) {
            double nA = 0.; // population on the atom.
            for (size_t iFn = pGroupOffsets[iGrp]; iFn != pGroupOffsets[iGrp+1]; ++ iFn)
               nA += pow2(CVec(iFn,iVec));
            L += std::pow(nA, (int)Opt.nLocExp);
         }
      L = std::pow(L, 1./double(Opt.nLocExp));  // <- easier to compare different powers that way.

      fVar2 = 0.;
      // loop over vector pairs.
      for (size_t iVec = 0; iVec < CVec.nCols; ++ iVec)
         for (size_t jVec = 0; jVec < iVec; ++ jVec) {
            double
               Aij = 0.0, // hessian for 2x2 rotation (with variable atan(phi/4) substituted)
               Bij = 0.0; // gradient for 2x2 rotation (with variable atan(phi/4) substituted)
            // loop over groups (normally atoms)
            for (size_t iGrp = 0; iGrp < nGroups; ++ iGrp) {
               // calculate the charge matrix elements
               //     Cii = <i|A|i>,
               //     Cjj = <j|A|j>,
               // and Cij = <i|A|j>.
               double
                  Cii = 0., Cij = 0., Cjj = 0.;
               for (size_t iFn = pGroupOffsets[iGrp]; iFn != pGroupOffsets[iGrp+1]; ++ iFn) {
                  Cii += CVec(iFn,iVec) * CVec(iFn,iVec);
                  Cjj += CVec(iFn,jVec) * CVec(iFn,jVec);
                  Cij += CVec(iFn,iVec) * CVec(iFn,jVec);
               }

               // see Supporting Information for `` Intrinsic atomic orbitals:
               // An unbiased bridge between quantum theory and chemical concepts''
               // for a description of what this does.
               if (Opt.nLocExp == 2 || Opt.nLocExp == 3) {
                  Aij += sqr(Cij) - .25*sqr(Cii - Cjj);
                  Bij += Cij*(Cii - Cjj);
               } else if (Opt.nLocExp == 4) {
                  Aij += -pow4(Cii) - pow4(Cjj) + 6.*(pow2(Cii) + pow2(Cjj))*pow2(Cij) + pow3(Cii)*Cjj + Cii*pow3(Cjj);
                  Bij += 4.*Cij*(pow3(Cii) - pow3(Cjj));
               }
            }

            // non-degenerate rotation? This condition is not quite right,
            // and in Python it was not required. However, the Fortran atan2
            // sometimes did mysterious things without it.
            if (pow2(Aij) + pow2(Bij) > Opt.ThrLoc * 1e-2) {
               double
                  phi = .25 * std::atan2(Bij,-Aij);
               // 2x2 rotate vectors iVec and jVec.
               Rot2x2(&CVec(0,iVec), &CVec(0,jVec), CVec.nRows, phi);
               fVar2 += pow2(phi);
            }
         }
      fVar2 = std::sqrt(fVar2 / pow2(CVec.nCols));

      if (Opt.pxout && Opt.Verbosity >= 2) {
         *Opt.pxout << boost::format("%6i     %12.6f      %8.2e\n") % (1+iIt) % L % fVar2;
      }

      if (fVar2 < Opt.ThrLoc)
         break;
   }

   if (Opt.pxout && Opt.Verbosity >= 1) {
      if (Opt.Verbosity >= 2)
         *Opt.pxout << "\n";
//       *Opt.pxout << boost::format(" Iterative localization: IB/PM, %i iter; Final gradient %8.2e") % iIt % fVar2 << std::endl;
      *Opt.pxout << boost::format(" Iterative localization: IB/PM, %i iter; Final gradient %8.2e") % iIt % fVar2;
      Opt.pxout->flush();
   }
}


// void LocalizeOrbitals(FBasisSetPtr pBasis, FMatrixView COrb, FLocalizeOptions const &Opt);


} // namespace ct


// // note: CIb is allocated on Mem.
// void MakeIaoBasis(FMatrixView &CIb, uint &nIb, FBasisSetPtr &pMinBasis,
//    FAtomSet const &Atoms, FBasisSet const *pOrbBasis, FMatrixView &COcc, FMemoryStack &Mem)
// {
//    pMinBasis = new FBasisSet(Atoms, BASIS_Guess);
//    uint
//       nBf = pOrbBasis->nFn(), // number of 'real basis' functions
//       nOcc = COcc.nCols;
//    nIb = pMinBasis->nFn(); // actual AOs in minimal basis
//    CIb = MakeStackMatrix(nBf, nIb, Mem);
//
//    FStackMatrix
//       S1(nBf, nBf, &Mem),
//       S1cd(nBf, nBf, &Mem);
//    MakeIntMatrix(S1, *pOrbBasis, *pOrbBasis, FKrn2i_Direct(&ir::g_IrOverlapKernel), Mem);
//    Move(S1cd,S1);
//    CalcCholeskyFactors(S1cd);
//
//    MakeIaoBasisNew(CIb, &*pMinBasis, Atoms, pOrbBasis, COcc, S1, S1cd, IAO_OrthSym|IAO_NormalizeInput, Mem);
// //    MakeIaoBasisNew(CIb, nIb, &*pMinBasis, Atoms, pOrbBasis, COcc, S1, S1cd, IAO_OrthSym, Mem);
// //    SymOrth(CIb, S1, Mem);
// }


double CalcMatrixRmsNorm(FMatrixView M)
{
   double d = 0;
   for (uint iCol = 0; iCol < M.nCols; ++ iCol)
      for (uint iRow = 0; iRow < M.nRows; ++ iRow)
         d += sqr(M(iRow,iCol));
   return std::sqrt(d);
}

// TODO: merge these back into IvDocument.cpp. It has a bunch of very similar
// functions in FOrbital::MakeFullDesc and FOrbital::MakeIaoCharges.
void MakeIaoChargesRaw(double *pOut, double fOcc, double const *pIaoCoeffs, ct::FBasisSet *pMinBasis, ct::FAtomSet *pAtoms)
{
   size_t
      iShOf = 0;
   memset(pOut, 0, sizeof(pOut[0]) * pAtoms->size());
   for (size_t iSh = 0; iSh < pMinBasis->Shells.size(); ++ iSh) {
      ct::FBasisShell
         &Sh = pMinBasis->Shells[iSh];
      uint
         nShFn = Sh.nFn();
      double
         fChg = 0;
      for (uint iFn = 0; iFn < nShFn; ++ iFn)
         fChg += sqr(pIaoCoeffs[iFn + iShOf]);
      assert(size_t(Sh.iCenter) < pAtoms->size());
      pOut[Sh.iCenter] += fOcc * fChg;
      iShOf += nShFn;
   }
   assert(iShOf == pMinBasis->nFn());
}

struct FPrintOrbitalOptions
{
   double
      ThrPrint;
   size_t
      nMaxAtoms;
   bool
      PrintHeader;
   size_t
      iOrbOffset;
   FPrintOrbitalOptions()
      : ThrPrint(0.02), nMaxAtoms(20), PrintHeader(true), iOrbOffset(0)
   {}
};


static void PrintOrbitalChargeComposition(FLog &Log, FMatrixView CIbOcc, FMatrixView Ew, ct::FAtomSet *pAtoms,
   FBasisSet *pMinBasis, std::string const sGrp, double fOcc1, FPrintOrbitalOptions &Opt,
   ct::FMemoryStack &Mem)
{
   size_t const
      nAt = pAtoms->size();
//       *pAtOffsets = &pMinBasis->pRawBasis->CenterOffsets[0];
   if (Opt.PrintHeader) {
      Log.Write(" Summary of localized orbital composition:   [THRPRINT={:6.3f}]", Opt.ThrPrint);
      Log.Write("\n   ORB  GRP  ORB.ENERGY      CENTERS/CHARGES");
      Opt.PrintHeader = false;
      // ^- that's for appending lines using multiple calls.
   }
   for (size_t iOrb = 0; iOrb < CIbOcc.nCols; ++ iOrb) {
      TMemoryLock<double>
         Charges(nAt, &Mem);
      MakeIaoChargesRaw(Charges, fOcc1, &CIbOcc(0,iOrb), pMinBasis, pAtoms);
      TMemoryLock<size_t>
         Indices(nAt, &Mem);
      ct::ArgSort1(&Indices[0], &Charges[0], 1, nAt, true); // last: reverse. Largest first.

      double
         fChgRest = 0.;
      size_t
         nAtomsEmitted = 0;
      fmt::MemoryWriter
         out;
      out.write("{:6} {:>3} {:12.6f}   ", iOrb+1, sGrp, Ew[iOrb]);
      for (size_t ii = 0; ii < nAt; ++ ii) {
         size_t
            iAt = Indices[ii];
         if (iAt >= nAt)
            continue; // not supposed to happen.
         double
            fAtChg = Charges[iAt];
         if (fAtChg < Opt.ThrPrint || nAtomsEmitted >= Opt.nMaxAtoms) {
            fChgRest += fAtChg;
            continue;
         }
         ct::FAtom
            &At = (*pAtoms)[iAt];
         out.write("  {:>2}{:3} {:6.3f}", At.GetElementName(), (1+iAt), Charges[iAt]);
         nAtomsEmitted += 1;
      }
      if (fChgRest > Opt.ThrPrint)
         out.write("   (other: {:6.3f})", fChgRest);
      Log.Write(out.str());
   }
   Opt.iOrbOffset += CIbOcc.nCols;
}

// predicate: sorts orbitals first by type, then by occupation, and finally
// by energy.
struct FSortOrbitalPred
{
   bool operator() (FDataSetPtr pA, FDataSetPtr pB) {
      FOrbital
         *pOrbA = dynamic_cast<FOrbital*>(pA.get()),
         *pOrbB = dynamic_cast<FOrbital*>(pB.get());
      if (pOrbA == 0 && pOrbB == 0)
         return false;
      if (pOrbA == 0)
         return true;
      if (pOrbB == 0)
         return false;
      // so... both are orbitals.
      if (pOrbA->info.Spin < pOrbB->info.Spin)
         return true;
      if (pOrbB->info.Spin < pOrbA->info.Spin)
         return false;
//       if (pOrbA->info.fOcc < pOrbB->info.fOcc)
//          return true;
//       if (pOrbB->info.fOcc < pOrbA->info.fOcc)
//          return false;
      return pOrbA->info.fEnergy < pOrbB->info.fEnergy;
   }
};


void FFrame::RunIaoAnalysis(FLog &Log, FWfOptions *pWfOptions, bool AllowLocalize, ct::FMemoryStack &Mem)
{
   if (!HaveOrbitals())
      // nothing to make coefficients of...
      return;

   Log.WriteProgramIntro("INTRINSIC BASIS BONDING ANALYSIS (IBBA)");
   FTimer
      tIbbaTotal;

   FWfType
      WfType = GetWfType();
   if (WfType != WFTYPE_Rhf && WfType != WFTYPE_Uhf) {
      Log.Write("Sorry, wave function type not supported. Can only do RHF/RKS and UHF/UKS at this moment.");
      return;
   }

   using namespace ct;
   FAtomSet
      *pAtoms = pGetAtoms();

   TMemoryLock<char>
      pBaseOfMemory(0, &Mem);

   FBasisSetPtr
      pBasis,
      pMinBasis = new FBasisSet(*pAtoms, BASIS_Guess);
   FMatrixView
      COccAC, // alpha and closed
      COccBC; // beta and closed
   {
      TArray<FOrbital*>
         RefOrbitalsAC,
         RefOrbitalsBC;
      MakeOrbitalMatrix(COccAC, pBasis, &RefOrbitalsAC, ORBMAT_OccupiedOnly | ORBMAT_AlphaAndClosedOnly, Mem);
      assert(RefOrbitalsAC.size() == COccAC.nCols);
      if (WfType == WFTYPE_Uhf) {
         MakeOrbitalMatrix(COccBC, pBasis, &RefOrbitalsBC, ORBMAT_OccupiedOnly | ORBMAT_BetaAndClosedOnly, Mem);
         assert(RefOrbitalsBC.size() == COccBC.nCols);
      } else
         COccBC = COccAC;
   }


   uint
      nBf = pBasis->nFn(),
      nIb = pMinBasis->nFn(),
      nOccAC = COccAC.nCols,
      nOccBC = COccBC.nCols;

   if (nOccAC > nIb || nOccBC > nIb) {
      Log.Write(" More occupied orbitals than valence orbitals. IAO analysis aborted.");
      return;
   }

   // form orthogonal IAO vectors.
   uint
      IaoFlags = IAO_NormalizeInput;
   char const
      *pOrthMethodName = "Unknown";
   if (pWfOptions->GetAoType() == "IAO (Sym Orth.)") {
      IaoFlags |= IAO_OrthSym;
      pOrthMethodName = "Symmetric (Loewdin)";
   } else if (pWfOptions->GetAoType() == "IAO (ZBD Orth.)") {
      IaoFlags |= IAO_OrthZbd;
      pOrthMethodName = "Zero-Bond-Dipole (Laikov)";
   } else {
      IaoFlags |= IAO_OrthSym;
      IvNotify(NOTIFY_Error, IvFmt("Intrinsic basis type '%1' not recognized. Using IAO (Sym Orth.) instead.", pWfOptions->GetAoType()));
      pOrthMethodName = "Symmetric (Loewdin)";
   }

   Log.Write(" {:<34}{}", "Polarized AO type:", "IAO/2014");
   Log.Write(" {:<34}{}", "Orthogonalization method:", pOrthMethodName);

   FMatrixView
      // intrinsic basis spanning the closed & alpha orbitals. In RHF, this includes the span of the beta orbitals.
      CIbAC = MakeStackMatrix(nBf, nIb, Mem),
      // intrinsic basis spanning closed & beta orbitals
      CIbBC;

   {
      FStackMatrix
         S(nBf, nBf, &Mem),
         Scd(nBf, nBf, &Mem);
      MakeIntMatrix(S, *pBasis, *pBasis, FKrn2i_Direct(&ir::g_IrOverlapKernel), Mem);
//       S.Print(std::cout, "MAIN BASIS OVERLAP MATRIX #24");
      if (1) {
         // apply a sanity check to the orbitals.
         FStackMatrix
            SMo(COccAC.nCols, COccAC.nCols, &Mem);
         ChainMxm(SMo, Transpose(COccAC), S, COccAC, Mem);
         double
            fRmsd = SMo.fRmsdFromIdentity();
         if (fRmsd > 1e-5) {
            fmt::MemoryWriter ss;
            ss.write("Sanity check for input orbitals failed! Rmsd from orthogonality: {:8.2e}", fRmsd);
            Log.EmitWarning(ss.str());
            Log.Write(" If the input orbitals came from a .molden file: Beware of differnt normalizations for"
                    "\n different programs. Options: (1) Check for a newer version of IboView, (2) If it"
                    "\n still happens with the current version, contact Gerald Knizia.");
         }
//          SMo.Print(std::cout, "MO BASIS OVERLAP MATRIX #22");
      }
      Move(Scd, S);
      CalcCholeskyFactors(Scd);
      MakeIaoBasisNew(CIbAC, &*pMinBasis, *pAtoms, &*pBasis, COccAC, S, Scd, IaoFlags, Mem);
      if (WfType == WFTYPE_Uhf) {
         CIbBC = MakeStackMatrix(nBf, nIb, Mem);
         MakeIaoBasisNew(CIbBC, &*pMinBasis, *pAtoms, &*pBasis, COccBC, S, Scd, IaoFlags, Mem);
      } else {
         CIbBC = CIbAC;
      }
   }


   if (0) {
      FStackMatrix
         Rdm(nBf, nBf, &Mem),
         S(nBf, nBf, &Mem),
         Op(nBf, nBf, &Mem);
      Mxm(Rdm, COccAC, Transpose(COccAC));
      Op.Clear();
      pAtoms->AddKineticMatrix(Op, *pBasis, *pBasis, Mem);
      MakeIntMatrix(S, *pBasis, *pBasis, FKrn2i_Direct(&ir::g_IrOverlapKernel), Mem);
//       MakeIntMatrix(Op, *pBasis, *pBasis, FKrn2i_Direct(&ir::g_IrOverlapKernel), Mem);
      Log.WriteResult("Nuclear repulsion energy", pAtoms->NuclearRepulsionEnergy());
      Log.WriteResult("Kinetic energy", 2*Dot(Rdm,Op));
      Log.WriteResult("Number of electrons", Dot(Rdm,S));
   }


//    if (1) {
//       // keep a copy of the IAO coefficients and the basis sets. May be needed to form hybrids.
//       // Not sure if this is a good idea; could be better to just make a new one whenever needed.
//       m_pOrbBasis = pBasis;
//       m_pMinBasis = pMinBasis;
//       m_CIaoData.resize(CIb.GetStridedSize());
//       memcpy(&m_CIaoData[0], &CIb[0], sizeof(CIb[0]) * CIb.GetStridedSize());
//       m_CIb = CIb;
//       m_CIb.pData = &m_CIaoData[0];
//    }

   // setup the localization options.
   FLocalizeOptions
      LocOpt;
   QString
      LocMethod = pWfOptions->GetLocMethod().toLower();
   if (AllowLocalize) {
      LocOpt.nLocExp = 2;
      if (LocMethod == "ibo (exponent 2)")
         LocOpt.nLocExp = 2;
      else if (LocMethod == "ibo (exponent 4)")
         LocOpt.nLocExp = 4;
      else if (LocMethod == "none")
         LocOpt.nLocExp = 0;
      else {
         LocOpt.nLocExp = 2;
         Log.EmitWarning(q2s(IvFmt("Localization type '%1' not recognized. Defaulting to IBO (exponent 2).", LocMethod)));
      }
      Log.Write(" {:<34}L = sum[A,i] (i|n_A|i)^{} -> max", "Localization functional:", LocOpt.nLocExp);
      Log.WriteLine();
   } else {
      LocOpt.nLocExp = 0;
      LocMethod = "none";
   }

   LocOpt.ThrLoc = std::min(double(1e-6), double(1e-2 * pWfOptions->GetThrGrad()));
   // ^- WfOptions::ThrGrad is actually the SCF gradient threshold.

   typedef std::set<FOrbital*>
      FOrbitalRefSet;
   FOrbitalRefSet
      ProcessedOrbitals;

   bool
      InputOrbitalsChanged = false;
   for (int iCase = 0; iCase < 2; ++ iCase)
   {
      // need to define:
      // - RefOrbitals set and COcc
      // - Coeffs of IB vectors which span this space
      // - factors for spin and total densities?

      QString
         SpaceDesc = "";
      TArray<FOrbital*>
         RefOrbitals; // orbitals in current subspace which are localized together (i.e., freely mixable w/o changing the wave function)
      FMatrixView
         COcc; // corresponding orbital matrix
      double
         fChargeFactor,
         fSpinFactor;
      FMatrixView
         CIb; // IB which spanns the orbitals of the current subspace.
      if (WfType == WFTYPE_Rhf) {
         CIb = CIbAC; // <- that's all the occupied orbtials in RHF.
         if (iCase == 0) {
            // closed-shell orbitals.
            SpaceDesc = "AB";
            fChargeFactor = 2.;
            fSpinFactor = 0.;
            MakeOrbitalMatrix(COcc, pBasis, &RefOrbitals, ORBMAT_OccupiedOnly | ORBMAT_ClosedOnly, Mem);
         } else if (iCase == 1) {
            // open-shell orbitals (alpha occupied).
            SpaceDesc = "A_";
            fChargeFactor = 1.;
            fSpinFactor = 1.;
            MakeOrbitalMatrix(COcc, pBasis, &RefOrbitals, ORBMAT_OccupiedOnly | ORBMAT_AlphaOnly, Mem);
         } else if (iCase == 2) {
            // external orbitals (unoccupied)
            SpaceDesc = "__";
            fChargeFactor = 1.;
            fSpinFactor = 1.;
            MakeOrbitalMatrix(COcc, pBasis, &RefOrbitals, ORBMAT_VirtualOnly, Mem);
            // ^- TODO: if I actually had all the virtual orbitals, would the LeastSquaresSolve be sufficient
            //    to get the span of the anti-bonds? (w/o explicitly constructing them).
            //    Would have to update the nOcc number, but if it is ONLY that...
         } else {
            continue;
         }
      } else if (WfType == WFTYPE_Uhf) {
         if (iCase == 0) {
            // alpha orbitals.
            CIb = CIbAC;
            SpaceDesc = "A_";
            fChargeFactor = 1.;
            fSpinFactor = 1.;
            MakeOrbitalMatrix(COcc, pBasis, &RefOrbitals, ORBMAT_OccupiedOnly | ORBMAT_AlphaOnly, Mem);
         } else {
            // beta orbitals.
            CIb = CIbBC;
            SpaceDesc = "_B";
            fChargeFactor = 1.;
            fSpinFactor = -1.;
            MakeOrbitalMatrix(COcc, pBasis, &RefOrbitals, ORBMAT_OccupiedOnly | ORBMAT_BetaOnly, Mem);
         }
      } else {
         Log.EmitError("Unrecognized wave function type in RunIaoAnalysis (inner).");
         continue;
      }
      IR_SUPPRESS_UNUSED_WARNING2(fSpinFactor, fChargeFactor);
      uint
         nOcc = COcc.nCols;
      if (nOcc == 0)
         continue; // no orbitals in this class.

      // re-express occupied vectors in terms of the IAO vectors. This is an exact
      // transformation.
      // (note: this probably can be done in a different way than via SVD... probably via
      //      Sib^{-1} CIb.T S1 COcc, where Sib = CIb.T S1 CIb
      // I just don't want to think about it for the moment)
      // Note: BEWARE OF THE ABSORBED OCCUPATION NUMBERS!!

      // solve CIb CIbOcc = COcc for CIbOcc
      FStackMatrix
         CIbOcc(nIb, nOcc, &Mem);
      LeastSquaresSolveSafe(CIb, CIbOcc, COcc, 1e-9, Mem);

      if (0) {
         // test if IB works.
         uint
            nBf = pBasis->nFn();
         FStackMatrix
            // COcc reconstructed from IB and CIbOcc.
            COccRec(nBf, nIb, &Mem);
         Mxm(COccRec, CIb, CIbOcc);
         Add(COccRec, COcc, -1.);
         std::cout << boost::format("Norm(COccRec - COcc) = %8.2e   [nBf=%i, nOcc=%i]\n") % CalcMatrixRmsNorm(COccRec) % nBf % COcc.nCols;
      }


      // make a fock matrix by assuming that the input orbitals are canonical.
      // (FIXME: need all orbitals for this. at least when making anti-bonds.)
      FStackMatrix
         Ew(nOcc, 1, &Mem);
      for (size_t iOcc = 0; iOcc < nOcc; ++ iOcc)
         Ew[iOcc] = RefOrbitals[iOcc]->info.fEnergy;


//       {
//          Log.Write("!INITIAL ORBITALS:");
//          FPrintOrbitalOptions PrintOpt;
//          PrintOrbitalChargeComposition(Log, CIbOcc, Ew, pAtoms, &*pMinBasis, "AB", 2., PrintOpt, Mem);
//       }


//       Log.Write("Run localize? {}  AllowLocalize? {}", LocOpt.nLocExp, (int)AllowLocalize);
      if (LocOpt.nLocExp != 0) {
         std::stringstream
            ss;
         LocOpt.Verbosity = 1;
         LocOpt.pxout = &ss;
         size_t
            nAt = pAtoms->size();
         TMemoryLock<size_t>
            pAtOffsets(nAt+1, &Mem);
            // ^- that is NOT the same as pRawBasis->CenterOffsets! The latter is in terms of shells!!!
         for (size_t iAt = 0; iAt <= nAt; ++ iAt)
            pAtOffsets[iAt] = pMinBasis->pRawBasis->iCenFn(iAt);
         if (nAt + 1 != pMinBasis->pRawBasis->CenterOffsets.size())
            Log.EmitError("Number of atoms in minimal basis inconsistent with frame main basis (?).");
         if (CIbOcc.nRows != pMinBasis->pRawBasis->nFn())
            Log.EmitError("Number of basis functions in minimal basis inconsistent with frame main basis (?).");

         if (0) {
            std::stringstream str;
            pMinBasis->Print(str);
            Log.Write(str.str());
            Log.Write("\ncenter offsets:");
            for (size_t i = 0; i < nAt+1; ++ i)
               Log.Write("  {:4}. {:6}", i, pAtOffsets[i]);
         }

         FStackMatrix
            PrevCIbOcc(nIb, nOcc, &Mem),
            PrevEw(nOcc, 1, &Mem);
         Move(PrevCIbOcc, CIbOcc);
         Move(PrevEw, Ew);

         if (iCase == 0)
            Log.WriteLine();
         LocalizeVectors(CIbOcc, pAtOffsets, nAt, LocOpt, Mem);
         Log.Write("{}; Class: {}", ss.str(), q2s(SpaceDesc));

         // make new set of vectors by transforming CIbOcc back to the
         // main basis.
         Mxm(COcc, CIb, CIbOcc);

         if (1) {
            // compute new orbital energies---assuming old orbitals were canonical.
            FStackMatrix
               U(nOcc, nOcc, &Mem);
            Mxm(U, Transpose(PrevCIbOcc), CIbOcc);
            for (size_t iOcc = 0; iOcc < nOcc; ++ iOcc) {
               Ew[iOcc] = 0;
               for (size_t iPrev = 0; iPrev < nOcc; ++ iPrev)
                  Ew[iOcc] += PrevEw[iPrev] * sqr(U(iPrev, iOcc));
            }
         }
         // fixme: fock matrix? occupation numbers?
         InputOrbitalsChanged = true;
      }

//       {
// //          Log.Write("!FINAL ORBITALS:");
//          FPrintOrbitalOptions PrintOpt;
//          PrintOpt.PrintHeader = iCase == 0;
//          PrintOrbitalChargeComposition(Log, CIbOcc, Ew, pAtoms, &*pMinBasis, q2s(SpaceDesc), fChargeFactor, PrintOpt, Mem);
//       }

      // store IAO coefficients of orbitals in the orbital data sets,
      // and also a link to the minimal basis (to distinguish what is what in the basis)
      for (uint iOcc = 0; iOcc < COcc.nCols; ++ iOcc) {
         FOrbital *pOrb = RefOrbitals[iOcc];
         if (ProcessedOrbitals.find(pOrb) != ProcessedOrbitals.end())
            IvNotify(NOTIFY_Warning, "An orbital occured in more than one localization group!");

         ProcessedOrbitals.insert(pOrb);
         pOrb->pMinBasis = pMinBasis;
         pOrb->pIaoCoeffs.clear();
         pOrb->pIaoCoeffs.insert(pOrb->pIaoCoeffs.begin(), &CIbOcc(0,iOcc), &CIbOcc(0,iOcc)+nIb);
         if (InputOrbitalsChanged) {
            RefOrbitals[iOcc]->info.fEnergy = Ew[iOcc];
            pOrb->pCoeffs.clear();
            pOrb->pCoeffs.insert(pOrb->pCoeffs.begin(), &COcc(0,iOcc), &COcc(0,iOcc)+nBf);
         }
      }
   }
   if (InputOrbitalsChanged) {
      // remove all orbitals which were not processed. This is used to delete either
      // all or non-valence virtual orbitals after the localization procedure.
      if (1) {
         FDataSetList
            NewData;
         for (size_t i = 0; i < m_Data.size(); ++ i) {
            FOrbital
               *pOrb = dynamic_cast<FOrbital*>(m_Data[i].get());
            if (pOrb == 0 || ProcessedOrbitals.find(pOrb) != ProcessedOrbitals.end())
               // copy pointers to all data sets which either are not orbitals,
               // or are orbitals which were processed.
               NewData.push_back(m_Data[i]);
         }
         m_Data.swap(NewData);
      }

      // sort new orbitals (by type and energy) and update their descriptions
      std::stable_sort(m_Data.begin(), m_Data.end(), FSortOrbitalPred());

      size_t iOrb = 0;
      for (size_t i = 0; i < m_Data.size(); ++ i) {
         FOrbital *pOrb = dynamic_cast<FOrbital*>(m_Data[i].get());
         if (pOrb) {
            pOrb->info.iOriginalIndex = i;
            pOrb->UpdateDescFromInfo(iOrb);
//             Log.Write("   pOrb[{:4}]  E = {:12.5f}  NewDesc = '{}'", i, pOrb->info.fEnergy, q2s(pOrb->GetDesc()));
            iOrb += 1;
         }
      }

      // print orbital composition in terms of atomic partial charges (with new order)
      double ThrPrint = 0.02;
      iOrb = 0;
      Log.Write("\n Summary of localized orbital composition:   [THRPRINT={:6.3f}]", ThrPrint);
      Log.Write("\n   ORB  GRP   ORB.ENERGY   CENTERS/CHARGES");
      for (size_t i = 0; i < m_Data.size(); ++ i) {
         FOrbital *pOrb = dynamic_cast<FOrbital*>(m_Data[i].get());
         if (pOrb) {
            Log.Write("{:6}  {:>3} {:12.6f} {}", iOrb+1, pOrbDescFromType(pOrb->info.Spin, pOrb->info.fOcc),
                  pOrb->info.fEnergy, q2s(pOrb->MakeFullDesc(ThrPrint, FOrbital::ORBDESC_ChargesOnly)));
            iOrb += 1;
         }
      }
   }

   if (1) {
      Log.WriteLine();
      RunChargeAnalysis(Log, 0);
   }

   Log.WriteLine();
   Log.WriteTiming("IBBA", (double)tIbbaTotal);
   Log.WriteLine();

//    GetBondOrderType
//    GetThrGrad
//    GetOrbDisplay
//    GetOrbDivision
//    GetLocMethod
}


ct::FMatrixView FFrame::MakeIaoBasis(ct::FAtomSet *pAtoms, ct::FMemoryStack &Mem)
{
   IR_SUPPRESS_UNUSED_WARNING2(pAtoms, Mem);
   return m_CIb;
}


void FDocument::MakeHybridsForSelectedAtomGroup()
{
}