Upgrade omeinsum to version 0.9.1

Project.toml

@@ -1,7 +1,7 @@
 name = "TensorInference"
 uuid = "c2297e78-99bd-40ad-871d-f50e56b81012"
 authors = ["Jin-Guo Liu", "Martin Roa Villescas"]
-version = "0.6.1"
+version = "0.6.2"
 
 [deps]
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
@@ -23,7 +23,7 @@ TensorInferenceCUDAExt = "CUDA"
 CUDA = "4, 5"
 DocStringExtensions = "0.8.6, 0.9"
 LinearAlgebra = "1"
-OMEinsum = "0.8.7"
+OMEinsum = "0.9.1"
 Pkg = "1"
 PrettyTables = "2"
 ProblemReductions = "0.3"

docs/src/api/public.md

@@ -34,6 +34,9 @@ TensorInference
 ```@docs
 GreedyMethod
 KaHyParBipartite
+HyperND
+TreeSASlicer
+ScoreFunction
 MergeGreedy
 MergeVectors
 SABipartite
@@ -73,4 +76,6 @@ update_temperature
 random_matrix_product_state
 random_matrix_product_uai
 random_tensor_train_uai
+save_tensor_network
+load_tensor_network
 ```

docs/src/performance-tips.jl

@@ -30,12 +30,12 @@ probability(tn)
 
 # For large scale applications, it is also possible to slice over certain degrees of freedom to reduce the space complexity, i.e.
 # loop and accumulate over certain degrees of freedom so that one can have a smaller tensor network inside the loop due to the removal of these degrees of freedom.
-# In the [`TreeSA`](@ref) optimizer, one can set `nslices` to a value larger than zero to turn on this feature.
-# As a comparison we slice over 5 degrees of freedom, which can reduce the space complexity by at most 5.
+# One can use the `slicer` keyword argument to reduce the space complexity by slicing over certain degrees of freedom.
+# In the following example, we use the `TreeSASlicer` to reduce the space complexity to `sc_target=10`.
 # In this application, the slicing achieves the largest possible space complexity reduction 5, while the time and read-write complexity are only increased by less than 1,
 # i.e. the peak memory usage is reduced by a factor ``32``, while the (theoretical) computing time is increased by at a factor ``< 2``.
-optimizer = TreeSA(ntrials = 1, niters = 5, βs = 0.1:0.3:100, nslices=5)
-tn = TensorNetworkModel(model; optimizer, evidence);
+optimizer = TreeSA(ntrials = 1, niters = 5, βs = 0.1:0.3:100)
+tn = TensorNetworkModel(model; optimizer, evidence, slicer=TreeSASlicer(score=ScoreFunction(sc_target=10)));
 contraction_complexity(tn)
 
 # ## Faster Tropical tensor contraction to speed up MAP and MMAP

src/Core.jl

@@ -49,7 +49,7 @@ Probabilistic modeling with a tensor network.
 * `code` is the tensor network contraction pattern.
 * `tensors` are the tensors fed into the tensor network, the leading tensors are unity tensors associated with `unity_tensors_labels`.
 * `evidence` is a dictionary used to specify degrees of freedom that are fixed to certain values.
-* `unity_tensors_idx` is a vector of indices of the unity tensors in the `tensors` array. Unity tensors are dummy tensors used to obtain the marginal probabilities.
+* `unity_tensors_idx` is a vector of indices pointing to the unity tensors in the `tensors` array. Unity tensors are dummy tensors with all entries equal to one, which are used to obtain the marginal probabilities.
 """
 struct TensorNetworkModel{ET, MT <: AbstractArray}
     nvars::Int
@@ -118,6 +118,7 @@ function TensorNetworkModel(
     evidence = Dict{Int,Int}(),
     optimizer = GreedyMethod(),
     simplifier = nothing,
+    slicer = nothing,
     unity_tensors_labels = [[i] for i=1:model.nvars]
 ) where {ET, FT}
     # `optimize_code` optimizes the contraction order of a raw tensor network without a contraction order specified.
@@ -127,7 +128,7 @@ function TensorNetworkModel(
     rawcode = EinCode([unity_tensors_labels..., [[factor.vars...] for factor in model.factors]...], collect(Int, openvars))  # labels for vertex tensors (unity tensors) and edge tensors
     tensors = Array{ET}[[ones(ET, [model.cards[i] for i in lb]...) for lb in unity_tensors_labels]..., [t.vals for t in model.factors]...]
     size_dict = OMEinsum.get_size_dict(getixsv(rawcode), tensors)
-    code = optimize_code(rawcode, size_dict, optimizer, simplifier)
+    code = optimize_code(rawcode, size_dict, optimizer; simplifier, slicer)
     return TensorNetworkModel(model.nvars, code, tensors, evidence, collect(Int, 1:length(unity_tensors_labels)))
 end
 

src/TensorInference.jl

@@ -9,6 +9,7 @@ module TensorInference
 
 using OMEinsum, LinearAlgebra
 using OMEinsum: CacheTree, cached_einsum
+using OMEinsum.OMEinsumContractionOrders.JSON
 using DocStringExtensions, TropicalNumbers
 # The Tropical GEMM support
 using StatsBase
@@ -19,7 +20,7 @@ import Pkg
 
 # reexport OMEinsum functions
 export RescaledArray
-export contraction_complexity, TreeSA, GreedyMethod, KaHyParBipartite, SABipartite, MergeGreedy, MergeVectors
+export contraction_complexity, TreeSA, GreedyMethod, KaHyParBipartite, HyperND, SABipartite, MergeGreedy, MergeVectors, TreeSASlicer, ScoreFunction
 
 # read and load uai files
 export read_model_file, read_td_file, read_evidence_file
@@ -44,6 +45,9 @@ export update_temperature
 # belief propagation
 export BeliefPropgation, belief_propagate
 
+# fileio
+export save_tensor_network, load_tensor_network
+
 # utils
 export random_matrix_product_state, random_tensor_train_uai, random_matrix_product_uai
 
@@ -56,5 +60,6 @@ include("mmap.jl")
 include("sampling.jl")
 include("cspmodels.jl")
 include("belief.jl")
+include("fileio.jl")
 
 end # module

src/fileio.jl

@@ -0,0 +1,147 @@
+"""
+    save_tensor_network(tn::TensorNetworkModel; folder::String)
+
+Save a tensor network model to a folder with separate files for code, tensors, and model metadata.
+The code is saved using `OMEinsum.writejson`, tensors as JSON, and model specifics in model.json.
+
+# Arguments
+- `tn::TensorNetworkModel`: The tensor network model to save
+- `folder::String`: The folder path to save the files
+
+# Files Created
+- `code.json`: Contains the einsum code using OMEinsum format
+- `tensors.json`: Contains the tensor data as JSON
+- `model.json`: Contains nvars, evidence, and unity_tensors_idx
+
+# Example
+```julia
+tn = TensorNetworkModel(...)  # create your model
+save_tensor_network(tn; folder="my_model")
+```
+"""
+function save_tensor_network(tn::TensorNetworkModel; folder::String)
+    !isdir(folder) && mkpath(folder)
+
+    # save code
+    OMEinsum.writejson(joinpath(folder, "code.json"), tn.code)
+
+    # save tensors
+    open(joinpath(folder, "tensors.json"), "w") do io
+        JSON.print(io, [tensor_to_dict(tensor) for tensor in tn.tensors], 2)
+    end
+
+    # save model metadata
+    open(joinpath(folder, "model.json"), "w") do io
+        JSON.print(io, Dict(
+            "nvars" => tn.nvars,
+            "evidence" => tn.evidence,
+            "unity_tensors_idx" => tn.unity_tensors_idx
+        ), 2)
+    end
+    return nothing
+end
+
+"""
+    load_tensor_network(folder::String)
+
+Load a tensor network model from a folder containing code, tensors, and model files.
+
+# Arguments
+- `folder::String`: The folder path containing the files
+
+# Returns
+- `TensorNetworkModel`: The loaded tensor network model
+
+# Required Files
+- `code.json`: Contains the einsum code using OMEinsum format
+- `tensors.json`: Contains the tensor data as JSON
+- `model.json`: Contains nvars, evidence, and unity_tensors_idx
+
+# Example
+```julia
+tn = load_tensor_network("my_model")
+```
+"""
+function load_tensor_network(folder::String)::TensorNetworkModel
+    !isdir(folder) && throw(SystemError("Folder not found: $folder"))
+
+    code_path = joinpath(folder, "code.json")
+    tensors_path = joinpath(folder, "tensors.json")
+    model_path = joinpath(folder, "model.json")
+    !isfile(code_path) && throw(SystemError("Code file not found: $code_path"))
+    !isfile(tensors_path) && throw(SystemError("Tensors file not found: $tensors_path"))
+    !isfile(model_path) && throw(SystemError("Model file not found: $model_path"))
+
+    code = OMEinsum.readjson(code_path)
+
+    tensors = [tensor_from_dict(t) for t in JSON.parsefile(tensors_path)]
+
+    model_dict = JSON.parsefile(model_path)
+
+    # Convert evidence keys to Int (JSON parses them as strings)
+    evidence = Dict{Int, Int}()
+    for (k, v) in model_dict["evidence"]
+        evidence[parse(Int, k)] = v
+    end
+
+    return TensorNetworkModel(
+        model_dict["nvars"],
+        code,
+        tensors,
+        evidence,
+        collect(Int, model_dict["unity_tensors_idx"])
+    )
+end
+
+"""
+    tensor_to_dict(tensor::AbstractArray{T}) where T
+
+Convert a tensor to a dictionary representation for JSON serialization.
+
+# Arguments
+- `tensor::AbstractArray{T}`: The tensor to convert
+
+# Returns
+- `Dict`: A dictionary containing tensor metadata and data
+
+# Dictionary Structure
+- `"size"`: The dimensions of the tensor
+- `"complex"`: Boolean indicating if the tensor contains complex numbers
+- `"data"`: The tensor data as a flat array of real numbers
+"""
+function tensor_to_dict(tensor::AbstractArray{T}) where T
+    d = Dict()
+    d["size"] = collect(size(tensor))
+    d["complex"] = T <: Complex
+    d["data"] = vec(reinterpret(real(T), tensor))
+    return d
+end
+
+"""
+    tensor_from_dict(dict::Dict)
+
+Convert a dictionary back to a tensor.
+
+# Arguments
+- `dict::Dict`: The dictionary representation of a tensor
+
+# Returns
+- `AbstractArray`: The reconstructed tensor
+
+# Dictionary Structure Expected
+- `"size"`: The dimensions of the tensor
+- `"complex"`: Boolean indicating if the tensor contains complex numbers
+- `"data"`: The tensor data as a flat array of real numbers
+"""
+function tensor_from_dict(dict::Dict)
+    size_vec = Tuple(dict["size"])
+    is_complex = dict["complex"]
+    data = collect(Float64, dict["data"])
+
+    if is_complex
+        complex_data = reinterpret(ComplexF64, data)
+        return reshape(complex_data, size_vec...)
+    else
+        return reshape(data, size_vec...)
+    end
+end 

src/mmap.jl

@@ -91,13 +91,13 @@ function MMAPModel(vars::AbstractVector{LT}, cards::AbstractVector{Int}, factors
         ixsi = all_ixs[cluster]
         vari = unique!(vcat(ixsi...))
         iyi = setdiff(vari, contracted)
-        codei = optimize_code(EinCode(ixsi, iyi), size_dict, marginalize_optimizer, marginalize_simplifier)
+        codei = optimize_code(EinCode(ixsi, iyi), size_dict, marginalize_optimizer; simplifier=marginalize_simplifier)
         push!(ixs, iyi)
         push!(clusters, Cluster(contracted, codei, ts))
     end
     rem_indices = setdiff(1:length(all_ixs), vcat([c.second for c in subsets]...))
     remaining_tensors = all_tensors[rem_indices]
-    code = optimize_code(EinCode([all_ixs[rem_indices]..., ixs...], iy), size_dict, optimizer, simplifier)
+    code = optimize_code(EinCode([all_ixs[rem_indices]..., ixs...], iy), size_dict, optimizer; simplifier)
     return MMAPModel(setdiff(vars, marginalized), code, remaining_tensors, clusters, evidence)
 end
 

src/sampling.jl

@@ -139,7 +139,7 @@ function generate_samples!(code::DynamicNestedEinsum, cache::CacheTree{T}, iy_en
             siblings = filter(x->x !== child, cache.siblings)
             siblings_ixs = filter(x->x !== ix, ixs)
             iy_subenv = batch_label ∈ ix ? ix : [ix..., batch_label]
-            envcode = optimize_code(EinCode([siblings_ixs..., iy_env], iy_subenv), size_dict, GreedyMethod(; nrepeat=1))
+            envcode = optimize_code(EinCode([siblings_ixs..., iy_env], iy_subenv), size_dict, GreedyMethod())
             subenv = einsum(envcode, (getfield.(siblings, :content)..., env), size_dict)
 
             # generate samples

test/fileio.jl

@@ -0,0 +1,91 @@
+using Test
+using TensorInference
+using Random
+
+@testset "TensorNetworkModel file I/O" begin
+    # Create a test model
+    n = 3
+    chi = 2
+    d = 2
+    tn = random_matrix_product_state(n, chi, d)
+
+    # Add evidence for testing
+    tn.evidence[1] = 0
+    tn.evidence[2] = 1
+
+    # Create temporary directory
+    test_dir = mktempdir()
+
+    # Test saving
+    @testset "Saving" begin
+        save_tensor_network(tn; folder=test_dir)
+        @test isfile(joinpath(test_dir, "code.json"))
+        @test isfile(joinpath(test_dir, "tensors.json"))
+        @test isfile(joinpath(test_dir, "model.json"))
+    end
+
+    # Test loading
+    @testset "Loading" begin
+        tn_loaded = load_tensor_network(test_dir)
+
+        # Verify basic properties
+        @test tn_loaded.nvars == tn.nvars
+        @test tn_loaded.evidence == tn.evidence
+        @test tn_loaded.unity_tensors_idx == tn.unity_tensors_idx
+
+        # Verify code structure
+        @test tn_loaded.code isa typeof(tn.code)
+
+        # Verify tensors
+        @test length(tn_loaded.tensors) == length(tn.tensors)
+        for (t_orig, t_loaded) in zip(tn.tensors, tn_loaded.tensors)
+            @test size(t_orig) == size(t_loaded)
+            @test eltype(t_orig) == eltype(t_loaded)
+            @test Array(t_orig) ≈ Array(t_loaded)
+        end
+
+        # Verify model functionality
+        @test probability(tn)[] ≈ probability(tn_loaded)[]
+    end
+
+    # Test error handling
+    @testset "Error handling" begin
+        # Invalid directory
+        @test_throws SystemError load_tensor_network("nonexistent_directory")
+
+        # Missing files
+        for file in ["code.json", "tensors.json", "model.json"]
+            bad_dir = mktempdir()
+            save_tensor_network(tn; folder=bad_dir)
+            rm(joinpath(bad_dir, file))
+            @test_throws SystemError load_tensor_network(bad_dir)
+            rm(bad_dir, recursive=true)
+        end
+    end
+
+    # Clean up
+    rm(test_dir, recursive=true)
+end
+
+@testset "Tensor serialization" begin
+    Random.seed!(42)
+
+    # Test real tensor
+    real_tensor = rand(2, 2)
+    dict_real = TensorInference.tensor_to_dict(real_tensor)
+    @test TensorInference.tensor_from_dict(dict_real) ≈ real_tensor
+
+    # Test complex tensor
+    complex_tensor = rand(ComplexF64, 2, 2)
+    dict_complex = TensorInference.tensor_to_dict(complex_tensor)
+    @test TensorInference.tensor_from_dict(dict_complex) ≈ complex_tensor
+
+    # Test higher-dimensional tensor
+    high_dim_tensor = rand(2, 3, 4)
+    dict_high_dim = TensorInference.tensor_to_dict(high_dim_tensor)
+    @test TensorInference.tensor_from_dict(dict_high_dim) ≈ high_dim_tensor
+
+    # Test invalid input
+    @test_throws KeyError TensorInference.tensor_from_dict(Dict())
+    @test_throws KeyError TensorInference.tensor_from_dict(Dict("size" => [2,2]))
+end 

test/runtests.jl

@@ -32,6 +32,10 @@ end
     include("belief.jl")
 end
 
+@testset "fileio" begin
+    include("fileio.jl")
+end
+
 using CUDA
 if CUDA.functional()
     include("cuda.jl")