Switch from Requires to extensions (#81)

* Create CUDA and GenericTensorNetworks extensions

* update project version

* fix document

* fix documentation

* prevent long vector

* fix the energy mode

Author: GiggleLiu

Project.toml

@@ -5,29 +5,32 @@ version = "0.4.2"
 
 [deps]
 Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
-GenericTensorNetworks = "3521c873-ad32-4bb4-b63d-f4f178f42b49"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OMEinsum = "ebe7aa44-baf0-506c-a96f-8464559b3922"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
+ProblemReductions = "899c297d-f7d2-4ebf-8815-a35996def416"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 TropicalNumbers = "b3a74e9c-7526-4576-a4eb-79c0d4c32334"
 
+[weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+
+[extensions]
+TensorInferenceCUDAExt = "CUDA"
+
 [compat]
 Artifacts = "1"
 CUDA = "4, 5"
 DocStringExtensions = "0.8.6, 0.9"
-GenericTensorNetworks = "2"
 LinearAlgebra = "1"
 OMEinsum = "0.8"
 Pkg = "1"
 PrecompileTools = "1"
 PrettyTables = "2"
-Requires = "1"
+ProblemReductions = "0.3"
 StatsBase = "0.34"
 TropicalNumbers = "0.5.4, 0.6"
-julia = "1.3"
+julia = "1.9"

benchmark/bench_map.jl

@@ -6,7 +6,7 @@ using Artifacts
 
 const SUITE = BenchmarkGroup()
 
-problem = problem_from_artifact("uai2014", "MAR" "Promedus", 14)
+problem = problem_from_artifact("uai2014", "MAR", "Promedus", 14)
 
 optimizer = TreeSA(ntrials = 1, niters = 2, βs = 1:0.1:40)
 tn = TensorNetworkModel(read_model(problem); optimizer, evidence=get_evidence(problem))

docs/Project.toml

@@ -1,7 +1,9 @@
 [deps]
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
-GenericTensorNetworks = "3521c873-ad32-4bb4-b63d-f4f178f42b49"
+Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 LiveServer = "16fef848-5104-11e9-1b77-fb7a48bbb589"
+LuxorGraphPlot = "1f49bdf2-22a7-4bc4-978b-948dc219fbbc"
+ProblemReductions = "899c297d-f7d2-4ebf-8815-a35996def416"
 TensorInference = "c2297e78-99bd-40ad-871d-f50e56b81012"
 TikzPictures = "37f6aa50-8035-52d0-81c2-5a1d08754b2d"

examples/hard-core-lattice-gas/main.jl

@@ -15,16 +15,17 @@
 
 a, b = (1, 0), (0.5, 0.5*sqrt(3))
 Na, Nb = 10, 10
-sites = vec([50 .* (a .* i .+ b .* j) for i=1:Na, j=1:Nb])
+sites = vec([50 .* (a .* i .+ b .* j) for i=1:Na, j=1:Nb]);
 
 # There exists blockade interactions between hard-core particles.
 # We connect two lattice sites within blockade radius by an edge.
 # Two ends of an edge can not both be occupied by particles.
-blockade_radius = 55
-using GenericTensorNetworks: show_graph, unit_disk_graph
-using GenericTensorNetworks.Graphs: edges, nv
-graph = unit_disk_graph(vec(sites), blockade_radius)
-show_graph(graph, sites; texts=fill("", length(sites)))
+blockade_radius = 55.0
+using LuxorGraphPlot: show_graph, GraphDisplayConfig
+using Graphs: edges, nv, SimpleGraph
+using TensorInference.ProblemReductions: UnitDiskGraph, IndependentSet
+graph = UnitDiskGraph(vec(sites), blockade_radius)
+show_graph(SimpleGraph(graph), sites; texts=fill("", length(sites)))
 
 # These constraints defines an independent set problem that characterized by the following energy based model.
 # Let $G = (V, E)$ be a graph, where $V$ is the set of vertices and $E$ is the set of edges.
@@ -38,7 +39,6 @@ show_graph(graph, sites; texts=fill("", length(sites)))
 # The solution space hard-core lattice gas is equivalent to that of an independent set problem.
 # The independent set problem involves finding a set of vertices in a graph such that no two vertices in the set are adjacent (i.e., there is no edge connecting them).
 # One can create a tensor network based modeling of an independent set problem with package [`GenericTensorNetworks.jl`](https://github.com/QuEraComputing/GenericTensorNetworks.jl).
-using GenericTensorNetworks
 problem = IndependentSet(graph)
 
 # There are plenty of discussions related to solution space properties in the `GenericTensorNetworks` [documentaion page](https://queracomputing.github.io/GenericTensorNetworks.jl/dev/generated/IndependentSet/).
@@ -59,14 +59,14 @@ partition_func[]
 
 # The marginal probabilities can be computed with the [`marginals`](@ref) function, which measures how likely a site is occupied.
 mars = marginals(pmodel)
-show_graph(graph, sites; vertex_colors=[(b = mars[[i]][2]; (1-b, 1-b, 1-b)) for i in 1:nv(graph)], texts=fill("", nv(graph)))
+show_graph(SimpleGraph(graph), sites; vertex_colors=[(b = mars[[i]][2]; (1-b, 1-b, 1-b)) for i in 1:nv(graph)], texts=fill("", nv(graph)))
 # The can see the sites at the corner is more likely to be occupied.
 # To obtain two-site correlations, one can set the variables to query marginal probabilities manually.
 pmodel2 = TensorNetworkModel(problem, β; mars=[[e.src, e.dst] for e in edges(graph)])
 mars = marginals(pmodel2);
 
 # We show the probability that both sites on an edge are not occupied
-show_graph(graph, sites; edge_colors=[(b = mars[[e.src, e.dst]][1, 1]; (1-b, 1-b, 1-b)) for e in edges(graph)], texts=fill("", nv(graph)), config=GraphDisplayConfig(; edge_line_width=5))
+show_graph(SimpleGraph(graph), sites; edge_colors=[(b = mars[[e.src, e.dst]][1, 1]; (1-b, 1-b, 1-b)) for e in edges(graph)], texts=fill("", nv(graph)), config=GraphDisplayConfig(; edge_line_width=5))
 
 # ## The most likely configuration
 # The MAP and MMAP can be used to get the most likely configuration given an evidence.
@@ -77,7 +77,7 @@ mars = marginals(pmodel3)
 logp, config = most_probable_config(pmodel3)
 
 # The log probability is 102. Let us visualize the configuration.
-show_graph(graph, sites; vertex_colors=[(1-b, 1-b, 1-b) for b in config], texts=fill("", nv(graph)))
+show_graph(SimpleGraph(graph), sites; vertex_colors=[(1-b, 1-b, 1-b) for b in config], texts=fill("", nv(graph)))
 # The number of particles is
 sum(config)
 
@@ -86,7 +86,7 @@ pmodel3 = TensorNetworkModel(problem, β; evidence=Dict(1=>0))
 logp2, config2 = most_probable_config(pmodel)
 
 # The log probability is 99, which is much smaller.
-show_graph(graph, sites; vertex_colors=[(1-b, 1-b, 1-b) for b in config2], texts=fill("", nv(graph)))
+show_graph(SimpleGraph(graph), sites; vertex_colors=[(1-b, 1-b, 1-b) for b in config2], texts=fill("", nv(graph)))
 # The number of particles is
 sum(config2)
 

src/cuda.jl renamed to ext/TensorInferenceCUDAExt.jl

@@ -1,4 +1,7 @@
-using .CUDA: CuArray
+module TensorInferenceCUDAExt
+using CUDA: CuArray
+import CUDA
+import TensorInference: match_arraytype, keep_only!, onehot_like, togpu
 
 function onehot_like(A::CuArray, j)
     mask = zero(A)
@@ -15,3 +18,7 @@ function keep_only!(x::CuArray{T}, j) where T
     CUDA.@allowscalar x[j] = hotvalue
     return x
 end
+
+togpu(x::AbstractArray) = CuArray(x)
+
+end

src/TensorInference.jl

@@ -12,6 +12,8 @@ using DocStringExtensions, TropicalNumbers
 # The Tropical GEMM support
 using StatsBase
 using PrettyTables
+using ProblemReductions
+
 import Pkg
 
 # reexport OMEinsum functions
@@ -35,6 +37,9 @@ export sample
 # MMAP
 export MMAPModel
 
+# for ProblemReductions
+export update_temperature
+
 # utils
 export random_matrix_product_state
 
@@ -45,12 +50,7 @@ include("mar.jl")
 include("map.jl")
 include("mmap.jl")
 include("sampling.jl")
-
-using Requires
-function __init__()
-    @require CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" include("cuda.jl")
-    @require GenericTensorNetworks = "3521c873-ad32-4bb4-b63d-f4f178f42b49" include("generictensornetworks.jl")
-end
+include("cspmodels.jl")
 
 # import PrecompileTools
 # PrecompileTools.@setup_workload begin

src/cspmodels.jl

@@ -0,0 +1,71 @@
+# generate the tensors for the constraint satisfaction problem
+function generate_tensors(β::T, problem::ConstraintSatisfactionProblem) where T <: Real
+    cons = ProblemReductions.constraints(problem)
+    objs = ProblemReductions.objectives(problem)
+    ixs = vcat([t.variables for t in cons], [t.variables for t in objs])
+	# generate tensors for x = e^β
+    x = energy_mode(problem) === LargerSizeIsBetter() ? exp(β) : exp(-β)
+    tensors = vcat(
+        Array{T}[reshape(map(s -> s ? one(x) : zero(x), t.specification), ntuple(i->num_flavors(problem), length(t.variables))) for t in cons],
+        Array{T}[reshape(map(s -> x^s, t.specification), ntuple(i->num_flavors(problem), length(t.variables))) for t in objs]
+    )
+    return tensors, ixs
+end
+
+"""
+$TYPEDSIGNATURES
+
+Convert a constraint satisfiability problem (or energy model) to a probabilistic model.
+
+### Arguments
+* `problem` is a `ConstraintSatisfactionProblem` instance in [`ProblemReductions`](https://github.com/GiggleLiu/ProblemReductions.jl).
+* `β` is the inverse temperature.
+
+### Keyword Arguments
+* `evidence` is a dictionary mapping variables to their values.
+* `optimizer` is the optimizer used to optimize the tensor network.
+* `openvars` is the list of variables to be marginalized.
+* `mars` is the list of variables to be marginalized.
+"""
+function TensorNetworkModel(problem::ConstraintSatisfactionProblem, β::T; evidence::Dict=Dict{Int,Int}(),
+        optimizer=GreedyMethod(), openvars=Int[], simplifier=nothing, mars=[[l] for l in variables(problem)]) where T <: Real
+    tensors, ixs = generate_tensors(β, problem)
+    factors = [Factor((ix...,), t) for (ix, t) in zip(ixs, tensors)]
+	return TensorNetworkModel(variables(problem), fill(num_flavors(problem), num_variables(problem)), factors; openvars, evidence, optimizer, simplifier, mars)
+end
+
+"""
+$TYPEDSIGNATURES
+
+Update the temperature of a tensor network model.
+The program will regenerate tensors from the problem, without repeated optimizing the contraction order.
+
+### Arguments
+- `tnet` is the [`TensorNetworkModel`](@ref) instance.
+- `problem` is the target constraint satisfiability problem.
+- `β` is the inverse temperature.
+"""
+function update_temperature(tnet::TensorNetworkModel, problem::ConstraintSatisfactionProblem, β::Real)
+	tensors, ixs = generate_tensors(β, problem)
+    alltensors = [tnet.tensors[1:length(tnet.mars)]..., tensors...]
+    return TensorNetworkModel(tnet.vars, tnet.code, alltensors, tnet.evidence, tnet.mars)
+end
+
+function MMAPModel(problem::ConstraintSatisfactionProblem, β::Real;
+            queryvars,
+            openvars = Int[],
+            evidence = Dict{Int, Int}(),
+            optimizer = GreedyMethod(), simplifier = nothing,
+            marginalize_optimizer = GreedyMethod(), marginalize_simplifier = nothing
+        )::MMAPModel
+	# generate tensors for x = e^β
+	tensors, ixs = generate_tensors(β, problem)
+    factors = [Factor((ix...,), t) for (ix, t) in zip(ixs, tensors)]
+    return MMAPModel(variables(problem), fill(num_flavors(problem), num_variables(problem)), factors; queryvars, openvars, evidence,
+        optimizer, simplifier,
+        marginalize_optimizer, marginalize_simplifier)
+end
+function update_temperature(tnet::MMAPModel, problem::ConstraintSatisfactionProblem, β::Real)
+    error("We haven't got time to implement setting temperatures for `MMAPModel`.
+It is about one or two hours of works. If you need it, please file an issue to let us know: https://github.com/TensorBFS/TensorInference.jl/issues")
+end

src/generictensornetworks.jl

@@ -7,7 +7,7 @@ function adapt_tensors(code, tensors, evidence; usecuda, rescale)
     map(tensors, ixs) do t, ix
         dims = map(ixi -> ixi ∉ keys(evidence) ? Colon() : ((evidence[ixi] + 1):(evidence[ixi] + 1)), ix)
         t2 = t[dims...]
-        t3 = usecuda ? CuArray(t2) : t2
+        t3 = usecuda ? togpu(t2) : t2
         rescale ? rescale_array(t3) : t3
     end
 end

src/mar.jl

@@ -306,6 +306,8 @@ function get_artifact_path(artifact_name::String)
     return Pkg.Artifacts.artifact_path(artifact_hash)
 end
 
+togpu(x) = error("You must import CUDA with `using CUDA` before using GPU!")
+
 """
 $TYPEDSIGNATURES
 

src/utils.jl

@@ -6,6 +6,7 @@ KaHyPar = "2a6221f6-aa48-11e9-3542-2d9e0ef01880"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OMEinsum = "ebe7aa44-baf0-506c-a96f-8464559b3922"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+ProblemReductions = "899c297d-f7d2-4ebf-8815-a35996def416"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/Project.toml

@@ -1,5 +1,6 @@
 using Test
-using GenericTensorNetworks, TensorInference
+using TensorInference, ProblemReductions.Graphs
+using GenericTensorNetworks
 
 @testset "marginals" begin
     # compute the probability
@@ -24,4 +25,10 @@ using GenericTensorNetworks, TensorInference
     model = MMAPModel(problem, β; queryvars=[1,4])
     logp, config = most_probable_config(model)
     @test config == [0, 0]
+
+    β = 1.0
+    problem = SpinGlass(g, -ones(Int, ne(g)), zeros(Int, nv(g)))
+    model = TensorNetworkModel(problem, β; mars=[[2, 3]])
+    samples = sample(model, 100)
+    @test sum(energy.(Ref(problem), samples))/100 <= -14
 end

test/generictensornetworks.jl renamed to test/cspmodels.jl

@@ -20,8 +20,8 @@ end
     include("sampling.jl")
 end
 
-@testset "generic tensor networks" begin
-    include("generictensornetworks.jl")
+@testset "cspmodels" begin
+    include("cspmodels.jl")
 end
 
 using CUDA