diff --git a/Project.toml b/Project.toml index 07229d4..f63d04e 100644 --- a/Project.toml +++ b/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" diff --git a/benchmark/bench_map.jl b/benchmark/bench_map.jl index 6d623ce..0ffda89 100644 --- a/benchmark/bench_map.jl +++ b/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)) diff --git a/docs/Project.toml b/docs/Project.toml index 868af75..20f7833 100644 --- a/docs/Project.toml +++ b/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" diff --git a/examples/hard-core-lattice-gas/main.jl b/examples/hard-core-lattice-gas/main.jl index 2b82f01..14cc289 100644 --- a/examples/hard-core-lattice-gas/main.jl +++ b/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) diff --git a/src/cuda.jl b/ext/TensorInferenceCUDAExt.jl similarity index 71% rename from src/cuda.jl rename to ext/TensorInferenceCUDAExt.jl index d2c6fc1..40204fc 100644 --- a/src/cuda.jl +++ b/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 \ No newline at end of file diff --git a/src/TensorInference.jl b/src/TensorInference.jl index bfa714c..19125ba 100644 --- a/src/TensorInference.jl +++ b/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 diff --git a/src/cspmodels.jl b/src/cspmodels.jl new file mode 100644 index 0000000..128f684 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/generictensornetworks.jl b/src/generictensornetworks.jl deleted file mode 100644 index 2fc89e0..0000000 --- a/src/generictensornetworks.jl +++ /dev/null @@ -1,66 +0,0 @@ -using .GenericTensorNetworks: generate_tensors, GraphProblem, flavors, labels - -# update models -export update_temperature - -""" -$TYPEDSIGNATURES - -Convert a constraint satisfiability problem (or energy model) to a probabilistic model. - -### Arguments -* `problem` is a `GraphProblem` instance in [`GenericTensorNetworks`](https://github.com/QuEraComputing/GenericTensorNetworks.jl). -* `β` is the inverse temperature. -""" -function TensorInference.TensorNetworkModel(problem::GraphProblem, β::Real; evidence::Dict=Dict{eltype(labels(problem)),Int}(), - optimizer=GreedyMethod(), openvars=empty(labels(problem)), simplifier=nothing, mars=[[l] for l in labels(problem)]) - ixs = [GenericTensorNetworks.energy_terms(problem)..., GenericTensorNetworks.extra_terms(problem)...] - lbs = labels(problem) - nflavors = length(flavors(problem)) - # generate tensors for x = e^β - tensors = generate_tensors(exp(β), problem) - factors = [Factor((ix...,), t) for (ix, t) in zip(ixs, tensors)] - return TensorNetworkModel(lbs, fill(nflavors, length(lbs)), 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::GraphProblem, β::Real) - tensors = generate_tensors(exp(β), problem) - alltensors = [tnet.tensors[1:end-length(tensors)]..., tensors...] - return TensorNetworkModel(tnet.vars, tnet.code, alltensors, tnet.evidence, tnet.mars) -end - -function TensorInference.MMAPModel(problem::GraphProblem, β::Real; - queryvars, - openvars = empty(labels(problem)), - evidence = Dict{eltype(labels(problem)), Int}(), - optimizer = GreedyMethod(), simplifier = nothing, - marginalize_optimizer = GreedyMethod(), marginalize_simplifier = nothing - )::MMAPModel - ixs = [GenericTensorNetworks.energy_terms(problem)..., GenericTensorNetworks.extra_terms(problem)...] - nflavors = length(flavors(problem)) - # generate tensors for x = e^β - tensors = generate_tensors(exp(β), problem) - factors = [Factor((ix...,), t) for (ix, t) in zip(ixs, tensors)] - lbs = labels(problem) - return MMAPModel(lbs, fill(nflavors, length(lbs)), factors; queryvars, openvars, evidence, - optimizer, simplifier, - marginalize_optimizer, marginalize_simplifier) -end -function update_temperature(tnet::MMAPModel, problem::GraphProblem, β::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 - -@info "`TensorInference` loaded `GenericTensorNetworks` extension successfully, -`TensorNetworkModel` and `MMAPModel` can be used for converting a `GraphProblem` to a probabilistic model now." \ No newline at end of file diff --git a/src/mar.jl b/src/mar.jl index b39c7ad..a436d3d 100644 --- a/src/mar.jl +++ b/src/mar.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 diff --git a/src/utils.jl b/src/utils.jl index 48072b0..ce90819 100644 --- a/src/utils.jl +++ b/src/utils.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 diff --git a/test/Project.toml b/test/Project.toml index aa8a432..97225a0 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -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" diff --git a/test/generictensornetworks.jl b/test/cspmodels.jl similarity index 73% rename from test/generictensornetworks.jl rename to test/cspmodels.jl index 4732542..c7b559b 100644 --- a/test/generictensornetworks.jl +++ b/test/cspmodels.jl @@ -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 \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl index c6a333e..6bd7da2 100644 --- a/test/runtests.jl +++ b/test/runtests.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