Trigonometric Sketching Matrices

A trigonometric sketch matrix can be applied to $\vec{A}$ in $O(nd\log(n))$ operations. A downside is that it is not easy to efficiently apply the sketch matrix to a sparse matrix $\vec{A}$. In addition, parallelizing the application of tigonometric transforms is not as straightforward as for Gaussian or sparse sketches.

In other words, with relative little overhead compared to CountSketch, SparseStack matrices substantially improve the embedding dimension.

Implementation¶

The implementation of a trigonometric sketch matrix is relatively straightforward. An efficient application requires sequentially applying the constituent operations, and hence we define a custom class.

class trig_sketch():
    
    def __init__(self,n,k,rng):

        self.shape = (k,n)
        self.pi = rng.permutation(n)
        self.E = 2*rng.randint(0,2,n)-1
        self.R = rng.choice(n,size=k,replace=False)

    def __matmul__(self,A):

        A_loc = A.A if sp.sparse.issparse(A) else A
        EPiA = self.E[:,None]*A_loc[self.pi] if A_loc.ndim == 2 else self.E*A_loc[self.pi]
        return np.sqrt(self.shape[1]/self.shape[0])*sp.fft.dct(EPiA,norm='ortho',axis=0)[self.R]