From e3d7148b30c501f3817f0ff8ba6d101b6424e36c Mon Sep 17 00:00:00 2001
From: annamariadziubyna <anna.maria.dziubyna@doctoral.uj.edu.pl>
Date: Fri, 19 Jan 2024 22:34:05 +0100
Subject: [PATCH] save BP results to file

---
 Project.toml             |  1 +
 src/SpinGlassNetworks.jl |  1 +
 src/truncate.jl          | 60 +++++++++++++++++++++++++++++++++-------
 3 files changed, 52 insertions(+), 10 deletions(-)

diff --git a/Project.toml b/Project.toml
index bf03767..38ed063 100644
--- a/Project.toml
+++ b/Project.toml
@@ -5,6 +5,7 @@ version = "1.0"
 
 [deps]
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
diff --git a/src/SpinGlassNetworks.jl b/src/SpinGlassNetworks.jl
index 12e010c..7f83c64 100644
--- a/src/SpinGlassNetworks.jl
+++ b/src/SpinGlassNetworks.jl
@@ -3,6 +3,7 @@ module SpinGlassNetworks
     using LightGraphs
     using MetaGraphs # TODO: to be replaced by MetaGraphsNext
     using CSV
+    using JLD2
     using DocStringExtensions
     using LinearAlgebra, MKL
     using Base.Cartesian
diff --git a/src/truncate.jl b/src/truncate.jl
index b434200..3ebbfe9 100644
--- a/src/truncate.jl
+++ b/src/truncate.jl
@@ -88,22 +88,62 @@ to keep. It computes the beliefs for all 2-site combinations and selects the sta
 # Returns:
 - `LabelledGraph{S, T}`: A truncated clustered Hamiltonian.
 """
+# function truncate_clustered_hamiltonian_2site_BP(
+#     cl_h::LabelledGraph{S, T}, 
+#     beliefs::Dict, 
+#     num_states::Int; 
+#     beta=1.0
+#     ) where {S, T}
+#     # TODO: name to be clean to make it consistent with square2 and squarestar2
+#     states = Dict()
+#     for node in vertices(cl_h)
+#         if node in keys(states) continue end
+#         i, j, _ = node
+#         sx = has_vertex(cl_h, (i, j, 1)) ? length(get_prop(cl_h, (i, j, 1), :spectrum).energies) : 1
+#         E = beliefs[(i, j)]
+#         ind1, ind2 = select_numstate_best(E, sx, num_states)
+#         push!(states, (i, j, 1) => ind1)
+#         push!(states, (i, j, 2) => ind2)
+#     end
+#     truncate_clustered_hamiltonian(cl_h, states)
+# end
+
+function load_file(filename)
+    if isfile(filename)
+        try 
+            load_object(string(filename))
+        catch e
+            return nothing
+        end
+    else
+        return nothing
+    end
+end
+
 function truncate_clustered_hamiltonian_2site_BP(
     cl_h::LabelledGraph{S, T}, 
     beliefs::Dict, 
-    num_states::Int; 
+    num_states::Int,
+    result_folder::String = "results_folder",
+    inst::String = "inst"; 
     beta=1.0
     ) where {S, T}
-    # TODO: name to be clean to make it consistent with square2 and squarestar2
     states = Dict()
-    for node in vertices(cl_h)
-        if node in keys(states) continue end
-        i, j, _ = node
-        sx = has_vertex(cl_h, (i, j, 1)) ? length(get_prop(cl_h, (i, j, 1), :spectrum).energies) : 1
-        E = beliefs[(i, j)]
-        ind1, ind2 = select_numstate_best(E, sx, num_states)
-        push!(states, (i, j, 1) => ind1)
-        push!(states, (i, j, 2) => ind2)
+
+    saved_states = load_file(joinpath(result_folder, "$(inst).jld2"))
+    if saved_states == nothing
+        for node in vertices(cl_h)
+            i, j, _ = node
+            sx = has_vertex(cl_h, (i, j, 1)) ? length(get_prop(cl_h, (i, j, 1), :spectrum).energies) : 1
+            E = beliefs[(i, j)]
+            ind1, ind2 = select_numstate_best(E, sx, num_states)
+            push!(states, (i, j, 1) => ind1)
+            push!(states, (i, j, 2) => ind2)
+        end
+        path = joinpath(result_folder, "$(inst).jld2")
+        save_object(string(path), states)
+    else
+        states = saved_states
     end
     truncate_clustered_hamiltonian(cl_h, states)
 end