Line data    Source code

       1             : // SPDX-FileCopyrightText: 2025 PairInteraction Developers
       2             : // SPDX-License-Identifier: LGPL-3.0-or-later
       3             : 
       4             : #include "pairinteraction/utils/tensor.hpp"
       5             : 
       6             : #include "pairinteraction/basis/BasisPair.hpp"
       7             : #include "pairinteraction/utils/eigen_assertion.hpp"
       8             : #include "pairinteraction/utils/traits.hpp"
       9             : 
      10             : #include <Eigen/SparseCore>
      11             : #include <algorithm>
      12             : #include <limits>
      13             : #include <memory>
      14             : #include <oneapi/tbb.h>
      15             : 
      16             : namespace pairinteraction::utils {
      17             : 
      18             : template <typename Scalar>
      19        2042 : Eigen::SparseMatrix<Scalar, Eigen::RowMajor> calculate_tensor_product_in_canonical_basis(
      20             :     const std::shared_ptr<const BasisPair<Scalar>> &basis_initial,
      21             :     const std::shared_ptr<const BasisPair<Scalar>> &basis_final,
      22             :     const Eigen::SparseMatrix<Scalar, Eigen::RowMajor> &matrix1,
      23             :     const Eigen::SparseMatrix<Scalar, Eigen::RowMajor> &matrix2) {
      24        2042 :     oneapi::tbb::concurrent_vector<Eigen::Triplet<Scalar>> triplets;
      25             : 
      26             :     // Loop over the rows of the first matrix in parallel (outer index == row)
      27        2035 :     oneapi::tbb::parallel_for(
      28      142957 :         oneapi::tbb::blocked_range<Eigen::Index>(0, matrix1.outerSize()), [&](const auto &range) {
      29      306539 :             for (Eigen::Index row1 = range.begin(); row1 != range.end(); ++row1) {
      30             : 
      31      165231 :                 const auto &range_row2 = basis_final->get_index_range(row1);
      32             : 
      33             :                 // Loop over the rows of the second matrix that are energetically allowed
      34     1549111 :                 for (auto row2 = static_cast<Eigen::Index>(range_row2.min());
      35     1549235 :                      row2 < static_cast<Eigen::Index>(range_row2.max()); ++row2) {
      36             : 
      37     1379975 :                     Eigen::Index row = basis_final->get_ket_index_from_tuple(row1, row2);
      38     1385370 :                     if (row < 0) {
      39      956977 :                         continue;
      40             :                     }
      41             : 
      42             :                     // Loop over the non-zero column elements of the first matrix
      43     6499226 :                     for (typename Eigen::SparseMatrix<Scalar, Eigen::RowMajor>::InnerIterator it1(
      44             :                              matrix1, row1);
      45     6476368 :                          it1; ++it1) {
      46             : 
      47     6070249 :                         Eigen::Index col1 = it1.col();
      48     6106689 :                         Scalar value1 = it1.value();
      49             : 
      50             :                         // Calculate the minimum and maximum values of the index pointer of the
      51             :                         // second matrix
      52     6099922 :                         Eigen::Index begin_idxptr2 = matrix2.outerIndexPtr()[row2];
      53     6094375 :                         Eigen::Index end_idxptr2 = matrix2.outerIndexPtr()[row2 + 1];
      54             : 
      55             :                         // The minimum value is chosen such that we start with an energetically
      56             :                         // allowed column
      57     6128681 :                         const auto &range_col2 = basis_initial->get_index_range(it1.index());
      58     6070614 :                         begin_idxptr2 +=
      59     6061310 :                             std::distance(matrix2.innerIndexPtr() + begin_idxptr2,
      60     6064473 :                                           std::lower_bound(matrix2.innerIndexPtr() + begin_idxptr2,
      61     6002985 :                                                            matrix2.innerIndexPtr() + end_idxptr2,
      62     5957093 :                                                            range_col2.min()));
      63             : 
      64             :                         // Loop over the non-zero column elements of the second matrix that are
      65             :                         // energetically allowed (we break the loop if the index pointer corresponds
      66             :                         // to a column that is not energetically allowed)
      67    11655693 :                         for (Eigen::Index idxptr2 = begin_idxptr2; idxptr2 < end_idxptr2;
      68             :                              ++idxptr2) {
      69             : 
      70    10560528 :                             Eigen::Index col2 = matrix2.innerIndexPtr()[idxptr2];
      71    10529518 :                             if (col2 >= static_cast<Eigen::Index>(range_col2.max())) {
      72     4975668 :                                 break;
      73             :                             }
      74             : 
      75     5455050 :                             Eigen::Index col = basis_initial->get_ket_index_from_tuple(col1, col2);
      76     5731464 :                             if (col < 0) {
      77       37248 :                                 continue;
      78             :                             }
      79             : 
      80     5694216 :                             Scalar value2 = matrix2.valuePtr()[idxptr2];
      81             : 
      82             :                             // Store the entry
      83     2850240 :                             Scalar value = value1 * value2;
      84     5712507 :                             triplets.emplace_back(row, col, value);
      85             :                         }
      86             :                     }
      87             :                 }
      88             :             }
      89             :         });
      90             : 
      91             :     // Construct the combined matrix from the triplets
      92        2029 :     Eigen::SparseMatrix<Scalar, Eigen::RowMajor> matrix(basis_final->get_number_of_kets(),
      93        2040 :                                                         basis_initial->get_number_of_kets());
      94        2035 :     matrix.setFromTriplets(triplets.begin(), triplets.end());
      95        2045 :     matrix.makeCompressed();
      96             : 
      97        4076 :     return matrix;
      98        2039 : }
      99             : 
     100             : // Explicit instantiations
     101             : template Eigen::SparseMatrix<double, Eigen::RowMajor>
     102             : calculate_tensor_product_in_canonical_basis(const std::shared_ptr<const BasisPair<double>> &,
     103             :                                             const std::shared_ptr<const BasisPair<double>> &,
     104             :                                             const Eigen::SparseMatrix<double, Eigen::RowMajor> &,
     105             :                                             const Eigen::SparseMatrix<double, Eigen::RowMajor> &);
     106             : template Eigen::SparseMatrix<std::complex<double>, Eigen::RowMajor>
     107             : calculate_tensor_product_in_canonical_basis(
     108             :     const std::shared_ptr<const BasisPair<std::complex<double>>> &,
     109             :     const std::shared_ptr<const BasisPair<std::complex<double>>> &,
     110             :     const Eigen::SparseMatrix<std::complex<double>, Eigen::RowMajor> &,
     111             :     const Eigen::SparseMatrix<std::complex<double>, Eigen::RowMajor> &);
     112             : 
     113             : } // namespace pairinteraction::utils