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             : Eigen::SparseMatrix<Scalar, Eigen::RowMajor>
      20          38 : calculate_tensor_product(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             :     using real_t = typename traits::NumTraits<Scalar>::real_t;
      25          38 :     constexpr real_t numerical_precision = 100 * std::numeric_limits<real_t>::epsilon();
      26             : 
      27          38 :     oneapi::tbb::concurrent_vector<Eigen::Triplet<Scalar>> triplets;
      28             : 
      29             :     // Loop over the rows of the first matrix in parallel (outer index == row)
      30          38 :     oneapi::tbb::parallel_for(
      31        3013 :         oneapi::tbb::blocked_range<Eigen::Index>(0, matrix1.outerSize()), [&](const auto &range) {
      32        5966 :             for (Eigen::Index row1 = range.begin(); row1 != range.end(); ++row1) {
      33             : 
      34        2984 :                 const auto &range_row2 = basis_final->get_index_range(row1);
      35             : 
      36             :                 // Loop over the rows of the second matrix that are energetically allowed
      37       14764 :                 for (auto row2 = static_cast<Eigen::Index>(range_row2.min());
      38       14766 :                      row2 < static_cast<Eigen::Index>(range_row2.max()); ++row2) {
      39             : 
      40       11754 :                     Eigen::Index row = basis_final->get_ket_index_from_tuple(row1, row2);
      41       11776 :                     if (row < 0) {
      42        2379 :                         continue;
      43             :                     }
      44             : 
      45             :                     // Loop over the non-zero column elements of the first matrix
      46      171928 :                     for (typename Eigen::SparseMatrix<Scalar, Eigen::RowMajor>::InnerIterator it1(
      47             :                              matrix1, row1);
      48      168784 :                          it1; ++it1) {
      49             : 
      50      160933 :                         Eigen::Index col1 = it1.col();
      51      160188 :                         Scalar value1 = it1.value();
      52             : 
      53             :                         // Calculate the minimum and maximum values of the index pointer of the
      54             :                         // second matrix
      55      159728 :                         Eigen::Index begin_idxptr2 = matrix2.outerIndexPtr()[row2];
      56      159723 :                         Eigen::Index end_idxptr2 = matrix2.outerIndexPtr()[row2 + 1];
      57             : 
      58             :                         // The minimum value is chosen such that we start with an energetically
      59             :                         // allowed column
      60      160887 :                         const auto &range_col2 = basis_initial->get_index_range(it1.index());
      61      159922 :                         begin_idxptr2 +=
      62      159724 :                             std::distance(matrix2.innerIndexPtr() + begin_idxptr2,
      63      157612 :                                           std::lower_bound(matrix2.innerIndexPtr() + begin_idxptr2,
      64      154072 :                                                            matrix2.innerIndexPtr() + end_idxptr2,
      65      151791 :                                                            range_col2.min()));
      66             : 
      67             :                         // Loop over the non-zero column elements of the second matrix that are
      68             :                         // energetically allowed (we break the loop if the index pointer corresponds
      69             :                         // to a column that is not energetically allowed)
      70      221309 :                         for (Eigen::Index idxptr2 = begin_idxptr2; idxptr2 < end_idxptr2;
      71             :                              ++idxptr2) {
      72             : 
      73      200625 :                             Eigen::Index col2 = matrix2.innerIndexPtr()[idxptr2];
      74      200831 :                             if (col2 >= static_cast<Eigen::Index>(range_col2.max())) {
      75      141847 :                                 break;
      76             :                             }
      77             : 
      78       59473 :                             Eigen::Index col = basis_initial->get_ket_index_from_tuple(col1, col2);
      79       61701 :                             if (col < 0) {
      80       37603 :                                 continue;
      81             :                             }
      82             : 
      83       24098 :                             Scalar value2 = matrix2.valuePtr()[idxptr2];
      84             : 
      85             :                             // Store the entry
      86           0 :                             Scalar value = value1 * value2;
      87       23964 :                             if (std::abs(value) > numerical_precision) {
      88       23243 :                                 triplets.emplace_back(row, col, value);
      89             :                             }
      90             :                         }
      91             :                     }
      92             :                 }
      93             :             }
      94             :         });
      95             : 
      96             :     // Construct the combined matrix from the triplets
      97          38 :     Eigen::SparseMatrix<Scalar, Eigen::RowMajor> matrix(basis_final->get_number_of_kets(),
      98          38 :                                                         basis_initial->get_number_of_kets());
      99          38 :     matrix.setFromTriplets(triplets.begin(), triplets.end());
     100          38 :     matrix.makeCompressed();
     101             : 
     102          76 :     return matrix;
     103          38 : }
     104             : 
     105             : // Explicit instantiations
     106             : template Eigen::SparseMatrix<double, Eigen::RowMajor>
     107             : calculate_tensor_product(const std::shared_ptr<const BasisPair<double>> &,
     108             :                          const std::shared_ptr<const BasisPair<double>> &,
     109             :                          const Eigen::SparseMatrix<double, Eigen::RowMajor> &,
     110             :                          const Eigen::SparseMatrix<double, Eigen::RowMajor> &);
     111             : template Eigen::SparseMatrix<std::complex<double>, Eigen::RowMajor>
     112             : calculate_tensor_product(const std::shared_ptr<const BasisPair<std::complex<double>>> &,
     113             :                          const std::shared_ptr<const BasisPair<std::complex<double>>> &,
     114             :                          const Eigen::SparseMatrix<std::complex<double>, Eigen::RowMajor> &,
     115             :                          const Eigen::SparseMatrix<std::complex<double>, Eigen::RowMajor> &);
     116             : 
     117             : } // namespace pairinteraction::utils