Tukukkk

I have implemented what is know as a countSketch in python (page 17: https://arxiv.org/pdf/1411.4357.pdf) but my implementation is currently lacking in performance. The algorithm is to compute the product SA where A is an n x d matrix, S is m x n matrix defined as follows: for every column of S uniformly at random select a row (hash bucket) from the m rows and for that given row, uniformly at random select +1 or -1. So S is a matrix with exactly one nonzero in every column and otherwise all zero.

My intention is to compute SA in a streaming fashion by reading the entries of A. The idea for my implementation is as follows: observe a sequence of triples (i,j,A_ij) and return a sequence (h(i), j, s(i)A_ij) where:
- h(i) is a hash bucket (row of matrix chosen uniformly at random from the m possible rows of S
- s(i) is the random sign function as described above.
I have assumed that the matrix is in row order so that the first row of A arrives in its entirety before the next row of A arrives because this limits the number of calls I need to select a random bucket or the need to use a hash library. I have also assumed that the number of nonzero entries (or the length of the input stream) is known so that I can efficiently encode the iteration.

My problem is that the matrix should compute (1+error)*||Ax||^2 <= ||SAx||^2 <= (1+error)*||Ax||^2 and also have the difference in frobenius norms between A^T S^T S A and A^T A being small. However, while my implementation for the first condition seems to be true, the latter is consistently too small. I was wondering if there is an obvious reason for this that I am missing because it seems to be underestimating the latter quantity.

I am open to feedback on changing the code if there are obvious improvements to be made.

nb. If you don't want to run using numba then just comment out the import and the function decorator and it will run in standard numpy/scipy.

import numpy as np

import numpy.random as npr

import scipy.sparse as sparse

from scipy.sparse import coo_matrix

import numba

from numba import jit





@jit(nopython=True) # comment this if want just numpy

def countSketch(input_rows, input_data,

                      input_nnz,

                      sketch_size, seed=None):

   '''

   input_rows: row indices for data (can be repeats)

   input_data: values seen in row location,

   input_nnz : number of nonzers in the data (can replace with

   len(input_data) but avoided here for speed)

   sketch_size: int

   seed=None : random seed

   '''

   hashed_rows = np.empty(input_rows.shape,dtype=np.int32)

   current_row = 0

   hash_val = npr.choice(sketch_size)

   sign_val = npr.choice(np.array([-1.0,1.0]))

   #print(hash_val)

   hashed_rows[0] = hash_val

   #print(hash_val)

   for idx in np.arange(input_nnz):

       print(idx)

       row_id = input_rows[idx]

       data_val = input_data[idx]

       if row_id == current_row:

           hashed_rows[idx] = hash_val

           input_data[idx] = sign_val*data_val

       else:

           # make new hashes

           hash_val = npr.choice(sketch_size)

           sign_val = npr.choice(np.array([-1.0,1.0]))

           hashed_rows[idx] = hash_val

           input_data[idx] = sign_val*data_val

   return hashed_rows, input_data



def sort_row_order(input_data):

   sorted_row_column = np.array((input_data.row,

                                 input_data.col,

                                 input_data.data))



   idx  = np.argsort(sorted_row_column[0])

   sorted_rows = np.array(sorted_row_column[0,idx], dtype=np.int32)

   sorted_cols = np.array(sorted_row_column[1,idx], dtype=np.int32)

   sorted_data = np.array(sorted_row_column[2,idx], dtype=np.float64)

   return sorted_rows, sorted_cols, sorted_data





if __name__=="__main__":

   import time

   from tabulate import tabulate



   matrix = sparse.random(1000, 50, 0.1)

   x = np.random.randn(matrix.shape[1])

   true_norm = np.linalg.norm(matrix@x,ord=2)**2

   tidy_data =  sort_row_order(matrix)



   sketch_size = 300

   start = time.time()

   hashed_rows, sketched_data = countSketch(tidy_data[0],

                                            tidy_data[2], matrix.nnz,sketch_size)

   duration_slow = time.time() - start

   S_A = sparse.coo_matrix((sketched_data, (hashed_rows,matrix.col)))

   approx_norm_slow = np.linalg.norm(S_A@x,ord=2)**2

   rel_error_slow = approx_norm_slow/true_norm

   #print("Sketch time: {}".format(duration_slow))

   start = time.time()

   hashed_rows, sketched_data = countSketch(tidy_data[0],

                                           tidy_data[2], matrix.nnz,sketch_size)

   duration = time.time() - start

   #print("Sketch time: {}".format(duration))

   S_A = sparse.coo_matrix((sketched_data, (hashed_rows,matrix.col)))

   approx_norm = np.linalg.norm(S_A@x,ord=2)**2

   rel_error = approx_norm/true_norm

   #print("Relative norms: {}".format(approx_norm/true_norm))

   print(tabulate([[duration_slow, rel_error_slow, 'Yes'],

                   [duration, rel_error, 'No']],

                   headers=['Sketch Time', 'Relative Error', 'Dry Run'],

                   tablefmt='orgtbl'))

as suggested by @SamOnela, code not working is off-topic. for your performance issue, you can group your calls to choice at the beginning of your function

   hash_vals = npr.choice(sketch_size, input_nnz)

   sign_vals = npr.choice(np.array([-1.0,1.0]), input_nnz)

and use it later in your code this way:

       hashed_rows[idx] = hash_vals[idx]

       input_data[idx] = sign_vals[idx]*data_val