Update wire order support for probabilities

.github/CHANGELOG.md

@@ -108,6 +108,9 @@
 
 ### Bug fixes
 
+* Fix wire order permutations when using `qml.probs` with out-of-order wires.
+  [(#707)](https://github.com/PennyLaneAI/pennylane-lightning/pull/707)
+
 * Lightning Qubit once again respects the wire order specified on device instantiation.
   [(#705)](https://github.com/PennyLaneAI/pennylane-lightning/pull/705)
 

.github/workflows/tests_linux_cpp.yml

@@ -142,7 +142,7 @@ jobs:
 
       - name: Install dependencies
         run: |
-          sudo apt-get update && sudo apt-get -y -q install cmake gcc-$GCC_VERSION  g++-$GCC_VERSION libopenblas-dev ninja-build gcovr lcov
+          sudo apt-get update && sudo apt-get -y -q install cmake gcc-$GCC_VERSION  g++-$GCC_VERSION libopenblas-base libopenblas-dev ninja-build gcovr lcov
           python -m pip install scipy
 
       - name: Build and run unit tests

.github/workflows/tests_lqcpu_python.yml

@@ -73,7 +73,7 @@ jobs:
 
       - name: Install dependencies
         run: |
-          sudo apt-get update && sudo apt-get -y -q install cmake gcc-$GCC_VERSION g++-$GCC_VERSION libopenblas-dev
+          sudo apt-get update && sudo apt-get -y -q install cmake gcc-$GCC_VERSION g++-$GCC_VERSION libopenblas-base libopenblas-dev
           python -m pip install scipy wheel
 
       - name: Get required Python packages

pennylane_lightning/core/src/measurements/tests/Test_MeasurementsBase.cpp

@@ -84,36 +84,41 @@ template <typename TypeList> void testProbabilities() {
         // Expected results calculated with Pennylane default.qubit:
         std::vector<std::pair<std::vector<size_t>, std::vector<PrecisionT>>>
             input = {
-#ifdef _ENABLE_PLGPU
+#if defined(_ENABLE_PLGPU)
                 // Bit index reodering conducted in the python layer
                 // for L-GPU. Also L-GPU backend doesn't support
                 // out of order wires for probability calculation
                 {{2, 1, 0},
                  {0.67078706, 0.03062806, 0.0870997, 0.00397696, 0.17564072,
                   0.00801973, 0.02280642, 0.00104134}}
 #else
-                {{0, 1, 2},
-                 {0.67078706, 0.03062806, 0.0870997, 0.00397696, 0.17564072,
-                  0.00801973, 0.02280642, 0.00104134}},
+#if defined(_ENABLE_PLQUBIT)
+                // LightningQubit currently supports arbitrary wire index
+                // ordering.
                 {{0, 2, 1},
                  {0.67078706, 0.0870997, 0.03062806, 0.00397696, 0.17564072,
                   0.02280642, 0.00801973, 0.00104134}},
                 {{1, 0, 2},
                  {0.67078706, 0.03062806, 0.17564072, 0.00801973, 0.0870997,
                   0.00397696, 0.02280642, 0.00104134}},
                 {{1, 2, 0},
-                 {0.67078706, 0.0870997, 0.17564072, 0.02280642, 0.03062806,
-                  0.00397696, 0.00801973, 0.00104134}},
-                {{2, 0, 1},
                  {0.67078706, 0.17564072, 0.03062806, 0.00801973, 0.0870997,
                   0.02280642, 0.00397696, 0.00104134}},
+                {{2, 0, 1},
+                 {0.67078706, 0.0870997, 0.17564072, 0.02280642, 0.03062806,
+                  0.00397696, 0.00801973, 0.00104134}},
                 {{2, 1, 0},
                  {0.67078706, 0.17564072, 0.0870997, 0.02280642, 0.03062806,
                   0.00801973, 0.00397696, 0.00104134}},
+                {{2, 1}, {0.84642778, 0.10990612, 0.0386478, 0.0050183}},
+
+#endif
+                {{0, 1, 2},
+                 {0.67078706, 0.03062806, 0.0870997, 0.00397696, 0.17564072,
+                  0.00801973, 0.02280642, 0.00104134}},
                 {{0, 1}, {0.70141512, 0.09107666, 0.18366045, 0.02384776}},
                 {{0, 2}, {0.75788676, 0.03460502, 0.19844714, 0.00906107}},
                 {{1, 2}, {0.84642778, 0.0386478, 0.10990612, 0.0050183}},
-                {{2, 1}, {0.84642778, 0.10990612, 0.0386478, 0.0050183}},
                 {{0}, {0.79249179, 0.20750821}},
                 {{1}, {0.88507558, 0.11492442}},
                 {{2}, {0.9563339, 0.0436661}}

pennylane_lightning/core/src/simulators/lightning_gpu/measurements/MeasurementsGPU.hpp

@@ -92,6 +92,10 @@ class Measurements final
      * @return std::vector<PrecisionT>
      */
     auto probs(const std::vector<size_t> &wires) -> std::vector<PrecisionT> {
+        PL_ABORT_IF_NOT(std::is_sorted(wires.cbegin(), wires.cend()),
+                        "LightningGPU does not currently support out-of-order "
+                        "wire indices with probability calculations");
+
         // Data return type fixed as double in custatevec function call
         std::vector<double> probabilities(Pennylane::Util::exp2(wires.size()));
         // this should be built upon by the wires not participating
@@ -193,6 +197,10 @@ class Measurements final
 
     std::vector<PrecisionT> probs(const std::vector<size_t> &wires,
                                   size_t num_shots) {
+        PL_ABORT_IF_NOT(std::is_sorted(wires.cbegin(), wires.cend()),
+                        "LightningGPU does not currently support out-of-order "
+                        "wire indices with probability calculations");
+
         return BaseType::probs(wires, num_shots);
     }
 

pennylane_lightning/core/src/simulators/lightning_kokkos/measurements/MeasurementsKokkos.hpp

@@ -506,6 +506,10 @@ class Measurements final
      * The basis columns are rearranged according to wires.
      */
     std::vector<PrecisionT> probs(const std::vector<size_t> &wires) {
+        PL_ABORT_IF_NOT(
+            std::is_sorted(wires.cbegin(), wires.cend()),
+            "LightningKokkos does not currently support out-of-order wire "
+            "indices with probability calculations");
         using MDPolicyType_2D =
             Kokkos::MDRangePolicy<Kokkos::Rank<2, Kokkos::Iterate::Left>>;
 
@@ -645,6 +649,11 @@ class Measurements final
 
     std::vector<PrecisionT> probs(const std::vector<size_t> &wires,
                                   size_t num_shots) {
+        PL_ABORT_IF_NOT(
+            std::is_sorted(wires.cbegin(), wires.cend()),
+            "LightningKokkos does not currently support out-of-order wire "
+            "indices with probability calculations");
+
         return BaseType::probs(wires, num_shots);
     }
 

pennylane_lightning/core/src/simulators/lightning_kokkos/measurements/tests/Test_StateVectorKokkos_Measure.cpp

@@ -226,25 +226,10 @@ TEMPLATE_TEST_CASE("Probabilities", "[Measures]", float, double) {
         {{0, 1, 2},
          {0.67078706, 0.03062806, 0.0870997, 0.00397696, 0.17564072, 0.00801973,
           0.02280642, 0.00104134}},
-        {{0, 2, 1},
-         {0.67078706, 0.0870997, 0.03062806, 0.00397696, 0.17564072, 0.02280642,
-          0.00801973, 0.00104134}},
-        {{1, 0, 2},
-         {0.67078706, 0.03062806, 0.17564072, 0.00801973, 0.0870997, 0.00397696,
-          0.02280642, 0.00104134}},
-        {{1, 2, 0},
-         {0.67078706, 0.0870997, 0.17564072, 0.02280642, 0.03062806, 0.00397696,
-          0.00801973, 0.00104134}},
-        {{2, 0, 1},
-         {0.67078706, 0.17564072, 0.03062806, 0.00801973, 0.0870997, 0.02280642,
-          0.00397696, 0.00104134}},
-        {{2, 1, 0},
-         {0.67078706, 0.17564072, 0.0870997, 0.02280642, 0.03062806, 0.00801973,
-          0.00397696, 0.00104134}},
+        // TODO: Fix LK out-of-order permutations
         {{0, 1}, {0.70141512, 0.09107666, 0.18366045, 0.02384776}},
         {{0, 2}, {0.75788676, 0.03460502, 0.19844714, 0.00906107}},
         {{1, 2}, {0.84642778, 0.0386478, 0.10990612, 0.0050183}},
-        {{2, 1}, {0.84642778, 0.10990612, 0.0386478, 0.0050183}},
         {{0}, {0.79249179, 0.20750821}},
         {{1}, {0.88507558, 0.11492442}},
         {{2}, {0.9563339, 0.0436661}}};

pennylane_lightning/core/src/simulators/lightning_qubit/measurements/MeasurementsLQubit.hpp

@@ -31,6 +31,7 @@
 
 #include "LinearAlgebra.hpp"
 #include "MeasurementsBase.hpp"
+#include "NDPermuter.hpp"
 #include "Observables.hpp"
 #include "SparseLinAlg.hpp"
 #include "StateVectorLQubitManaged.hpp"
@@ -44,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
 
@@ -100,16 +102,17 @@ class Measurements final
         // Determining index that would sort the vector.
         // This information is needed later.
         const auto sorted_ind_wires = Pennylane::Util::sorting_indices(wires);
+
         // Sorting wires.
         std::vector<size_t> sorted_wires(wires.size());
         for (size_t pos = 0; pos < wires.size(); pos++) {
             sorted_wires[pos] = wires[sorted_ind_wires[pos]];
         }
+
         // Determining probabilities for the sorted wires.
         const ComplexT *arr_data = this->_statevector.getData();
 
         size_t num_qubits = this->_statevector.getNumQubits();
-
         const std::vector<size_t> all_indices =
             Gates::generateBitPatterns(sorted_wires, num_qubits);
         const std::vector<size_t> all_offsets = Gates::generateBitPatterns(
@@ -125,12 +128,29 @@ class Measurements final
             }
             ind_probs++;
         }
-        // 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;
         }
+
         return probabilities;
     }
 

pennylane_lightning/core/src/simulators/lightning_qubit/measurements/tests/CMakeLists.txt

@@ -17,6 +17,7 @@ FetchAndIncludeCatch()
 add_library(lightning_qubit_measurements_tests INTERFACE)
 target_link_libraries(lightning_qubit_measurements_tests INTERFACE  Catch2::Catch2
                                                                     lightning_measurements
+                                                                    lightning_qubit_observables
                                                                     lightning_qubit_measurements
                                                                     )