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          38 :     oneapi::tbb::concurrent_vector<Eigen::Triplet<Scalar>> triplets;
      25             : 
      26             :     // Loop over the rows of the first matrix in parallel (outer index == row)
      27          38 :     oneapi::tbb::parallel_for(
      28        3014 :         oneapi::tbb::blocked_range<Eigen::Index>(0, matrix1.outerSize()), [&](const auto &range) {
      29        5969 :             for (Eigen::Index row1 = range.begin(); row1 != range.end(); ++row1) {
      30             : 
      31        2984 :                 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       14759 :                 for (auto row2 = static_cast<Eigen::Index>(range_row2.min());
      35       14762 :                      row2 < static_cast<Eigen::Index>(range_row2.max()); ++row2) {
      36             : 
      37       11766 :                     Eigen::Index row = basis_final->get_ket_index_from_tuple(row1, row2);
      38       11758 :                     if (row < 0) {
      39        2368 :                         continue;
      40             :                     }
      41             : 
      42             :                     // Loop over the non-zero column elements of the first matrix
      43      172670 :                     for (typename Eigen::SparseMatrix<Scalar, Eigen::RowMajor>::InnerIterator it1(
      44             :                              matrix1, row1);
      45      169853 :                          it1; ++it1) {
      46             : 
      47      161416 :                         Eigen::Index col1 = it1.col();
      48      161403 :                         Scalar value1 = it1.value();
      49             : 
      50             :                         // Calculate the minimum and maximum values of the index pointer of the
      51             :                         // second matrix
      52      160982 :                         Eigen::Index begin_idxptr2 = matrix2.outerIndexPtr()[row2];
      53      160857 :                         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      162261 :                         const auto &range_col2 = basis_initial->get_index_range(it1.index());
      58      161142 :                         begin_idxptr2 +=
      59      160999 :                             std::distance(matrix2.innerIndexPtr() + begin_idxptr2,
      60      159633 :                                           std::lower_bound(matrix2.innerIndexPtr() + begin_idxptr2,
      61      156064 :                                                            matrix2.innerIndexPtr() + end_idxptr2,
      62      157078 :                                                            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      222719 :                         for (Eigen::Index idxptr2 = begin_idxptr2; idxptr2 < end_idxptr2;
      68             :                              ++idxptr2) {
      69             : 
      70      202458 :                             Eigen::Index col2 = matrix2.innerIndexPtr()[idxptr2];
      71      201874 :                             if (col2 >= static_cast<Eigen::Index>(range_col2.max())) {
      72      143019 :                                 break;
      73             :                             }
      74             : 
      75       61505 :                             Eigen::Index col = basis_initial->get_ket_index_from_tuple(col1, col2);
      76       61628 :                             if (col < 0) {
      77       37571 :                                 continue;
      78             :                             }
      79             : 
      80       24057 :                             Scalar value2 = matrix2.valuePtr()[idxptr2];
      81             : 
      82             :                             // Store the entry
      83           0 :                             Scalar value = value1 * value2;
      84       24046 :                             triplets.emplace_back(row, col, value);
      85             :                         }
      86             :                     }
      87             :                 }
      88             :             }
      89             :         });
      90             : 
      91             :     // Construct the combined matrix from the triplets
      92          38 :     Eigen::SparseMatrix<Scalar, Eigen::RowMajor> matrix(basis_final->get_number_of_kets(),
      93          38 :                                                         basis_initial->get_number_of_kets());
      94          38 :     matrix.setFromTriplets(triplets.begin(), triplets.end());
      95          38 :     matrix.makeCompressed();
      96             : 
      97          76 :     return matrix;
      98          38 : }
      99             : 
     100             : // Explicit instantiations
     101             : template Eigen::SparseMatrix<double, Eigen::RowMajor>
     102             : calculate_tensor_product(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(const std::shared_ptr<const BasisPair<std::complex<double>>> &,
     108             :                          const std::shared_ptr<const BasisPair<std::complex<double>>> &,
     109             :                          const Eigen::SparseMatrix<std::complex<double>, Eigen::RowMajor> &,
     110             :                          const Eigen::SparseMatrix<std::complex<double>, Eigen::RowMajor> &);
     111             : 
     112             : } // namespace pairinteraction::utils