Expand quantum simulator gate set and partial measurement

README.md

@@ -14,11 +14,13 @@ ten famous unsolved problems in mathematics:
 
 1. **Quantum Computing Simulation** – a minimal yet functional simulator for quantum
    circuits written in pure Python (found in `quantum_simulator.py`). The simulator
-   supports a small set of single-qubit gates (Hadamard and Pauli-X), a
+   supports a broad set of single-qubit gates (Hadamard, Pauli-X/Y/Z, S, T and
-   two-qubit controlled-NOT (CNOT) gate and measurement. You can create
+   arbitrary rotations about the X/Y/Z axes), controlled operations such as CNOT
-   circuits, apply gates, and measure qubits to observe the probabilistic outcomes
+   and controlled-Z, custom statevector initialisation and both full and partial
-   expected from quantum mechanics. The code uses the vector–state model implemented
+   measurement. You can create circuits, apply gates, and measure qubits to observe
-   directly with Python lists and complex numbers, so it has no third-party dependencies.
+   the probabilistic outcomes expected from quantum mechanics. The code uses the
+   vector–state model implemented directly with Python lists and complex numbers,
+   so it has no third-party dependencies.
 
 2. **Energy and Particle Simulation** – a simple set of utilities (in
    `energy_simulator.py`) that model energy generation and consumption as well as
@@ -59,19 +61,22 @@ You can run the modules directly or import the functions into your own scripts.
 For example, to create a simple quantum circuit:
 
 ```python
+import math
 from native_ai_quantum_energy.quantum_simulator import QuantumCircuit
 
 # Create a two-qubit circuit
 qc = QuantumCircuit(2)
 
-# Put the first qubit into superposition and entangle with the second qubit
+# Put the first qubit into superposition, rotate the second qubit and entangle them
 qc.apply_hadamard(0)
-qc.apply_cnot(0, 1)
+qc.apply_ry(1, math.pi / 4)
+qc.apply_cz(0, 1)
 
-# Measure both qubits
+# Measure only the control qubit while leaving the target unmeasured
-qc.measure_all()
+subset_result = qc.measure_subset([0])
 
-print("Measurement results:", qc.measurements)
+print("Measured control qubit:", subset_result)
+print("Statevector after measurement:", qc.statevector())
 ```
 
 Similarly, to simulate energy production:

native_ai_quantum_energy/quantum_simulator.py

@@ -2,9 +2,12 @@
 
 This module implements a tiny quantum circuit simulator based on the
 state-vector formalism.  It supports constructing a register of qubits,
-applying single-qubit gates (Hadamard and Pauli-X) and a controlled-NOT (CNOT)
+applying a range of single-qubit gates (Hadamard, the Pauli set, phase gates
-gate, and performing measurements.  The simulator is intentionally small and
+and arbitrary rotations) and multi-qubit controlled operations such as CNOT and
-focused on clarity rather than performance.
+controlled-Z.  The simulator can also be initialised from an arbitrary state
+vector, and supports measuring either all qubits or only a subset of them.
+The simulator is intentionally small and focused on clarity rather than
+performance.
 
 Example
 -------
@@ -46,6 +49,13 @@ class QuantumCircuit:
         (1 / math.sqrt(2) + 0j, -(1 / math.sqrt(2)) + 0j),
     )
     X_GATE: GateMatrix = ((0 + 0j, 1 + 0j), (1 + 0j, 0 + 0j))
+    Y_GATE: GateMatrix = ((0 + 0j, -1j), (1j, 0 + 0j))
+    Z_GATE: GateMatrix = ((1 + 0j, 0 + 0j), (0 + 0j, -1 + 0j))
+    S_GATE: GateMatrix = ((1 + 0j, 0 + 0j), (0 + 0j, 1j))
+    T_GATE: GateMatrix = (
+        (1 + 0j, 0 + 0j),
+        (0 + 0j, math.cos(math.pi / 4) + 1j * math.sin(math.pi / 4)),
+    )
 
     def __init__(self, num_qubits: int) -> None:
         if num_qubits < 1:
@@ -55,7 +65,7 @@ class QuantumCircuit:
         dim = 2 ** num_qubits
         self.state: List[complex] = [0j] * dim
         self.state[0] = 1.0 + 0j
-        # Store measurement results
+        # Store the outcomes of the most recent measurement call
         self.measurements: List[int] = []
 
     def _apply_single_qubit_gate(self, gate: GateMatrix, qubit: int) -> None:
@@ -91,6 +101,62 @@ class QuantumCircuit:
 
         self._apply_single_qubit_gate(self.X_GATE, qubit)
 
+    def apply_pauli_y(self, qubit: int) -> None:
+        """Apply a Pauli-Y gate to one qubit."""
+
+        self._apply_single_qubit_gate(self.Y_GATE, qubit)
+
+    def apply_pauli_z(self, qubit: int) -> None:
+        """Apply a Pauli-Z gate to one qubit."""
+
+        self._apply_single_qubit_gate(self.Z_GATE, qubit)
+
+    def apply_s(self, qubit: int) -> None:
+        """Apply the phase (S) gate to one qubit."""
+
+        self._apply_single_qubit_gate(self.S_GATE, qubit)
+
+    def apply_t(self, qubit: int) -> None:
+        """Apply the T (π/8) gate to one qubit."""
+
+        self._apply_single_qubit_gate(self.T_GATE, qubit)
+
+    def apply_rx(self, qubit: int, angle: float) -> None:
+        """Apply a rotation around the X axis by ``angle`` radians."""
+
+        half = angle / 2.0
+        cos = math.cos(half)
+        sin = math.sin(half)
+        gate: GateMatrix = (
+            (cos + 0j, -1j * sin),
+            (-1j * sin, cos + 0j),
+        )
+        self._apply_single_qubit_gate(gate, qubit)
+
+    def apply_ry(self, qubit: int, angle: float) -> None:
+        """Apply a rotation around the Y axis by ``angle`` radians."""
+
+        half = angle / 2.0
+        cos = math.cos(half)
+        sin = math.sin(half)
+        gate: GateMatrix = (
+            (cos + 0j, -sin + 0j),
+            (sin + 0j, cos + 0j),
+        )
+        self._apply_single_qubit_gate(gate, qubit)
+
+    def apply_rz(self, qubit: int, angle: float) -> None:
+        """Apply a rotation around the Z axis by ``angle`` radians."""
+
+        half = angle / 2.0
+        phase_neg = math.cos(half) - 1j * math.sin(half)
+        phase_pos = math.cos(half) + 1j * math.sin(half)
+        gate: GateMatrix = (
+            (phase_neg, 0 + 0j),
+            (0 + 0j, phase_pos),
+        )
+        self._apply_single_qubit_gate(gate, qubit)
+
     def apply_cnot(self, control: int, target: int) -> None:
         """Apply a controlled-NOT (CNOT) gate.
 
@@ -114,6 +180,22 @@ class QuantumCircuit:
             new_state[new_index] += amplitude
         self.state = new_state
 
+    def apply_cz(self, control: int, target: int) -> None:
+        """Apply a controlled-Z (CZ) gate."""
+
+        if control == target:
+            raise ValueError("Control and target must be different for CZ")
+        if any(q < 0 or q >= self.num_qubits for q in (control, target)):
+            raise IndexError("Qubit index out of range")
+
+        control_mask = 1 << (self.num_qubits - 1 - control)
+        target_mask = 1 << (self.num_qubits - 1 - target)
+        new_state = self.state.copy()
+        for index, amplitude in enumerate(self.state):
+            if (index & control_mask) and (index & target_mask):
+                new_state[index] = -amplitude
+        self.state = new_state
+
     def measure(self, qubit: int) -> int:
         """Measure a single qubit and collapse the state.
 
@@ -122,33 +204,10 @@ class QuantumCircuit:
         observed result.
         """
 
-        if qubit < 0 or qubit >= self.num_qubits:
+        outcome, new_state = self._collapse_qubits([qubit])
-            raise IndexError("Qubit index out of range")
-
-        prob0 = 0.0
-        prob1 = 0.0
-        mask = 1 << (self.num_qubits - 1 - qubit)
-        for idx, amp in enumerate(self.state):
-            if idx & mask:
-                prob1 += abs(amp) ** 2
-            else:
-                prob0 += abs(amp) ** 2
-
-        rand = random.random()
-        outcome = 1 if rand < prob1 else 0
-
-        new_state = [0j] * len(self.state)
-        for idx, amp in enumerate(self.state):
-            bit = 1 if idx & mask else 0
-            if bit == outcome:
-                new_state[idx] = amp
-
-        norm = math.sqrt(prob1 if outcome == 1 else prob0)
-        if norm > 0:
-            new_state = [amp / norm for amp in new_state]
-
         self.state = new_state
-        return outcome
+        self.measurements = [int(outcome)] if outcome else []
+        return int(outcome) if outcome else 0
 
     def measure_all(self) -> str:
         """Measure all qubits sequentially, returning a bitstring.
@@ -157,14 +216,31 @@ class QuantumCircuit:
         ``measurements`` and returned as a concatenated string (qubit 0 first).
         """
 
-        self.measurements = []
         bits: List[int] = []
+        self.measurements = []
         for q in range(self.num_qubits):
-            bit = self.measure(q)
+            outcome, new_state = self._collapse_qubits([q])
+            bit = int(outcome) if outcome else 0
             bits.append(bit)
             self.measurements.append(bit)
+            self.state = new_state
         return "".join(str(b) for b in bits)
 
+    def measure_subset(self, qubits: Sequence[int]) -> str:
+        """Measure only the specified ``qubits`` and collapse the state.
+
+        The ``qubits`` sequence is interpreted in the provided order, and the
+        returned bitstring follows the same ordering.  The amplitudes of basis
+        states incompatible with the sampled outcome are set to zero while the
+        remaining amplitudes are renormalised.  Unmeasured qubits retain their
+        relative amplitudes, preserving entanglement when possible.
+        """
+
+        outcome, new_state = self._collapse_qubits(qubits)
+        self.state = new_state
+        self.measurements = [int(bit) for bit in outcome]
+        return outcome
+
     def statevector(self) -> List[complex]:
         """Return a copy of the current state vector."""
 
@@ -179,3 +255,57 @@ class QuantumCircuit:
         """Return a tuple view of the current amplitudes."""
 
         return tuple(self.state)
+
+    def initialize_statevector(self, amplitudes: Sequence[complex]) -> None:
+        """Initialise the circuit with a custom ``amplitudes`` state vector."""
+
+        expected = 2 ** self.num_qubits
+        if len(amplitudes) != expected:
+            raise ValueError(
+                f"State vector must have length {expected}, got {len(amplitudes)}"
+            )
+        norm = sum(abs(amp) ** 2 for amp in amplitudes)
+        if not math.isclose(norm, 1.0, rel_tol=1e-9, abs_tol=1e-9):
+            raise ValueError("State vector must be normalised to 1.0")
+        self.state = [complex(amp) for amp in amplitudes]
+
+    def _collapse_qubits(self, qubits: Sequence[int]) -> Tuple[str, List[complex]]:
+        """Return outcome and collapsed state for measuring ``qubits``."""
+
+        if any(q < 0 or q >= self.num_qubits for q in qubits):
+            raise IndexError("Qubit index out of range")
+        if len(set(qubits)) != len(qubits):
+            raise ValueError("Qubit indices must be unique for measurement")
+        if not qubits:
+            return "", self.state.copy()
+
+        outcome_probs: dict[Tuple[int, ...], float] = {}
+        for index, amplitude in enumerate(self.state):
+            bits = tuple((index >> (self.num_qubits - 1 - q)) & 1 for q in qubits)
+            outcome_probs[bits] = outcome_probs.get(bits, 0.0) + abs(amplitude) ** 2
+
+        rand = random.random()
+        cumulative = 0.0
+        chosen_outcome: Tuple[int, ...] | None = None
+        sorted_outcomes = sorted(outcome_probs.items(), key=lambda item: item[0], reverse=True)
+        for bits, prob in sorted_outcomes:
+            cumulative += prob
+            if rand < cumulative:
+                chosen_outcome = bits
+                break
+        if chosen_outcome is None:
+            # Fallback in case of floating-point accumulation error.
+            chosen_outcome = sorted_outcomes[-1][0]
+
+        new_state = [0j] * len(self.state)
+        for index, amplitude in enumerate(self.state):
+            bits = tuple((index >> (self.num_qubits - 1 - q)) & 1 for q in qubits)
+            if bits == chosen_outcome:
+                new_state[index] = amplitude
+
+        prob = outcome_probs[chosen_outcome]
+        if prob > 0:
+            norm = math.sqrt(prob)
+            new_state = [amp / norm for amp in new_state]
+
+        return "".join(str(bit) for bit in chosen_outcome), new_state

tests/test_quantum_simulator.py

@@ -27,6 +27,37 @@ def test_pauli_x_flips_qubit():
     assert_complex_approx(state[1], 1.0)
 
 
+def test_pauli_y_adds_phase_to_one_state():
+    qc = QuantumCircuit(1)
+    qc.apply_pauli_y(0)
+    state = qc.statevector()
+    assert_complex_approx(state[0], 0.0)
+    assert_complex_approx(state[1], 0.0, 1.0)
+
+
+def test_pauli_z_flips_phase_of_one_state():
+    qc = QuantumCircuit(1)
+    qc.apply_pauli_x(0)
+    qc.apply_pauli_z(0)
+    state = qc.statevector()
+    assert_complex_approx(state[0], 0.0)
+    assert_complex_approx(state[1], -1.0)
+
+
+def test_phase_gates_apply_expected_phases():
+    qc = QuantumCircuit(1)
+    qc.apply_pauli_x(0)
+    qc.apply_s(0)
+    state = qc.statevector()
+    assert_complex_approx(state[1], 0.0, 1.0)
+
+    qc.apply_t(0)
+    state = qc.statevector()
+    expected = math.cos(math.pi / 4)
+    expected_imag = math.sin(math.pi / 4)
+    assert_complex_approx(state[1], -expected_imag, expected)
+
+
 def test_cnot_entangles_qubits():
     qc = QuantumCircuit(2)
     qc.apply_hadamard(0)
@@ -38,6 +69,17 @@ def test_cnot_entangles_qubits():
     assert probs[2] == pytest.approx(0.0, abs=1e-12)
 
 
+def test_controlled_z_adds_phase_to_11_state():
+    qc = QuantumCircuit(2)
+    qc.apply_hadamard(0)
+    qc.apply_hadamard(1)
+    qc.apply_cz(0, 1)
+    state = qc.statevector()
+    amplitude_11 = state[3]
+    assert amplitude_11.real == pytest.approx(-0.5, rel=1e-6)
+    assert amplitude_11.imag == pytest.approx(0.0, abs=1e-12)
+
+
 def test_measure_collapses_state(monkeypatch):
     qc = QuantumCircuit(1)
     qc.apply_hadamard(0)
@@ -59,3 +101,96 @@ def test_cnot_requires_distinct_qubits():
     qc = QuantumCircuit(2)
     with pytest.raises(ValueError):
         qc.apply_cnot(0, 0)
+
+
+def test_rotation_gates_match_expected_unitaries():
+    qc = QuantumCircuit(1)
+    initial = [1 / math.sqrt(2), 1 / math.sqrt(2)]
+    qc.initialize_statevector(initial)
+
+    angle = math.pi / 3
+    qc.apply_rx(0, angle)
+    cos = math.cos(angle / 2)
+    sin = math.sin(angle / 2)
+    expected_rx = [
+        cos * initial[0] - 1j * sin * initial[1],
+        -1j * sin * initial[0] + cos * initial[1],
+    ]
+    for actual, expected in zip(qc.statevector(), expected_rx):
+        assert actual == pytest.approx(expected)
+
+    qc.initialize_statevector(initial)
+    qc.apply_ry(0, angle)
+    expected_ry = [
+        math.cos(angle / 2) * initial[0] - math.sin(angle / 2) * initial[1],
+        math.sin(angle / 2) * initial[0] + math.cos(angle / 2) * initial[1],
+    ]
+    for actual, expected in zip(qc.statevector(), expected_ry):
+        assert actual == pytest.approx(expected)
+
+    qc.initialize_statevector(initial)
+    qc.apply_rz(0, angle)
+    phase_neg = math.cos(angle / 2) - 1j * math.sin(angle / 2)
+    phase_pos = math.cos(angle / 2) + 1j * math.sin(angle / 2)
+    expected_rz = [phase_neg * initial[0], phase_pos * initial[1]]
+    for actual, expected in zip(qc.statevector(), expected_rz):
+        assert actual == pytest.approx(expected)
+
+
+def test_rotation_identity_cases():
+    qc = QuantumCircuit(1)
+    qc.apply_rx(0, 0.0)
+    assert qc.statevector()[0] == pytest.approx(1.0)
+    qc.apply_ry(0, 0.0)
+    assert qc.statevector()[0] == pytest.approx(1.0)
+    qc.apply_rz(0, 0.0)
+    assert qc.probabilities() == pytest.approx([1.0, 0.0])
+
+    qc.apply_hadamard(0)
+    before = qc.probabilities()
+    qc.apply_rx(0, 2 * math.pi)
+    qc.apply_ry(0, 2 * math.pi)
+    qc.apply_rz(0, 2 * math.pi)
+    after = qc.probabilities()
+    assert after == pytest.approx(before)
+
+
+def test_measure_subset_collapses_only_requested_qubits(monkeypatch):
+    qc = QuantumCircuit(2)
+    qc.apply_hadamard(0)
+    monkeypatch.setattr(random, "random", lambda: 0.2)
+    outcome = qc.measure_subset([1])
+    assert outcome == "0"
+    state = qc.statevector()
+    assert state[0] == pytest.approx(1 / math.sqrt(2))
+    assert state[2] == pytest.approx(1 / math.sqrt(2))
+    assert state[1] == pytest.approx(0.0)
+    assert state[3] == pytest.approx(0.0)
+
+
+def test_measure_subset_on_entangled_pair(monkeypatch):
+    qc = QuantumCircuit(2)
+    qc.apply_hadamard(0)
+    qc.apply_cnot(0, 1)
+    monkeypatch.setattr(random, "random", lambda: 0.25)
+    outcome = qc.measure_subset([0])
+    assert outcome == "1"
+    state = qc.statevector()
+    assert state[3] == pytest.approx(1.0 + 0j)
+    for idx in (0, 1, 2):
+        assert state[idx] == pytest.approx(0.0 + 0j)
+
+
+def test_initialize_statevector_validates_input():
+    qc = QuantumCircuit(2)
+    with pytest.raises(ValueError):
+        qc.initialize_statevector([1.0, 0.0])
+
+    with pytest.raises(ValueError):
+        qc.initialize_statevector([0.5, 0.5, 0.5, 0.2])
+
+    amplitudes = [0.5 + 0j, 0.5j, -0.5 + 0j, 0.5 - 0.5j]
+    norm = math.sqrt(sum(abs(a) ** 2 for a in amplitudes))
+    amplitudes = [a / norm for a in amplitudes]
+    qc.initialize_statevector(amplitudes)
+    assert qc.statevector() == pytest.approx(amplitudes)