From 5a3a2862f1480f8b450163210c2e6b5917b0ed8c Mon Sep 17 00:00:00 2001
From: "Lee J. O'Riordan" <lee@xanadu.ai>
Date: Fri, 3 May 2024 08:44:17 -0400
Subject: [PATCH] Fix permutation ordering

---
 .../measurements/MeasurementsLQubit.hpp       | 58 ++++++++-----------
 .../core/src/utils/NDPermuter.hpp             |  4 +-
 2 files changed, 25 insertions(+), 37 deletions(-)

diff --git a/pennylane_lightning/core/src/simulators/lightning_qubit/measurements/MeasurementsLQubit.hpp b/pennylane_lightning/core/src/simulators/lightning_qubit/measurements/MeasurementsLQubit.hpp
index b8626d24c0..6c2276df4e 100644
--- a/pennylane_lightning/core/src/simulators/lightning_qubit/measurements/MeasurementsLQubit.hpp
+++ b/pennylane_lightning/core/src/simulators/lightning_qubit/measurements/MeasurementsLQubit.hpp
@@ -31,7 +31,7 @@
 
 #include "LinearAlgebra.hpp"
 #include "MeasurementsBase.hpp"
-#include "NDPermuter.hpp" //transpose_state_tensor, sorting_indices
+#include "NDPermuter.hpp"
 #include "Observables.hpp"
 #include "SparseLinAlg.hpp"
 #include "StateVectorLQubitManaged.hpp"
@@ -45,6 +45,7 @@ using namespace Pennylane::Measures;
 using namespace Pennylane::Observables;
 using Pennylane::LightningQubit::StateVectorLQubitManaged;
 using Pennylane::LightningQubit::Util::innerProdC;
+namespace PUtil = Pennylane::Util;
 } // namespace
 /// @endcond
 
@@ -108,10 +109,6 @@ class Measurements final
             sorted_wires[pos] = wires[sorted_ind_wires[pos]];
         }
 
-        std::cout << "wires=" << wires << "\n";
-        std::cout << "sorted_ind_wires=" << sorted_ind_wires << "\n";
-        std::cout << "sorted_wires=" << sorted_wires << "\n";
-
         // Determining probabilities for the sorted wires.
         const ComplexT *arr_data = this->_statevector.getData();
 
@@ -122,9 +119,6 @@ class Measurements final
             Gates::getIndicesAfterExclusion(sorted_wires, num_qubits),
             num_qubits);
 
-        std::cout << "all_indices=" << all_indices << "\n";
-        std::cout << "all_offsets=" << all_offsets << "\n";
-
         std::vector<PrecisionT> probabilities(all_indices.size(), 0);
 
         size_t ind_probs = 0;
@@ -134,36 +128,30 @@ class Measurements final
             }
             ind_probs++;
         }
-        std::cout << "probabilities=" << probabilities << "\n";
 
-        // Transposing the probabilities tensor with the indices determined
-        // at the beginning.
+        // Permute the data according to the required wire ordering
         if (wires != sorted_wires) {
-            probabilities = Pennylane::Util::transpose_state_tensor(
-                probabilities, sorted_ind_wires);
+            static constexpr std::size_t CACHE_SIZE = 8;
+            PUtil::Permuter<PUtil::DefaultPermuter<CACHE_SIZE>> p;
+            std::vector<std::size_t> shape(wires.size(), 2);
+            std::vector<std::string> wire_labels_old(sorted_wires.size(), "");
+            std::vector<std::string> wire_labels_new(wires.size(), "");
+
+            std::transform(sorted_wires.begin(), sorted_wires.end(),
+                           wire_labels_old.begin(), [](std::size_t index) {
+                               return std::to_string(index);
+                           });
+            std::transform(
+                wires.begin(), wires.end(), wire_labels_new.begin(),
+                [](std::size_t index) { return std::to_string(index); });
+
+            auto probs_sorted = probabilities;
+            p.Transpose(probabilities, shape, probs_sorted, wire_labels_old,
+                        wire_labels_new);
+            return probs_sorted;
         }
-        std::cout << "probabilities+1=" << probabilities << "\n";
-
-        // return probabilities;
-
-        Pennylane::Util::Permuter<Pennylane::Util::DefaultPermuter<1024>> p;
-        std::vector<std::size_t> shape(num_qubits, 2);
-        std::vector<std::string> wire_labels_old(sorted_ind_wires.size(), "");
-        std::vector<std::string> wire_labels_new(sorted_ind_wires.size(), "");
-
-        std::transform(sorted_ind_wires.begin(), sorted_ind_wires.end(),
-                       wire_labels_old.begin(),
-                       [](std::size_t index) { return std::to_string(index); });
-        std::transform(wires.begin(), wires.end(), wire_labels_new.begin(),
-                       [](std::size_t index) { return std::to_string(index); });
-        // for ()
-
-        auto probs_sorted = probabilities;
-        p.Transpose(probabilities, shape, probs_sorted, wire_labels_old,
-                    wire_labels_new);
-        std::cout << "probabilities+2=" << probs_sorted << "\n";
-        return probs_sorted;
-        // return probabilities;
+
+        return probabilities;
     }
 
     /**
diff --git a/pennylane_lightning/core/src/utils/NDPermuter.hpp b/pennylane_lightning/core/src/utils/NDPermuter.hpp
index 356995b821..6aac4a7eb6 100644
--- a/pennylane_lightning/core/src/utils/NDPermuter.hpp
+++ b/pennylane_lightning/core/src/utils/NDPermuter.hpp
@@ -54,8 +54,8 @@ template <class PermuterBackend> class Permuter {
         const std::set<std::string> idx_old(current_order.begin(),
                                             current_order.end());
         const std::set<std::string> idx_new(new_order.begin(), new_order.end());
-        const auto data_size = std::accumulate(shape.begin(), shape.end(), 1,
-                                               std::multiplies<std::size_t>());
+        const std::size_t data_size = std::accumulate(
+            shape.begin(), shape.end(), 1, std::multiplies<std::size_t>());
         PL_ABORT_IF_NOT(idx_old.size() == current_order.size(),
                         "Duplicate existing indices found. Please ensure "
                         "indices are unique.");