- Implemented partial localization for different occupations

- Fixed values for tests - some cleanup
MRChemSoft · Jun 10, 2024 · deeda7a · deeda7a
1 parent b612183
commit deeda7a
Show file tree

Hide file tree

Showing 9 changed files with 1,052 additions and 965 deletions.
diff --git a/src/driver.cpp b/src/driver.cpp
@@ -511,7 +511,7 @@ bool driver::scf::guess_energy(const json &json_guess, Molecule &mol, FockBuilde
     auto &F_mat = mol.getFockMatrix();
 
     F_mat = ComplexMatrix::Zero(Phi.size(), Phi.size());
-    if (localize && rotate) orbital::localize(prec, Phi, F_mat);
+    if (localize && rotate) orbital::localize(prec, Phi, F_mat, true);
 
     F.setup(prec);
     F_mat = F(Phi, Phi);

diff --git a/src/qmfunctions/orbital_utils.cpp b/src/qmfunctions/orbital_utils.cpp
@@ -47,7 +47,7 @@ namespace mrchem {
 extern mrcpp::MultiResolutionAnalysis<3> *MRA; // Global MRA
 
 namespace orbital {
-ComplexMatrix localize(double prec, OrbitalVector &Phi, int spin);
+ComplexMatrix localize(double prec, OrbitalVector &Phi, int spin, bool partial = false);
 ComplexMatrix calc_localization_matrix(double prec, OrbitalVector &Phi);
 
 /* POD struct for orbital meta data. Used for simple MPI communication. */
@@ -388,6 +388,34 @@ OrbitalVector orbital::disjoin(OrbitalVector &Phi, int spin) {
     return out;
 }
 
+ /** @brief Erase orbitals from a vector
+ *
+ * Orbitals are copied in a new vector, except the ones which are defined as to_erase.
+ * The number of orbitals in the new vector will be smaller if to_erase is not empty.
+ * The indices to be erased are marked with 1, for example
+ * to_erase[4] = 1; to_erase[5] = 1; will erase orbitals with indices 4 and 5
+ *
+ */
+OrbitalVector orbital::deep_copy_erase(OrbitalVector &Phi, std::map<int,int> to_erase) {
+    OrbitalVector out;
+    for (auto &i : Phi) {
+        if (to_erase[i.getRank()] == 0) {
+            Orbital out_i(i.spin(), i.occ(), i.getRank());
+            if (i.getRank() % mrcpp::mpi::wrk_size != out.size() % mrcpp::mpi::wrk_size) {
+                // need to send orbital from owner to new owner
+                if (mrcpp::mpi::my_orb(i)) { mrcpp::mpi::send_function(i, out.size() % mrcpp::mpi::wrk_size, i.getRank(), mrcpp::mpi::comm_wrk); }
+                if (mrcpp::mpi::my_orb(out.size())) { mrcpp::mpi::recv_function(out_i, i.getRank() % mrcpp::mpi::wrk_size, i.getRank(), mrcpp::mpi::comm_wrk); }
+            } else {
+                // owner of old and new orbital. Just copy
+                mrcpp::cplxfunc::deep_copy(out_i, i);
+            }
+            out_i.setRank(out.size());
+            out.push_back(out_i);
+        }
+    }
+    return out;
+}
+
 /** @brief Write orbitals to disk
  *
  * @param Phi: orbitals to save
@@ -700,7 +728,7 @@ ComplexMatrix orbital::calc_lowdin_matrix(OrbitalVector &Phi) {
     return S_m12;
 }
 
-ComplexMatrix orbital::localize(double prec, OrbitalVector &Phi, ComplexMatrix &F) {
+ComplexMatrix orbital::localize(double prec, OrbitalVector &Phi, ComplexMatrix &F, bool partial) {
     Timer t_tot;
     auto plevel = Printer::getPrintLevel();
     mrcpp::print::header(2, "Localizing orbitals");
@@ -713,9 +741,9 @@ ComplexMatrix orbital::localize(double prec, OrbitalVector &Phi, ComplexMatrix &
     int nA = size_alpha(Phi);
     int nB = size_beta(Phi);
     ComplexMatrix U = ComplexMatrix::Identity(nO, nO);
-    if (nP > 0) U.block(0, 0, nP, nP) = localize(prec, Phi, SPIN::Paired);
-    if (nA > 0) U.block(nP, nP, nA, nA) = localize(prec, Phi, SPIN::Alpha);
-    if (nB > 0) U.block(nP + nA, nP + nA, nB, nB) = localize(prec, Phi, SPIN::Beta);
+    if (nP > 0) U.block(0, 0, nP, nP) = localize(prec, Phi, SPIN::Paired, partial);
+    if (nA > 0) U.block(nP, nP, nA, nA) = localize(prec, Phi, SPIN::Alpha, partial);
+    if (nB > 0) U.block(nP + nA, nP + nA, nB, nB) = localize(prec, Phi, SPIN::Beta, partial);
 
     // Transform Fock matrix
     F = U.adjoint() * F * U;
@@ -733,9 +761,43 @@ Localization is done for each set of spins separately (we don't want to mix spin
 The localization matrix is returned for further processing.
 
 */
-ComplexMatrix orbital::localize(double prec, OrbitalVector &Phi, int spin) {
+ComplexMatrix orbital::localize(double prec, OrbitalVector &Phi, int spin, bool partial) {
     OrbitalVector Phi_s = orbital::disjoin(Phi, spin);
-    ComplexMatrix U = calc_localization_matrix(prec, Phi_s);
+    ComplexMatrix U;
+    if (partial) {
+        DoubleVector occ = orbital::get_occupations(Phi_s);
+        int mainOcc = (spin == SPIN::Paired) ? 2 : 1;
+        std::map<int, int> to_erase;
+        for (int i = 0; i < Phi_s.size(); i++) {
+            if (occ(i) == mainOcc) {
+                to_erase[i] = 0;
+            }
+            else {
+                to_erase[i] = 1;
+            }
+        }
+        OrbitalVector Phi_s_small = orbital::deep_copy_erase(Phi_s, to_erase);
+        ComplexMatrix U_small = calc_localization_matrix(prec, Phi_s_small);
+        U = ComplexMatrix::Identity(Phi_s.size(), Phi_s.size());
+        unsigned int j = 0;
+        for (int i = 0; i < Phi_s.size(); i++) {
+            if (to_erase[i] == 0) {
+                unsigned int l = 0;
+                for (int k = 0; k < Phi_s.size(); k++) {
+                    if (to_erase[k] == 0) {
+                        U(i, k) = U_small(j, l);
+                        l++;
+                    }
+                }
+                j++;
+            }
+        }
+        print_utils::matrix(2, "Loc matrix small", U_small.real());
+        print_utils::matrix(2, "Localization matrix", U.real());
+    }
+    else {
+        U = calc_localization_matrix(prec, Phi_s);
+    }
     Timer rot_t;
     mrcpp::mpifuncvec::rotate(Phi_s, U, prec);
     Phi = orbital::adjoin(Phi, Phi_s);

diff --git a/src/qmfunctions/orbital_utils.h b/src/qmfunctions/orbital_utils.h
@@ -53,6 +53,8 @@ OrbitalVector param_copy(const OrbitalVector &Phi);
 OrbitalVector adjoin(OrbitalVector &Phi_a, OrbitalVector &Phi_b);
 OrbitalVector disjoin(OrbitalVector &Phi, int spin);
 
+OrbitalVector deep_copy_erase(OrbitalVector &Phi, std::map<int,int> to_erase);
+
 void save_orbitals(OrbitalVector &Phi, const std::string &file, int spin = -1);
 OrbitalVector load_orbitals(const std::string &file, int n_orbs = -1);
 
@@ -67,7 +69,7 @@ ComplexMatrix calc_overlap_matrix(OrbitalVector &BraKet);
 ComplexMatrix calc_overlap_matrix(OrbitalVector &Bra, OrbitalVector &Ket);
 DoubleMatrix calc_norm_overlap_matrix(OrbitalVector &BraKet);
 
-ComplexMatrix localize(double prec, OrbitalVector &Phi, ComplexMatrix &F);
+ComplexMatrix localize(double prec, OrbitalVector &Phi, ComplexMatrix &F, bool partial = false);
 ComplexMatrix diagonalize(double prec, OrbitalVector &Phi, ComplexMatrix &F);
 ComplexMatrix orthonormalize(double prec, OrbitalVector &Phi, ComplexMatrix &F);
 

diff --git a/src/scf_solver/GroundStateSolver.cpp b/src/scf_solver/GroundStateSolver.cpp
@@ -325,23 +325,26 @@ json GroundStateSolver::optimize(Molecule &mol, FockBuilder &F, OrbitalVector &P
             if (restricted) {
                 DoubleVector occNew = getNewOccupations(Phi_n, Phi_mom);
                 orbital::set_occupations(Phi_n, occNew);
+
+                // orbital::print(Phi_n);
+                // mol.calculateOrbitalPositions();
+                // mol.printOrbitalPositions();
             }
             else {
                 // in case of unrestricted calculation, get the new occupation for alpha and beta spins independently
-                OrbitalVector Phi_n_copy = orbital::deep_copy(Phi_n);
-                OrbitalVector Phi_mom_copy = orbital::deep_copy(Phi_mom);
-                auto Phi_n_a = orbital::disjoin(Phi_n_copy, SPIN::Alpha);
-                auto Phi_mom_a = orbital::disjoin(Phi_mom_copy, SPIN::Alpha);
+                OrbitalVector Phi_n_a = orbital::disjoin(Phi_n, SPIN::Alpha);
+                OrbitalVector Phi_mom_a = orbital::disjoin(Phi_mom, SPIN::Alpha);
                 DoubleVector occAlpha = getNewOccupations(Phi_n_a, Phi_mom_a);
-                DoubleVector occBeta = getNewOccupations(Phi_n_copy, Phi_mom_copy);
+                DoubleVector occBeta = getNewOccupations(Phi_n, Phi_mom);
                 DoubleVector occNew(occAlpha.size() + occBeta.size());
                 occNew << occAlpha, occBeta;
                 orbital::set_occupations(Phi_n, occNew);
-            }
-            if (plevel >= 2) { // TODO: whats the right print level here?
-                orbital::print(Phi_n);
-                mol.calculateOrbitalPositions();
-                mol.printOrbitalPositions();
+                orbital::adjoin(Phi_n, Phi_n_a);
+                orbital::adjoin(Phi_mom, Phi_mom_a);
+
+                // orbital::print(Phi_n);
+                // mol.calculateOrbitalPositions();
+                // mol.printOrbitalPositions();
             }
         }
 
@@ -350,8 +353,6 @@ json GroundStateSolver::optimize(Molecule &mol, FockBuilder &F, OrbitalVector &P
             Phi_mom = orbital::deep_copy(Phi_n);
         }
 
-        // TODO: keep MOM or just use IMOM?
-
         orbital::orthonormalize(orb_prec, Phi_n, F_mat);
 
         // Compute Fock matrix and energy
@@ -373,7 +374,7 @@ json GroundStateSolver::optimize(Molecule &mol, FockBuilder &F, OrbitalVector &P
 
         // Rotate orbitals
         if (needLocalization(nIter, converged)) {
-            ComplexMatrix U_mat = orbital::localize(orb_prec, Phi_n, F_mat);
+            ComplexMatrix U_mat = orbital::localize(orb_prec, Phi_n, F_mat, true);
             F.rotate(U_mat);
             kain.clear();
         } else if (needDiagonalization(nIter, converged)) {

diff --git a/tests/li2h_imom/li2h_gs.inp b/tests/li2h_imom/li2h_gs.inp
@@ -22,7 +22,6 @@ WaveFunction {
 
 SCF {
   run = false
-  localize = true                   # Use canonical of localized orbitals
   write_orbitals = true
   path_orbitals = initial_guess
 }