cov.jl

using Tullio, Statistics

function tcov(x::AbstractMatrix; corrected=true, dims)
  @assert dims==2
  pre = 1/(size(x,2)-corrected)
  xbar=vec(mean(x, dims=2))
  @tullio Q[i,j] :=pre * (x[i,k] - xbar[i]) * (x[j,k] - xbar[j]) # verbose=true  # sum(k)
end

mat = 100randn(4,5) # .+ im.*rand.();  # wrong for complex
mat2 = 100randn(4,5)
cov(mat; dims=2)
tcov(mat; dims=2)
mcov(mat; dims=2)

function mcov(x::AbstractMatrix; corrected=true, dims)
    @assert dims==2
    pre = 1/(size(x,2)-corrected)
    xcorr = x .- mean(x, dims=2)
    z = pre * xcorr * adjoint(xcorr)  # right for complex
end

function mcov(x::AbstractMatrix, y::AbstractMatrix; corrected=true, dims)
    @assert dims==2
    pre = 1/(size(x,2)-corrected)
    xcorr = x .- mean(x, dims=2)
    ycorr = y .- mean(y, dims=2)
    z = pre * xcorr * adjoint(ycorr)
end

mat = 100randn(4,5) .+ im.*rand.();
mat2 = 100randn(4,5) .+ im.*rand.();
cov(mat; dims=2)
mcov(mat; dims=2)

cov(mat, mat2; dims=2)
mcov(mat, mat2; dims=2)

using ChainRulesCore, Diffractor, ForwardDiff

function ChainRulesCore.rrule(::typeof(mcov), x::AbstractMatrix{T}; corrected=true, dims) where {T}
    @assert dims==2
    pre = one(float(T))/(size(x,2) - corrected)
    xcorr = x .- mean(x, dims=2)
    z = pre * xcorr * adjoint(xcorr)
    z, dz -> (NoTangent(), Hermitian(unthunk(dz)) * xcorr * 2pre)  # not sure about complex case
end

ForwardDiff.gradient(x -> real(sum(sin, mcov(x; dims=2))), real(mat))
Diffractor.gradient(x -> real(sum(sin, mcov(x; dims=2))), real(mat))[1]

function ChainRulesCore.rrule(::typeof(mcov), x::AbstractMatrix{T}, y::AbstractMatrix; corrected=true, dims) where {T}
    @assert dims==2
    pre = one(float(T))/(size(x,2) - corrected)
    xcorr = x .- mean(x, dims=2)
    ycorr = y .- mean(y, dims=2)
    z = pre * xcorr * adjoint(ycorr)
    z, dz -> (NoTangent(), dz * ycorr * pre, adjoint(dz) * xcorr * pre)
end

ForwardDiff.gradient(x -> real(sum(sin, mcov(x, real(mat2); dims=2))), real(mat))
Diffractor.gradient((x,y) -> real(sum(sin, mcov(x, y; dims=2))), real(mat), real(mat2))[1]

ForwardDiff.gradient(y -> real(sum(sin, mcov(real(mat), y; dims=2))), real(mat2))
Diffractor.gradient((x,y) -> real(sum(sin, mcov(x, y; dims=2))), real(mat), real(mat2))[2]