PoolOfProjectors in clustered_hamiltonian

euro-hpc-pl · Sep 20, 2023 · 39a095f · 39a095f
1 parent 5ace384
commit 39a095f
Show file tree

Hide file tree

Showing 5 changed files with 100 additions and 47 deletions.
diff --git a/src/SpinGlassNetworks.jl b/src/SpinGlassNetworks.jl
@@ -8,7 +8,6 @@ module SpinGlassNetworks
     using Base.Cartesian
     using SparseArrays
     using CUDA, CUDA.CUSPARSE
-
     import Base.Prehashed
 
     export unique_neighbors

diff --git a/src/bp.jl b/src/bp.jl
@@ -73,13 +73,13 @@ function get_neighbors(cl_h::LabelledGraph{S, T}, vertex::NTuple) where {S, T}
             en = get_prop(cl_h, src_node, dst_node, :en)
             pv = get_prop(cl_h, src_node, dst_node, :pl)
             # idx_pv = get_prop(cl_h, src_node, dst_node, :ipl)
-            # pv = get_projector2!(get_prop(cl_h, :pool_of_projectors), idx_pv, :CPU)
+            # pv = get_projector!(get_prop(cl_h, :pool_of_projectors), idx_pv, :CPU)
             push!(neighbors, (dst_node, pv, en))
         elseif dst_node == vertex
             en = get_prop(cl_h, src_node, dst_node, :en)'
             pv = get_prop(cl_h, src_node, dst_node, :pr)
             # idx_pv = get_prop(cl_h, src_node, dst_node, :ipr)
-            # pv = get_projector2!(get_prop(cl_h, :pool_of_projectors), idx_pv, :CPU)
+            # pv = get_projector!(get_prop(cl_h, :pool_of_projectors), idx_pv, :CPU)
             push!(neighbors, (src_node, pv, en))
         end
     end
@@ -235,10 +235,10 @@ end
 function projector(cl_h::LabelledGraph{S, T}, v::NTuple{N, Int64}, w::NTuple{N, Int64}) where {S, T, N}
     if has_edge(cl_h, w, v)
         idx_p = get_prop(cl_h, w, v, :ipr)
-        p = get_projector2!(get_prop(cl_h, :pool_of_projectors), idx_p, :CPU)
+        p = get_projector!(get_prop(cl_h, :pool_of_projectors), idx_p, :CPU)
     elseif has_edge(cl_h, v, w)
         idx_p = get_prop(cl_h, v, w, :ipl)
-        p = get_projector2!(get_prop(cl_h, :pool_of_projectors), idx_p, :CPU)
+        p = get_projector!(get_prop(cl_h, :pool_of_projectors), idx_p, :CPU)
     else
         p = ones(Int, v ∈ vertices(cl_h) ? length(get_prop(cl_h, v, :spectrum).energies) : 1)
     end

diff --git a/src/clustered_hamiltonian.jl b/src/clustered_hamiltonian.jl
@@ -40,7 +40,7 @@ end
 # ) where T
 #     L = maximum(values(cluster_assignment_rule))
 #     cl_h = LabelledGraph{MetaDiGraph}(sort(unique(values(cluster_assignment_rule))))
-#     lp = PoolOfProjectors2{Int}()
+#     lp = PoolOfProjectors{Int}()
 
 #     for (v, cl) ∈ split_into_clusters(ig, cluster_assignment_rule)
 #         sp = spectrum(cl, num_states=get(num_states_cl, v, basis_size(cl)))
@@ -64,7 +64,7 @@ end
 #             pl, unique_states_v = rank_reveal([s[ind1] for s ∈ states_v], :PE)
 #             pr, unique_states_w = rank_reveal([s[ind2] for s ∈ states_w], :PE)
 #             en = inter_cluster_energy(unique_states_v, JJ, unique_states_w)
-#             add_projector2!(lp, pl)
+#             add_projector!(lp, pl)
 #             add_edge!(cl_h, v, w)
 #             set_props!(
 #                 cl_h, v, w, Dict(:outer_edges => outer_edges, :pl => pl, :en => en, :pr => pr)
@@ -83,7 +83,7 @@ function clustered_hamiltonian(
 ) where T
     L = maximum(values(cluster_assignment_rule))
     cl_h = LabelledGraph{MetaDiGraph}(sort(unique(values(cluster_assignment_rule))))
-    lp = PoolOfProjectors2{Int}()
+    lp = PoolOfProjectors{Int}()
 
     for (v, cl) ∈ split_into_clusters(ig, cluster_assignment_rule)
         sp = spectrum(cl, num_states=get(num_states_cl, v, basis_size(cl)))
@@ -107,8 +107,8 @@ function clustered_hamiltonian(
             pl, unique_states_v = rank_reveal([s[ind1] for s ∈ states_v], :PE)
             pr, unique_states_w = rank_reveal([s[ind2] for s ∈ states_w], :PE)
             en = inter_cluster_energy(unique_states_v, JJ, unique_states_w)
-            ipl = add_projector2!(lp, pl)
-            ipr = add_projector2!(lp, pr)
+            ipl = add_projector!(lp, pl)
+            ipr = add_projector!(lp, pr)
 
             add_edge!(cl_h, v, w)
             set_props!(
@@ -180,9 +180,9 @@ function energy(cl_h::LabelledGraph{S, T}, σ::Dict{T, Int}) where {S, T}
     for v ∈ vertices(cl_h) en_cl_h += get_prop(cl_h, v, :spectrum).energies[σ[v]] end
     for edge ∈ edges(cl_h)
         idx_pl = get_prop(cl_h, edge, :ipl)
-        pl = get_projector2!(get_prop(cl_h, :pool_of_projectors), idx_pl, :CPU)
+        pl = get_projector!(get_prop(cl_h, :pool_of_projectors), idx_pl, :CPU)
         idx_pr = get_prop(cl_h, edge, :ipr)
-        pr = get_projector2!(get_prop(cl_h, :pool_of_projectors), idx_pr, :CPU)
+        pr = get_projector!(get_prop(cl_h, :pool_of_projectors), idx_pr, :CPU)
         en = get_prop(cl_h, edge, :en)
         en_cl_h += en[pl[σ[src(edge)]], pr[σ[dst(edge)]]]
     end
@@ -195,16 +195,16 @@ function energy_2site(cl_h::LabelledGraph{S, T}, i::Int, j::Int) where {S, T}
     if has_edge(cl_h, (i, j, 1), (i, j, 2))
         en12 = copy(get_prop(cl_h, (i, j, 1), (i, j, 2), :en))
         idx_pl = get_prop(cl_h, (i, j, 1), (i, j, 2), :ipl)
-        pl = copy(get_projector2!(get_prop(cl_h, :pool_of_projectors), idx_pl, :CPU))
+        pl = copy(get_projector!(get_prop(cl_h, :pool_of_projectors), idx_pl, :CPU))
         idx_pr = get_prop(cl_h, (i, j, 1), (i, j, 2), :ipr)
-        pr = copy(get_projector2!(get_prop(cl_h, :pool_of_projectors), idx_pr, :CPU))
+        pr = copy(get_projector!(get_prop(cl_h, :pool_of_projectors), idx_pr, :CPU))
         int_eng = en12[pl, pr]
     elseif has_edge(cl_h, (i, j, 2), (i, j, 1))
         en21 = copy(get_prop(cl_h, (i, j, 2), (i, j, 1), :en))
         idx_pl = get_prop(cl_h, (i, j, 2), (i, j, 1), :ipl)
-        pl = copy(get_projector2!(get_prop(cl_h, :pool_of_projectors), idx_pl, :CPU))
+        pl = copy(get_projector!(get_prop(cl_h, :pool_of_projectors), idx_pl, :CPU))
         idx_pr = get_prop(cl_h, (i, j, 2), (i, j, 1), :ipr)
-        pr = copy(get_projector2!(get_prop(cl_h, :pool_of_projectors), idx_pr, :CPU))
+        pr = copy(get_projector!(get_prop(cl_h, :pool_of_projectors), idx_pr, :CPU))
         int_eng = en21[pl, pr]'
     else
         int_eng = zeros(1, 1)
@@ -217,15 +217,15 @@ function bond_energy(cl_h::LabelledGraph{S, T}, cl_h_u::NTuple{N, Int64}, cl_h_v
         ipu, en, ipv = get_prop.(
                         Ref(cl_h), Ref(cl_h_u), Ref(cl_h_v), (:ipl, :en, :ipr)
                     )
-        pu = get_projector2!(get_prop(cl_h, :pool_of_projectors), ipu, :CPU)
-        pv = get_projector2!(get_prop(cl_h, :pool_of_projectors), ipv, :CPU)
+        pu = get_projector!(get_prop(cl_h, :pool_of_projectors), ipu, :CPU)
+        pv = get_projector!(get_prop(cl_h, :pool_of_projectors), ipv, :CPU)
         @inbounds energies = en[pu, pv[σ]]
     elseif has_edge(cl_h, cl_h_v, cl_h_u)
         ipv, en, ipu = get_prop.(
                         Ref(cl_h), Ref(cl_h_v), Ref(cl_h_u), (:ipl, :en, :ipr)
                     )
-        pu = get_projector2!(get_prop(cl_h, :pool_of_projectors), ipu, :CPU)
-        pv = get_projector2!(get_prop(cl_h, :pool_of_projectors), ipv, :CPU)
+        pu = get_projector!(get_prop(cl_h, :pool_of_projectors), ipu, :CPU)
+        pv = get_projector!(get_prop(cl_h, :pool_of_projectors), ipv, :CPU)
         @inbounds energies = en[pv[σ], pu]
     else
         energies = zeros(cluster_size(cl_h, cl_h_u))
@@ -249,7 +249,7 @@ function exact_cond_prob(clustered_hamiltonian::LabelledGraph{S, T}, beta, targe
 function truncate_clustered_hamiltonian(cl_h::LabelledGraph{S, T}, states::Dict) where {S, T}
 
     new_cl_h = LabelledGraph{MetaDiGraph}(vertices(cl_h))
-    new_lp = PoolOfProjectors2{Int}()
+    new_lp = PoolOfProjectors{Int}()
 
     for v ∈ vertices(new_cl_h)
         cl = get_prop(cl_h, v, :cluster)
@@ -267,17 +267,17 @@ function truncate_clustered_hamiltonian(cl_h::LabelledGraph{S, T}, states::Dict)
         add_edge!(new_cl_h, v, w)
         outer_edges = get_prop(cl_h, v, w, :outer_edges)
         ipl = get_prop(cl_h, v, w, :ipl)
-        pl = get_projector2!(get_prop(cl_h, :pool_of_projectors), ipl, :CPU)
+        pl = get_projector!(get_prop(cl_h, :pool_of_projectors), ipl, :CPU)
         ipr = get_prop(cl_h, v, w, :ipr)
-        pr = get_projector2!(get_prop(cl_h, :pool_of_projectors), ipr, :CPU)
+        pr = get_projector!(get_prop(cl_h, :pool_of_projectors), ipr, :CPU)
         en = get_prop(cl_h, v, w, :en)
         pl = pl[states[v]]
         pr = pr[states[w]]
         pl_transition, pl_unique = rank_reveal(pl, :PE)
         pr_transition, pr_unique = rank_reveal(pr, :PE)
         en = en[pl_unique, pr_unique]
-        ipl = add_projector2!(new_lp, pl_transition)
-        ipr = add_projector2!(new_lp, pr_transition)
+        ipl = add_projector!(new_lp, pl_transition)
+        ipr = add_projector!(new_lp, pr_transition)
         set_props!(
                   new_cl_h, v, w, Dict(:outer_edges => outer_edges, :ipl => ipl, :en => en, :ipr => ipr)
               )

diff --git a/src/projectors.jl b/src/projectors.jl
@@ -1,41 +1,41 @@
 export
-    PoolOfProjectors2,
-    get_projector2!,
-    add_projector2!
+    PoolOfProjectors,
+    get_projector!,
+    add_projector!
 
-const Proj{T} = Union{Vector{T}, Array{T, 1}}
+const Proj{T} = Union{Vector{T}, CuArray{T, 1}}
 
-struct PoolOfProjectors2{T <: Integer}
+struct PoolOfProjectors{T <: Integer}
     data::Dict{Symbol, Dict{Int, Proj{T}}}
     default_device::Symbol
     sizes::Dict{Int, Int}
 
-    PoolOfProjectors2(data::Dict{Int, Dict{Int, Vector{T}}}) where T = new{T}(Dict(:CPU => data),
+    PoolOfProjectors(data::Dict{Int, Dict{Int, Vector{T}}}) where T = new{T}(Dict(:CPU => data),
                                                                             :CPU,
                                                                             Dict{Int, Int}(k => maximum(v) for (k, v) ∈ data))
-    PoolOfProjectors2{T}() where T = new{T}(Dict(:CPU => Dict{Int, Proj{T}}()),
+    PoolOfProjectors{T}() where T = new{T}(Dict(:CPU => Dict{Int, Proj{T}}()),
                                            :CPU,
                                            Dict{Int, Int}())
 end
 
 
-Base.eltype(lp::PoolOfProjectors2{T}) where T = T
-Base.length(lp::PoolOfProjectors2) = length(lp.data[lp.default_device])
-Base.length(lp::PoolOfProjectors2, device::Symbol) = length(lp.data[device])
+Base.eltype(lp::PoolOfProjectors{T}) where T = T
+Base.length(lp::PoolOfProjectors) = length(lp.data[lp.default_device])
+Base.length(lp::PoolOfProjectors, device::Symbol) = length(lp.data[device])
 
-function Base.empty!(lp::PoolOfProjectors2, device::Symbol) 
+function Base.empty!(lp::PoolOfProjectors, device::Symbol) 
     if device ∈ keys(lp.data)
         empty!(lp.data[device])
     end
 end
 
-Base.length(lp::PoolOfProjectors2, index::Int) = length(lp.data[lp.default_device][index])
-Base.size(lp::PoolOfProjectors2, index::Int) = lp.sizes[index]
+Base.length(lp::PoolOfProjectors, index::Int) = length(lp.data[lp.default_device][index])
+Base.size(lp::PoolOfProjectors, index::Int) = lp.sizes[index]
 
-get_projector2!(lp::PoolOfProjectors2, index::Int) = get_projector2!(lp, index, lp.default_device)
+get_projector!(lp::PoolOfProjectors, index::Int) = get_projector!(lp, index, lp.default_device)
 
 # TODO This is version for only one GPU
-function get_projector2!(lp::PoolOfProjectors2{T}, index::Int, device::Symbol) where T <: Integer
+function get_projector!(lp::PoolOfProjectors{T}, index::Int, device::Symbol) where T <: Integer
     if device ∉ keys(lp.data)
         push!(lp.data, device => Dict{Int, Proj{T}}())
     end
@@ -53,7 +53,7 @@ function get_projector2!(lp::PoolOfProjectors2{T}, index::Int, device::Symbol) w
     lp.data[device][index]
 end
 
-function add_projector2!(lp::PoolOfProjectors2{T}, p::Proj) where T <: Integer
+function add_projector!(lp::PoolOfProjectors{T}, p::Proj) where T <: Integer
     if lp.default_device == :CPU
         p = Array{T}(p)
     elseif lp.default_device == :GPU
@@ -76,3 +76,57 @@ function add_projector2!(lp::PoolOfProjectors2{T}, p::Proj) where T <: Integer
     end
     key
 end
+
+# get_projector!(lp::PoolOfProjectors, index::Int) = get_projector!(lp, index, lp.default_device)
+
+# # TODO This is version for only one GPU
+# function get_projector!(lp::PoolOfProjectors{T}, index::Int, ::Val{:CPU}) where T <: Integer
+#     device = :CPU
+#     if device ∉ keys(lp.data)
+#         push!(lp.data, device => Dict{Int, Proj{T}}())
+#     end
+
+#     if index ∉ keys(lp.data[device])
+#         if device == :GPU
+#             p = CuArray{T}(lp.data[lp.default_device][index])
+#         elseif device == :CPU
+#             p = Array{T}(lp.data[lp.default_device][index])
+#         else
+#             throw(ArgumentError("device should be :CPU or :GPU"))
+#         end
+#         push!(lp.data[device], index => p)
+#     end
+#     lp.data[device][index]
+# end
+
+
+# add_projector!(lp::PoolOfProjectors, p::Proj) = add_projector!(lp::PoolOfProjectors, p::Proj, lp.default_device)
+
+# function add_projector!(lp::PoolOfProjectors{T}, p::Proj, device::Symbol) where T <: Integer
+#     if device == :CPU
+#         return add_projector!(lp::PoolOfProjectors, p::Proj, Val(device))
+#     elseif device == :GPU
+#         return SpinGlassTensors.add_projector!(lp::PoolOfProjectors, p::Proj, Val(device))
+#     else
+#         throw(ArgumentError("default_device should be :CPU or :GPU"))
+#     end
+# end
+
+# function add_projector!(lp::PoolOfProjectors{T}, p::Proj, ::Val{:CPU}) where T <: Integer
+#     p = Array{T}(p)
+
+#     if p in values(lp.data[lp.default_device])
+#         key = -1
+#         for guess in keys(lp.data[lp.default_device])
+#             if lp.data[lp.default_device][guess] == p
+#                 key = guess
+#                 break
+#             end
+#         end
+#     else
+#         key = length(lp.data[lp.default_device]) + 1
+#         push!(lp.data[lp.default_device], key => p)
+#         push!(lp.sizes, key => maximum(p))
+#     end
+#     key
+# end
diff --git a/test/clustered_hamiltonian.jl b/test/clustered_hamiltonian.jl
@@ -88,16 +88,16 @@ end
 
    for (bd, e) in zip(bond_dimensions, edges(cl_h))
       ipl, en, ipr = get_prop(cl_h, e, :ipl), get_prop(cl_h, e, :en), get_prop(cl_h, e, :ipr)
-      pl = get_projector2!(get_prop(cl_h, :pool_of_projectors), ipl, :CPU)
-      pr = get_projector2!(get_prop(cl_h, :pool_of_projectors), ipr, :CPU)
+      pl = get_projector!(get_prop(cl_h, :pool_of_projectors), ipl, :CPU)
+      pr = get_projector!(get_prop(cl_h, :pool_of_projectors), ipr, :CPU)
 
       @test minimum(size(en)) == bd
    end
 
    for ((i, j), cedge) ∈ cedges
       ipl, en, ipr = get_prop(cl_h, i, j, :ipl), get_prop(cl_h, i, j, :en), get_prop(cl_h, i, j, :ipr)
-      pl = get_projector2!(get_prop(cl_h, :pool_of_projectors), ipl, :CPU)
-      pr = get_projector2!(get_prop(cl_h, :pool_of_projectors), ipr, :CPU)
+      pl = get_projector!(get_prop(cl_h, :pool_of_projectors), ipl, :CPU)
+      pr = get_projector!(get_prop(cl_h, :pool_of_projectors), ipr, :CPU)
       base_i = all_states(rank[i])
       base_j = all_states(rank[j])
 
@@ -242,8 +242,8 @@ function clustered_hamiltonian_energy(cl_h, state)
    for edge ∈ edges(cl_h)
       i, j = cl_h.reverse_label_map[src(edge)], cl_h.reverse_label_map[dst(edge)]
       ipl, en, ipr = get_prop(cl_h, edge, :ipl), get_prop(cl_h, edge, :en), get_prop(cl_h, edge, :ipr)
-      pl = get_projector2!(get_prop(cl_h, :pool_of_projectors), ipl, :CPU)
-      pr = get_projector2!(get_prop(cl_h, :pool_of_projectors), ipr, :CPU)
+      pl = get_projector!(get_prop(cl_h, :pool_of_projectors), ipl, :CPU)
+      pr = get_projector!(get_prop(cl_h, :pool_of_projectors), ipr, :CPU)
       edge_energy = en[pl, pr]
       total_en += edge_energy[state[i], state[j]]
    end