Remove type restrictions that prevent PCA on Float32 and GPU arrays

[brendanjohnharris]

Currently, parts of the PCA calculation force concrete Vector and Matrix types. This prevents both fit(::PCA, x::Matrix{Float32}) or fit(::PCA, x::CuArray), both of which can be significantly faster for large arrays (besides not having to offload or convert these large Float32/CUDA arrays); a comparison is below.

This PR removes these concrete type annotations. It also revises choose_pcadim() to avoid scalar indexing, but has (?) to calculate the cumulative sum over the full explained variance (this would perform slightly worse than the current implementation, which breaks the sum after the optimal dimensionality is found).

[brendanjohnharris]

[brendanjohnharris]

using CUDA
using LinearAlgebra
using CairoMakie
using Distributions
using MultivariateStats

function gencov(N, n; σ = 5, η = 0.01)
    _Σ² = zeros(N, N)
    [_Σ²[i, i] = rand(1:σ) for i in 1:n]
    _Σ² = _Σ² .+ Diagonal(abs.(randn(N)) .* η)
    Q, R = qr(randn(N, N))
    R = Q * Diagonal(sign.(diag(R))) # Random rotation matrix
    Σ² = Q * _Σ² * Q'
    if η == 0
        @assert all(real.(eigvals(Σ²)[(N - n + 1):N]) .> 0.5)
        @assert all(.≈(real.(eigvals(Σ²)[1:(N - n)]), 0, atol = 1e-10))
    end
    if eltype(eigvals(Σ²)) <: Complex
        @warn "gencov produced complex eigenvalues, take care"
    end
    return round.(Σ², sigdigits = 10) # Round to avoid numerical issues
end

if CUDA.functional()
    N = 50
    n = 5 # The latent dimensionality
    method = :svd
    Σ² = Symmetric(gencov(N, n))
    μ = randn(size(Σ², 1))
    D = MvNormal(μ, Σ²)

    Ns = round.(Int, exp10.(1.5:0.5:6.5))
    ts = map(Ns) do N
        GC.gc()
        @info N
        x = rand(D, N)'
        t = @timed fit(PCA, x; method, mean = 0)
        x = Float32.(x)
        t32 = @timed fit(PCA, x; method, mean = 0)
        x = CuArray(x)
        tgpu = @timed CUDA.@sync fit(PCA, x; method, mean = 0)
        atol = 1e-3
        q = map(t.value.proj |> collect |> eachcol,
                t32.value.proj |> collect |> eachcol,
                tgpu.value.proj |> collect |> eachcol) do a, b, c
            (≈(a, b; atol) || ≈(a, -b; atol)) && (≈(a, c; atol) || ≈(a, -c; atol))
        end
        return t.time, t32.time, tgpu.time, all(q)
    end
    t = getindex.(ts, 1)[2:end]
    t32 = getindex.(ts, 2)[2:end]
    tgpu = getindex.(ts, 3)[2:end]
    q = getindex.(ts, 4)
    @assert all(q) == true # Checks that the fit PCAs are similar

    f = Figure()
    ax = Axis(f[1, 1]; yscale = log10, xscale = log10, xlabel = "N",
              ylabel = "Time (s)")
    scatterlines!(ax, Ns[2:end], t, label = "CPU (Float64)", color = :crimson)
    scatterlines!(ax, Ns[2:end], t32, label = "CPU (Float32)", color = :cornflowerblue)
    scatterlines!(ax, Ns[2:end], tgpu, label = "GPU (Float32)", color = :green)
    axislegend(ax; position = :lt)
    f
    save("./test/testfigures/cuda.pdf", f)
end