def sketch()

in sketching.py [0:0]

• 20 lines of code
• 10 McCabe index (conditional complexity)

    def sketch(self, r):
        """
        Computes S^T A, S^T A S and S^T r for S being the count sketch matrix
        """
        self.set_sketch()
        if self.build_matrix:
            SA = (self.A @ self.S).T  # A @ S --> S should be CSC
            SAS = csc_matrix.dot(SA, self.S)  # SA @ S --> S should be CSC
            rs = csr_matrix.dot(r.T, self.S).T  # r.T @ S --> S should be CSC
            # rs = csr_matrix.dot(self.S.T, r)  # S.T @ r --> S.T should be CSC, which is the case if S is CSR
        else:
            # Computational issue: two loops over m
            SA = np.zeros((self.sketch_size, self.m))
            rs = np.zeros((self.sketch_size, 1))
            for i in range(self.m):  # TODO: remove loop or vectorize this step
                SA[self.cols[i], :] += (
                    self.signs[i] * self.A._A[i, :]
                )  # could be efficient if A were CSR sparse matrix
                # SA[self.cols[i],:] += self.signs[i] * self.A.get_rows(i) # if A is an AOperator
                rs[self.cols[i], :] += self.signs[i] * r[i]
            SAS = np.zeros(
                (self.sketch_size, self.sketch_size)
            )  # need for SA to be built before computing SAS
            for i in range(self.m):  # TODO: remove loop or vectorize this step
                SAS[:, self.cols[i]] += self.signs[i] * SA[:, i]
        return SA, SAS, rs