CKKS-encode-encrypt.py

import numpy as np
from numpy.polynomial import Polynomial

# Set the parameters
M = 8
N = M // 2
scale = 64
xi = np.exp(2 * np.pi * 1j / M)

def vandermonde(xi: np.complex128, M: int) -> np.array:
    N = M // 2
    matrix = []
    for i in range(N):
        root = xi ** (2 * i + 1)
        row = [root ** j for j in range(N)]
        matrix.append(row)
    return matrix

def sigma_inverse(xi, M, b: np.array) -> Polynomial:
    A = vandermonde(xi, M)
    coeffs = np.linalg.solve(A, b)
    p = Polynomial(coeffs)
    return p

def sigma(xi, M, p: Polynomial) -> np.array:
    outputs = []
    N = M // 2
    for i in range(N):
        root = xi ** (2 * i + 1)
        output = p(root)
        outputs.append(output)
    return np.array(outputs)

def create_sigma_R_basis(xi, M):
    return np.array(vandermonde(xi, M)).T

def sigma_R_discretization(xi, M, z):
    sigma_R_basis = create_sigma_R_basis(xi, M)
    coordinates = compute_basis_coordinates(sigma_R_basis, z)
    rounded_coordinates = coordinate_wise_random_rounding(coordinates)
    y = np.matmul(sigma_R_basis.T, rounded_coordinates)
    return y

def compute_basis_coordinates(sigma_R_basis, z):
    return np.array([np.real(np.vdot(z, b) / np.vdot(b,b)) for b in sigma_R_basis])

def round_coordinates(coordinates):
    return coordinates - np.floor(coordinates)

def coordinate_wise_random_rounding(coordinates):
    r = round_coordinates(coordinates)
    f = np.array([np.random.choice([c, c-1], 1, p=[1-c, c]) for c in r]).reshape(-1)
    rounded_coordinates = coordinates - f
    return [int(coeff) for coeff in rounded_coordinates]

def pi(M, z: np.array) -> np.array:
    N = M // 4
    return z[:N]

def pi_inverse(z: np.array) -> np.array:
    z_conjugate = [np.conjugate(x) for x in z[::-1]]
    return np.concatenate([z, z_conjugate])

def encode(xi, M, scale, z: np.array) -> Polynomial:
    pi_z = pi_inverse(z)
    scaled_pi_z = scale * pi_z
    rounded_scale_pi_zi = sigma_R_discretization(xi, M, scaled_pi_z)
    p = sigma_inverse(xi, M, rounded_scale_pi_zi)
    coef = np.round(np.real(p.coef)).astype(int)
    return Polynomial(coef)

def decode(xi, M, scale, p: Polynomial) -> np.array:
    rescaled_p = p / scale
    z = sigma(xi, M, rescaled_p)
    return pi(M, z)

# Example usage
z = np.array([3 + 4j, 2 - 1j])
p = encode(xi, M, scale, z)
decoded_z = decode(xi, M, scale, p)
print(decoded_z)