An online windowed approximation for auto covariance in Julia

onlinewindowedapproximateautocovariance.jl

# to run
# install Julia: https://julialang.org/
# create gist.jl in dir
# $ cd dir
# $ julia gist.jl
# tested on macOS 12.2 and Julia 1.7.2

struct OnlineCov{T <: AbstractFloat}
    mx::Vector{T}
    my::Vector{T}
    C::Vector{T}
    n::Vector{Int}
end

OnlineCov(T, D) = OnlineCov(zeros(T, D),
                         zeros(T, D),
                         zeros(T, D),
                         zeros(Int, 1))

OnlineCov(D) = OnlineCov(Float64, D)

function accumulate!(oc::OnlineCov, x::AbstractVector, y::AbstractVector)
    oc.n[1] += 1
    for d in eachindex(x, y)
        oc.mx[d] += (x[d] - oc.mx[d]) / oc.n[1]
        dy = y[d] - oc.my[d]
        oc.my[d] += dy / oc.n[1]
        oc.C[d] += (x[d] - oc.mx[d]) * dy
    end
end

function onlinecov(oc::OnlineCov)
    return oc.C ./ oc.n[1]
end

struct OnlineWindowedApproximateAutoCovariance{T <: AbstractFloat}
    cov::Vector{OnlineCov{T}}
    buffer::Matrix{T}
    N::Vector{Int}
    W::Int
end

function OnlineWindowedApproximateAutoCovariance(T, W, D)
    return OnlineWindowedApproximateAutoCovariance(
        [OnlineCov(T, D) for w in 1:W],
        zeros(T, W, D),
        zeros(Int, 1),
        W
    )
end

OnlineWindowedApproximateAutoCovariance(W, D = 1) = OnlineWindowedApproximateAutoCovariance(Float64, W, D)

function accumulate!(owaac::OnlineWindowedApproximateAutoCovariance,
                     x::AbstractVector)
    owaac.N[1] += 1
    n = owaac.N[1]
    W = n
    if n < owaac.W
        owaac.buffer[n, :] .= x
    else
        W = owaac.W
        for w in 1:(W-1)
            owaac.buffer[w, :] .= owaac.buffer[w + 1, :]
        end
        owaac.buffer[W, :] .= x
    end
    for w in 1:W
        for i in 1:(W-w+1)
            accumulate!(owaac.cov[w],
                        owaac.buffer[i, :],
                        owaac.buffer[i + w - 1, :])
        end
    end
end

function onlineautocov(owaac::OnlineWindowedApproximateAutoCovariance)
    oac = zero(owaac.buffer)
    for w in axes(oac, 1)
        oac[w, :] .= onlinecov(owaac.cov[w])
    end
    return oac
end

# O(N^2) Autocovariance used as baseline
function autocov(x)
    N = length(x)
    xc = x .- (sum(x) / N);
    ac = zero(xc)
    for n in 1:N
        for i in 1:N-n+1
            ac[n] += xc[i] * xc[i + n - 1]
        end
    end
    return ac ./ N
end

#################################### TESTS #####################################

W = 10;
N = 1_000;
D = 1;
x = randn(N, D);
wac = OnlineWindowedApproximateAutoCovariance(W, D);

# Round 1
println("# Number of Samples $N, Window size $W\n")


## run time
autocov(x);                     # compile and then time
tac = @elapsed ac = autocov(x);

println("## Time (seconds) for O(N^2) Autocovariance")
println("     $tac")

accumulate!(wac, x[1, :])       # compile and then time
wac = OnlineWindowedApproximateAutoCovariance(W);

otac = @elapsed begin for n in axes(x, 1)
    accumulate!(wac, x[n, :])
end
end

println("## Time (seconds) for Online Windowed Approximate Autocovariance")
println("    $otac \n")


## side by side comparison
oac = onlineautocov(wac);
println("## This    O(N^2) Autocovariance")
display("text/plain", round.([oac first(ac, W)], digits = 4))


## errors in w
println("\n")
println("## Error Increases in W")
for w in 1:W
    error = round(sqrt(sum( (oac[1:w] .- first(ac, w)) .^ 2)), digits = 4)
    println("error in $w/$W estimates: $error")
end


# Round 2
println()
N = 50_000;
println("# Error decreases in N: new number of samples $N, window size $W\n")
x = randn(N, D);


## run time
tac = @elapsed ac = autocov(x);

println("## Time (seconds) for O(N^2) Autocovariance")
println("     $tac")

wac = OnlineWindowedApproximateAutoCovariance(W, D);

otac = @elapsed begin for n in axes(x, 1)
    accumulate!(wac, x[n, :])
end
end

println("## Time (seconds) for Online Windowed Approximate Autocovariance")
println("    $otac")
println("scales as O(N * W^2)\n")


## side by side comparison
oac = onlineautocov(wac);
println("## This    O(N^2) Autocovariance")
display("text/plain", round.([oac first(ac, W)], digits = 4))


## overall error
println("\n")
error = sqrt(sum( (oac .- first(ac, W)) .^ 2))
println("## Error in Estimates: $error")