Random matrix sampling on the GPU

Project.toml

@@ -3,20 +3,22 @@ uuid = "3c0b384b-479c-5684-b2ef-9d7a46dd931e"
 authors = ["Dariusz Kurzyk <dkurzyk@iitis.pl>", "Łukasz Pawela <lpawela@iitis.pl>", "Piotr Gawron <gawron@iitis.pl>", "Marcin Przewięźlikowski <m.przewie@gmail.com>"]
 version = "0.4.6"
 
-[compat]
-julia = "0.7, 1.0, 1.1, 1.2"
-
 [deps]
 Convex = "f65535da-76fb-5f13-bab9-19810c17039a"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SCS = "c946c3f1-0d1f-5ce8-9dea-7daa1f7e2d13"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 
+[compat]
+julia = "0.7, 1.0, 1.1, 1.2"
+
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-        
+
 [targets]
 test = ["Test"]

REQUIRE

@@ -4,3 +4,5 @@ DocStringExtensions
 TensorOperations 
 Convex 
 SCS 
+Pkg
+Requires

curandommatrices/src/CuRandomMatrices.jl

@@ -0,0 +1,15 @@
+module CuRandomMatrices
+export curand
+
+using LinearAlgebra
+using ..CuArrays
+using CUDAnative
+
+include("../../randommatrices/src/RandomMatrices.jl")
+using ..RandomMatrices
+
+include("ginibre.jl")
+include("circular.jl")
+include("wigner.jl")
+include("wishart.jl")
+end

curandommatrices/src/circular.jl

@@ -0,0 +1,50 @@
+# FIXME: this can be accelerated
+function cplx_phase!(a)
+    i = (blockIdx().x-1) * blockDim().x + threadIdx().x
+    if i <= length(a)
+        @inbounds a[i] = a[i] / sqrt(real(a[i])^2 + imag(a[i])^2)
+    end
+    return nothing
+end
+
+# FIXME: use blocks to support larger than 1024
+
+function _qr_fix!(z::CuMatrix)
+    q, r = CuArrays.qr!(z)
+    ph = diag(r)
+    len = min(length(ph), 1024) #hack, warpsize() segfaults
+    @cuda threads=len blocks=16 cplx_phase!(ph)
+    q = CuMatrix(q)
+    idim = size(r, 1)
+    for i = 1:idim
+        q[:, i] .*= ph[i]
+    end
+    q[:, 1:idim]
+end
+
+function _qr_fix(z::CuMatrix)
+    a = copy(z)
+    _qr_fix!(a)
+end
+
+function curand(c::COE)
+    z = curand(c.g)
+    u = _qr_fix!(z)
+    transpose(u)*u
+end
+
+function curand(c::CSE)
+    z = curand(c.g)
+    u = _qr_fix!(z)
+    ur = cat([CuMatrix{Float32}([0 -1; 1 0]) for _=1:c.d÷2]..., dims=[1,2])
+    ur*u*ur'*transpose(u)
+end
+
+for T in (CUE, CircularRealEnsemble, CircularQuaternionEnsemble, HaarIsometry)
+    @eval begin
+        function curand(c::$T)
+            z = curand(c.g)
+            _qr_fix!(z)
+        end
+    end
+end

curandommatrices/src/ginibre.jl

@@ -0,0 +1,10 @@
+curand(g::GinibreEnsemble{1}) = CuArrays.randn(g.m, g.n)
+curand(g::GinibreEnsemble{2}) = CuArrays.randn(g.m, g.n)+1im*CuArrays.randn(g.m, g.n)
+
+function curand(g::GinibreEnsemble{4})
+    q0=CuArrays.randn(g.m, g.n)
+    q1=CuArrays.randn(g.m, g.n)
+    q2=CuArrays.randn(g.m, g.n)
+    q3=CuArrays.randn(g.m, g.n)
+    [q0+1im*q1 q2+1im*q3; -q2+1im*q3 q0-1im*q1]
+end

curandommatrices/src/wigner.jl

@@ -0,0 +1,4 @@
+function curand(w::WignerEnsemble{β}) where β
+    z = curand(w.g)
+    (z + z') / 2sqrt(2β * w.d)
+end

curandommatrices/src/wishart.jl

@@ -0,0 +1,4 @@
+function curand(w::WishartEnsemble{β, K}) where {β, K}
+    z = curand(w.g)/sqrt(2β * w.d)
+    z*z'
+end

curandommatrices/test/circular.jl

@@ -0,0 +1,74 @@
+Random.seed!(42)
+
+@testset "CircularEnsemble" begin
+    @testset "CUE" begin
+        n = 10
+        c = CUE(n)
+        u = CuRandomMatrices.curand(c)
+        @test norm(u*u' - I) ≈ 0 atol=1e-5
+
+        n = 100
+        c = CUE(n)
+        steps = 100
+        r = zeros(steps, n)
+
+        for i=1:steps
+            u = collect(CuRandomMatrices.curand(c))
+            r[i, :] = angle.(eigvals(u))
+        end
+        r = vec(r)
+        h = normalize(fit(Histogram, r, weights(ones(size(r))), -π:0.1π:π, closed=:left))
+        @test all(isapprox.(h.weights, 1/2π, atol=0.01))
+    end
+
+    @testset "COE" begin
+        n = 10
+        c = COE(n)
+        o = CuRandomMatrices.curand(c)
+        @test norm(o*o' - I) ≈ 0 atol=1e-5
+
+        n = 100
+        c = COE(n)
+        steps = 100
+        r = zeros(steps, n)
+        for i=1:steps
+            o = collect(CuRandomMatrices.curand(c))
+            r[i, :] = angle.(eigvals(o))
+        end
+        r = vec(r)
+        h = normalize(fit(Histogram, r, weights(ones(size(r))), -π:0.1π:π, closed=:left))
+        @test all(isapprox.(h.weights, 1/2π, atol=0.1))
+    end
+
+    @testset "CircularRealEnsemble" begin
+        c = CircularRealEnsemble(10)
+        o = CuRandomMatrices.curand(c)
+        @test size(o) == (10, 10)
+        @test eltype(o) <: Real
+    end
+end
+
+    @testset "HaarIsometry" begin
+        idim = 2
+        odim = 3
+        c = HaarIsometry(idim, odim)
+        u = CuRandomMatrices.curand(c)
+        @test size(u) == (odim, idim)
+        @test isapprox(norm(u'*u - I), 0, atol=1e-6)
+        @test_throws ArgumentError HaarIsometry(odim, idim)
+
+    @testset "CSE" begin
+        n = 10
+        c = CSE(n)
+        o = CuRandomMatrices.curand(c)
+        @test norm(o*o' - I) ≈ 0 atol=1e-5
+        @test size(o) == (n, n)
+    end
+
+    @testset "Circular quaternion ensemble" begin
+        c = CircularQuaternionEnsemble(10)
+        u = CuRandomMatrices.curand(c)
+        @test size(u) == (20, 20)
+        @test isapprox(norm(u'*u - I), 0, atol=1e-5)
+    end
+end

curandommatrices/test/ginibre.jl

@@ -0,0 +1,23 @@
+Random.seed!(42)
+
+@testset "GinibreEnsemble" begin
+    g = GinibreEnsemble{1}(10, 20)
+    z = rand(g)
+    @test eltype(z) <: Real
+    @test size(z) == (10, 20)
+    @test GinibreEnsemble(10, 20) == GinibreEnsemble{2}(10, 20)
+    @test GinibreEnsemble(10) == GinibreEnsemble{2}(10, 10)
+
+    @test_throws ArgumentError GinibreEnsemble{4}(11, 21)
+    g = GinibreEnsemble{4}(10, 20)
+    z = rand(g)
+    @test size(z) == (20, 40)
+    @test eltype(z) == ComplexF64
+
+@testset "_qr_fix" begin
+    a = rand(2, 2)
+    u1 = RandomMatrices._qr_fix(a)
+    u2 = RandomMatrices._qr_fix!(a)
+    @test norm(u1 - u2) ≈ 0
+end
+end

curandommatrices/test/runtests.jl

@@ -0,0 +1,17 @@
+using Test
+using Random
+using StatsBase
+using LinearAlgebra
+using CuArrays
+
+include("../../randommatrices/src/RandomMatrices.jl")
+using ..RandomMatrices
+
+include("../src/CuRandomMatrices.jl")
+using .CuRandomMatrices
+
+my_tests = ["circular.jl", "ginibre.jl", "wigner.jl", "wishart.jl"]
+
+for my_test in my_tests
+    include(my_test)
+end

curandommatrices/test/wigner.jl

@@ -0,0 +1,11 @@
+Random.seed!(42)
+
+@testset "WignerEnsemble" begin
+    w = WignerEnsemble{1}(10)
+    z = rand(w)
+
+    @test size(z) == (10, 10)
+    @test eltype(z) <: Real
+    @test ishermitian(z)
+    @test WignerEnsemble(10) == WignerEnsemble{2}(10)
+end

curandommatrices/test/wishart.jl

@@ -0,0 +1,15 @@
+Random.seed!(42)
+
+@testset "WishartEnsemble" begin
+    w = WishartEnsemble{1, 0.1}(10)
+    z = rand(w)
+    ev = eigvals(z)
+
+    @test size(z) == (10, 10)
+    @test eltype(z) <: Real
+    @test length(ev[ev.>1e-5]) == 1
+    @test all(ev[ev.>1e-5] .> 0)
+
+    @test WishartEnsemble{1}(10) == WishartEnsemble{1, 1}(10)
+    @test WishartEnsemble(10) == WishartEnsemble{2, 1}(10)
+end

randommatrices/src/ginibre.jl

@@ -19,9 +19,9 @@ rand(rng::AbstractRNG, g::GinibreEnsemble{2}) = randn(rng, g.m, g.n)+1im*randn(r
 
 function rand(rng::AbstractRNG, g::GinibreEnsemble{4})
     # TODO: fix dimensions of blocks
-    q0=randn(rng, g.m, g.n)
-    q1=randn(rng, g.m, g.n)
-    q2=randn(rng, g.m, g.n)
-    q3=randn(rng, g.m, g.n)
+    q0 = randn(rng, g.m, g.n)
+    q1 = randn(rng, g.m, g.n)
+    q2 = randn(rng, g.m, g.n)
+    q3 = randn(rng, g.m, g.n)
     [q0+1im*q1 q2+1im*q3; -q2+1im*q3 q0-1im*q1]
 end

src/QuantumInformation.jl

@@ -8,7 +8,7 @@ using TensorOperations
 using Convex, SCS
 using Random: AbstractRNG, GLOBAL_RNG
 
-import Base: convert, size, length, kron, *, rand, show
+import Base: convert, size, length, kron, *, show, rand
 
 const ⊗ = kron
 
@@ -18,6 +18,11 @@ include("../randommatrices/src/RandomMatrices.jl")
 using .RandomMatrices
 eval(Expr(:export, names(RandomMatrices)...))
 
+using Requires
+@init @require CuArrays = "3a865a2d-5b23-5a0f-bc46-62713ec82fae" include("../curandommatrices/src/CuRandomMatrices.jl")
+@init @require CuArrays = "3a865a2d-5b23-5a0f-bc46-62713ec82fae" using ..CuRandomMatrices
+@init @require CuArrays = "3a865a2d-5b23-5a0f-bc46-62713ec82fae" export curand
+
 include("base.jl")
 include("randomqobjects.jl")
 include("gates.jl")