bb84/run_bb84_full.py at main · codr-void/bb84 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
import hashlib
import numpy as np

from braket.circuits import noises
from braket.devices import LocalSimulator

from utils.bb84 import initialize_protocol, encode_qubits, measure_qubits, filter_qubits
from utils.golay_code import GolayCode

# -------------------------
# BB84 full-pipeline config
# -------------------------
BIT_FLIP_PROBABILITY = 0.10
NUMBER_OF_QUBITS_PER_ROUND = 12
TARGET_SIFTED_BITS = 12  # one Golay message block


def privacy_amplification_sha256(key_bits: np.ndarray, out_bits: int = 128) -> str:
    """Hash-based privacy amplification demo (truncate SHA-256)."""
    b = "".join(str(int(x)) for x in key_bits.astype(int)).encode("utf-8")
    digest = hashlib.sha256(b).digest()
    out_bytes = (out_bits + 7) // 8
    return digest[:out_bytes].hex()


def main():
    golay = GolayCode()

    alice_sifted_all = np.array([], dtype=np.uint8)
    bob_sifted_all = np.array([], dtype=np.uint8)
    rounds = 0

    # Use density-matrix sim when noise is enabled; otherwise statevector is fine.
    device = LocalSimulator("braket_dm") if BIT_FLIP_PROBABILITY > 0 else LocalSimulator("default")

    while alice_sifted_all.size < TARGET_SIFTED_BITS:
        rounds += 1

        encoding_basis_A, states_A, _ = initialize_protocol(NUMBER_OF_QUBITS_PER_ROUND)
        _, _, measurement_basis_B = initialize_protocol(NUMBER_OF_QUBITS_PER_ROUND)

        circuit = encode_qubits(NUMBER_OF_QUBITS_PER_ROUND, states_A, encoding_basis_A)

        if BIT_FLIP_PROBABILITY > 0:
            noise = noises.BitFlip(probability=BIT_FLIP_PROBABILITY)
            circuit.apply_gate_noise(noise)

        circuit = measure_qubits(circuit, measurement_basis_B)
        circuit.measure(list(range(NUMBER_OF_QUBITS_PER_ROUND)))

        result = device.run(circuit, shots=1).result()
        measured_bits = np.array(result.measurements[0], dtype=np.uint8)

        a_keep = filter_qubits(states_A, encoding_basis_A, measurement_basis_B)
        b_keep = filter_qubits(measured_bits, encoding_basis_A, measurement_basis_B)

        alice_sifted_all = np.concatenate([alice_sifted_all, a_keep.astype(np.uint8)])
        bob_sifted_all = np.concatenate([bob_sifted_all, b_keep.astype(np.uint8)])

    alice_raw = alice_sifted_all[:TARGET_SIFTED_BITS]
    bob_raw = bob_sifted_all[:TARGET_SIFTED_BITS]

    qber = float(np.mean(alice_raw != bob_raw))

    print("\n=== SIFTED (RAW) KEYS ===")
    print("Rounds used:", rounds)
    print("Alice raw (12):", alice_raw.astype(int).tolist())
    print("Bob raw   (12):", bob_raw.astype(int).tolist())
    print(f"QBER before error correction: {qber:.3f} ({int(np.sum(alice_raw != bob_raw))}/{TARGET_SIFTED_BITS})")

    # --- Golay reconciliation (Alice sends parity bits) ---
    codeword_A = golay.encode(alice_raw)
    parity_bits = codeword_A[12:]

    received_at_bob = np.concatenate([bob_raw, parity_bits]).astype(np.uint8)
    corrected_bob = golay.decode_to_message(received_at_bob)

    print("\n=== GOLAY RECONCILIATION ===")
    print("Parity bits sent (12):", parity_bits.astype(int).tolist())
    print("Bob corrected (12):   ", corrected_bob.astype(int).tolist())
    print("Correction success?:  ", bool(np.all(corrected_bob == alice_raw)))

    # --- Privacy amplification (demo) ---
    final_key_hex = privacy_amplification_sha256(alice_raw, out_bits=128)
    print("\n=== PRIVACY AMPLIFICATION (DEMO) ===")
    print("Final 128-bit key (hex):", final_key_hex)


if __name__ == "__main__":
    main()