unitary_quantization.py

import numpy as np
import sys

# Parameters: (w_i, theta_i)

n_theta_vals = 6

if len(sys.argv) > 1:
    n_theta_vals = int(sys.argv[1])

codebook = [
    (0, 0),
    (np.pi / 2, 0)
    ] + list(map(lambda x : (np.pi / 4, x), (np.arange(n_theta_vals) * 2 * np.pi / n_theta_vals).tolist()))

def flag_distance(a, b):
    return np.sin(b[0] - a[0])**2 + np.sin(2 * a[0]) * np.sin(2 * b[0]) * np.sin((b[1] - a[1]) / 2)**2

def parameterize_unitary_matrix(V):
    D = np.diag(np.exp(-1j * np.angle(np.diag(V))))
    V = V @ D
    w = np.arccos(np.real(V[0,0]))
    theta = np.angle(V[0,1])
    return (w, theta)

err = 0
for i in range(1000):
    H = np.random.randn(2,2) + 1j * np.random.randn(2,2)
    U, S, Vh = np.linalg.svd(H)
    V = Vh.conj().T
    D = np.diag(np.exp(-1j * np.angle(np.diag(V))))
    V = V @ D
    U = U @ D
    assert(np.max(np.abs(U @ np.diag(S) @ V.conj().T - H)) < 1e-10)
    w = np.arccos(np.real(V[0,0]))
    theta = np.angle(V[0,1])
    # Find the closest codebook element
    distances = np.array(list(map(lambda x : flag_distance((w, theta), x), codebook)))
    minidx = np.argmin(distances)
    minerr = np.min(distances)
    err = err + minerr
print(err / 1000)