OMP-ed denominator related lines in spin-adapted MRDSRG (#426)

* openmp denominator related lines

* also for t1

* another version

Author: lcyyork

forte/mrdsrg-spin-adapted/sa_dsrgpt.cc

@@ -120,26 +120,12 @@ void SA_DSRGPT::compute_t2_full() {
     std::vector<std::string> T2blocks(T2_.block_labels());
     if (ccvv_source_ == "ZERO") {
         T2blocks.erase(std::remove(T2blocks.begin(), T2blocks.end(), "ccvv"), T2blocks.end());
-        T2_.block("ccvv").iterate([&](const std::vector<size_t>& i, double& value) {
-            size_t i0 = core_mos_[i[0]];
-            size_t i1 = core_mos_[i[1]];
-            size_t i2 = virt_mos_[i[2]];
-            size_t i3 = virt_mos_[i[3]];
-            value /= Fdiag_[i0] + Fdiag_[i1] - Fdiag_[i2] - Fdiag_[i3];
-        });
+        apply_denominator(T2_, {"ccvv"}, [&](double d) { return 1.0 / d; });
     }
 
     // build T2
-    for (const std::string& block : T2blocks) {
-        T2_.block(block).iterate([&](const std::vector<size_t>& i, double& value) {
-            size_t i0 = label_to_spacemo_[block[0]][i[0]];
-            size_t i1 = label_to_spacemo_[block[1]][i[1]];
-            size_t i2 = label_to_spacemo_[block[2]][i[2]];
-            size_t i3 = label_to_spacemo_[block[3]][i[3]];
-            double denom = Fdiag_[i0] + Fdiag_[i1] - Fdiag_[i2] - Fdiag_[i3];
-            value *= dsrg_source_->compute_renormalized_denominator(denom);
-        });
-    }
+    apply_denominator(T2_, T2blocks,
+                      [&](double v) { return dsrg_source_->compute_renormalized_denominator(v); });
 
     // transform back to non-canonical basis
     if (!semi_canonical_) {
@@ -183,21 +169,12 @@ void SA_DSRGPT::compute_t1() {
     std::vector<std::string> T1blocks(T1_.block_labels());
     if (ccvv_source_ == "ZERO") {
         T1blocks.erase(std::remove(T1blocks.begin(), T1blocks.end(), "cv"), T1blocks.end());
-        T1_.block("cv").iterate([&](const std::vector<size_t>& i, double& value) {
-            size_t i0 = core_mos_[i[0]];
-            size_t i1 = virt_mos_[i[1]];
-            value /= Fdiag_[i0] - Fdiag_[i1];
-        });
+        apply_denominator(T1_, {"cv"}, [](double d) { return 1.0 / d; });
     }
 
     // build T1
-    for (const std::string& block : T1blocks) {
-        T1_.block(block).iterate([&](const std::vector<size_t>& i, double& value) {
-            size_t i0 = label_to_spacemo_[block[0]][i[0]];
-            size_t i1 = label_to_spacemo_[block[1]][i[1]];
-            value *= dsrg_source_->compute_renormalized_denominator(Fdiag_[i0] - Fdiag_[i1]);
-        });
-    }
+    apply_denominator(T1_, T1blocks,
+                      [&](double d) { return dsrg_source_->compute_renormalized_denominator(d); });
 
     // transform back to non-canonical basis
     if (!semi_canonical_) {
@@ -230,27 +207,11 @@ void SA_DSRGPT::renormalize_integrals(bool add) {
     }
 
     if (add) {
-        for (const std::string& block : Vblocks) {
-            V_.block(block).iterate([&](const std::vector<size_t>& i, double& value) {
-                size_t i0 = label_to_spacemo_[block[0]][i[0]];
-                size_t i1 = label_to_spacemo_[block[1]][i[1]];
-                size_t i2 = label_to_spacemo_[block[2]][i[2]];
-                size_t i3 = label_to_spacemo_[block[3]][i[3]];
-                double denom = Fdiag_[i0] + Fdiag_[i1] - Fdiag_[i2] - Fdiag_[i3];
-                value *= 1.0 + dsrg_source_->compute_renormalized(denom);
-            });
-        }
+        apply_denominator(V_, Vblocks,
+                          [&](double d) { return 1.0 + dsrg_source_->compute_renormalized(d); });
     } else {
-        for (const std::string& block : Vblocks) {
-            V_.block(block).iterate([&](const std::vector<size_t>& i, double& value) {
-                size_t i0 = label_to_spacemo_[block[0]][i[0]];
-                size_t i1 = label_to_spacemo_[block[1]][i[1]];
-                size_t i2 = label_to_spacemo_[block[2]][i[2]];
-                size_t i3 = label_to_spacemo_[block[3]][i[3]];
-                double denom = Fdiag_[i0] + Fdiag_[i1] - Fdiag_[i2] - Fdiag_[i3];
-                value *= dsrg_source_->compute_renormalized(denom);
-            });
-        }
+        apply_denominator(V_, Vblocks,
+                          [&](double d) { return dsrg_source_->compute_renormalized(d); });
     }
 
     // transform back if necessary
@@ -283,10 +244,8 @@ void SA_DSRGPT::renormalize_integrals(bool add) {
     }
 
     // scale by exp(-s * D^2)
-    temp.iterate([&](const std::vector<size_t>& i, const std::vector<SpinType>&, double& value) {
-        double denom = Fdiag_[i[0]] - Fdiag_[i[1]];
-        value *= dsrg_source_->compute_renormalized(denom);
-    });
+    apply_denominator(temp, temp.block_labels(),
+                      [&](double d) { return dsrg_source_->compute_renormalized(d); });
 
     // transform back if necessary
     if (!semi_canonical_) {

forte/mrdsrg-spin-adapted/sa_mrdsrg_amps.cc

@@ -95,26 +95,11 @@ void SA_MRDSRG::guess_t2_impl(BlockedTensor& T2) {
     std::vector<std::string> T2blocks(T2.block_labels());
     if (ccvv_source_ == "ZERO") {
         T2blocks.erase(std::remove(T2blocks.begin(), T2blocks.end(), "ccvv"), T2blocks.end());
-        T2.block("ccvv").iterate([&](const std::vector<size_t>& i, double& value) {
-            size_t i0 = core_mos_[i[0]];
-            size_t i1 = core_mos_[i[1]];
-            size_t i2 = virt_mos_[i[2]];
-            size_t i3 = virt_mos_[i[3]];
-
-            value /= Fdiag_[i0] + Fdiag_[i1] - Fdiag_[i2] - Fdiag_[i3];
-        });
+        apply_denominator(T2, {"ccvv"}, [](double d) { return 1.0 / d; });
     }
 
-    for (const std::string& block : T2blocks) {
-        T2.block(block).iterate([&](const std::vector<size_t>& i, double& value) {
-            size_t i0 = label_to_spacemo_[block[0]][i[0]];
-            size_t i1 = label_to_spacemo_[block[1]][i[1]];
-            size_t i2 = label_to_spacemo_[block[2]][i[2]];
-            size_t i3 = label_to_spacemo_[block[3]][i[3]];
-            double denom = Fdiag_[i0] + Fdiag_[i1] - Fdiag_[i2] - Fdiag_[i3];
-            value *= dsrg_source_->compute_renormalized_denominator(denom);
-        });
-    }
+    apply_denominator(T2, T2blocks,
+                      [&](double d) { return dsrg_source_->compute_renormalized_denominator(d); });
 
     // transform back to non-canonical basis
     if (!semi_canonical_) {
@@ -161,22 +146,12 @@ void SA_MRDSRG::guess_t1(BlockedTensor& F, BlockedTensor& T2, BlockedTensor& T1)
             std::vector<std::string> T1blocks(T1.block_labels());
             if (ccvv_source_ == "ZERO") {
                 T1blocks.erase(std::remove(T1blocks.begin(), T1blocks.end(), "cv"), T1blocks.end());
-                T1.block("cv").iterate([&](const std::vector<size_t>& i, double& value) {
-                    size_t i0 = core_mos_[i[0]];
-                    size_t i1 = virt_mos_[i[1]];
-
-                    value /= Fdiag_[i0] - Fdiag_[i1];
-                });
+                apply_denominator(T1, {"cv"}, [](double d) { return 1.0 / d; });
             }
 
-            for (const std::string& block : T1blocks) {
-                T1.block(block).iterate([&](const std::vector<size_t>& i, double& value) {
-                    size_t i0 = label_to_spacemo_[block[0]][i[0]];
-                    size_t i1 = label_to_spacemo_[block[1]][i[1]];
-                    value *=
-                        dsrg_source_->compute_renormalized_denominator(Fdiag_[i0] - Fdiag_[i1]);
-                });
-            }
+            apply_denominator(T1, T1blocks, [&](double d) {
+                return dsrg_source_->compute_renormalized_denominator(d);
+            });
 
             // transform back to non-canonical basis
             if (!semi_canonical_) {
@@ -230,25 +205,11 @@ void SA_MRDSRG::update_t2() {
     timer t2("scale Hbar2 by renormalized denominator");
     // scale Hbar2 by renormalized denominator
     if (ccvv_source_ == "ZERO") {
-        DT2_.block("ccvv").iterate([&](const std::vector<size_t>& i, double& value) {
-            size_t i0 = core_mos_[i[0]];
-            size_t i1 = core_mos_[i[1]];
-            size_t i2 = virt_mos_[i[2]];
-            size_t i3 = virt_mos_[i[3]];
-            value /= Fdiag_[i0] + Fdiag_[i1] - Fdiag_[i2] - Fdiag_[i3];
-        });
+        apply_denominator(DT2_, {"ccvv"}, [](double d) { return 1.0 / d; });
     }
 
-    for (const std::string& block : T2blocks) {
-        DT2_.block(block).iterate([&](const std::vector<size_t>& i, double& value) {
-            size_t i0 = label_to_spacemo_[block[0]][i[0]];
-            size_t i1 = label_to_spacemo_[block[1]][i[1]];
-            size_t i2 = label_to_spacemo_[block[2]][i[2]];
-            size_t i3 = label_to_spacemo_[block[3]][i[3]];
-            double denom = Fdiag_[i0] + Fdiag_[i1] - Fdiag_[i2] - Fdiag_[i3];
-            value *= dsrg_source_->compute_renormalized_denominator(denom);
-        });
-    }
+    apply_denominator(DT2_, T2blocks,
+                      [&](double d) { return dsrg_source_->compute_renormalized_denominator(d); });
     t2.stop();
 
     // Step 2: work on T2 where Hbar2 is treated as intermediate
@@ -269,16 +230,8 @@ void SA_MRDSRG::update_t2() {
 
     timer t6("scale T2 by delta exponential");
     // scale T2 by delta exponential
-    for (const std::string& block : T2blocks) {
-        T2_.block(block).iterate([&](const std::vector<size_t>& i, double& value) {
-            size_t i0 = label_to_spacemo_[block[0]][i[0]];
-            size_t i1 = label_to_spacemo_[block[1]][i[1]];
-            size_t i2 = label_to_spacemo_[block[2]][i[2]];
-            size_t i3 = label_to_spacemo_[block[3]][i[3]];
-            double denom = Fdiag_[i0] + Fdiag_[i1] - Fdiag_[i2] - Fdiag_[i3];
-            value *= dsrg_source_->compute_renormalized(denom);
-        });
-    }
+    apply_denominator(T2_, T2blocks,
+                      [&](double d) { return dsrg_source_->compute_renormalized(d); });
     if (ccvv_source_ == "ZERO") {
         T2_.block("ccvv").zero();
     }
@@ -349,21 +302,11 @@ void SA_MRDSRG::update_t1() {
 
     // scale Hbar1 by renormalized denominator
     if (ccvv_source_ == "ZERO") {
-        DT1_.block("cv").iterate([&](const std::vector<size_t>& i, double& value) {
-            size_t i0 = core_mos_[i[0]];
-            size_t i1 = virt_mos_[i[1]];
-            value /= Fdiag_[i0] - Fdiag_[i1];
-        });
+        apply_denominator(DT1_, {"cv"}, [](double d) { return 1.0 / d; });
     }
 
-    for (const std::string& block : T1blocks) {
-        DT1_.block(block).iterate([&](const std::vector<size_t>& i, double& value) {
-            size_t i0 = label_to_spacemo_[block[0]][i[0]];
-            size_t i1 = label_to_spacemo_[block[1]][i[1]];
-            double denom = Fdiag_[i0] - Fdiag_[i1];
-            value *= dsrg_source_->compute_renormalized_denominator(denom);
-        });
-    }
+    apply_denominator(DT1_, T1blocks,
+                      [&](double d) { return dsrg_source_->compute_renormalized_denominator(d); });
 
     // Step 2: work on T1 where Hbar1 is treated as intermediate
 
@@ -379,14 +322,9 @@ void SA_MRDSRG::update_t1() {
     }
 
     // scale T1 by delta exponential
-    for (const std::string& block : T1blocks) {
-        T1_.block(block).iterate([&](const std::vector<size_t>& i, double& value) {
-            size_t i0 = label_to_spacemo_[block[0]][i[0]];
-            size_t i1 = label_to_spacemo_[block[1]][i[1]];
-            double denom = Fdiag_[i0] - Fdiag_[i1];
-            value *= dsrg_source_->compute_renormalized(denom);
-        });
-    }
+    apply_denominator(T1_, T1blocks,
+                      [&](double d) { return dsrg_source_->compute_renormalized(d); });
+
     if (ccvv_source_ == "ZERO") {
         T1_.block("cv").zero();
     }

forte/mrdsrg-spin-adapted/sa_mrpt2.cc

@@ -377,10 +377,8 @@ void SA_MRPT2::compute_t2_df_minimal() {
     }
 
     // build T2
-    T2_.iterate([&](const std::vector<size_t>& i, const std::vector<SpinType>&, double& value) {
-        double denom = Fdiag_[i[0]] + Fdiag_[i[1]] - Fdiag_[i[2]] - Fdiag_[i[3]];
-        value *= dsrg_source_->compute_renormalized_denominator(denom);
-    });
+    apply_denominator(T2_, T2_.block_labels(),
+                      [&](double v) { return dsrg_source_->compute_renormalized_denominator(v); });
 
     // transform back to non-canonical basis
     if (!semi_canonical_) {

forte/mrdsrg-spin-adapted/sadsrg.cc

@@ -1210,6 +1210,53 @@ ambit::Tensor SADSRG::read_Bcanonical(const std::string& block,
     return T;
 }
 
+void SADSRG::apply_denominator(ambit::BlockedTensor& T, const std::vector<std::string>& Tblocks,
+                               std::function<double(double)> func) {
+    for (const std::string& block : Tblocks) {
+        auto _bsize = block.size();
+        if (!T.is_block(block) or (_bsize != 2 and _bsize != 4))
+            continue;
+
+        auto size = T.block(block).numel();
+        size_t _size = 1;
+        for (char i : block) {
+            if (label_to_spacemo_.find(i) != label_to_spacemo_.end())
+                _size *= label_to_spacemo_[i].size();
+        }
+        if (size != _size)
+            continue;
+
+        auto& data = T.block(block).data();
+        int nthreads = omp_get_max_threads();
+        if (size < static_cast<size_t>(nthreads))
+            nthreads = size;
+
+        if (_bsize == 2) {
+            auto s1 = label_to_spacemo_[block[1]].size();
+#pragma omp parallel for num_threads(nthreads)
+            for (size_t i = 0; i < size; ++i) {
+                auto n0 = label_to_spacemo_[block[0]][i / s1];
+                auto n1 = label_to_spacemo_[block[1]][i % s1];
+                data[i] *= func(Fdiag_[n0] - Fdiag_[n1]);
+            }
+        } else if (_bsize == 4) {
+            auto s3 = label_to_spacemo_[block[3]].size();
+            auto s2 = label_to_spacemo_[block[2]].size();
+            auto s1 = label_to_spacemo_[block[1]].size();
+            auto s23 = s2 * s3;
+            auto s123 = s1 * s23;
+#pragma omp parallel for num_threads(nthreads)
+            for (size_t i = 0; i < size; ++i) {
+                auto n0 = label_to_spacemo_[block[0]][i / s123];
+                auto n1 = label_to_spacemo_[block[1]][(i / s23) % s1];
+                auto n2 = label_to_spacemo_[block[2]][(i % s23) / s3];
+                auto n3 = label_to_spacemo_[block[3]][i % s3];
+                data[i] *= func(Fdiag_[n0] + Fdiag_[n1] - Fdiag_[n2] - Fdiag_[n3]);
+            }
+        }
+    }
+}
+
 void SADSRG::print_contents(const std::string& str, size_t size) {
     if (str.size() + 4 > size)
         size = str.size() + 4;

forte/mrdsrg-spin-adapted/sadsrg.h

@@ -70,7 +70,7 @@ class SADSRG : public DynamicCorrelationSolver {
     void set_Uactv(ambit::Tensor& U);
 
     /// If the amplitudes are converged or not
-    bool converged() {return converged_; }
+    bool converged() { return converged_; }
 
   protected:
     /// Startup function called in constructor
@@ -380,9 +380,10 @@ class SADSRG : public DynamicCorrelationSolver {
     double H2_T1_C0(BlockedTensor& H2, BlockedTensor& T1, const double& alpha, double& C0);
     /// Compute zero-body term of commutator [H2, T2], S2[ijab] = 2 * T[ijab] - T[ijba]
     std::vector<double> H2_T2_C0(BlockedTensor& H2, BlockedTensor& T2, BlockedTensor& S2,
-                                 const double& alpha, double& C0, bool load_mps=false);
+                                 const double& alpha, double& C0, bool load_mps = false);
     /// Compute zero-body term of commutator [H2, T2], T2 and S2 contain at least two active indices
-    std::vector<double> H2_T2_C0_T2small(BlockedTensor& H2, BlockedTensor& T2, BlockedTensor& S2, bool load_mps=false);
+    std::vector<double> H2_T2_C0_T2small(BlockedTensor& H2, BlockedTensor& T2, BlockedTensor& S2,
+                                         bool load_mps = false);
 
     /// Compute one-body term of commutator [H1, T1]
     void H1_T1_C1(BlockedTensor& H1, BlockedTensor& T1, const double& alpha, BlockedTensor& C1);
@@ -406,7 +407,7 @@ class SADSRG : public DynamicCorrelationSolver {
     void V_T1_C0_DF(BlockedTensor& B, BlockedTensor& T1, const double& alpha, double& C0);
     /// Compute zero-body term of commutator [V, T2], V is constructed from B (DF/CD)
     std::vector<double> V_T2_C0_DF(BlockedTensor& B, BlockedTensor& T1, BlockedTensor& S2,
-                                   const double& alpha, double& C0, bool load_mps=false);
+                                   const double& alpha, double& C0, bool load_mps = false);
 
     /// Compute one-body term of commutator [V, T1], V is constructed from B (DF/CD)
     void V_T1_C1_DF(BlockedTensor& B, BlockedTensor& T1, const double& alpha, BlockedTensor& C1);
@@ -478,6 +479,10 @@ class SADSRG : public DynamicCorrelationSolver {
 
     // ==> common amplitudes analysis and printing <==
 
+    /// Apply denominator to BlockedTensor T (only 1- and 2-body)
+    void apply_denominator(ambit::BlockedTensor& T, const std::vector<std::string>& Tblocks,
+                           std::function<double(double)> func);
+
     /// Prune internal amplitudes for T1
     void internal_amps_T1(BlockedTensor& T1);
     /// Prune internal amplitudes for T2