diff --git a/docs/src/api/solvers.md b/docs/src/api/solvers.md index 05f9c1a1..ca98abc3 100644 --- a/docs/src/api/solvers.md +++ b/docs/src/api/solvers.md @@ -12,7 +12,9 @@ SolverRecipe ## Solver Structures ```@docs +IHS +IHSRecipe ``` ## Exported Functions diff --git a/docs/src/refs.bib b/docs/src/refs.bib index b7297ec8..e6be4075 100644 --- a/docs/src/refs.bib +++ b/docs/src/refs.bib @@ -34,8 +34,30 @@ @misc{martinsson2020randomized year = {2020}, publisher = {arXiv}, doi = {10.48550/ARXIV.2002.01387}, - copyright = {arXiv.org perpetual, non-exclusive license}, - keywords = {FOS: Mathematics,Numerical Analysis (math.NA)} +} + +@article{needell2014paved, + title = {Paved with Good Intentions: {{Analysis}} of a Randomized Block {{Kaczmarz}} Method}, + shorttitle = {Paved with Good Intentions}, + author = {Needell, Deanna and Tropp, Joel A.}, + year = {2014}, + month = jan, + journal = {Linear Algebra and its Applications}, + volume = {441}, + pages = {199--221}, + issn = {00243795}, + doi = {10.1016/j.laa.2012.12.022}, + langid = {english}, +} + +@misc{pilanci2014iterative, + title = {Iterative {{Hessian}} Sketch: {{Fast}} and Accurate Solution Approximation for Constrained Least-Squares}, + shorttitle = {Iterative {{Hessian}} Sketch}, + author = {Pilanci, Mert and Wainwright, Martin J.}, + year = {2014}, + publisher = {arXiv}, + doi = {10.48550/ARXIV.1411.0347}, + keywords = {FOS: Computer and information sciences,FOS: Mathematics,Information Theory (cs.IT),Machine Learning (cs.LG),Machine Learning (stat.ML),Optimization and Control (math.OC)} } @article{pritchard2023practical, @@ -64,3 +86,26 @@ @article{pritchard2024solving issn = {0266-5611, 1361-6420}, doi = {10.1088/1361-6420/ad5583}, } + +@article{strohmer2009randomized, + title = {A {{Randomized Kaczmarz Algorithm}} with {{Exponential Convergence}}}, + author = {Strohmer, Thomas and Vershynin, Roman}, + year = {2009}, + month = apr, + journal = {Journal of Fourier Analysis and Applications}, + volume = {15}, + number = {2}, + pages = {262--278}, + issn = {1069-5869, 1531-5851}, + doi = {10.1007/s00041-008-9030-4}, + langid = {english} +} +@misc{pilanci2014iterative, + title = {Iterative {{Hessian}} Sketch: {{Fast}} and Accurate Solution Approximation for Constrained Least-Squares}, + shorttitle = {Iterative {{Hessian}} Sketch}, + author = {Pilanci, Mert and Wainwright, Martin J.}, + year = {2014}, + publisher = {arXiv}, + doi = {10.48550/ARXIV.1411.0347}, +} + diff --git a/src/RLinearAlgebra.jl b/src/RLinearAlgebra.jl index cfd4324f..8816e23b 100644 --- a/src/RLinearAlgebra.jl +++ b/src/RLinearAlgebra.jl @@ -1,7 +1,8 @@ module RLinearAlgebra import Base.:* import Base: transpose, adjoint -import LinearAlgebra: Adjoint, axpby!, dot, ldiv!, lmul!, lq!, lq, LQ, mul!, norm, qr! +import LinearAlgebra: Adjoint, axpby!, axpy!, dot, ldiv!, lmul!, lq!, lq, LQ, mul!, norm +import LinearAlgebra: qr!, UpperTriangular import StatsBase: sample!, ProbabilityWeights, wsample! import Random: bitrand, rand!, randn! import SparseArrays: SparseMatrixCSC, sprandn @@ -33,6 +34,7 @@ export Uniform, UniformRecipe # Export Solver types and functions export Solver, SolverRecipe export complete_solver, update_solver!, rsolve! +export IHS, IHSRecipe # Export Logger types and functions export Logger, LoggerRecipe diff --git a/src/Solvers.jl b/src/Solvers.jl index ffe0aaa1..4242c0e4 100644 --- a/src/Solvers.jl +++ b/src/Solvers.jl @@ -150,3 +150,4 @@ include("Solvers/ErrorMethods.jl") ############################# # The Solver Routine Files ############################ +include("Solvers/ihs.jl") diff --git a/src/Solvers/ihs.jl b/src/Solvers/ihs.jl new file mode 100644 index 00000000..1616edb2 --- /dev/null +++ b/src/Solvers/ihs.jl @@ -0,0 +1,234 @@ +""" + IHS <: Solver + +An implementation of the Iterative Hessian Sketch solver for solving over determined +least squares problems (@cite)[pilanci2014iterative]. + +# Mathematical Description +Let ``A \\in \\mathbb{R}^{m \\times n}`` and consider the least square problem ``\\min_x +\\|Ax - b \\|_2^2``. If we let ``S \\in \\mathbb{R}^{s \\times m}`` be a compression matrix, then +Iterative Hessian Sketch iteratively finds a solution to this problem +by repeatedly updating ``x_{k+1} = x_k + \\alpha u_k``where ``u_k`` is the solution to the +convex optimization problem, +``u_k = \\min_u \\{\\|S_k Au\\|_2^2 - \\langle A, b - Ax_k \\rangle \\}.`` This method +has been to shown to converge geometrically at a rate ``\\rho \\in (0, 1/2]``, typically the +required compression dimension needs to be 4-8 times the size of n for the algorithm to +perform successfully. + +# Fields +- `alpha::Float64`, a step size parameter. +- `compressor::Compressor`, a technique for forming the compressed linear system. +- `log::Logger`, a technique for logging the progress of the solver. +- `error::SolverError`, a method for estimating the progress of the solver. + +# Constructor + function IHS(; + compressor::Compressor = SparseSign(cardinality = Left()), + log::Logger = BasicLogger(), + error::SolverError = FullResidual(), + alpha::Float64 = 1.0 + ) +## Keywords +- `compressor::Compressor`, a technique for forming the compressed linear system. +- `log::Logger`, a technique for logging the progress of the solver. +- `error::SolverError', a method for estimating the progress of the solver. +- `alpha::Float64`, a step size parameter. + +# Returns +- A `IHS` object. +""" +mutable struct IHS <: Solver + alpha::Float64 + log::Logger + compressor::Compressor + error::SolverError + function IHS(alpha, log, compressor, error) + if typeof(compressor.cardinality) != Left + @warn "Compressor has cardinality `Right` but IHS compresses from the `Left`." + end + + new(alpha, log, compressor, error) + end + +end + +function IHS(; + compressor::Compressor = SparseSign(cardinality = Left()), + log::Logger = BasicLogger(), + error::SolverError = FullResidual(), + alpha::Float64 = 1.0 +) + return IHS( + alpha, + log, + compressor, + error + ) +end + +""" + IHSRecipe{ + Type<:Number, + LR<:LoggerRecipe, + CR<:CompressorRecipe, + ER<:ErrorRecipe, + M<:AbstractArray, + MV<:SubArray, + V<:AbstractVector + } <: SolverRecip + +A mutable structure containing all information relevant to the Iterative Hessian Sketch +solver. It is formed by calling the function `complete_solver` on a `IHS` solver, which +includes all the user controlled parameters, the linear system `A`, and the constant +vector `b`. + +# Fields +- `compressor::CompressorRecipe`, a technique for compressing the matrix ``A``. +- `logger::LoggerRecipe`, a technique for logging the progress of the solver. +- `error::SolverErrorRecipe`, a technique for estimating the progress of the solver. +- `compressed_mat::AbstractMatrix`, a buffer for storing the compressed matrix. +- `mat_view::SubArray`, a container for storing a view of the compressed matrix buffer. +- `residual_vec::AbstractVector`, a vector that contains the residual of the linear system + ``Ax-b``. +- `gradient_vec::AbstractVector`, a vector that contains the gradient of the least squares + problem, ``A^\\top(b-Ax)``. +- `buffer_vec::AbstractVector`, a buffer vector for storing intermediate linear system solves. +- `solution_vec::AbstractVector`, a vector storing the current IHS solution. +- `R::UpperTriangular`, a container for storing the upper triangular portion of the R + factor from a QR factorization of `mat_view`. This is used to solve the IHS sub-problem. +""" +mutable struct IHSRecipe{ + Type<:Number, + LR<:LoggerRecipe, + CR<:CompressorRecipe, + ER<:SolverErrorRecipe, + M<:AbstractArray, + MV<:SubArray, + V<:AbstractVector +} <: SolverRecipe + log::LR + compressor::CR + error::ER + alpha::Float64 + compressed_mat::M + mat_view::MV + residual_vec::V + gradient_vec::V + buffer_vec::V + solution_vec::V + R::UpperTriangular{Type, M} +end + +function complete_solver( + ingredients::IHS, + x::AbstractVector, + A::AbstractMatrix, + b::AbstractVector +) + compressor = complete_compressor(ingredients.compressor, x, A, b) + logger = complete_logger(ingredients.log) + error = complete_error(ingredients.error, ingredients, A, b) + sample_size::Int64 = compressor.n_rows + rows_a, cols_a = size(A) + # Check that required fields are in the types + if !isdefined(error, :residual) + throw( + ArgumentError( + "ErrorRecipe $(typeof(error)) does not contain the \ + field 'residual' and is not valid for an IHS solver." + ) + ) + end + + if !isdefined(logger, :converged) + throw( + ArgumentError( + "LoggerRecipe $(typeof(logger)) does not contain \ + the field 'converged' and is not valid for an IHS solver." + ) + ) + end + + # Check that the sketch size is larger than the column dimension and return a warning + # otherwise + if cols_a > sample_size + throw( + ArgumentError( + "Compression dimension not larger than column dimension this will lead to \ + singular QR decompositions, which cannot be inverted." + ) + ) + end + + compressed_mat = zeros(eltype(A), sample_size, cols_a) + res = zeros(eltype(A), rows_a) + grad = zeros(eltype(A), cols_a) + buffer_vec = zeros(eltype(A), cols_a) + solution_vec = x + mat_view = view(compressed_mat, 1:sample_size, :) + R = UpperTriangular(mat_view[1:cols_a, :]) + + return IHSRecipe{ + eltype(compressed_mat), + typeof(logger), + typeof(compressor), + typeof(error), + typeof(compressed_mat), + typeof(mat_view), + typeof(buffer_vec) + }( + logger, + compressor, + error, + ingredients.alpha, + compressed_mat, + mat_view, + res, + grad, + buffer_vec, + solution_vec, + R + ) +end + +function rsolve!(solver::IHSRecipe, x::AbstractVector, A::AbstractMatrix, b::AbstractVector) + reset_logger!(solver.log) + solver.solution_vec = x + err = 0.0 + copyto!(solver.residual_vec, b) + # compute the initial residual r = b - Ax + mul!(solver.residual_vec, A, solver.solution_vec, -1.0, 1.0) + for i in 1:solver.log.max_it + # compute the gradient A'r + mul!(solver.gradient_vec, A', solver.residual_vec) + err = compute_error(solver.error, solver, A, b) + update_logger!(solver.log, err, i) + if solver.log.converged + return nothing + end + + # generate a new compressor + update_compressor!(solver.compressor, x, A, b) + # Based on the size of the compressor update views of the matrix + rows_s, cols_s = size(solver.compressor) + solver.mat_view = view(solver.compressed_mat, 1:rows_s, :) + # Compress the matrix + mul!(solver.mat_view, solver.compressor, A) + # Update the subsolver + # This is the only piece of allocating code + solver.R = UpperTriangular(qr!(solver.mat_view).R) + # Compute first R' solver R'R x = g + ldiv!(solver.buffer_vec, solver.R', solver.gradient_vec) + # Compute second R Solve Rx = (R')^(-1)g will be stored in gradient_vec + ldiv!(solver.gradient_vec, solver.R, solver.buffer_vec) + # update the solution + # solver.solution_vec = solver.solution_vec + alpha * solver.gradient_vec + axpy!(solver.alpha, solver.gradient_vec, solver.solution_vec) + # compute the fast update of r = r - A * gradient_vec + # note: in this case gradient vec stores the update + mul!(solver.residual_vec, A, solver.gradient_vec, -1.0, 1.0) + end + + return nothing + +end diff --git a/test/Solvers/ihs.jl b/test/Solvers/ihs.jl new file mode 100644 index 00000000..1eee8d20 --- /dev/null +++ b/test/Solvers/ihs.jl @@ -0,0 +1,448 @@ +module IHSTest +using Test, RLinearAlgebra, LinearAlgebra +import RLinearAlgebra: complete_compressor +import LinearAlgebra: mul!, norm +import Random: randn!, seed! +using ..FieldTest +using ..ApproxTol + +# Define the test structures +########################## +# Compressors +########################## +mutable struct ITestCompressor <: Compressor + cardinality::Cardinality + compression_dim::Int64 +end + +ITestCompressor() = ITestCompressor(Left(), 5) + +mutable struct ITestCompressorRecipe <: CompressorRecipe + cardinality::Cardinality + n_rows::Int64 + n_cols::Int64 + op::AbstractMatrix +end + +function RLinearAlgebra.complete_compressor( + comp::ITestCompressor, + A::AbstractMatrix, + b::AbstractVector +) + n_rows = comp.compression_dim + n_cols = size(A, 1) + # Make a gaussian compressor + op = randn(n_rows, n_cols) ./ sqrt(n_cols) + return ITestCompressorRecipe(comp.cardinality, n_rows, n_cols, op) +end + +function RLinearAlgebra.update_compressor!( + comp::ITestCompressorRecipe, + x::AbstractVector, + A::AbstractMatrix, + b::AbstractVector +) + randn!(comp.op) + comp.op ./= sqrt(comp.n_cols) +end + +# Define a mul function for the test compressor +function RLinearAlgebra.mul!( + C::AbstractArray, + S::Main.IHSTest.ITestCompressorRecipe, + A::AbstractArray, + alpha::Number, + beta::Number +) + mul!(C, S.op, A, alpha, beta) +end + +########################## +# Error Method +########################## +mutable struct ITestError <: RLinearAlgebra.SolverError + g::Real +end + +mutable struct ITestErrorRecipe <: RLinearAlgebra.SolverErrorRecipe + residual::Vector{Number} +end + +ITestError() = ITestError(1.0) + +function RLinearAlgebra.complete_error( + error::ITestError, + solver::IHS, + A::AbstractMatrix, + b::AbstractVector +) + return ITestErrorRecipe(zeros(typeof(error.g), size(A, 1))) +end + +function RLinearAlgebra.compute_error(error::ITestErrorRecipe, solver, A, b) + error.residual = A * solver.solution_vec - b + return norm(error.residual) +end + +############################## +# Residual-less Error Recipe +############################## +mutable struct ITestErrorNoRes <: RLinearAlgebra.SolverError end + +mutable struct ITestErrorRecipeNoRes <: RLinearAlgebra.SolverErrorRecipe end + +function RLinearAlgebra.complete_error( + error::ITestErrorNoRes, + solver::IHS, + A::AbstractMatrix, + b::AbstractVector +) + return ITestErrorRecipeNoRes() +end + +############################ +# Loggers +############################ +mutable struct ITestLog <: Logger + max_it::Int64 + g::Number +end + +ITestLog() = ITestLog(5, 1.0) + +ITestLog(max_it) = ITestLog(max_it, 1.0) + +mutable struct ITestLogRecipe <: LoggerRecipe + max_it::Int64 + hist::Vector{Real} + thresh::Float64 + converged::Bool +end + +function RLinearAlgebra.complete_logger(logger::ITestLog) + return ITestLogRecipe( + logger.max_it, + zeros(typeof(logger.g), logger.max_it), + logger.g, + false + ) +end + +function RLinearAlgebra.update_logger!(logger::ITestLogRecipe, err::Real, i::Int64) + logger.hist[i] = err + logger.converged = err < logger.thresh ? true : false +end + +function RLinearAlgebra.reset_logger!(logger::ITestLogRecipe) + fill!(logger.hist, 0.0) +end + +############################## +# Converged-less Logger +############################## +mutable struct ITestLogNoCov <: Logger end + +mutable struct ITestLogRecipeNoCov <: LoggerRecipe end + +function RLinearAlgebra.complete_logger(logger::ITestLogNoCov) + return ITestLogRecipeNoCov() +end + + +@testset "IHS" begin + seed!(12312) + n_rows = 20 + n_cols = 2 + A = rand(n_rows, n_cols) + xsol = rand(n_cols) + b = A * xsol + + @testset "IHS" begin + @test supertype(IHS) == Solver + + # test fieldnames and types + @test fieldnames(IHS) == (:alpha, :log, :compressor, :error) + @test fieldtypes(IHS) == (Float64, Logger, Compressor, SolverError) + + # test default constructor + + let solver = IHS() + @test solver.alpha == 1.0 + @test typeof(solver.compressor) == SparseSign + @test typeof(solver.compressor.cardinality) == Left + @test typeof(solver.log) == BasicLogger + @test typeof(solver.error) == FullResidual + end + + # test constructor + let solver = IHS( + alpha = 2.0, + compressor = ITestCompressor(), + log = ITestLog(), + error = ITestError(), + ) + + @test solver.alpha == 2.0 + @test typeof(solver.compressor) == ITestCompressor + @test typeof(solver.log) == ITestLog + @test typeof(solver.error) == ITestError + end + + # Test that error gets returned with right compressor + @test_logs (:warn, + "Compressor has cardinality `Right` but IHS compresses from the `Left`." + ) IHS( + alpha = 2.0, + compressor = ITestCompressor(Right(), 5), + log = ITestLog(), + error = ITestError(), + ) + end + + @testset "IHSRecipe" begin + @test supertype(IHSRecipe) == SolverRecipe + + # test fieldnames and types + @test fieldnames(IHSRecipe) == ( + :log, + :compressor, + :error, + :alpha, + :compressed_mat, + :mat_view, + :residual_vec, + :gradient_vec, + :buffer_vec, + :solution_vec, + :R + ) + @test fieldtypes(IHSRecipe) == ( + LoggerRecipe, + CompressorRecipe, + SolverErrorRecipe, + Float64, + AbstractArray, + SubArray, + AbstractVector, + AbstractVector, + AbstractVector, + AbstractVector, + UpperTriangular{Type, M} where {Type<:Number, M<:AbstractArray} + ) + end + + @testset "IHS: Complete Solver" begin + # test error method with no residual error + let A = A, + xsol = xsol, + b = b, + comp_dim = 2 * size(A, 2), + alpha = 1.0, + n_rows = size(A, 1), + n_cols = size(A, 2), + x = zeros(n_cols) + + comp = ITestCompressor(Left(), comp_dim) + log = ITestLog() + err = ITestErrorNoRes() + solver = IHS( + log = log, + compressor = comp, + error = err, + alpha = alpha + ) + + @test_throws ArgumentError( + "ErrorRecipe $(typeof(ITestErrorRecipeNoRes())) does not contain the \ + field 'residual' and is not valid for an IHS solver." + ) complete_solver(solver, x, A, b) + end + + # test logger method with no converged field + let A = A, + xsol = xsol, + b = b, + comp_dim = 2 * size(A, 2), + alpha = 1.0, + n_rows = size(A, 1), + n_cols = size(A, 2), + x = zeros(n_cols) + + comp = ITestCompressor(Left(), comp_dim) + log = ITestLogNoCov() + err = ITestError() + solver = IHS( + compressor = comp, + log = log, + error = err, + alpha = alpha + ) + + @test_throws ArgumentError( + "LoggerRecipe $(typeof(ITestLogRecipeNoCov())) does not contain \ + the field 'converged' and is not valid for an IHS solver." + ) complete_solver(solver, x, A, b) + end + + # Test the error message about not large enough compression_dim + let A = A, + xsol = xsol, + b = b, + comp_dim = 1, + alpha = 1.0, + n_rows = size(A, 1), + n_cols = size(A, 2), + x = zeros(n_cols) + + comp = ITestCompressor(Left(), comp_dim) + log = ITestLog() + err = ITestError() + solver = IHS( + compressor = comp, + log = log, + error = err, + alpha = alpha + ) + + @test_throws ArgumentError( + "Compression dimension not larger than column dimension this will lead to \ + singular QR decompositions, which cannot be inverted." + ) complete_solver(solver, x, A, b) + end + + let A = A, + xsol = xsol, + b = b, + comp_dim = 8 * size(A, 2), + alpha = 1.0, + n_rows = size(A, 1), + n_cols = size(A, 2), + x = zeros(n_cols) + + comp = ITestCompressor(Left(), comp_dim) + log = ITestLog() + err = ITestError() + solver = IHS( + compressor = comp, + log = log, + error = err, + alpha = alpha + ) + + solver_rec = complete_solver(solver, x, A, b) + + # test types of the contents of the solver + @test typeof(solver_rec) == IHSRecipe{ + Float64, + Main.IHSTest.ITestLogRecipe, + Main.IHSTest.ITestCompressorRecipe, + Main.IHSTest.ITestErrorRecipe, + Matrix{Float64}, + SubArray{ + Float64, + 2, + Matrix{Float64}, + Tuple{UnitRange{Int64}, Base.Slice{Base.OneTo{Int64}}}, + false + }, + Vector{Float64}, + } + @test typeof(solver_rec.compressor) == ITestCompressorRecipe + @test typeof(solver_rec.log) == ITestLogRecipe + @test typeof(solver_rec.error) == ITestErrorRecipe + @test typeof(solver_rec.alpha) == Float64 + @test typeof(solver_rec.compressed_mat) == Matrix{Float64} + @test typeof(solver_rec.solution_vec) == Vector{Float64} + @test typeof(solver_rec.buffer_vec) == Vector{Float64} + @test typeof(solver_rec.gradient_vec) == Vector{Float64} + @test typeof(solver_rec.residual_vec) == Vector{Float64} + @test typeof(solver_rec.mat_view) <: SubArray + @test typeof(solver_rec.R) <: UpperTriangular + # Test sizes of vectors and matrices + @test size(solver_rec.compressor) == (comp_dim, n_rows) + @test size(solver_rec.compressed_mat) == (comp_dim, n_cols) + @test size(solver_rec.buffer_vec) == (n_cols,) + @test size(solver_rec.solution_vec) == (n_cols,) + @test size(solver_rec.gradient_vec) == (n_cols,) + @test size(solver_rec.residual_vec) == (n_rows,) + + # test values of entries + solver_rec.alpha == alpha + solver_rec.solution_vec == x + solver_rec.buffer_vec == zeros(n_cols) + solver_rec.gradient_vec == zeros(n_cols) + solver_rec.residual_vec == zeros(n_rows) + end + + end + + @testset "IHS: rsolve!" begin + for type in [Float16, Float32, Float64, ComplexF32, ComplexF64] + # test maxit stop + let A = Array(qr(rand(type, n_rows, n_cols)).Q), + xsol = ones(type, n_cols), + b = A * xsol, + # need to choose compression dim to be large enough + comp_dim = 9 * size(A, 2), + alpha = 1.0, + n_rows = size(A, 1), + n_cols = size(A, 2), + x = zeros(type, n_cols) + x_st = deepcopy(x) + + comp = ITestCompressor(Left(), comp_dim) + #check 10 iterations + log = ITestLog(10, 0.0) + err = ITestError() + solver = IHS( + compressor = comp, + log = log, + error = err, + alpha = alpha + ) + + solver_rec = complete_solver(solver, x, A, b) + + result = rsolve!(solver_rec, x, A, b) + #test that the residual decreases + @test norm(A * x_st - b) > norm(A * x - b) + end + + # using threshold stop + let A = Array(qr(rand(type, n_rows, n_cols)).Q), + xsol = ones(type, n_cols), + b = A * xsol, + # need to choose compression dim to be large enough + comp_dim = 9 * size(A, 2), + alpha = 1.0, + n_rows = size(A, 1), + n_cols = size(A, 2), + x = zeros(type, n_cols) + x_st = deepcopy(x) + + comp = ITestCompressor(Left(), comp_dim) + #check 40 iterations we make a 1% improvement + log = ITestLog(40, norm(b) * .99) + err = ITestError() + solver = IHS( + compressor = comp, + log = log, + error = err, + alpha = alpha + ) + + solver_rec = complete_solver(solver, x, A, b) + + result = rsolve!(solver_rec, x, A, b) + #test that the error decreases + @test norm(A * x_st - b) > norm(A * x - b) + @test solver_rec.log.converged + end + + end + + end + +end + +end