revise interface for ising energy levels (#20)

lanl-ansi · Feb 16, 2022 · 99f4183 · 99f4183
1 parent 95c0ea0
commit 99f4183
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 10 deletions.
diff --git a/src/ising.jl b/src/ising.jl
@@ -46,7 +46,9 @@ function compute_ising_energy_levels(ising_model::Dict)
         push!(energy_levels[energy], state_id)
     end
 
-    return energy_levels
+    energies = sort(collect(keys(energy_levels)))
+
+    return [(energy=e, states=energy_levels[e]) for e in energies]
 end
 
 """
@@ -59,13 +61,9 @@ function print_ising_energy_levels(ising_model::Dict; limit=3)
 
     energy_levels = compute_ising_energy_levels(ising_model)
 
-    energies = sort(collect(keys(energy_levels)))
-
-    for (i,energy) in enumerate(energies)
-        state_ids = energy_levels[energy]
-
-        println("\033[1menergy: $(energy)\033[0m")
-        for state_id in sort(collect(state_ids))
+    for (i,energy_level) in enumerate(energy_levels)
+        println("\033[1menergy: $(energy_level.energy)\033[0m")
+        for state_id in sort(collect(energy_level.states))
             state = binary2spin(int2binary(state_id, pad=n))
             state_string = spin2braket(state)
             println("  $(state_string)")

diff --git a/test/base.jl b/test/base.jl
@@ -100,8 +100,10 @@ end
 
     @testset "two qubit, energy levels" begin
         energy_levels = compute_ising_energy_levels(Dict((1,2) => -1))
-        @test energy_levels[-1.0] == Set([0,3])
-        @test energy_levels[1.0] == Set([1,2])
+        @test energy_levels[1].energy == -1.0
+        @test energy_levels[1].states == Set([0,3])
+        @test energy_levels[2].energy == 1.0
+        @test energy_levels[2].states == Set([1,2])
     end
 
     @testset "two qubit, print energy levels" begin