Fix: Kaczmarz + Sampling on Sparse Matrix

src/Compressors/sampling.jl

@@ -215,4 +215,44 @@ function mul!(
     mul!(C_sub, A, I, alpha, 1)
 
     return nothing
-end
+end
+
+###############################################################################
+# Binary Operator Compressor-Array Multiplications for sparse matrices/vectors
+###############################################################################
+# S * A 
+function (*)(S::SamplingRecipe, A::Union{SparseMatrixCSC, SparseVector})
+    s_rows = size(S, 1)
+    a_cols = size(A, 2)
+    C = a_cols == 1 ? spzeros(eltype(A), s_rows) : spzeros(eltype(A), s_rows, a_cols)
+    mul!(C, S, A)
+    return C
+end
+
+# A * S 
+function (*)(A::Union{SparseMatrixCSC, SparseVector}, S::SamplingRecipe)
+    s_cols = size(S, 2)
+    a_rows = size(A, 1)
+    C = a_rows == 1 ? spzeros(eltype(A), s_cols)' : spzeros(eltype(A), a_rows, s_cols)
+    mul!(C, A, S)
+    return C
+end
+
+# S' * A
+function (*)(S::CompressorAdjoint{<:SamplingRecipe}, A::Union{SparseMatrixCSC, SparseVector})
+    s_rows = size(S, 1)
+    a_cols = size(A, 2)
+    C = a_cols == 1 ? spzeros(eltype(A), s_rows) : spzeros(eltype(A), s_rows, a_cols)
+    mul!(C, S, A)
+    return C
+end
+
+# A * S'
+function (*)(A::Union{SparseMatrixCSC, SparseVector}, S::CompressorAdjoint{<:SamplingRecipe})
+    s_cols = size(S, 2)
+    a_rows = size(A, 1)
+    C = a_rows == 1 ? spzeros(eltype(A), s_cols)' : spzeros(eltype(A), a_rows, s_cols)
+    mul!(C, A, S)
+    return C
+end
+

src/RLinearAlgebra.jl

@@ -4,7 +4,7 @@ import Base: transpose, adjoint
 import LinearAlgebra: Adjoint, axpby!, dot, I, ldiv!, lmul!, lq!, lq, LQ, mul!, norm, qr!, svd
 import StatsBase: sample, sample!, ProbabilityWeights, wsample!
 import Random: bitrand, rand!, randn!
-import SparseArrays: SparseMatrixCSC, sprandn, sparse
+import SparseArrays: SparseMatrixCSC, SparseVector, spzeros, sprandn, sparse
 
 # Include the files correspoding to the top-level techniques
 include("Compressors.jl")

src/Solvers/SubSolvers/lq.jl

@@ -39,6 +39,8 @@ function ldiv!(
 )
     fill!(x, zero(eltype(b)))
     # this will modify B in place so you cannot use it again
-    ldiv!(x, lq!(solver.A), b)
+    # using qr here on the transpose of the matrix will work for sparse and dense matrices
+    # while the lq would have only worked for dense matrices
+    ldiv!(x, qr!(solver.A')', b)
     return nothing
 end

src/Solvers/kaczmarz.jl

@@ -15,9 +15,9 @@ Let ``A`` be an ``m \\times n`` matrix and consider the consistent linear system
     this interesection is to iteratively project some abritrary point, ``x`` from one 
     hyperplane to the next, through 
     ``
-    x_{+} = x + \\alpha \\frac{b_i - \\lange A_{i\\cdot}, x\\rangle}{\\| A_{i\\cdot}.
+    x_{+} = x + \\alpha \\frac{b_i -  A_{i\\cdot}^\\top x}{\\| A_{i\\cdot} \\|_2^2}
     ``
-    Doing this with random permutation of ``i`` can lead to a geometric convergence 
+     Doing this with random permutation of ``i`` can lead to a geometric convergence 
     [strohmer2009randomized](@cite).
     Here ``\\alpha`` is viewed as an over-relaxation parameter and can improve convergence. 
     One can also generalize this procedure to blocks by considering the ``S`` being a 
@@ -158,8 +158,8 @@ mutable struct KaczmarzRecipe{
     alpha::Float64
     compressed_mat::M
     compressed_vec::V
-    solution_vec::V
-    update_vec::V
+    solution_vec::Vector{T}
+    update_vec::Vector{T}
     mat_view::MV
     vec_view::VV
 end
@@ -198,9 +198,22 @@ function complete_solver(
     # We assume the user is using compressors to only decrease dimension
     sample_size::Int64 = compressor.n_rows
     cols_a = size(A, 2)
-    # Allocate the information in the buffer using the types of A and b
-    compressed_mat = zeros(eltype(A), sample_size, cols_a)
-    compressed_vec = zeros(eltype(b), sample_size) 
+    # because sampling recipe use subsets sparse matrices will be problematic
+    if typeof(compressor) <: SamplingRecipe && typeof(A) <: SparseMatrixCSC
+        compressed_mat = spzeros(eltype(A), sample_size, cols_a)
+    else
+        # Allocate the information in the buffer using the types of A
+        compressed_mat = zeros(eltype(A), sample_size, cols_a)
+    end
+
+    # because sampling recipe use subsets sparse vectors will be problematic
+    if typeof(compressor) <: SamplingRecipe && typeof(b) <: SparseVector
+        compressed_vec = spzeros(eltype(b), sample_size)
+    else
+        # Allocate the information in the buffer using the type of b
+        compressed_vec = zeros(eltype(b), sample_size)
+    end
+
     # Since sub_solver is applied to compressed matrices use here
     sub_solver = complete_sub_solver(ingredients.sub_solver, compressed_mat, compressed_vec)
     mat_view = view(compressed_mat, 1:sample_size, :)
@@ -250,12 +263,15 @@ function kaczmarz_update!(solver::KaczmarzRecipe)
     # when the constant vector is a zero dimensional subArray we know that we should perform
     # the one dimension kaczmarz update
 
-    # Compute the projection scaling (bi - dot(ai,x)) / ||ai||^2
-    scaling = solver.alpha * (dotu(solver.mat_view, solver.solution_vec) 
-        - solver.vec_view[1]) 
-    scaling /= dot(solver.mat_view, solver.mat_view)
-    # udpate the solution computes solution_vec = solution_vec - scaling * mat_view'
-    axpby!(-scaling, solver.mat_view', 1.0, solver.solution_vec)
+    # check that the row is non-zero
+    if !iszero(solver.mat_view)
+        # Compute the projection scaling (bi - dot(ai,x)) / ||ai||^2
+        scaling = solver.alpha * (dotu(solver.mat_view, solver.solution_vec) 
+            - solver.vec_view[1]) 
+        scaling /= dot(solver.mat_view, solver.mat_view)
+        # udpate the solution computes solution_vec = solution_vec - scaling * mat_view'
+        axpby!(-scaling, solver.mat_view', 1.0, solver.solution_vec)
+    end
     return nothing
 end
 
@@ -278,16 +294,19 @@ A function that performs the kaczmarz update when the compression dim is greater
 function kaczmarz_update_block!(solver::KaczmarzRecipe)
     # when the constant vector is a one dimensional subArray we know that we should perform
     # the one dimension kaczmarz update
-    # sub-solver needs to designed for new compressed matrix
-    update_sub_solver!(solver.sub_solver, solver.mat_view)
-    # Compute the block residual 
-    # (computes solver.vec_view - solver.mat_view * solver.solution_vec)
-    mul!(solver.vec_view, solver.mat_view, solver.solution_vec, -1.0, 1.0)
-    # use sub-solver to find update the solution (solves min ||tilde A - tilde b|| and 
-    # stores in update_vec)
-    ldiv!(solver.update_vec, solver.sub_solver, solver.vec_view)
-    # computes solver.solution_vec = solver.solution_vec + alpha * solver.update_vec
-    axpby!(solver.alpha, solver.update_vec, 1.0, solver.solution_vec)
+    if !iszero(solver.mat_view)
+        # sub-solver needs to designed for new compressed matrix
+        update_sub_solver!(solver.sub_solver, solver.mat_view)
+        # Compute the block residual 
+        # (computes solver.vec_view - solver.mat_view * solver.solution_vec)
+        mul!(solver.vec_view, solver.mat_view, solver.solution_vec, -1.0, 1.0)
+        # use sub-solver to find update the solution (solves min ||tilde A - tilde b|| and 
+        # stores in update_vec)
+        ldiv!(solver.update_vec, solver.sub_solver, solver.vec_view)
+        # computes solver.solution_vec = solver.solution_vec + alpha * solver.update_vec
+        axpby!(solver.alpha, solver.update_vec, 1.0, solver.solution_vec)
+    end
+
     return nothing
 end
 

test/Compressors/sampling.jl

@@ -2,6 +2,7 @@ module Sampling_compressor
 using Test, RLinearAlgebra, Random
 using StatsBase: ProbabilityWeights, sample
 import LinearAlgebra: mul!, Adjoint
+import SparseArrays: sprandn
 using ..FieldTest
 using ..ApproxTol
 
@@ -499,8 +500,139 @@ Random.seed!(2131)
                 mul!(x, S', yc, alpha, beta)
                 @test x ≈ alpha * Sty_exact + beta * xc
             end
+
+        end
+
+        @testset "Left Cardinality Sparse" begin
+            let a_matrix_rows = 20,
+                a_matrix_cols = 12,
+                comp_dim = 7,
+                alpha = 2.5,
+                beta = 1.5
+
+                # Setup matrices and vectors 
+                A = sprandn(a_matrix_rows, a_matrix_cols, .8)
+                B = sprandn(comp_dim, a_matrix_cols, .8)
+                C1 = sprandn(comp_dim, a_matrix_cols, .8)
+                C2 = sprandn(a_matrix_rows, a_matrix_cols, .8)
+                x = sprandn(a_matrix_rows, .8)
+                y = sprandn(comp_dim, .8)
+
+                # Keep copies for 5-argument mul! verification
+                C1c = deepcopy(C1)
+                C2c = deepcopy(C2)
+                yc = deepcopy(y)
+                xc = deepcopy(x)
+
+                # Setup the Sampling compressor recipe
+                S_info = Sampling(
+                    cardinality=Left(),
+                    compression_dim=comp_dim,
+                    distribution=Uniform(cardinality=Left(), replace=false)
+                )
+                S = complete_compressor(S_info, A)
+
+                # Calculate all ground truth results
+                SA_exact = A[S.idx, :]
+                StB_exact = zeros(a_matrix_rows, a_matrix_cols); for i in 1:comp_dim; StB_exact[S.idx[i], :] = B[i, :]; end
+                Sx_exact = x[S.idx]
+                Sty_exact = zeros(a_matrix_rows); for i in 1:comp_dim; Sty_exact[S.idx[i]] = y[i]; end
+
+                # Test '*' operations by comparing to ground truths
+                @test S * A ≈ SA_exact
+                @test S' * B ≈ StB_exact
+                @test A' * S' ≈ SA_exact'
+                @test B' * S ≈ StB_exact'
+                @test S * x ≈ Sx_exact
+                @test x' * S' ≈ Sx_exact'
+                @test S' * y ≈ Sty_exact
+                @test y' * S ≈ Sty_exact'
+
+                # Test the 5-argument mul!
+                mul!(C1, S, A, alpha, beta)
+                @test C1 ≈ alpha * SA_exact + beta * C1c
+
+                mul!(C2, S', B, alpha, beta)
+                @test C2 ≈ alpha * StB_exact + beta * C2c
+
+                mul!(y, S, xc, alpha, beta)
+                @test y ≈ alpha * Sx_exact + beta * yc
+
+                mul!(x, S', yc, alpha, beta)
+                @test x ≈ alpha * Sty_exact + beta * xc
+            end
+        end
+
+        # Test multiplications with right compressors
+        @testset "Right Cardinality" begin
+            let n = 20,         
+                comp_dim = 10,     
+                alpha = 2.0,
+                beta = 2.0
+
+                # Setup matrices and vectors with dimensions
+                A = sprandn(n, comp_dim, .8)
+                B = sprandn(n, n, .8)
+                # C1 is for S'*A, C2 is for B*S
+                C1 = sprandn(n, n, .8)
+                C2 = sprandn(n, comp_dim, .8)
+                x = sprandn(comp_dim, .8)
+                y = sprandn(n, .8)
+
+                # Keep copies for 5-argument mul! verification
+                C1c = deepcopy(C1)
+                C2c = deepcopy(C2)
+                yc = deepcopy(y)
+                xc = deepcopy(x)
+
+                # Setup the Sampling compressor recipe. It's created from B, an n x n matrix.
+                # The operator S will have conceptual dimensions (n x comp_dim).
+                S_info = Sampling(
+                    cardinality=Right(),
+                    compression_dim=comp_dim,
+                    distribution=Uniform(cardinality=Right(), replace=false)
+                )
+                S = complete_compressor(S_info, B)
+
+                # Calculate all ground truth results based on direct indexing/operations
+                StA_exact = A[S.idx, :]
+                BS_exact = B[:, S.idx]
+                BtS_exact = B'[:, S.idx]
+                ASt_exact = zeros(n, n); for i in 1:comp_dim; ASt_exact[:, S.idx[i]] = A[:, i]; end
+                Sx_exact = zeros(n); for i in 1:comp_dim; Sx_exact[S.idx[i]] = x[i]; end
+                Sty_exact = y[S.idx]
+
+                # Test '*' operations by comparing to ground truths
+                @test S' * A ≈ StA_exact
+                @test A' * S ≈ StA_exact'
+                @test B * S ≈ BS_exact
+                @test S' * B' ≈ BS_exact'
+                @test B' * S ≈ BtS_exact
+                @test S' * B ≈ BtS_exact'
+                @test A * S' ≈ ASt_exact
+                @test S * A' ≈ ASt_exact'
+                @test S * x ≈ Sx_exact
+                @test x' * S' ≈ Sx_exact'
+                @test y' * S ≈ Sty_exact'
+                @test S' * y ≈ Sty_exact
+
+                # Test the 5-argument mul!
+                mul!(C1, A, S', alpha, beta)
+                @test C1 ≈ alpha * ASt_exact + beta * C1c
+
+                mul!(C2, B, S, alpha, beta)
+                @test C2 ≈ alpha * BS_exact + beta * C2c
+
+                mul!(y, S, xc, alpha, beta)
+                @test y ≈ alpha * Sx_exact + beta * yc
+
+                mul!(x, S', yc, alpha, beta)
+                @test x ≈ alpha * Sty_exact + beta * xc
+            end
+
         end
     end
+
 end
 
 end