Azam/feature/hardcoded biorthogonal coefficients

.github/workflows/cmake.yml

@@ -125,7 +125,7 @@ jobs:
       working-directory: ${{github.workspace}}/build
       shell: bash
       # Execute the build.  You can specify a specific target with "--target <NAME>"
-      run: ccache -p && ccache -z && cmake --build . && ccache -s
+      run: ccache -p && ccache -z && cmake --build . -j 2 && ccache -s
 
     - name: Test
       if: ${{ !matrix.valgrind }}

CMakeLists.txt

@@ -367,6 +367,9 @@ set(SeQuant_src
         SeQuant/domain/mbpt/antisymmetrizer.hpp
         SeQuant/domain/mbpt/biorthogonalization.cpp
         SeQuant/domain/mbpt/biorthogonalization.hpp
+        SeQuant/domain/mbpt/biorthogonalization_hardcoded.ipp
+        SeQuant/domain/mbpt/biorthogonalization_hardcoded.cpp
+        SeQuant/domain/mbpt/biorthogonalization_hardcoded.hpp
         SeQuant/domain/mbpt/context.cpp
         SeQuant/domain/mbpt/context.hpp
         SeQuant/domain/mbpt/convention.cpp

SeQuant/core/eval/result.hpp

@@ -8,6 +8,8 @@
 #include <SeQuant/core/index.hpp>
 #include <SeQuant/core/logger.hpp>
 #include <SeQuant/core/utility/macros.hpp>
+#include <SeQuant/domain/mbpt/biorthogonalization.hpp>
+#include <SeQuant/domain/mbpt/biorthogonalization_hardcoded.hpp>
 
 #include <TiledArray/einsum/tiledarray.h>
 #include <btas/btas.h>
@@ -361,11 +363,23 @@ auto biorthogonal_nns_project_ta(TA::DistArray<Args...> const& arr,
 
   using numeric_type = typename TA::DistArray<Args...>::numeric_type;
 
-  size_t factorial_ket = 1;
-  for (size_t i = 2; i <= ket_rank; ++i) {
-    factorial_ket *= i;
+  std::vector<numeric_type> nns_p_coeffs;
+
+  constexpr std::size_t max_rank_hardcoded_nns_projector = 5;
+  if (ket_rank >= 3) {
+    if (ket_rank <= max_rank_hardcoded_nns_projector) {
+      nns_p_coeffs = hardcoded_nns_projector<numeric_type>(ket_rank);
+    } else {
+      Eigen::MatrixXd nns_matrix = compute_nns_p_matrix(ket_rank);
+      size_t num_perms = nns_matrix.rows();
+
+      nns_p_coeffs.reserve(num_perms);
+      for (size_t i = 0; i < num_perms; ++i) {
+        nns_p_coeffs.push_back(
+            static_cast<numeric_type>(nns_matrix(num_perms - 1, i)));
+      }
+    }
   }
-  numeric_type norm_factor = numeric_type(1) / numeric_type(factorial_ket);
 
   TA::DistArray<Args...> result;
 
@@ -375,43 +389,28 @@ auto biorthogonal_nns_project_ta(TA::DistArray<Args...> const& arr,
 
   const auto lannot = ords_to_annot(perm);
 
-  auto process_permutations = [&lannot](const TA::DistArray<Args...>& input_arr,
-                                        size_t range_rank, perm_t range_perm,
-                                        const std::string& other_annot,
-                                        bool is_bra) -> TA::DistArray<Args...> {
-    if (range_rank <= 1) return input_arr;
-    TA::DistArray<Args...> result;
+  if (ket_rank > 2 && !nns_p_coeffs.empty()) {
+    const auto bra_annot = bra_rank == 0 ? "" : ords_to_annot(bra_perm);
 
-    auto callback = [&]([[maybe_unused]] int parity) {
-      const auto range_annot = ords_to_annot(range_perm);
-      const auto annot = other_annot.empty()
-                             ? range_annot
-                             : (is_bra ? range_annot + "," + other_annot
-                                       : other_annot + "," + range_annot);
+    size_t num_perms = nns_p_coeffs.size();
+    for (size_t perm_rank = 0; perm_rank < num_perms; ++perm_rank) {
+      perm_t permuted_ket =
+          compute_permuted_indices(ket_perm, perm_rank, ket_rank);
+
+      numeric_type coeff = nns_p_coeffs[perm_rank];
+
+      const auto ket_annot = ords_to_annot(permuted_ket);
+      const auto annot =
+          bra_annot.empty() ? ket_annot : bra_annot + "," + ket_annot;
 
-      // ignore parity, all permutations get same coefficient
-      numeric_type p_ = 1;
       if (result.is_initialized()) {
-        result(lannot) += p_ * input_arr(annot);
+        result(lannot) += coeff * arr(annot);
       } else {
-        result(lannot) = p_ * input_arr(annot);
+        result(lannot) = coeff * arr(annot);
       }
-    };
-    antisymmetric_permutation(ParticleRange{range_perm.begin(), range_rank},
-                              callback);
-    return result;
-  };
-
-  // identity term with coefficient +1
-  result(lannot) = arr(lannot);
-
-  // process only ket permutations with coefficient norm_factor
-  if (ket_rank > 1) {
-    const auto bra_annot = bra_rank == 0 ? "" : ords_to_annot(bra_perm);
-    auto ket_result =
-        process_permutations(arr, ket_rank, ket_perm, bra_annot, false);
-
-    result(lannot) -= norm_factor * ket_result(lannot);
+    }
+  } else {
+    result(lannot) = arr(lannot);
   }
 
   TA::DistArray<Args...>::wait_for_lazy_cleanup(result.world());
@@ -428,7 +427,6 @@ auto biorthogonal_nns_project_ta(TA::DistArray<Args...> const& arr,
 template <typename... Args>
 auto biorthogonal_nns_project_btas(btas::Tensor<Args...> const& arr,
                                    size_t bra_rank) {
-  using ranges::views::concat;
   using ranges::views::iota;
   size_t const rank = arr.rank();
   SEQUANT_ASSERT(bra_rank <= rank);
@@ -439,53 +437,57 @@ auto biorthogonal_nns_project_btas(btas::Tensor<Args...> const& arr,
   }
 
   using numeric_type = typename btas::Tensor<Args...>::numeric_type;
+  std::vector<numeric_type> nns_p_coeffs;
+  constexpr std::size_t max_rank_hardcoded_nns_projector = 5;
 
-  size_t factorial_ket = 1;
-  for (size_t i = 2; i <= ket_rank; ++i) {
-    factorial_ket *= i;
+  if (ket_rank >= 3) {
+    if (ket_rank <= max_rank_hardcoded_nns_projector) {
+      nns_p_coeffs = hardcoded_nns_projector<numeric_type>(ket_rank);
+    } else {
+      Eigen::MatrixXd nns_matrix = compute_nns_p_matrix(ket_rank);
+      size_t num_perms = nns_matrix.rows();
+
+      nns_p_coeffs.reserve(num_perms);
+      for (size_t i = 0; i < num_perms; ++i) {
+        nns_p_coeffs.push_back(
+            static_cast<numeric_type>(nns_matrix(num_perms - 1, i)));
+      }
+    }
   }
-  numeric_type norm_factor = numeric_type(1) / numeric_type(factorial_ket);
 
+  btas::Tensor<Args...> result;
+
+  perm_t perm = iota(size_t{0}, rank) | ranges::to<perm_t>;
   perm_t bra_perm = iota(size_t{0}, bra_rank) | ranges::to<perm_t>;
   perm_t ket_perm = iota(bra_rank, rank) | ranges::to<perm_t>;
-  const auto lannot = iota(size_t{0}, rank) | ranges::to<perm_t>;
 
-  auto process_permutations = [&lannot](const btas::Tensor<Args...>& input_arr,
-                                        size_t range_rank, perm_t range_perm,
-                                        const perm_t& other_perm, bool is_bra) {
-    if (range_rank <= 1) return input_arr;
-    btas::Tensor<Args...> result{input_arr.range()};
-    result.fill(0);
+  if (ket_rank > 2 && !nns_p_coeffs.empty()) {
+    bool result_initialized = false;
 
-    auto callback = [&]([[maybe_unused]] int parity) {
-      const auto annot =
-          is_bra ? concat(range_perm, other_perm) | ranges::to<perm_t>()
-                 : concat(other_perm, range_perm) | ranges::to<perm_t>();
+    size_t num_perms = nns_p_coeffs.size();
+    for (size_t perm_rank = 0; perm_rank < num_perms; ++perm_rank) {
+      perm_t permuted_ket =
+          compute_permuted_indices(ket_perm, perm_rank, ket_rank);
 
-      // ignore parity, all permutations get same coefficient
-      numeric_type p_ = 1;
-      btas::Tensor<Args...> temp;
-      btas::permute(input_arr, lannot, temp, annot);
-      btas::scal(p_, temp);
-      result += temp;
-    };
+      numeric_type coeff = nns_p_coeffs[perm_rank];
 
-    antisymmetric_permutation(ParticleRange{range_perm.begin(), range_rank},
-                              callback);
-    return result;
-  };
+      perm_t annot = bra_perm;
+      annot.insert(annot.end(), permuted_ket.begin(), permuted_ket.end());
 
-  // identity term with coefficient +1
-  auto result = arr;
+      btas::Tensor<Args...> temp;
+      btas::permute(arr, annot, temp, perm);
+      btas::scal(coeff, temp);
 
-  // process only ket permutations with coefficient norm_factor
-  if (ket_rank > 1) {
-    const auto bra_annot = bra_rank == 0 ? perm_t{} : bra_perm;
-    auto ket_result =
-        process_permutations(arr, ket_rank, ket_perm, bra_annot, false);
+      if (result_initialized) {
+        result += temp;
+      } else {
+        result = temp;
+        result_initialized = true;
+      }
+    }
 
-    btas::scal(norm_factor, ket_result);
-    result -= ket_result;
+  } else {
+    result = arr;
   }
 
   return result;
@@ -903,8 +905,14 @@ class ResultTensorTA final : public Result {
 
   [[nodiscard]] ResultPtr biorthogonal_nns_project(
       size_t bra_rank) const override {
-    return eval_result<this_type>(
-        biorthogonal_nns_project_ta(get<ArrayT>(), bra_rank));
+    if constexpr (std::is_integral_v<numeric_type>) {
+      SEQUANT_ABORT(
+          "biorthogonal_nns_project is not supported for integral numeric "
+          "types");
+    } else {
+      return eval_result<this_type>(
+          biorthogonal_nns_project_ta(get<ArrayT>(), bra_rank));
+    }
   }
 
  private:
@@ -1166,8 +1174,14 @@ class ResultTensorBTAS final : public Result {
 
   [[nodiscard]] ResultPtr biorthogonal_nns_project(
       [[maybe_unused]] size_t bra_rank) const override {
-    return eval_result<ResultTensorBTAS<T>>(
-        biorthogonal_nns_project_btas(get<T>(), bra_rank));
+    if constexpr (std::is_integral_v<numeric_type>) {
+      SEQUANT_ABORT(
+          "biorthogonal_nns_project is not supported for integral numeric "
+          "types");
+    } else {
+      return eval_result<ResultTensorBTAS<T>>(
+          biorthogonal_nns_project_btas(get<T>(), bra_rank));
+    }
   }
 
  private:

SeQuant/domain/mbpt/biorthogonalization.cpp

@@ -1,4 +1,5 @@
 #include <SeQuant/domain/mbpt/biorthogonalization.hpp>
+#include <SeQuant/domain/mbpt/biorthogonalization_hardcoded.hpp>
 
 #include <SeQuant/core/container.hpp>
 #include <SeQuant/core/expr.hpp>
@@ -9,6 +10,7 @@
 #include <SeQuant/core/utility/macros.hpp>
 #include <SeQuant/core/utility/permutation.hpp>
 
+#include <Eigen/Core>
 #include <Eigen/Eigenvalues>
 
 #include <libperm/Permutation.hpp>
@@ -129,6 +131,28 @@ Eigen::MatrixXd compute_biorth_coeffs(std::size_t n_particles,
   return pinv;
 }
 
+Eigen::MatrixXd compute_nns_p_matrix(std::size_t n_particles,
+                                     double threshold) {
+  auto perm_ovlp_mat = permutational_overlap_matrix(n_particles);
+  auto normalized_pinv = compute_biorth_coeffs(n_particles, threshold);
+
+  auto nns = perm_ovlp_mat * normalized_pinv;
+
+  return nns;
+}
+
+container::svector<size_t> compute_permuted_indices(
+    const container::svector<size_t>& indices, size_t perm_rank,
+    size_t n_particles) {
+  perm::Permutation perm_obj = perm::unrank(perm_rank, n_particles);
+
+  container::svector<size_t> permuted_indices(n_particles);
+  for (size_t i = 0; i < n_particles; ++i) {
+    permuted_indices[i] = indices[perm_obj[i]];
+  }
+  return permuted_indices;
+}
+
 void sort_pairings(ParticlePairings& pairing) {
   std::stable_sort(pairing.begin(), pairing.end(),
                    compare_first_less<IndexPair>{});
@@ -309,7 +333,7 @@ void biorthogonal_transform(container::svector<ResultExpr>& result_exprs,
   // like R^{IJ}_{AB} and the index pairing of the result is what determines
   // the required symmetrization. Hence, the symmetrization operator must not
   // be changed when transforming from one representation into the other.
-  assert(std::all_of(
+  SEQUANT_ASSERT(std::all_of(
       result_exprs.begin(), result_exprs.end(), [](const ResultExpr& res) {
         bool found = false;
         res.expression()->visit(
@@ -336,19 +360,27 @@ void biorthogonal_transform(container::svector<ResultExpr>& result_exprs,
                ranges::to<container::svector<std::size_t>>();
 
   const std::size_t n_particles = externals.front().size();
-
-  Eigen::MatrixXd coefficients = compute_biorth_coeffs(n_particles, threshold);
-
   auto num_perms = factorial(n_particles);
-  SEQUANT_ASSERT(num_perms == coefficients.rows());
-  SEQUANT_ASSERT(num_perms == coefficients.cols());
 
   auto original_exprs = result_exprs |
                         ranges::views::transform([](const ResultExpr& res) {
                           return res.expression();
                         }) |
                         ranges::to<container::svector<ExprPtr>>();
 
+  Eigen::Matrix<sequant::rational, Eigen::Dynamic, Eigen::Dynamic>
+      hardcoded_coefficients;
+  Eigen::MatrixXd computed_coefficients;
+  constexpr std::size_t max_rank_hardcoded_biorth_coeffs = 6;
+
+  if (n_particles <= max_rank_hardcoded_biorth_coeffs) {
+    hardcoded_coefficients = hardcoded_biorth_coeffs_matrix(n_particles);
+  } else {
+    computed_coefficients = compute_biorth_coeffs(n_particles, threshold);
+    SEQUANT_ASSERT(num_perms == computed_coefficients.rows());
+    SEQUANT_ASSERT(num_perms == computed_coefficients.cols());
+  }
+
   for (std::size_t i = 0; i < result_exprs.size(); ++i) {
     result_exprs.at(i).expression() = ex<Constant>(0);
     perm::Permutation reference = perm::unrank(ranks.at(i), n_particles);
@@ -358,9 +390,14 @@ void biorthogonal_transform(container::svector<ResultExpr>& result_exprs,
       perm::Permutation perm = perm::unrank(rank, n_particles);
       perm->postMultiply(reference);
 
+      sequant::rational coeff =
+          (n_particles <= max_rank_hardcoded_biorth_coeffs)
+              ? hardcoded_coefficients(ranks.at(i), rank)
+              : to_rational(computed_coefficients(ranks.at(i), rank),
+                            threshold);
+
       result_exprs.at(i).expression() +=
-          ex<Constant>(
-              to_rational(coefficients(ranks.at(i), rank), threshold)) *
+          ex<Constant>(coeff) *
           create_expr_for(externals.at(i), perm, externals, original_exprs);
     }