use ideal counts in measure

README.md

@@ -45,10 +45,12 @@ simulator.u(2, 0.3, 0.4, 0.5)  # Generic gate
 Measure the state of the qubits:
 
 ```python
-simulator.run(shots=100)
+print(simulator.run(shots=100))
 ```
 
-> {'000': 51, '101': 41, '001': 6, '100': 2}
+```plaintext
+{'000': 46, '001': 4, '100': 4, '101': 46}
+```
 
 ### Circuit Representation
 
@@ -74,6 +76,8 @@ Show the amplitude and phase of all quantum states:
 simulator.plot_wavefunction()
 ```
 
+> Note: This method requires `matplotlib`
+
 ![Wavefunction Scatter Plot](https://github.com/splch/qubit-simulator/assets/25377399/de3242ef-9c14-44be-b49b-656e9727c618)
 
 ## Testing

qubit_simulator/simulator.py

@@ -182,7 +182,7 @@ def cu(
     def measure(self, shots: int = 1, basis: Optional[np.ndarray] = None) -> List[str]:
         """
         Measures the state of the qubits.
-
+    
         :param shots: Number of measurements.
         :param basis: Optional basis transformation.
         :return: List of measurement results.
@@ -192,10 +192,21 @@ def measure(self, shots: int = 1, basis: Optional[np.ndarray] = None) -> List[st
             raise ValueError("Number of shots must be non-negative.")
         state_vector = self.state_vector
         if basis is not None:
-            state_vector = basis @ self.state_vector
-        probabilities = np.abs(state_vector) ** 2
-        result = np.random.choice(2**self.num_qubits, p=probabilities, size=shots)
-        return [format(r, f"0{self.num_qubits}b") for r in result]
+            state_vector = basis @ state_vector
+        probabilities = np.abs(state_vector)**2
+        ideal_counts = np.round(probabilities * shots).astype(int)
+        unique_states = np.array([format(i, f"0{self.num_qubits}b") for i in range(len(ideal_counts))])
+        results = np.repeat(unique_states, ideal_counts)
+        total_counts = np.sum(ideal_counts)
+        diff = total_counts - shots
+        if diff > 0:
+            remove_indices = np.random.choice(len(results), size=diff, replace=False)
+            results = np.delete(results, remove_indices)
+        elif diff < 0:
+            add_indices = np.random.choice(len(ideal_counts), size=-diff, p=probabilities)
+            additional_results = unique_states[add_indices]
+            results = np.concatenate([results, additional_results])
+        return results.tolist()
 
     def run(
         self, shots: int = 100, basis: Optional[np.ndarray] = None

requirements.txt

@@ -1 +1,2 @@
-numpy>=1.16.0
+numpy
+matplotlib

setup.py

@@ -9,8 +9,7 @@
     author="Spencer Churchill",
     author_email="[email protected]",
     packages=find_packages(),
-    install_requires=["numpy"],
-    extras_require={"all": ["matplotlib"]},
+    install_requires=["numpy", "matplotlib"],
     classifiers=[
         "Development Status :: 5 - Production/Stable",
         "Intended Audience :: Science/Research",

tests/test_simulator.py

@@ -361,15 +361,6 @@ def apply_qft(simulator):
 
 
 def test_getsize():
-    simulator = QubitSimulator(3)
-    # Apply some gates to make the instance more complex
-    simulator.h(0)
-    simulator.u(1, np.pi / 4, np.pi / 4, np.pi / 2)
-    simulator.cx(1, 2)
-    assert simulator.__getsize__() == 412
-
-
-def test_getsize_relative():
     simulator = QubitSimulator(2)
     initial_size = simulator.__getsize__()
     simulator.h(0)