Ensuring performance of sketching/streaming algorithm (countSketch)

Question

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 nonzeroes 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'))

Answer 1

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