Feature request: Add randomized Gram-Schmidt alongside Classical and Modified Variants

This repository currently supports Classical and Modified Gram-Schmidt variants for building Krylov subspaces. I’d like to request the addition of a Randomized Gram-Schmidt variant, which uses randomized sketching to accelerate orthogonalization.

I have found this to be mainly useful in large-scale settings, this paper would agree.

An implementation of this papers code in matlab can be found here.

The core request would "just" be to have something like this in KrylovKit, of course having additional niceties (like in the matlab implementation), such as the flags for reducing the precision, to improve time in some steps.

Further, this would probably require some adaptions in the rest of KrylovKit as, for example, randomized Gram-Schmidt variants seem to profit from having a larger Krylov dimension (they choose 500, opposed to 30 in KrylovKit with MGS and CGS variants).

Sadly, I currently lack the time to work on something like this, but I still wanted suggest it here, maybe for later.

[Jutho]

Unfortunately I also don't have the bandwidth for this now. I would need to look into this approach more (as I am not familiar with it), but I am not even sure if the randomized approach and the sketching matrix work well in the KrylovKit context, where try to make minimal assumptions about what vector are (in particular, we cannot index or slice them).

[observant-owl]

An answer to @Jutho, if am not disturbing: Below is a very rudimentary proof-of-concept of Randomized Gram–Schmidt in Julia. AbstractVector or AbstractMatrix should work out in my opinion.

Since RGS uses a non-standard inner product ⟨x,y⟩_Θ := (Θx)ᵀ(Θy), InnerProductVec(x, (x,y) -> dot(Θ*x, Θ*y)) can already do the sketching. The problem is that each inner(vᵢ, w) call would re-sketch vᵢ from scratch, which is (too) expensive.

A better way could be a dedicated RandomizedGramSchmidt <: Orthogonalizer carrying Θ and maintaining S = ΘV incrementally, with orthogonalize!! solving the small k × j LS against S and appending Θq in lockstep. But as existing Orthogonalizer subtypes are immutable, S would need to live in a mutable struct or be threaded through the orthogonalize!! argument list.

Of course, there are probably additional problems arising, which I just didn't think of so far, since I only briefly looked at the code in KrylovKit and there are further adjustments for RGS like mixed precision, more/custom sketch matrices, a Block-RGS implementation, and so on...😅

So, before I perhaps do anything, I just wanted to ask whether with these possibly larger changes in mind, would this even still be in-scope for KrylovKit?
Also probably take my assessment with a grain of salt, since I am still in school, doing this for fun, and very much not a professional😅

Edit:
Okay, I looked a bit deeper, and there'd be more changes than I stated above: The statefulness of RGS would seep into every factorization.

This cache needed by RGS could either be stored in the Orthogonalizer or in each Factorization. When storing it in the Orthogonalizer, it is not really elegant because as far as I can do it, I have to store that cache also in the Factorization for restarts.
So maybe it'd be a bit nicer to store it in the Factorizations, but then one would need to refactor every Factorization to have a cache::C <: OrthCache, which would either be dispatched with NoOrthCache or RGSOrthCache. This cache would need to be maintained with something like on_expand!/on_shrink!/on_rotate! hooks for example.
Else there is mainly implementation details, and watching out that every norm or inner product is the sketched alternative for RGS.

---

Oh and my small piece of code I wrote, if it is relevant:

# Randomized Gram-Schmidt (RGS) - proof-of-concept


using LinearAlgebra
using Random

"""
    rademacher(rng, ::Type{T}, k, n) -> Matrix{T}

Rademacher sketching matrix with entries `±1/√k`, suitable as `Θ` for `rgs`.
"""
rademacher(rng::AbstractRNG, ::Type{T}, k::Integer, n::Integer) where {T} =
    (T(2) .* rand(rng, Bool, k, n) .- one(T)) ./ sqrt(T(k))

"""
    rgs(W; sketch_ratio=4, rng=Random.default_rng()) -> (Q, R)
    rgs(W, Θ) -> (Q, R)

Randomized Gram-Schmidt QR factorization (Balabanov & Grigori, Algorithm 1).

Given an `n × m` matrix `W` (n ≫ m) and a `k × n` sketching matrix `Θ` with
`m ≤ k ≪ n`, produces `Q (n × m)` and upper-triangular `R (m × m)` such that
`W ≈ Q*R` and `(Θ*Q)' * (Θ*Q) ≈ I`.
"""
function rgs(W::AbstractMatrix;
             sketch_ratio::Real = 4,
             rng::AbstractRNG = Random.default_rng())
    T = float(eltype(W))
    n, m = size(W)
    k = clamp(ceil(Int, sketch_ratio * m), m, n)
    return rgs(W, rademacher(rng, T, k, n))
end

function rgs(W::AbstractMatrix, Θ::AbstractMatrix)
    T = float(eltype(W))
    n, m = size(W)
    k = size(Θ, 1)
    size(Θ, 2) == n || throw(DimensionMismatch(
        "Θ has $(size(Θ, 2)) columns but W has $n rows"))
    m ≤ k ≤ n || throw(ArgumentError(
        "sketch dimension k=$k must satisfy m=$m ≤ k ≤ n=$n"))
    eltype(Θ) == T || throw(ArgumentError(
        "eltype(Θ) ($(eltype(Θ))) must match float(eltype(W)) ($T)"))
    Q = Matrix{T}(undef, n, m)
    R = zeros(T, m, m)
    S = Matrix{T}(undef, k, m)
    p = Vector{T}(undef, k)
    q = Vector{T}(undef, n)
    s = Vector{T}(undef, k)
    for i in 1:m
        wi = @view W[:, i]
        mul!(p, Θ, wi)
        copyto!(q, wi)
        if i > 1
            Si = @view S[:, 1:i-1]
            Q1 = @view Q[:, 1:i-1]
            Rcol = @view R[1:i-1, i]
            ldiv!(Rcol, qr(Si), p)
            mul!(q, Q1, Rcol, -one(T), one(T))
        end
        mul!(s, Θ, q)
        rii = norm(s)
        iszero(rii) && throw(ArgumentError("RGS breakdown at column $i"))
        R[i, i] = rii
        @. S[:, i] = s / rii
        @. Q[:, i] = q / rii
    end
    return Q, R
end