Merge pull request #353 from JuliaGPU/tb/statistics

Port statistics functions from CUDA.jl.

Author: maleadt

Manifest.toml

@@ -30,5 +30,13 @@ uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 [[Serialization]]
 uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 
+[[SparseArrays]]
+deps = ["LinearAlgebra", "Random"]
+uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+
+[[Statistics]]
+deps = ["LinearAlgebra", "SparseArrays"]
+uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+
 [[Unicode]]
 uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"

Project.toml

@@ -9,6 +9,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
 AbstractFFTs = "0.4, 0.5, 1.0"

src/GPUArrays.jl

@@ -33,6 +33,7 @@ include("host/math.jl")
 include("host/random.jl")
 include("host/quirks.jl")
 include("host/uniformscaling.jl")
+include("host/statistics.jl")
 
 include("deprecated.jl")
 

src/host/statistics.jl

@@ -0,0 +1,43 @@
+using Statistics
+
+Statistics.varm(A::AbstractGPUArray{<:Real},m::AbstractArray{<:Real}; dims, corrected::Bool=true) =
+    sum((A .- m).^2, dims=dims)/(prod(size(A)[[dims...]])::Int-corrected)
+
+Statistics.stdm(A::AbstractGPUArray{<:Real},m::AbstractArray{<:Real}, dim::Int; corrected::Bool=true) =
+    sqrt.(varm(A,m;dims=dim,corrected=corrected))
+
+Statistics._std(A::AbstractGPUArray, corrected::Bool, mean, dims) =
+    sqrt.(Statistics.var(A; corrected=corrected, mean=mean, dims=dims))
+
+Statistics._std(A::AbstractGPUArray, corrected::Bool, mean, ::Colon) =
+    sqrt.(Statistics.var(A; corrected=corrected, mean=mean, dims=:))
+
+# Revert https://github.com/JuliaLang/Statistics.jl/pull/25
+Statistics._mean(A::AbstractGPUArray, ::Colon)    = sum(A) / length(A)
+Statistics._mean(f, A::AbstractGPUArray, ::Colon) = sum(f, A) / length(A)
+Statistics._mean(A::AbstractGPUArray, dims)    = mean!(Base.reducedim_init(t -> t/2, +, A, dims), A)
+Statistics._mean(f, A::AbstractGPUArray, dims) = sum(f, A, dims=dims) / mapreduce(i -> size(A, i), *, unique(dims); init=1)
+
+function Statistics.covzm(x::AbstractGPUMatrix, vardim::Int=1; corrected::Bool=true)
+    C = Statistics.unscaled_covzm(x, vardim)
+    T = promote_type(typeof(one(eltype(C)) / 1), eltype(C))
+    A = convert(AbstractArray{T}, C)
+    b = 1//(size(x, vardim) - corrected)
+    A .*= b
+    return A
+end
+
+function Statistics.cov2cor!(C::AbstractGPUMatrix{T}, xsd::AbstractGPUArray) where T
+    nx = length(xsd)
+    size(C) == (nx, nx) || throw(DimensionMismatch("inconsistent dimensions"))
+    tril!(C, -1)
+    C += adjoint(C)
+    C = Statistics.clampcor.(C ./ (xsd * xsd'))
+    C[diagind(C)] .= oneunit(T)
+    return C
+end
+
+function Statistics.corzm(x::AbstractGPUMatrix, vardim::Int=1)
+    c = Statistics.unscaled_covzm(x, vardim)
+    return Statistics.cov2cor!(c, sqrt.(diag(c)))
+end

test/testsuite.jl

@@ -73,6 +73,7 @@ include("testsuite/linalg.jl")
 include("testsuite/math.jl")
 include("testsuite/random.jl")
 include("testsuite/uniformscaling.jl")
+include("testsuite/statistics.jl")
 
 """
 Runs the entire GPUArrays test suite on array type `AT`

test/testsuite/statistics.jl

@@ -0,0 +1,52 @@
+using Statistics
+
+@testsuite "statistics" AT->begin
+    @testset "std" begin
+        @test compare(std, AT, rand(10))
+        @test compare(std, AT, rand(10,1,2))
+        @test compare(std, AT, rand(10,1,2); corrected=true)
+        @test compare(std, AT, rand(10,1,2); dims=1)
+    end
+
+    @testset "var" begin
+        @test compare(var, AT, rand(10))
+        @test compare(var, AT, rand(10,1,2))
+        @test compare(var, AT, rand(10,1,2); corrected=true)
+        @test compare(var, AT, rand(10,1,2); dims=1)
+        @test compare(var, AT, rand(10,1,2); dims=[1])
+        @test compare(var, AT, rand(10,1,2); dims=(1,))
+        @test compare(var, AT, rand(10,1,2); dims=[2,3])
+        @test compare(var, AT, rand(10,1,2); dims=(2,3))
+    end
+
+    @testset "mean" begin
+        @test compare(mean, AT, rand(2,2))
+        @test compare(mean, AT, rand(2,2); dims=2)
+        @test compare(mean, AT, rand(2,2,2); dims=[1,3])
+        @test compare(x->mean(sin, x), AT, rand(2,2))
+        @test compare(x->mean(sin, x; dims=2), AT, rand(2,2))
+        @test compare(x->mean(sin, x; dims=[1,3]), AT, rand(2,2,2))
+    end
+
+    @testset "cov" begin
+        s = 100
+        @test compare(cov, AT, rand(s))
+        @test compare(cov, AT, rand(Complex{Float64}, s))
+        @test compare(cov, AT, rand(s, 2))
+        @test compare(cov, AT, rand(Complex{Float64}, s, 2))
+        @test compare(cov, AT, rand(s, 2); dims=2)
+        @test compare(cov, AT, rand(Complex{Float64}, s, 2); dims=2)
+        @test compare(cov, AT, rand(1:100, s))
+    end
+
+    @testset "cor" begin
+        s = 100
+        @test compare(cor, AT, rand(s))
+        @test compare(cor, AT, rand(Complex{Float64}, s))
+        @test compare(cor, AT, rand(s, 2))
+        @test compare(cor, AT, rand(Complex{Float64}, s, 2))
+        @test compare(cor, AT, rand(s, 2); dims=2)
+        @test compare(cor, AT, rand(Complex{Float64}, s, 2); dims=2)
+        @test compare(cor, AT, rand(1:100, s))
+    end
+end