qqz.discrete_log

  1from typing import Optional
  2import math
  3
  4import numpy as np
  5import matplotlib.pyplot as plt
  6from sympy import gcdex
  7from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, Aer, transpile, assemble
  8from qiskit.visualization import plot_histogram
  9from sympy import Rational, gcdex
 10from sympy.ntheory.continued_fraction import continued_fraction, continued_fraction_convergents
 11
 12from qft import qft
 13from elementary import ax_modM
 14from order_finding import order_finding
 15from classical_utils import decode_bin
 16
 17def discrete_log(a: int, b: int, p: int, show_hist: Optional[bool] = False) -> int:
 18    """Shor's discrete log algorithm: given $a,b,p\in\mathbb{Z}$, it finds $s$ such that $a^s\equiv b\pmod p$.
 19
 20    Args:
 21        a (int): $a$
 22        b (int): $b$
 23        p (int): $p$
 24    """
 25
 26    r = order_finding(x=a, N=p, show_hist=False)
 27    t = int(np.ceil(np.log2(p)))
 28
 29    first_register = QuantumRegister(t)
 30    second_register = QuantumRegister(t)
 31    third_register = QuantumRegister(2 * t)
 32    auxiliary_register_mid = QuantumRegister(t)
 33    auxiliary_register_end = QuantumRegister(6 * t - 2)
 34    classical_register = ClassicalRegister(2 * t)
 35
 36    qc = QuantumCircuit(
 37        first_register,
 38        second_register,
 39        third_register,
 40        auxiliary_register_mid,
 41        auxiliary_register_end,
 42        classical_register,
 43    )
 44
 45    qc.h(first_register)
 46    qc.h(second_register)
 47
 48    qc.append(ax_modM(a=b, M=p, N_len=t), list(first_register) + list(auxiliary_register_mid) + list(third_register) + list(auxiliary_register_end))
 49    qc.append(ax_modM(a=a, M=p, N_len=t, x_0_at_first=False), list(second_register) + list(auxiliary_register_mid) + list(third_register) + list(auxiliary_register_end))
 50
 51    qc.append(qft(n=t).inverse(), first_register)
 52    qc.append(qft(n=t).inverse(), second_register)
 53
 54    qc.measure(list(first_register) + list(second_register), classical_register)
 55    #qc.measure(third_register, classical_register)
 56
 57    backend = Aer.get_backend('aer_simulator_matrix_product_state')#('aer_simulator')
 58    qc = transpile(qc, backend)
 59    job = backend.run(qc, shots=10000)
 60    hist = job.result().get_counts()
 61
 62    if show_hist:
 63        figsize_x = max(7 * (len(hist) // 8), 7)
 64        plot_histogram(hist, figsize=(figsize_x, 5))
 65        plt.savefig(f'img/discrete_log_a{a}_b{b}_p{p}_r{r}.png', bbox_inches='tight')
 66
 67    for measured_key, _ in sorted(hist.items(), key=lambda x: x[1], reverse=True):
 68        tilde_l_per_r = Rational(decode_bin(measured_key[:t]), 2 ** t) # decoded from second register: $\widetilde{l/r}$
 69        if tilde_l_per_r == 0:
 70            continue
 71
 72        l = None
 73        for fraction in continued_fraction_convergents(continued_fraction(tilde_l_per_r)):
 74            if fraction.denominator == r:
 75                l = fraction.numerator # get correct $l$
 76                break
 77
 78        if l is None:
 79            continue
 80
 81        tilde_beta_per_r = Rational(decode_bin(measured_key[-t:]), 2 ** t) # decoded from first register: $\widetilde{\beta/r}$
 82        if tilde_beta_per_r == 0:
 83            continue
 84
 85        beta = None
 86        for fraction in continued_fraction_convergents(continued_fraction(tilde_beta_per_r)):
 87            if fraction.denominator == r:
 88                beta = fraction.numerator # get correct $\beta$
 89                break
 90
 91        if beta is None:
 92            continue
 93
 94        s, alpha, _ = gcdex(l, -r)
 95        s = int(s * beta)
 96        if pow(a, s, p) == b:
 97            return s
 98
 99    raise Exception('s is NOT found!')
100
101
102if __name__ == '__main__':
103    print(discrete_log(a=2, b=4, p=7, show_hist=True))