Implementation of the CVP algorithm for lattice of the Voronoi first kind from [Mckilliam 2014] in Python

mckilliam-cvp.py

import numpy as np
import graph_tool.all as gt

# B is the obtuse superbasis of a n-dimensional gadget lattice
# p is a (n+1)-dimensional vector
# returns t \in \{0,1\}^{n+1} such that ||B(p-t)||_2 is minimized
def next_step(B, p):
    # construct adjacency matrix
    n = B.shape[0]
    Q = B.T @ B
    s = -2 * (Q @ p)
    W = np.hstack((np.zeros((n+1, 1)), -Q, np.zeros((n+1, 1))))
    W = np.vstack((np.zeros((1, n+3)), W, np.zeros((1, n+3))))
    for i in range(1, n+2):
        s_i = s[i-1]
        if s_i >= 0:
            W[i][n+2] = s_i
            W[n+2][i] = s_i
            W[0][i] = 0
            W[i][0] = 0
        else:
            W[i][n+2] = 0
            W[n+2][i] = 0
            W[0][i] = -s_i
            W[i][0] = -s_i

    # construct graph
    g = gt.Graph()
    g.add_vertex(n+3)
    weight_prop = g.new_edge_property("double")
    g.ep.weight = weight_prop
    for i in range(n+3):
        for j in range(n+3):
            if W[i][j] > 0:
                edge = g.add_edge(g.vertex(i), g.vertex(j))
                g.ep.weight[edge] = W[i][j]

    # find min cut
    residual = gt.edmonds_karp_max_flow(g, g.vertex(0), g.vertex(n+2), g.ep.weight)
    part = gt.min_st_cut(g, g.vertex(0), g.ep.weight, residual)
    
    # find t
    t = np.zeros(n+1)
    for i in range(1, n+2):
        in_src_part = part[g.vertex(i)]
        if in_src_part:
            t[i-1] = 1
        else:
            t[i-1] = 0
    return t

# get the n-dimensional basis for the base-b gadget lattice
def get_gadget_superbasis(n, b):
    def get_base_entry(r, c):
        if r == c:
            return b
        elif r == c + 1:
            return -1
        else:
            return 0
    B_raw = np.array([[get_base_entry(r, c) for c in range(n)] for r in range(n)])
    # convert B into an obtuse superbasis
    B = np.hstack((B_raw, -np.sum(B_raw, axis=1).reshape((n, 1))))
    return B

# compute inverse(B) @ t
# B is the base-b gadget lattice basis
# faster than computing matrix inverse due to niceness of inverse(B)
def b_inv_mul(b, t):
    n = len(t)

    t1 = np.zeros(n)
    prev = 0
    for i in range(n):
        t1[i] = (prev + t[i]) / b
        prev = t1[i]

    return t1

# B is the obtuse superbasis of the lattice, y is the target point
# finds the closest lattice point to y
def cvp(B, y, next_step_func, epsilon=1e-9):
    # find z s.t. y = Bz
    n = B.shape[0]
    b = B[0][0]
    B_raw = B[:, :n]
    z = b_inv_mul(b, y)
    z = np.hstack((z, [0]))

    # main loop of the algorithm
    u = np.floor(z)
    dist = np.linalg.norm(B @ (z - u))
    for i in range(n):
        t = next_step_func(B, z - u)
        new_u = u + t
        new_dist = np.linalg.norm(B @ (z - new_u))
        if np.abs(new_dist - dist) <= epsilon:
            break
        u = new_u
        dist = new_dist
    return (B @ u, i)

# experiment
n = 6
b = 2
B = get_gadget_superbasis(n, b)
num_tests = 1000000
for i in range(num_tests):
    y = (2 * np.random.random(n) - 1) * 1e2
    x, num_iter = cvp(B, y, next_step)
    if num_iter > 1:
        print(y, x, num_iter)