Fix apply gate + more tests

README.md

@@ -61,7 +61,7 @@ simulator.measure()
 ## Tests
 
 ```shell
-pytest
+python -m pytest
 ```
 
 ## License

qubit_simulator/simulator.py

@@ -5,48 +5,37 @@
 class QubitSimulator:
     def __init__(self, num_qubits):
         self.num_qubits = num_qubits
-        self.state_vector = np.zeros(2**num_qubits)
+        self.state_vector = np.zeros(2 ** num_qubits, dtype=complex)
         self.state_vector[0] = 1
 
-    def _apply_gate(self, gate, target_qubit, control_qubit=None):
-        if control_qubit is not None:  # Handle controlled gates
-            # Create the controlled gate matrix
-            control_gate = np.kron(
-                gates.I
-                - np.outer(gates.I[:, control_qubit], gates.I[control_qubit, :]),
-                gate,
-            )
-            control_gate += np.kron(
-                np.outer(gates.I[:, control_qubit], gates.I[control_qubit, :]), gates.I
-            )
-            # Create the operator using Kronecker products
-            operator = 1
-            for qubit in range(self.num_qubits):
-                if qubit == control_qubit or qubit == target_qubit:
-                    continue
-                operator = np.kron(operator, gates.I)
-            operator = np.kron(operator, control_gate)
-        else:  # Handle single qubit gates
-            operator = 1
-            for qubit in range(self.num_qubits):
-                if qubit == target_qubit:
-                    operator = np.kron(operator, gate)
-                else:
-                    operator = np.kron(operator, gates.I)
-        # Apply the operator to the state vector
-        self.state_vector = np.dot(operator, self.state_vector)
+    def _apply_gate(self, gate, target_qubits, control_qubit=None):
+        op_matrix = 1
+        target_idx = 0
+        for qubit in range(self.num_qubits):
+            current_gate = gate[target_idx] if qubit in target_qubits else gates.I
+            op_matrix = np.kron(op_matrix, current_gate)
+            if qubit in target_qubits:
+                target_idx += 1
+        if control_qubit is not None:
+            controlled_op_matrix = np.eye(2 ** self.num_qubits, dtype=complex)
+            for state in range(2 ** self.num_qubits):
+                binary_state = format(state, f"0{self.num_qubits}b")
+                if all(binary_state[c] == '1' for c in control_qubit):
+                    controlled_op_matrix[state, :] = np.dot(op_matrix, controlled_op_matrix[state, :])
+            op_matrix = controlled_op_matrix
+        self.state_vector = np.dot(op_matrix, self.state_vector)
 
     def H(self, target_qubit):
-        self._apply_gate(gates.H, target_qubit)
+        self._apply_gate([gates.H], [target_qubit])
 
     def T(self, target_qubit):
-        self._apply_gate(gates.T, target_qubit)
+        self._apply_gate([gates.T], [target_qubit])
 
     def X(self, target_qubit):
-        self._apply_gate(gates.X, target_qubit)
+        self._apply_gate([gates.X], [target_qubit])
 
     def CNOT(self, control_qubit, target_qubit):
-        self._apply_gate(gates.X, target_qubit, control_qubit)
+        self._apply_gate([gates.X], [target_qubit], [control_qubit])
 
     def Measure(self):
         probabilities = np.abs(self.state_vector) ** 2

tests/test_simulator.py

@@ -1,6 +1,7 @@
 import pytest
 import numpy as np
 from qubit_simulator import QubitSimulator
+from qubit_simulator import gates
 
 
 def test_hadamard_gate():
@@ -26,6 +27,15 @@ def test_cnot_gate():
     assert np.allclose(simulator.state_vector, [0, 0, 1, 0])
 
 
+def test_target_control():
+    simulator = QubitSimulator(3)
+    simulator.X(0)
+    simulator.X(2)  # Set the initial state to |101>
+    simulator.CNOT(control_qubit=2, target_qubit=0)
+    # After applying the CNOT gate, the state should be |001>
+    assert np.allclose(simulator.state_vector, [0, 1, 0, 0, 0, 0, 0, 0])
+
+
 def test_measure():
     simulator = QubitSimulator(1)
     simulator.X(0)
@@ -47,3 +57,31 @@ def test_bell_state():
     simulator.CNOT(0, 1)
     # After applying the Hadamard and CNOT gates, the state should be a Bell state
     assert np.allclose(simulator.state_vector, [0.70710678, 0, 0, 0.70710678])
+
+
+def test_ghz_state():
+    simulator = QubitSimulator(3)
+    simulator.H(0)
+    simulator.CNOT(0, 1)
+    simulator.CNOT(0, 2)
+    # The state should be a GHZ state
+    assert np.allclose(simulator.state_vector, [0.70710678, 0, 0, 0, 0, 0, 0, 0.70710678])
+
+
+def test_custom_2qubit_gate():
+    custom_gate = np.kron(gates.X, gates.H)
+    simulator = QubitSimulator(2)
+    simulator._apply_gate([gates.X, gates.H], [0, 1])
+    assert np.allclose(simulator.state_vector, np.dot(custom_gate, [1, 0, 0, 0]))
+
+
+def test_qft():
+    simulator = QubitSimulator(3)
+    # Applying QFT manually
+    for k in range(3):
+        simulator.H(k)
+        for j in range(k + 1, 3):
+            phase = np.pi / (2 ** (j - k))
+            phase_shift = np.array([[1, 0], [0, np.exp(1j * phase)]])
+            simulator._apply_gate([phase_shift], [j], [k])
+    assert np.allclose(simulator.state_vector, [1 / np.sqrt(8)] * 8)