From 465e5853ef93ba4f9c69ad7e8e18c0cca5d8a484 Mon Sep 17 00:00:00 2001
From: Sergei Isakov <54642992+sergeisakov@users.noreply.github.com>
Date: Tue, 26 Mar 2024 14:58:12 +0100
Subject: [PATCH] Fix sampling.

---
 lib/statespace_avx.h    | 12 ++++++++----
 lib/statespace_avx512.h | 12 ++++++++----
 lib/statespace_basic.h  | 12 ++++++++----
 lib/statespace_sse.h    | 12 ++++++++----
 4 files changed, 32 insertions(+), 16 deletions(-)

diff --git a/lib/statespace_avx.h b/lib/statespace_avx.h
index dbdbb6c0..876058bb 100644
--- a/lib/statespace_avx.h
+++ b/lib/statespace_avx.h
@@ -382,8 +382,8 @@ class StateSpaceAVX :
 
       for (uint64_t k = 0; k < size; ++k) {
         for (unsigned j = 0; j < 8; ++j) {
-          auto re = p[16 * k + j];
-          auto im = p[16 * k + 8 + j];
+          double re = p[16 * k + j];
+          double im = p[16 * k + 8 + j];
           norm += re * re + im * im;
         }
       }
@@ -396,8 +396,8 @@ class StateSpaceAVX :
 
       for (uint64_t k = 0; k < size; ++k) {
         for (unsigned j = 0; j < 8; ++j) {
-          auto re = p[16 * k + j];
-          auto im = p[16 * k + 8 + j];
+          double re = p[16 * k + j];
+          double im = p[16 * k + 8 + j];
           csum += re * re + im * im;
           while (rs[m] < csum && m < num_samples) {
             bitstrings.emplace_back(8 * k + j);
@@ -405,6 +405,10 @@ class StateSpaceAVX :
           }
         }
       }
+
+      for (; m < num_samples; ++m) {
+        bitstrings.emplace_back((uint64_t{1} << state.num_qubits()) - 1);
+      }
     }
 
     return bitstrings;
diff --git a/lib/statespace_avx512.h b/lib/statespace_avx512.h
index 4865b3e2..879fd89d 100644
--- a/lib/statespace_avx512.h
+++ b/lib/statespace_avx512.h
@@ -333,8 +333,8 @@ class StateSpaceAVX512 :
 
       for (uint64_t k = 0; k < size; ++k) {
         for (unsigned j = 0; j < 16; ++j) {
-          auto re = p[32 * k + j];
-          auto im = p[32 * k + 16 + j];
+          double re = p[32 * k + j];
+          double im = p[32 * k + 16 + j];
           norm += re * re + im * im;
         }
       }
@@ -347,8 +347,8 @@ class StateSpaceAVX512 :
 
       for (uint64_t k = 0; k < size; ++k) {
         for (unsigned j = 0; j < 16; ++j) {
-          auto re = p[32 * k + j];
-          auto im = p[32 * k + 16 + j];
+          double re = p[32 * k + j];
+          double im = p[32 * k + 16 + j];
           csum += re * re + im * im;
           while (rs[m] < csum && m < num_samples) {
             bitstrings.emplace_back(16 * k + j);
@@ -356,6 +356,10 @@ class StateSpaceAVX512 :
           }
         }
       }
+
+      for (; m < num_samples; ++m) {
+        bitstrings.emplace_back((uint64_t{1} << state.num_qubits()) - 1);
+      }
     }
 
     return bitstrings;
diff --git a/lib/statespace_basic.h b/lib/statespace_basic.h
index 604745fe..6468483a 100644
--- a/lib/statespace_basic.h
+++ b/lib/statespace_basic.h
@@ -203,8 +203,8 @@ class StateSpaceBasic :
       const fp_type* p = state.get();
 
       for (uint64_t k = 0; k < size; ++k) {
-        auto re = p[2 * k];
-        auto im = p[2 * k + 1];
+        double re = p[2 * k];
+        double im = p[2 * k + 1];
         norm += re * re + im * im;
       }
 
@@ -215,14 +215,18 @@ class StateSpaceBasic :
       bitstrings.reserve(num_samples);
 
       for (uint64_t k = 0; k < size; ++k) {
-        auto re = p[2 * k];
-        auto im = p[2 * k + 1];
+        double re = p[2 * k];
+        double im = p[2 * k + 1];
         csum += re * re + im * im;
         while (rs[m] < csum && m < num_samples) {
           bitstrings.emplace_back(k);
           ++m;
         }
       }
+
+      for (; m < num_samples; ++m) {
+        bitstrings.emplace_back((uint64_t{1} << state.num_qubits()) - 1);
+      }
     }
 
     return bitstrings;
diff --git a/lib/statespace_sse.h b/lib/statespace_sse.h
index f164b9b4..cf41a09f 100644
--- a/lib/statespace_sse.h
+++ b/lib/statespace_sse.h
@@ -342,8 +342,8 @@ class StateSpaceSSE :
 
       for (uint64_t k = 0; k < size; ++k) {
         for (unsigned j = 0; j < 4; ++j) {
-          auto re = p[8 * k + j];
-          auto im = p[8 * k + 4 + j];
+          double re = p[8 * k + j];
+          double im = p[8 * k + 4 + j];
           norm += re * re + im * im;
         }
       }
@@ -356,8 +356,8 @@ class StateSpaceSSE :
 
       for (uint64_t k = 0; k < size; ++k) {
         for (unsigned j = 0; j < 4; ++j) {
-          auto re = p[8 * k + j];
-          auto im = p[8 * k + 4 + j];
+          double re = p[8 * k + j];
+          double im = p[8 * k + 4 + j];
           csum += re * re + im * im;
           while (rs[m] < csum && m < num_samples) {
             bitstrings.emplace_back(4 * k + j);
@@ -365,6 +365,10 @@ class StateSpaceSSE :
           }
         }
       }
+
+      for (; m < num_samples; ++m) {
+        bitstrings.emplace_back((uint64_t{1} << state.num_qubits()) - 1);
+      }
     }
 
     return bitstrings;