Commit few changes in apply methods of PauliString to get James opinion

qognitive · Jul 25, 2024 · 80bf7d5 · 80bf7d5
1 parent 5f38586
commit 80bf7d5
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 35 deletions.
diff --git a/fast_pauli/cpp/include/__pauli_string.hpp b/fast_pauli/cpp/include/__pauli_string.hpp
@@ -197,9 +197,7 @@ struct PauliString {
   }
 
   /**
-   * @brief Apply the PauliString (using the sparse representation) to a vector.
-   * This performs following matrix-vector multiplication \f$ \mathcal{\hat{P}}
-   * \ket{\psi} \f$
+   * @brief @copybrief PauliString::apply(std::mdspan)
    *
    * @tparam T The floating point base to use for all the complex numbers
    * @param v The input vector to apply the PauliString to. Must be the same
@@ -210,26 +208,14 @@ struct PauliString {
   template <std::floating_point T>
   std::vector<std::complex<T>>
   apply(std::vector<std::complex<T>> const &v) const {
-    // Input check
-    if (v.size() != dims()) {
-      throw std::invalid_argument(
-          "Input vector size must match the number of qubits");
-    }
-
-    std::vector<size_t> k;
-    std::vector<std::complex<T>> m;
-    get_sparse_repr(k, m);
-
-    std::vector<std::complex<T>> result(v.size(), 0);
-    for (size_t i = 0; i < k.size(); ++i) {
-      result[i] += m[i] * v[k[i]];
-    }
-
-    return result;
+    // route this to implementation we have for mdspan specialization
+    return this->apply(std::mdspan(v.data(), v.size()));
   }
 
   /**
-   * @brief @copybrief PauliString::apply(std::vector<std::complex<T>>)
+   * @brief Apply the PauliString (using the sparse representation) to a vector.
+   * This performs following matrix-vector multiplication \f$ \mathcal{\hat{P}}
+   * \ket{\psi} \f$
    *
    * @tparam T The floating point base to use for all the complex numbers
    * @param v The input vector to apply the PauliString to. Must be the same
@@ -239,7 +225,7 @@ struct PauliString {
    */
   template <std::floating_point T>
   std::vector<std::complex<T>>
-  apply(std::mdspan<std::complex<T>, std::dextents<size_t, 1>> v) const {
+  apply(std::mdspan<const std::complex<T>, std::dextents<size_t, 1>> v) const {
     // Input check
     if (v.size() != dims()) {
       throw std::invalid_argument(
@@ -270,9 +256,9 @@ struct PauliString {
    *
    * @tparam T The floating point base to use for all the complex numbers
    * @param new_states_T The output states after applying the PauliString
-   * (n_data x n_dim)
-   * @param states_T THe original states to apply the PauliString to (n_data x
-   * n_dim)
+   * (n_dim x n_states)
+   * @param states_T THe original states to apply the PauliString to
+   * (n_dim x n_states)
    * @param c Multiplication factor to apply to the PauliString
    */
   template <std::floating_point T>
@@ -300,16 +286,12 @@ struct PauliString {
     std::vector<std::complex<T>> m;
     get_sparse_repr(k, m);
 
-    // TODO we have bad memory access patterns
-    // std::vector<std::complex<T>> col(states_T.extent(1), 0);
     for (size_t i = 0; i < states_T.extent(0); ++i) {
-      std::complex<T> c_m_i = c * m[i];
-      size_t k_i = k[i];
-      std::memcpy(&new_states_T(i, 0), &states_T(k_i, 0),
-                  states_T.extent(1) * sizeof(std::complex<T>));
-
+      // std::memcpy(&new_states_T(i, 0), &states_T(k_i, 0),
+      //             states_T.extent(1) * sizeof(std::complex<T>));
+      const std::complex<T> c_m_i = c * m[i];
       for (size_t t = 0; t < states_T.extent(1); ++t) {
-        new_states_T(i, t) *= c_m_i; // * states_T(k_i, t);
+        new_states_T(i, t) = c_m_i * states_T(k[i], t);
       }
     }
   }

diff --git a/fast_pauli/cpp/src/fast_pauli.cpp b/fast_pauli/cpp/src/fast_pauli.cpp
@@ -74,8 +74,8 @@ PYBIND11_MODULE(_fast_pauli, m) {
               if (vec.size() != n_states)
                 throw std::invalid_argument("Bad shape of states array");
               else
-                std::copy(vec.begin(), vec.end(),
-                          std::back_inserter(flat_inputs));
+                std::ranges::copy(vec.begin(), vec.end(),
+                                  std::back_inserter(flat_inputs));
 
             std::vector<std::complex<double>> flat_outputs(flat_inputs.size(),
                                                            0);
@@ -90,7 +90,8 @@ PYBIND11_MODULE(_fast_pauli, m) {
                 inputs.size());
             for (size_t i = 0; i < inputs.size(); i++) {
               auto it = flat_outputs.begin() + i * n_states;
-              std::copy(it, it + n_states, std::back_inserter(results[i]));
+              std::ranges::copy(it, it + n_states,
+                                std::back_inserter(results[i]));
             }
             return results;
           },