Residual Vector Quantization – Scott H. Hawley

def find_closest_centroids(data, codebook):
    distances = np.sum((data[:, np.newaxis] - codebook) ** 2, axis=2)
    closest_centroids = np.argmin(distances, axis=1)
    return closest_centroids # indices not coordinates

def compute_error(data, codebook):
    n_grid = int(np.sqrt(codebook.shape[0]))
    h = 1/n_grid * (DATA_MAX - DATA_MIN) 
    quantized_ind = find_closest_centroids(data, codebook)
    quantized_xy = codebook[quantized_ind]       
    error = np.sqrt( np.sum( (data - quantized_xy)**2 ) )
    return error 

n_points2 = 100
data2 = DATA_MIN + (DATA_MAX - DATA_MIN)*np.random.rand(n_points2, 2)

errors = []
grids = np.array([5,10,25,100,200])
for n_grid2 in grids:
    codebook = generate_codebook(n_grid2)
    error = compute_error(data2, codebook)               
    errors.append(error)
    
fig, ax = plt.subplots(1,2, figsize=(10,4))    
ax[0].plot(grids, errors, 'o-')
ax[1].loglog(grids, errors, 'o-')
for i in range(2):
    ax[i].set_ylabel('Error')
    ax[i].set_xlabel('n_grid')
plt.show()
print(f"lowest error (for n_grid={grids[-1]}) = ",errors[-1])

from matplotlib import collections  as mc


def get_region_membership(data:np.array, h=0.2, codebook=None):
    "Tells which region each point is in. TBD: this is slow but it works! ;-) "
    memb = np.zeros(data.shape[0], dtype=np.int32)
    if codebook is None:  # just assume basic squares
        for di, p in enumerate(data):
            i = (p[0]-DATA_MIN) // h
            j = (p[1]-DATA_MIN) // h
            ind = i + j*n_grid
            memb[di] = ind
    else: 
        memb = calc_cluster_membership(data, codebook)
    return memb


def plot_data_grid_with_resids(data, n_grid=5, hide_tick_labels=True, show_indices=False, show_centroids=False, 
                   show_next_level_grid=False, show_resids=True, codebook=None, show_grid=True):
    "big ol' workhorse plotting routine that we'll progressively make use of as the lesson proceeds"
    fig, ax = plt.subplots(figsize=(3,3))
    h = 1.0 / n_grid
    ax.set_xlim(DATA_MIN, DATA_MAX)
    ax.set_ylim(DATA_MIN, DATA_MAX)

    if show_grid:
        for i in range(n_grid + 1):  
            ax.axhline(DATA_MIN + i*h, color='black')
            ax.axvline(DATA_MIN + i*h, color='black')
        
    if show_next_level_grid: # draws lines in the middle
        x_start = 2*h
        y_start = -h/2
        for i in range(n_grid): # horizontal lines
            y = y_start + i*h/n_grid
            ax.axhline(y, xmin=x_start, xmax=x_start+h, color='black')
        y_start, x_start = x_start, y_start
        for j in range(n_grid): # horizontal lines
            x = x_start + j*h/n_grid
            ax.axvline(x, ymin=y_start, ymax=y_start+h, color='black')
        
    if show_indices:
        index = 0
        for j in range(n_grid):
            for i in range(n_grid):
                x = DATA_MIN + (i + 0.5) / n_grid
                y = DATA_MIN + 1 - (j + 0.5) / n_grid
                ax.text(x, y, str(index), ha='center', va='center',fontsize=14)
                index += 1

    if show_centroids: 
        for j in range(n_grid):
            for i in range(n_grid):
                x = DATA_MIN + (i + 0.5)*h
                y = DATA_MIN + (j + 0.5)*h
                ax.plot(x, y, 'ro', markersize=6)
                
    if show_resids and codebook is not None:
        memb = get_region_membership(data, codebook=codebook)
        resids = data*0
        lines = []
        for i, p in enumerate(data):
            #resids[i] = p - codebook[memb[i]] # don't actually need to compute resids for this
            lines.append([p, codebook[memb[i]] ] )
        lc = mc.LineCollection(lines, colors=(1,0,1,1), linewidths=2)
        ax.add_collection(lc)
        
    if hide_tick_labels:
        ax.set_xticks([])
        ax.set_yticks([])

    ax.set_aspect('equal')
    if data is not None:
        plt.scatter(data[:, 0], data[:, 1], s=16)
    plt.show()
    
    
n_grid=5
codebook = generate_codebook(n_grid)
plot_data_grid_with_resids(data, n_grid=n_grid, show_next_level_grid=True, show_centroids=True, 
                           hide_tick_labels=False, codebook=codebook, show_resids=False)

plot_data_grid_with_resids(data, n_grid=n_grid, show_next_level_grid=True, show_centroids=True, 
                           hide_tick_labels=True, codebook=codebook, show_resids=True)

plot_data_grid_with_resids(data, n_grid=n_grid, show_next_level_grid=False, show_centroids=False, 
                           hide_tick_labels=False, codebook=np.zeros((data.shape[0],2)), show_resids=True, show_grid=False)

K = n_codebooks
fig, ax_all = plt.subplots(nrows=1, ncols=K, figsize=(10+(K-3),3))
recon = data*0
for axi in range(K):
    ax = ax_all[axi]
    ax.scatter(data[:, 0], data[:, 1], s=16)
    recon += codebooks[axi][indices[axi]]
    ax.scatter(recon[:, 0], recon[:, 1], s=16, color='orange')
    ax.set_xticks([])
    ax.set_yticks([])
    hpos = 1/K*0.85 + 1/K*0.81*axi if K==4 else .25 + .27*axi
    fig.text(hpos, .05, f"{axi+1} codebook{'s' if axi>0 else ''}", ha='center')

    error = ((recon - data)**2).mean()
    fig.text(hpos, .0001, f"Error = {error:.2e}", ha='center')
    
plt.show()

#### label: fig-rvq-recon-highdim2
#### fig-cap: "Error for high-dimensional datasets using various levels of RVQ. 'cost savings factor' refers to the ratio of using regular VQ (at uniform resolution) vs RVQ"

d_choices = [2, 3, 4, 6]  # we can't go much higher with 5x5 uniform grids!
K_choices = [1,2,3,4]     # variable numbers of codebooks

npoints_hd = 1000  # points in high-dim spaces

print("Here we show the error for high-dimensional datasets using various levels of RVQ.")
print("'cost savings factor' refers to the ratio of using regular VQ (at uniform resolution)\nvs RVQ.")

for d in d_choices:
    print(f"\nd = {d}:")
    np.random.seed(1)
    data_hd = DATA_MIN + (DATA_MAX-DATA_MIN)*np.random.rand(npoints_hd, d)
    codebook0 = generate_codebook(n_grid, n_dim=d)
    codebooks = [codebook0/n_grid**level for level in range(max(K_choices))] # lets get this over with rather than in the loop
    for K in K_choices: 
        indices = quantizer(data_hd, codebooks)
        recon = data_hd*0
        for lil_k in range(K):   # reconstruct using all codebooks
            recon += codebooks[lil_k][indices[lil_k]]
        error = ((recon - data_hd)**2).mean()
        grid_0_points = n_grid**(d)
        rvq_points = grid_0_points*K 
        uni_res = grid_0_points**K # comparable uniform resolution
        savings = uni_res/rvq_points
        print(f"  K = {K}, error = {error:.2e}, cost savings factor = {savings:.1f}")
        pass

import torch
from vector_quantize_pytorch import ResidualVQ
from functools import partial 

def make_pretty(d,styler):
    styler.set_caption(f"d = {d}:")
    styler.format(lambda e: f'{10**e:0.1e}')   # undo the log op
    styler.background_gradient(axis=None, vmin=-12, vmax=1, cmap="plasma")
    return styler

def show_colored_data(results_d, d, cb_lengths, K_choices):
    "displays a dataframe table but colors the background according to log10 of the numbers"
    rdf = pd.DataFrame(np.log10(results_d.numpy()), 
                       columns=[f'K = {k}' for k in K_choices], 
                       index=[f'cb_len = {c}' for c in cb_lengths])
    ourfunc = partial(make_pretty,d)
    display(rdf.style.pipe(ourfunc))
    return

n_dim = 256  # number of dimensions
cb_len = 256 # codebook length
K = 4        # number of codebooks
npoints_hd = 4096  # number of data points in high-dim space

d_choices = [2, 3, 6, 8, 16, 32, 64, 128, 256, 512]  # dimensions to try
cb_lengths = [25, 64, 256, 1024, 2048]                 # codebook lengths
K_choices = [1,2,3,4,6,8,10]#,16]                           # variable numbers of codebooks

results = torch.empty( ( len(d_choices), len(cb_lengths), len(K_choices) )).cpu()
for q1, n_dim in enumerate(d_choices):
    #print(f"\nd={n_dim:4}:       {' '.join([f'      K={i}' for i in K_choices])}")
    for q2, cb_len in enumerate(cb_lengths):
        #print(f"   cb_len={cb_len:4}: ",end="")
        for q3, K in enumerate(K_choices):
            
            residual_vq = ResidualVQ(
                dim = n_dim,
                codebook_size = cb_len,
                num_quantizers = K,
                kmeans_init = True,   # set to True
                kmeans_iters = 10     # number of kmeans iterations to calculate the centroids for the codebook on init
            )
            torch.manual_seed(0)
            x = torch.randn(1, npoints_hd, n_dim)
            quantized, indices, commit_loss = residual_vq(x)
            error = ((quantized - x)**2).mean()
            #print(f"  {error:.1e} ", end="")
            results[q1,q2,q3] = error
        #print()
    show_colored_data(results[q1], n_dim, cb_lengths, K_choices)     
        

import numpy as np
import plotly.graph_objects as go
import plotly.io as pio
pio.renderers.default='notebook'
import plotly.express as px 
from plotly.offline import plot, iplot, init_notebook_mode
import plotly.graph_objs as go
init_notebook_mode(connected=True)
import plotly.graph_objects as go

def plot_results(results, d_choices, cb_lengths, K_choices, log_xy=True, log_z=False, vary_K=False,
                camera=dict( up=dict(x=0, y=0, z=1), center=dict(x=0, y=0, z=-0.2), eye=dict(x=1.3, y=1.3, z=1.25)),
                width=750,):
    y, ylabel = np.array(d_choices), 'd'
    if vary_K:
        x, xlabel = np.array(K_choices), 'K'
        z = results[:,-1,:] 
    else:
        x, xlabel = np.array(cb_lengths), 'cb_len'
        z = results[:,:,3] 
    z, zlabel = z.numpy(), 'error'

    if log_xy:
        if not vary_K: 
            x, xlabel = np.log10(x), f"log10( {xlabel} )"
        y, ylabel = np.log10(y), f"log10( {ylabel} )"
    if log_z:
        z, zlabel = np.log10(z), f"log10( {zlabel} )"

    X, Y = np.meshgrid(x, y)

    fig = go.Figure(data=[go.Surface(x=X, y=Y, z=z, )], )
    fig.update_layout(scene = dict(xaxis_title=xlabel,yaxis_title=ylabel,zaxis_title=zlabel))
    fig.update_layout(template='plotly_dark',
        autosize=False,
        width=width,
        height=400,
        margin=dict(l=20,r=20,b=10,t=20,pad=4),
        #paper_bgcolor="LightSteelBlue",
        scene_camera = camera,
    )
    return fig

fig = plot_results(results, d_choices, cb_lengths, K_choices, log_xy=True, log_z=False)
fig.show()

camera=dict( up=dict(x=0, y=0, z=1), center=dict(x=-0.15, y=0, z=-0.2), eye=dict(x=1.2, y=-1.4, z=1))
fig = plot_results(results, d_choices, cb_lengths, K_choices, log_xy=True, log_z=True, camera=camera)
fig.show()

camera=dict( up=dict(x=0, y=0, z=1), center=dict(x=-0.1, y=0, z=-0.22), eye=dict(x=0.3, y=-1.8, z=1))
fig = plot_results(results, d_choices, cb_lengths, K_choices, vary_K=True, log_xy=True, log_z=True, camera=camera)
fig.show()

fig, ax = plt.subplots(figsize=(6.5,4))
ourdata = results[:,-1,:] # max cb_len
for q1, d in enumerate(d_choices):
    x = np.array(K_choices)
    y = ourdata[q1]
    ax.semilogy(x,y,'o-',label=f"d = {d}")
    
box = ax.get_position()
ax.set_position([box.x0, box.y0, box.width * 0.8, box.height])
ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))
ax.set_xlabel("K")
ax.set_ylabel("error")
plt.show()