save and load configurations

Author: GiggleLiu

.gitignore

@@ -1,3 +1,5 @@
 /Manifest.toml
 *.swp
-_*
+_*
+_test.bin
+_test.txt

Project.toml

@@ -6,6 +6,7 @@ version = "0.1.0"
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Compose = "a81c6b42-2e10-5240-aca2-a61377ecd94b"
+DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 LightGraphs = "093fc24a-ae57-5d10-9952-331d41423f4d"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/arithematics.jl

@@ -197,4 +197,4 @@ function onehotv(::Type{CountingTropical{TV,BS}}, x, v) where {TV,BS}
     CountingTropical{TV,BS}(one(TV), onehotv(BS, x, v))
 end
 onehotv(::Type{ConfigEnumerator{N,S,C}}, i::Integer, v) where {N,S,C} = ConfigEnumerator([onehotv(StaticElementVector{N,S,C}, i, v)])
-onehotv(::Type{ConfigSampler{N,S,C}}, i::Integer, v) where {N,S,C} = ConfigSampler(onehotv(StaticElementVector{N,S,C}, i, v))
+onehotv(::Type{ConfigSampler{N,S,C}}, i::Integer, v) where {N,S,C} = ConfigSampler(onehotv(StaticElementVector{N,S,C}, i, v))

src/bitvector.jl

@@ -65,10 +65,10 @@ end
 ##### BitVectors
 const StaticBitVector{N,C} = StaticElementVector{N,1,C}
 @inline function Base.getindex(x::StaticBitVector{N,C}, i::Integer) where {N,C}
-    @boundscheck i <= N || throw(BoundsError(x, i))  # TODO: make this @boundscheck work.
+    @boundscheck (i <= N || throw(BoundsError(x, i)))  # NOTE: still checks bounds in global scope, why?
     i -= 1
     ii = i ÷ 64
-    @inbounds (x.data[ii+1] >> (i-ii*64)) & 1
+    return @inbounds (x.data[ii+1] >> (i-ii*64)) & 1
 end
 
 function StaticBitVector(x::AbstractVector)

src/configurations.jl

@@ -89,5 +89,4 @@ end
 for GP in [:Independence, :Matching, :MaximalIndependence, :Coloring]
     @eval symbols(gp::$GP) = labels(gp.code)
 end
-symbols(gp::MaxCut) = collect(OMEinsum.getixs(OMEinsum.flatten(gp.code)))
-# TODO: coloring
+symbols(gp::MaxCut) = collect(OMEinsum.getixs(OMEinsum.flatten(gp.code)))

src/cuda.jl

@@ -23,3 +23,4 @@ for TT in [:(Tropical{<:NativeTypes}), :TropicalTypes]
        end
    end
 end
+

src/interfaces.jl

@@ -51,3 +51,75 @@ function solve(gp::GraphProblem, task; usecuda=false, kwargs...)
         error("unknown task $task.")
     end
 end
+
+export save_configs, load_configs
+using DelimitedFiles
+function save_configs(filename, data::ConfigEnumerator{N,S,C}; format::Symbol=:binary) where {N,S,C}
+    if format == :binary
+        write(filename, raw_matrix(data))
+    elseif format == :text
+        writedlm(filename, plain_matrix(data))
+    else
+        error("format must be `:binary` or `:text`, got `:$format`")
+    end
+end
+function load_configs(filename; len=nothing, format::Symbol=:binary, nflavors=2)
+    if format == :binary
+        len === nothing && error("you need to specify `len` for reading configurations from binary files.")
+        S = ceil(Int, log2(nflavors))
+        C = _nints(len, S)
+        return _from_raw_matrix(StaticElementVector{len,S,C}, reshape(reinterpret(UInt64, read(filename)),C,:))
+    elseif format == :text
+        return from_plain_matrix(readdlm(filename); nflavors=nflavors)
+    else
+        error("format must be `:binary` or `:text`, got `:$format`")
+    end
+end
+
+function raw_matrix(x::ConfigEnumerator{N,S,C}) where {N,S,C}
+    m = zeros(UInt64, C, length(x))
+    @inbounds for i=1:length(x), j=1:C
+        m[j,i] = x.data[i].data[j]
+    end
+    return m
+end
+function plain_matrix(x::ConfigEnumerator{N,S,C}) where {N,S,C}
+    m = zeros(UInt8, N, length(x))
+    @inbounds for i=1:length(x), j=1:N
+        m[j,i] = x.data[i][j]
+    end
+    return m
+end
+
+function from_raw_matrix(m; len, nflavors=2)
+    S = ceil(Int,log2(nflavors))
+    C = size(m, 1)
+    T = StaticElementVector{len,S,C}
+    @assert len*S <= C*64
+    _from_raw_matrix(T, m)
+end
+function _from_raw_matrix(::Type{StaticElementVector{N,S,C}}, m::AbstractMatrix) where {N,S,C}
+    data = zeros(StaticElementVector{N,S,C}, size(m, 2))
+    @inbounds for i=1:size(m, 2)
+        data[i] = StaticElementVector{N,S,C}(NTuple{C,UInt64}(view(m,:,i)))
+    end
+    return ConfigEnumerator(data)
+end
+function from_plain_matrix(m::Matrix; nflavors=2)
+    S = ceil(Int,log2(nflavors))
+    N = size(m, 1)
+    C = _nints(N, S)
+    T = StaticElementVector{N,S,C}
+    _from_plain_matrix(T, m)
+end
+function _from_plain_matrix(::Type{StaticElementVector{N,S,C}}, m::AbstractMatrix) where {N,S,C}
+    data = zeros(StaticElementVector{N,S,C}, size(m, 2))
+    @inbounds for i=1:size(m, 2)
+        data[i] = convert(StaticElementVector{N,S,C}, view(m, :, i))
+    end
+    return ConfigEnumerator(data)
+end
+
+# convert to Matrix
+Base.Matrix(ce::ConfigEnumerator) = plain_matrix(ce)
+Base.Vector(ce::StaticElementVector) = collect(ce)

test/cuda.jl

@@ -0,0 +1,21 @@
+using CUDA, Random
+using LinearAlgebra: mul!
+using GraphTensorNetworks, Test
+
+@testset "cuda patch" begin
+    for T in [Tropical{Float64}, CountingTropical{Float64,Float64}]
+        a = T.(CUDA.randn(4, 4))
+        b = T.(CUDA.randn(4))
+        for A in [transpose(a), a, transpose(b)]
+            for B in [transpose(a), a, b]
+                if !(size(A) == (1,4) && size(B) == (4,))
+                    res0 = Array(A) * Array(B)
+                    res1 = A * B
+                    res2 = mul!(CUDA.zeros(T, size(res0)...), A, B, true, false)
+                    @test Array(res1) ≈ res0
+                    @test Array(res2) ≈ res0
+                end
+            end
+        end
+    end
+end

test/interfaces.jl

@@ -30,4 +30,43 @@ using LightGraphs, Test
     @test res11 == res5
     @test res12.c.data ∈ res13.c.data
     @test res13.c.data == res7.c.data
+end
+
+@testset "save load" begin
+    M = 10
+    m = ConfigEnumerator([StaticBitVector(rand(Bool, 300)) for i=1:M])
+    bm = GraphTensorNetworks.plain_matrix(m)
+    rm = GraphTensorNetworks.raw_matrix(m)
+    m1 = GraphTensorNetworks.from_raw_matrix(rm; len=300, nflavors=2)
+    m2 = GraphTensorNetworks.from_plain_matrix(bm; nflavors=2)
+    @test m1 == m
+    @test m2 == m
+    save_configs("_test.bin", m; format=:binary)
+    @test_throws ErrorException load_configs("_test.bin"; format=:binary)
+    ma = load_configs("_test.bin"; format=:binary, len=300, nflavors=2)
+    @test ma == m
+
+    save_configs("_test.txt", m; format=:text)
+    mb = load_configs("_test.txt"; format=:text, nflavors=2)
+    @test mb == m
+
+    M = 10
+    m = ConfigEnumerator([StaticElementVector(3, rand(1:3, 300)) for i=1:M])
+    bm = GraphTensorNetworks.plain_matrix(m)
+    rm = GraphTensorNetworks.raw_matrix(m)
+    m1 = GraphTensorNetworks.from_raw_matrix(rm; len=300, nflavors=3)
+    m2 = GraphTensorNetworks.from_plain_matrix(bm; nflavors=3)
+    @test m1 == m
+    @test m2 == m
+    @test Matrix(m) == bm
+    @test Vector(m.data[1]) == bm[:,1]
+
+    save_configs("_test.bin", m; format=:binary)
+    @test_throws ErrorException load_configs("_test.bin"; format=:binary)
+    ma = load_configs("_test.bin"; format=:binary, len=300, nflavors=3)
+    @test ma == m
+
+    save_configs("_test.txt", m; format=:text)
+    mb = load_configs("_test.txt"; format=:text, nflavors=3)
+    @test mb == m
 end