Support artifact downloading and reading (#50)

Author: mroavi

test/Artifacts.toml renamed to Artifacts.toml

@@ -4,10 +4,12 @@ authors = ["Jin-Guo Liu", "Martin Roa Villescas"]
 version = "0.1.0"
 
 [deps]
+Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OMEinsum = "ebe7aa44-baf0-506c-a96f-8464559b3922"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"

Project.toml

@@ -57,6 +57,7 @@ most_probable_config
 probability
 read_evidence_file
 read_instance
+read_instance_from_artifact
 read_model_file
 read_solution_file
 read_td_file

docs/src/api/public.md

@@ -11,13 +11,15 @@ using OMEinsum, LinearAlgebra
 using DocStringExtensions, TropicalNumbers
 # The Tropical GEMM support
 using StatsBase
+import Pkg
 
 # reexport OMEinsum functions
 export RescaledArray
 export contraction_complexity, TreeSA, GreedyMethod, KaHyParBipartite, SABipartite, MergeGreedy, MergeVectors
 
 # read and load uai files
-export read_model_file, read_td_file, read_evidence_file, read_solution_file, read_instance, UAIInstance
+export read_model_file, read_td_file, read_evidence_file, read_solution_file
+export read_instance, read_instance_from_artifact, UAIInstance
 export set_evidence!, set_query!
 
 # marginals

src/TensorInference.jl

@@ -2,12 +2,12 @@
 $(TYPEDSIGNATURES)
 
 Parse the problem instance found in `model_filepath` defined in the UAI model
-format.
+format. If the provided file path is empty, return `nothing`.
 
 The UAI file formats are defined in:
 https://uaicompetition.github.io/uci-2022/file-formats/
 """
-function read_model_file(model_filepath; factor_eltype = Float64)
+function read_model_file(model_filepath::AbstractString; factor_eltype = Float64)
     # Read the uai file into an array of lines
     str = open(model_filepath) do file
         read(file, String)
@@ -244,10 +244,79 @@ function read_instance(
     return UAIInstance(nvars, ncliques, cards, factors, obsvars, obsvals, queryvars, reference_solution)
 end
 
-function read_instance_from_string(uai::AbstractString; eltype = Float64)::UAIInstance
-    nvars, cards, ncliques, factors = read_model_string(uai; factor_eltype = eltype)
+function read_instance_from_string(model::AbstractString; eltype = Float64)::UAIInstance
+    nvars, cards, ncliques, factors = read_model_string(model; factor_eltype = eltype)
     return UAIInstance(nvars, ncliques, cards, factors, Int[], Int[], Int[], nothing)
 end
 
 # patch to get content by broadcasting into array, while keep array size unchanged.
 broadcasted_content(x) = asarray(content.(x), x)
+
+"""
+$(TYPEDSIGNATURES)
+
+Load a specific instance, identified by `task/name`, from the provided
+artifact.
+"""
+function read_instance_from_artifact(
+    artifact_name::AbstractString,
+    problem_name::AbstractString,
+    task::AbstractString;
+    eltype = Float64
+)
+    model_filepath, evidence_filepath, query_filepath, solution_filepath = get_instance_filepaths(
+        artifact_name,
+        problem_name,
+        task
+    )
+    read_instance(model_filepath; evidence_filepath, query_filepath, solution_filepath, eltype)
+end
+
+"""
+Helper function to obtain the filepaths for an instance's model, evidence, and
+solution files within the provided artifact, corresponding to the provided
+task. If a given file path does not exist in the instance's directory, an
+empty file path is returned.
+"""
+function get_instance_filepaths(
+    artifact_name::AbstractString,
+    problem_name::AbstractString,
+    task::AbstractString
+)
+    artifact_toml = pkgdir(TensorInference, "Artifacts.toml")
+    Pkg.ensure_artifact_installed(artifact_name, artifact_toml)
+    artifact_hash = Pkg.Artifacts.artifact_hash(artifact_name, artifact_toml)
+    artifact_path = Pkg.Artifacts.artifact_path(artifact_hash)
+    model_filepath = joinpath(artifact_path, task, problem_name * ".uai") |> topath
+    evidence_filepath = joinpath(artifact_path, task, problem_name * ".uai.evid") |> topath
+    query_filepath = joinpath(artifact_path, task, problem_name * ".uai.query") |> topath
+    solution_filepath = joinpath(artifact_path, task, problem_name * ".uai." * task) |> topath
+    return model_filepath, evidence_filepath, query_filepath, solution_filepath
+end
+
+"""
+This function returns the provided file path if it exists on disk, otherwise it
+  returns an empty string.
+"""
+topath(filepath::AbstractString) = isfile(filepath) ? filepath : ""
+
+"""
+Helper function that captures the problem names that belong to `problem_set`
+for the given task.
+"""
+function get_problem_names(
+    artifact_name::AbstractString,
+    problem_set::AbstractString,
+    task::AbstractString
+)
+    artifact_toml = pkgdir(TensorInference, "Artifacts.toml")
+    Pkg.ensure_artifact_installed(artifact_name, artifact_toml)
+    artifact_hash = Pkg.Artifacts.artifact_hash(artifact_name, artifact_toml)
+    artifact_path = Pkg.Artifacts.artifact_path(artifact_hash)
+
+    regex = Regex("($(problem_set)_\\d*)(\\.uai)\$")
+    return readdir(joinpath(artifact_path, task); sort = false) |>
+           x -> map(y -> match(regex, y), x) |> # apply regex
+                x -> filter(!isnothing, x) |> # filter out `nothing` values
+                     x -> map(first, x) # get the first capture of each element
+end

src/utils.jl

@@ -1,5 +1,4 @@
 [deps]
-Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 KaHyPar = "2a6221f6-aa48-11e9-3542-2d9e0ef01880"

test/Project.toml

@@ -4,8 +4,7 @@ using TensorInference, CUDA
 CUDA.allowscalar(false)
 
 @testset "gradient-based tensor network solvers" begin
-    model_filepath, evidence_filepath, _, solution_filepath = get_instance_filepaths("Promedus_14", "MAR")
-    instance = read_instance(model_filepath; evidence_filepath, solution_filepath)
+    instance = read_instance_from_artifact("uai2014", "Promedus_14", "MAR")
 
     # does not optimize over open vertices
     tn = TensorNetworkModel(instance; optimizer = TreeSA(ntrials = 1, niters = 2, βs = 1:0.1:40))
@@ -21,8 +20,7 @@ CUDA.allowscalar(false)
 end
 
 @testset "map" begin
-    model_filepath, evidence_filepath, _, solution_filepath = get_instance_filepaths("Promedus_14", "MAR")
-    instance = read_instance(model_filepath; evidence_filepath, solution_filepath)
+    instance = read_instance_from_artifact("uai2014", "Promedus_14", "MAR")
 
     # does not optimize over open vertices
     tn = TensorNetworkModel(instance; optimizer = TreeSA(ntrials = 1, niters = 2, βs = 1:0.1:40))
@@ -36,8 +34,7 @@ end
 end
 
 @testset "mmap" begin
-    model_filepath, evidence_filepath, _, solution_filepath = get_instance_filepaths("Promedus_14", "MAR")
-    instance = read_instance(model_filepath; evidence_filepath, solution_filepath)
+    instance = read_instance_from_artifact("uai2014", "Promedus_14", "MAR")
 
     optimizer = TreeSA(ntrials = 1, niters = 2, βs = 1:0.1:40)
     tn_ref = TensorNetworkModel(instance; optimizer)

test/cuda.jl

@@ -3,8 +3,7 @@ using OMEinsum
 using TensorInference
 
 @testset "gradient-based tensor network solvers" begin
-    model_filepath, evidence_filepath, _, solution_filepath = get_instance_filepaths("Promedus_14", "MAR")
-    instance = read_instance(model_filepath; evidence_filepath, solution_filepath)
+    instance = read_instance_from_artifact("uai2014", "Promedus_14", "MAR")
 
     # does not optimize over open vertices
     tn = TensorNetworkModel(instance; optimizer = TreeSA(ntrials = 3, niters = 2, βs = 1:0.1:80))
@@ -18,8 +17,7 @@ end
 @testset "UAI Reference Solution Comparison" begin
     problem_name = "Promedas_70"
     @info "Testing: $problem_name"
-    model_filepath, evidence_filepath, _, solution_filepath = get_instance_filepaths(problem_name, "MAP")
-    instance = read_instance(model_filepath; evidence_filepath, solution_filepath)
+    instance = read_instance_from_artifact("uai2014", problem_name, "MAP")
     tn = TensorNetworkModel(instance; optimizer = TreeSA(ntrials = 1, niters = 5, βs = 0.1:0.1:100))
     _, solution = most_probable_config(tn)
     @test solution == instance.reference_solution

test/map.jl

@@ -1,4 +1,4 @@
-using Test, Artifacts
+using Test
 using OMEinsum
 using KaHyPar
 using TensorInference
@@ -12,8 +12,7 @@ using TensorInference
 end
 
 @testset "cached, rescaled contract" begin
-    model_filepath, evidence_filepath, _, solution_filepath = get_instance_filepaths("Promedus_14", "MAR")
-    instance = read_instance(model_filepath; evidence_filepath, solution_filepath)
+    instance = read_instance_from_artifact("uai2014", "Promedus_14", "MAR")
     ref_sol = instance.reference_solution
     optimizer = TreeSA(ntrials = 1, niters = 5, βs = 0.1:0.1:100)
     tn = TensorNetworkModel(instance; optimizer)
@@ -50,12 +49,11 @@ end
     for (problem_set, optimizer) in problem_sets
         @testset "$(problem_set) problem set" begin
             # Capture the problem names that belong to the current problem set
-            problem_names = get_problem_names(problem_set, "MAR")
+            problem_names = TensorInference.get_problem_names("uai2014", problem_set, "MAR")
             for problem_name in problem_names
                 @info "Testing: $problem_name"
                 @testset "$(problem_name)" begin
-                    model_filepath, evidence_filepath, _, solution_filepath = get_instance_filepaths(problem_name, "MAR")
-                    instance = read_instance(model_filepath; evidence_filepath, solution_filepath)
+                    instance = read_instance_from_artifact("uai2014", problem_name, "MAR")
                     ref_sol = instance.reference_solution
                     obsvars = instance.obsvars
 

test/mar.jl

@@ -8,8 +8,7 @@ using TensorInference
 end
 
 @testset "gradient-based tensor network solvers" begin
-    model_filepath, evidence_filepath, _, solution_filepath = get_instance_filepaths("Promedus_14", "MAR")
-    instance = read_instance(model_filepath; evidence_filepath, solution_filepath)
+    instance = read_instance_from_artifact("uai2014", "Promedus_14", "MAR")
 
     optimizer = TreeSA(ntrials = 1, niters = 2, βs = 1:0.1:40)
     tn_ref = TensorNetworkModel(instance; optimizer)
@@ -43,10 +42,9 @@ end
     ]
     for (problem_name, optimizer) in problems
       @info "Testing: $problem_name"
-      model_filepath, evidence_filepath, query_filepath, solution_filepath = get_instance_filepaths(problem_name, "MMAP")
-      instance = read_instance(model_filepath; evidence_filepath, query_filepath, solution_filepath)
+      instance = read_instance_from_artifact("uai2014", problem_name, "MMAP")
       model = MMAPModel(instance; optimizer)
       _, solution = most_probable_config(model)
       @test solution == instance.reference_solution
     end
-end
+end