Expose additional PT2 options, refactor sort_and_accumulate, add coeffs into append_xyz_contributions

wavefunction91 · wavefunction91 · commit d868a9e64bcd · 2023-11-14T15:41:00.000-08:00
diff --git a/include/macis/asci/determinant_contributions.hpp b/include/macis/asci/determinant_contributions.hpp
@@ -20,6 +20,7 @@ struct asci_contrib {
   double h_diag;
 
   auto rv() const { return c_times_matel / h_diag; }
+  auto pt2() const { return rv() * c_times_matel; }
 };
 
 template <typename WfnT>
@@ -48,7 +49,7 @@ void append_singles_asci_contributions(
       for(auto p : occ_othr) h_el += V_ov[p];
 
       // Early Exit
-      if(std::abs(h_el) < h_el_tol) continue;
+      if(std::abs(coeff * h_el) < h_el_tol) continue;
 
       // Calculate Excited Determinant
       auto ex_det = wfn_traits::template single_excitation_no_check<Sigma>(
@@ -97,7 +98,7 @@ void append_ss_doubles_asci_contributions(
           const auto jb = b + j * LDG;
           const auto G_aibj = G_ai[jb];
 
-          if(std::abs(G_aibj) < h_el_tol) continue;
+          if(std::abs(coeff * G_aibj) < h_el_tol) continue;
 
 #if 0
           // Calculate excited determinant string (spin)
@@ -165,7 +166,7 @@ void append_os_doubles_asci_contributions(
           const auto jb = b + j * LDV;
           const auto V_aibj = V_ai[jb * LDV2];
 
-          if(std::abs(V_aibj) < h_el_tol) continue;
+          if(std::abs(coeff * V_aibj) < h_el_tol) continue;
 
           double sign_beta = single_excitation_sign(state_beta, b, j);
           double sign = sign_alpha * sign_beta;
diff --git a/include/macis/asci/determinant_search.hpp b/include/macis/asci/determinant_search.hpp
@@ -39,6 +39,12 @@ struct ASCISettings {
   double h_el_tol = 1e-8;
   double rv_prune_tol = 1e-8;
   size_t pair_size_max = 5e8;
+
+  double pt2_tol = 1e-16;
+  size_t pt2_reserve_count = 70000000;
+  bool   pt2_prune = false;
+  bool   pt2_precompute_eps = false;
+
   bool just_singles = false;
   size_t grow_factor = 8;
   size_t max_refine_iter = 6;
@@ -49,6 +55,7 @@ struct ASCISettings {
 
   // bool dist_triplet_random = false;
   int constraint_level = 2;  // Up To Quints
+  int pt2_constraint_level = 5;
 };
 
 template <size_t N>
diff --git a/include/macis/asci/determinant_sort.hpp b/include/macis/asci/determinant_sort.hpp
@@ -11,6 +11,7 @@
 #if __has_include(<boost/sort/pdqsort/pdqsort.hpp>)
 #define MACIS_USE_BOOST_SORT
 #include <boost/sort/pdqsort/pdqsort.hpp>
+#include <boost/sort/sort.hpp>
 #endif
 
 namespace macis {
@@ -61,6 +62,31 @@ void reorder_ci_on_alpha(WfnIterator begin, WfnIterator end, double* C) {
   std::copy(reorder_C.begin(), reorder_C.end(), C);
 }
 
+
+template <typename PairIterator>
+PairIterator accumulate_asci_pairs(PairIterator pairs_begin,
+                                   PairIterator pairs_end) {
+
+  // Accumulate the ASCI scores into first instance of unique bitstrings
+  auto cur_it = pairs_begin;
+  for(auto it = cur_it + 1; it != pairs_end; ++it) {
+    // If iterate is not the one being tracked, update the iterator
+    if(it->state != cur_it->state) {
+      cur_it = it;
+    }
+
+    // Accumulate
+    else {
+      cur_it->c_times_matel += it->c_times_matel;
+      it->c_times_matel = NAN;  // Zero out to expose potential bugs
+    }
+  }
+
+  // Remote duplicate bitstrings
+  return std::unique(pairs_begin, pairs_end,
+                     [](auto x, auto y) { return x.state == y.state; });
+}
+
 template <typename PairIterator>
 PairIterator sort_and_accumulate_asci_pairs(PairIterator pairs_begin,
                                             PairIterator pairs_end) {
@@ -80,24 +106,29 @@ PairIterator sort_and_accumulate_asci_pairs(PairIterator pairs_begin,
 #endif
       (pairs_begin, pairs_end, comparator);
 
-  // Accumulate the ASCI scores into first instance of unique bitstrings
-  auto cur_it = pairs_begin;
-  for(auto it = cur_it + 1; it != pairs_end; ++it) {
-    // If iterate is not the one being tracked, update the iterator
-    if(it->state != cur_it->state) {
-      cur_it = it;
-    }
+  return accumulate_asci_pairs(pairs_begin, pairs_end);
+}
 
-    // Accumulate
-    else {
-      cur_it->c_times_matel += it->c_times_matel;
-      it->c_times_matel = NAN;  // Zero out to expose potential bugs
-    }
-  }
+template <typename PairIterator>
+PairIterator stable_sort_and_accumulate_asci_pairs(PairIterator pairs_begin,
+                                                   PairIterator pairs_end) {
+  const size_t npairs = std::distance(pairs_begin, pairs_end);
 
-  // Remote duplicate bitstrings
-  return std::unique(pairs_begin, pairs_end,
-                     [](auto x, auto y) { return x.state == y.state; });
+  if(!npairs) return pairs_end;
+
+  auto comparator = [](const auto& x, const auto& y) {
+    return bitset_less(x.state, y.state);
+  };
+
+// Sort by bitstring
+#ifdef MACIS_USE_BOOST_SORT
+  boost::sort::flat_stable_sort
+#else
+  std::stable_sort
+#endif
+      (pairs_begin, pairs_end, comparator);
+
+  return accumulate_asci_pairs(pairs_begin, pairs_end);
 }
 
 template <typename WfnT>
diff --git a/include/macis/asci/mask_constraints.hpp b/include/macis/asci/mask_constraints.hpp
@@ -630,7 +630,7 @@ auto dist_constraint_general(size_t nlevels, size_t norb, size_t ns_othr,
 template <typename WfnType, typename ContainerType>
 auto gen_constraints_general(size_t nlevels, size_t norb, size_t ns_othr,
                              size_t nd_othr, const ContainerType& unique_alpha,
-                             int world_size) {
+                             int world_size, size_t nlevel_min = 0) {
   using wfn_traits = wavefunction_traits<WfnType>;
   using constraint_type = alpha_constraint<wfn_traits>;
   using string_type = typename constraint_type::constraint_type;
@@ -671,6 +671,21 @@ auto gen_constraints_general(size_t nlevels, size_t norb, size_t ns_othr,
         auto constraint = constraint_type::make_triplet(t_i, t_j, t_k);
         constraint_sizes.emplace_back(constraint, 0ul);
   }
+  // Build up higher-order constraints as base if requested
+  for(size_t ilevel = 0; ilevel < nlevel_min; ++ilevel) {
+    std::vector cur_constraints = constraint_sizes;
+    for(auto [c,nw] : cur_constraints) {
+      const auto C_min = c.C_min();
+      for(auto q_l = 0; q_l < C_min; ++q_l) {
+        // Generate masks / counts
+        string_type cn_C = c.C();
+        cn_C.flip(q_l);
+        string_type cn_B = c.B() >> (C_min - q_l);
+        constraint_type c_next(cn_C, cn_B, q_l);
+        constraint_sizes.emplace_back(c_next, 0ul);
+      }
+    }
+  }
 
   struct atomic_wrapper {
     std::atomic<size_t> value;
@@ -686,8 +701,10 @@ auto gen_constraints_general(size_t nlevels, size_t norb, size_t ns_othr,
   // Compute histogram
   const auto ntrip_full = constraint_sizes.size(); 
   std::vector<atomic_wrapper> constraint_work(ntrip_full, 0ul);
+  int world_rank = comm_rank(MPI_COMM_WORLD);
   #pragma omp parallel for schedule(dynamic)
-  for(auto i_trip = 0; i_trip < ntrip_full; ++i_trip) {
+  for(auto i_trip = 0ul; i_trip < ntrip_full; ++i_trip) {
+    if(!world_rank and !(i_trip%1000)) printf("cgen %lu / %lu\n", i_trip, ntrip_full);
     auto& [constraint, __nw] = constraint_sizes[i_trip];
     auto& c_nw = constraint_work[i_trip];
     size_t nw = 0;
diff --git a/include/macis/asci/pt2.hpp b/include/macis/asci/pt2.hpp
@@ -8,13 +8,16 @@
 
 #pragma once
 #include <deque>
+#include <fstream>
+#include <sstream>
 #include <macis/asci/determinant_contributions.hpp>
 #include <macis/asci/determinant_sort.hpp>
 
 namespace macis {
 
 template <size_t N>
-double asci_pt2_constraint(wavefunction_iterator_t<N> cdets_begin,
+double asci_pt2_constraint(ASCISettings asci_settings,
+                           wavefunction_iterator_t<N> cdets_begin,
                            wavefunction_iterator_t<N> cdets_end,
                            const double E_ASCI, const std::vector<double>& C,
                            size_t norb, const double* T_pq, const double* G_red,
@@ -39,26 +42,34 @@ double asci_pt2_constraint(wavefunction_iterator_t<N> cdets_begin,
                         : spdlog::stdout_color_mt("asci_pt2");
 
   const size_t ncdets = std::distance(cdets_begin, cdets_end);
+  logger->info("[ASCI PT2 Settings]");
+  logger->info("  * NDETS              = {}", ncdets);
+  logger->info("  * PT2_TOL            = {}", asci_settings.pt2_tol);
+  logger->info("  * PT2_RESERVE_COUNT  = {}", asci_settings.pt2_reserve_count);
+  logger->info("  * PT2_CONSTRAINT_LVL = {}", asci_settings.pt2_constraint_level);
+  logger->info("  * PT2_PRUNE          = {}", asci_settings.pt2_prune);
+  logger->info("  * PT2_PRECMP_EPS     = {}", asci_settings.pt2_precompute_eps);
+  logger->info("");
 
   // For each unique alpha, create a list of beta string and store metadata
   struct beta_coeff_data {
     spin_wfn_type beta_string;
-    std::vector<uint32_t> occ_beta;
-    std::vector<uint32_t> vir_beta;
+    std::vector<uint8_t> occ_beta;
+    std::vector<uint8_t> vir_beta;
     std::vector<double> orb_ens_alpha;
     std::vector<double> orb_ens_beta;
     double coeff;
     double h_diag;
 
     size_t mem() const {
       return sizeof(spin_wfn_type) + 
-             (occ_beta.capacity() + vir_beta.capacity()) * sizeof(uint32_t) +
+             (occ_beta.capacity() + vir_beta.capacity()) * sizeof(uint8_t) +
              (2 + orb_ens_alpha.capacity() + orb_ens_beta.capacity()) * sizeof(double); 
     }
 
     beta_coeff_data(double c, size_t norb,
                     const std::vector<uint32_t>& occ_alpha, wfn_t<N> w,
-                    const HamiltonianGenerator<wfn_t<N>>& ham_gen) {
+                    const HamiltonianGenerator<wfn_t<N>>& ham_gen, bool pce) {
       coeff = c;
 
       beta_string = wfn_traits::beta_string(w);
@@ -67,16 +78,24 @@ double asci_pt2_constraint(wavefunction_iterator_t<N> cdets_begin,
       h_diag = ham_gen.matrix_element(w, w);
 
       // Compute occ/vir for beta string
-      spin_wfn_traits::state_to_occ_vir(norb, beta_string, occ_beta, vir_beta);
+      std::vector<uint32_t> o_32, v_32;
+      spin_wfn_traits::state_to_occ_vir(norb, beta_string, o_32, v_32);
+      occ_beta.resize(o_32.size()); 
+      std::copy(o_32.begin(), o_32.end(), occ_beta.begin());
+      vir_beta.resize(v_32.size());
+      std::copy(v_32.begin(), v_32.end(), vir_beta.begin());
 
       // Precompute orbital energies
-      //orb_ens_alpha = ham_gen.single_orbital_ens(norb, occ_alpha, occ_beta);
-      //orb_ens_beta = ham_gen.single_orbital_ens(norb, occ_beta, occ_alpha);
+      if(pce) {
+        orb_ens_alpha = ham_gen.single_orbital_ens(norb, occ_alpha, o_32);
+        orb_ens_beta = ham_gen.single_orbital_ens(norb, o_32, occ_alpha);
+      }
     }
   };
 
   auto uniq_alpha = get_unique_alpha(cdets_begin, cdets_end);
   const size_t nuniq_alpha = uniq_alpha.size();
+  logger->info("  * NUNIQ_ALPHA = {}", nuniq_alpha);
 
   using unique_alpha_data = std::vector<beta_coeff_data>;
   std::vector<unique_alpha_data> uad(nuniq_alpha);
@@ -89,7 +108,7 @@ double asci_pt2_constraint(wavefunction_iterator_t<N> cdets_begin,
     uad[i].reserve(nbeta);
     for(auto j = 0; j < nbeta; ++j, ++iw) {
       const auto& w = *(cdets_begin + iw);
-      uad[i].emplace_back(C[iw], norb, occ_alpha, w, ham_gen);
+      uad[i].emplace_back(C[iw], norb, occ_alpha, w, ham_gen,asci_settings.pt2_precompute_eps);
     }
   }
 
@@ -105,7 +124,7 @@ double asci_pt2_constraint(wavefunction_iterator_t<N> cdets_begin,
       }
     }
     printf("MEM REQ ALPH = %.2e\n", mem_alpha / gib);
-    printf("MEM REQ CONT = %.2e\n", 70000000 * sizeof(asci_contrib<wfn_t<N>>)/ 1024./1024./1024);
+    printf("MEM REQ CONT = %.2e\n", asci_settings.pt2_reserve_count * sizeof(asci_contrib<wfn_t<N>>)/ gib);
   }
   MPI_Barrier(comm);
 
@@ -125,11 +144,29 @@ double asci_pt2_constraint(wavefunction_iterator_t<N> cdets_begin,
   // auto constraints = dist_constraint_general<wfn_t<N>>(
   //     5, norb, n_sing_beta, n_doub_beta, uniq_alpha, comm);
   auto constraints = gen_constraints_general<wfn_t<N>>(
-      10, norb, n_sing_beta, n_doub_beta, uniq_alpha,
-      world_size * omp_get_max_threads());
+      asci_settings.pt2_constraint_level, norb, n_sing_beta, 
+      n_doub_beta, uniq_alpha, world_size * omp_get_max_threads(), 0);
   auto gen_c_en = clock_type::now();
   duration_type gen_c_dur = gen_c_en - gen_c_st;
   logger->info("  * GEN_DUR = {:.2e} ms", gen_c_dur.count());
+  //if(!world_rank) {
+  //  std::ofstream c_file("constraint_work.txt");
+  //  std::stringstream ss; 
+  //  for(auto [c,s] : constraints) {
+  //    ss << c.C() << " " << s << std::endl;
+  //  }
+  //  auto str = ss.str();
+  //  c_file.write(str.c_str(), str.size());
+  //}
+  //if(!world_rank) {
+  //  std::ofstream c_file("unique_alpha.txt");
+  //  std::stringstream ss; 
+  //  for(size_t i = 0; i < nuniq_alpha; ++i) {
+  //    ss << uniq_alpha[i].first << " " << uniq_alpha[i].second << std::endl;
+  //  }
+  //  auto str = ss.str();
+  //  c_file.write(str.c_str(), str.size());
+  //}
 
   double EPT2 = 0.0;
   size_t NPT2 = 0;
@@ -138,14 +175,14 @@ double asci_pt2_constraint(wavefunction_iterator_t<N> cdets_begin,
 
   // Global atomic task-id counter
   global_atomic<size_t> nxtval(comm);
-  const double h_el_tol = 1e-6;
+  const double h_el_tol = asci_settings.pt2_tol;
 
   auto pt2_st = clock_type::now();
 #pragma omp parallel reduction(+ : EPT2) reduction(+ : NPT2)
   {
     // Process ASCI pair contributions for each constraint
     asci_contrib_container<wfn_t<N>> asci_pairs;
-    asci_pairs.reserve(70000000ul);
+    //asci_pairs.reserve(asci_settings.pt2_reserve_count);
     size_t ic = 0;
     while(ic < ncon_total) {
       // Atomically get the next task ID and increment for other
@@ -173,12 +210,22 @@ double asci_pt2_constraint(wavefunction_iterator_t<N> cdets_begin,
             const auto c = C[iw];
             const auto& beta_det = bcd[j_beta].beta_string;
             const auto h_diag = bcd[j_beta].h_diag;
-            const auto& occ_beta = bcd[j_beta].occ_beta;
-            const auto& vir_beta = bcd[j_beta].vir_beta;
-            //const auto& orb_ens_alpha = bcd[j_beta].orb_ens_alpha;
-            //const auto& orb_ens_beta = bcd[j_beta].orb_ens_beta;
-            auto orb_ens_alpha = ham_gen.single_orbital_ens(norb, occ_alpha, occ_beta);
-            auto orb_ens_beta = ham_gen.single_orbital_ens(norb, occ_beta, occ_alpha);
+
+            // TODO: These copies are slow
+            const auto& occ_beta_8 = bcd[j_beta].occ_beta;
+            const auto& vir_beta_8 = bcd[j_beta].vir_beta;
+            std::vector<uint32_t> occ_beta(occ_beta_8.size()), vir_beta(vir_beta_8.size());
+            std::copy(occ_beta_8.begin(), occ_beta_8.end(), occ_beta.begin());
+            std::copy(vir_beta_8.begin(), vir_beta_8.end(), vir_beta.begin());
+
+            std::vector<double> orb_ens_alpha, orb_ens_beta;
+            if(asci_settings.pt2_precompute_eps) {
+              orb_ens_alpha = bcd[j_beta].orb_ens_alpha;
+              orb_ens_beta  = bcd[j_beta].orb_ens_beta;
+            } else {
+              orb_ens_alpha = ham_gen.single_orbital_ens(norb, occ_alpha, occ_beta);
+              orb_ens_beta  = ham_gen.single_orbital_ens(norb, occ_beta, occ_alpha);
+            }
 
             // AA excitations
             generate_constraint_singles_contributions_ss(
@@ -215,10 +262,12 @@ double asci_pt2_constraint(wavefunction_iterator_t<N> cdets_begin,
                   {w, std::numeric_limits<double>::infinity(), 1.0});
             }
           }
-          if(asci_pairs.size() > 70000000 and asci_pairs.size() != old_pair_size) {
+          if(asci_settings.pt2_prune and asci_pairs.size() > asci_settings.pt2_reserve_count and asci_pairs.size() != old_pair_size) {
           // Cleanup
-          auto uit = sort_and_accumulate_asci_pairs(asci_pairs.begin(),
+          auto uit = stable_sort_and_accumulate_asci_pairs(asci_pairs.begin(),
                                                     asci_pairs.end());
+          asci_pairs.erase(uit, asci_pairs.end());
+          uit = std::stable_partition(asci_pairs.begin(), asci_pairs.end(), [&](const auto& p){ return std::abs(p.pt2()) > h_el_tol; });
           asci_pairs.erase(uit, asci_pairs.end());
             printf("[rank %4d tid:%4d] IC = %lu / %lu IA = %lu / %lu SZ = %lu\n", world_rank,
                    omp_get_thread_num(), ic, ncon_total, i_alpha,
@@ -235,8 +284,7 @@ double asci_pt2_constraint(wavefunction_iterator_t<N> cdets_begin,
                                                     asci_pairs.end());
           for(auto it = asci_pairs.begin(); it != uit; ++it) {
             if(!std::isinf(it->c_times_matel)) {
-              EPT2_local +=
-                  (it->c_times_matel * it->c_times_matel) / it->h_diag;
+              EPT2_local += it->pt2();
               NPT2_local++;
             }
           }
diff --git a/tests/standalone_driver.cxx b/tests/standalone_driver.cxx
