From 2fbf9493b23ca6bb212e16c695e587df9575c516 Mon Sep 17 00:00:00 2001
From: fernandopalafox <ferpalafoxe@gmail.com>
Date: Thu, 1 Feb 2024 10:57:47 -0600
Subject: [PATCH 01/11] added surface heatmap

---
 experiments/tower_defense.jl | 49 ++++++++++++++++++++++++++++++------
 1 file changed, 42 insertions(+), 7 deletions(-)

diff --git a/experiments/tower_defense.jl b/experiments/tower_defense.jl
index 460e37f..18a771f 100644
--- a/experiments/tower_defense.jl
+++ b/experiments/tower_defense.jl
@@ -2,7 +2,7 @@ using GamesVoI
 using BlockArrays
 using LinearAlgebra: norm_sqr, norm
 using Zygote  
-using GLMakie: Figure, Axis, Colorbar, heatmap!, text!
+using GLMakie: Figure, Axis, Colorbar, heatmap!, text!, surface!, Axis3
 
 """ Nomenclature
     N                         : Number of worlds (=3)
@@ -76,11 +76,11 @@ end
 """
 Temp. script to calculate and plot heatmap of Stage 1 cost function 
 """
-function run_heatmap()
+function run_visualization(; type = "heatmap")
     ps = [1/3, 1/3, 1/3]
     βs = [[2, 1, 1], [1, 2, 1], [1, 1, 2]]
-    Ks = calculate_stage_1_heatmap(ps, βs)
-    fig = display_heatmap(ps, Ks)
+    Ks = calculate_stage_1_costs(ps, βs)
+    fig = type == "heatmap" ? display_heatmap(ps, Ks) : display_surface(ps, Ks)
     fig
 end
 
@@ -93,9 +93,9 @@ Inputs:
     βs: vector containing P2's cost parameters for each world. vector of nx1 vectors
     dr: step size for r
 Outputs:
-    Ks: 2D Matrix of stage 1's objective function values for each r in the simplex.
+    Ks: 2D Matrix of stage 1's objective function values for each r in the simplex. Normalized, by default 
 """
-function calculate_stage_1_heatmap(ps, βs; dr = 0.05)
+function calculate_stage_1_costs(ps, βs; dr=0.05, normalize=true)
     @assert sum(ps) ≈ 1.0 "Prior distribution ps must be a probability distribution"
     game, _ = build_stage_2(ps, βs)
     rs = 0:dr:1
@@ -130,13 +130,48 @@ function display_heatmap(ps, Ks)
     ax = Axis(fig[1, 1]; xlabel="r₁", ylabel="r₂", 
     title="Stage 1 cost as a function of r \n priors = $(round.(ps, digits=2))", aspect=1)
     hmap = heatmap!(ax, rs, rs, Ks)
-    Colorbar(fig[1, 2], hmap; label = "K", width = 15, ticksize = 15, tickalign = 1)
+    Colorbar(fig[1, 2], hmap; label = "K̃", width = 15, ticksize = 15, tickalign = 1)
     text!(ax, "$(round(ps[1], digits=2))", position = (0.9, 0.15), font = "Bold")
     text!(ax, "$(round(ps[2], digits=2))", position = (0.1, 0.95), font = "Bold")
     text!(ax, "$(round(ps[3], digits=2))", position = (0.1, 0.1), font = "Bold")
     fig
 end
 
+"""
+Display surface of Stage 1's objective function. Assumes number of worlds is 3.
+
+Input: 
+    Ks: 2D Matrix of stage 1's objective function values for each r in the simplex
+Output: 
+    fig: Figure with simplex heatmap
+"""
+function display_surface(ps, Ks)
+    rs = 0:1/(size(Ks)[1] - 1):1
+    fig = Figure(size = (600, 400))  
+    ax = Axis3(
+        fig[1, 1],
+        aspect = (1, 1, 1),
+        perspectiveness = 0.5,
+        elevation = pi / 20,
+        azimuth = -π * (1 / 2 + 1 / 4),
+        zgridcolor = :grey,
+        ygridcolor = :grey,
+        xgridcolor = :grey;
+        xlabel = "r₁",
+        ylabel = "r₂",
+        zlabel = "K",
+        title = "Stage 1 cost as a function of r \n priors = $(round.(ps, digits=2))",
+    )
+    Ks_min = minimum(filter(!isnan, Ks))
+    hmap = surface!(ax, rs, rs, Ks)
+    Colorbar(fig[1, 2], hmap; label = "K", width = 15, ticksize = 15, tickalign = 1)
+    text!(ax, "$(round(ps[1], digits=2))", position = (0.9, 0.4, Ks_min), font = "Bold")
+    text!(ax, "$(round(ps[2], digits=2))", position = (0.1, 0.95, Ks_min), font = "Bold")
+    text!(ax, "$(round(ps[3], digits=2))", position = (0.2, 0.1, Ks_min), font = "Bold")
+    fig
+end
+
+
 """
 Project onto simplex using Fig. 1 Duchi 2008
 """

From e42f4032c60e33322799324f103ff82562b4d6f7 Mon Sep 17 00:00:00 2001
From: fernandopalafox <ferpalafoxe@gmail.com>
Date: Thu, 1 Feb 2024 10:58:06 -0600
Subject: [PATCH 02/11] added normalization arg

---
 experiments/tower_defense.jl | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/experiments/tower_defense.jl b/experiments/tower_defense.jl
index 18a771f..381ad2c 100644
--- a/experiments/tower_defense.jl
+++ b/experiments/tower_defense.jl
@@ -113,6 +113,13 @@ function calculate_stage_1_costs(ps, βs; dr=0.05, normalize=true)
         end
     end
 
+    if !normalize
+        return Ks
+    end
+
+    max_value = maximum(filter(!isnan, Ks))
+    Ks = [isnan(K) ? NaN : K / max_value for K in Ks]
+
     return Ks
 end
 

From 4062db192a17cdb30f6d620b92fe6f4f4c7e9b52 Mon Sep 17 00:00:00 2001
From: fernandopalafox <ferpalafoxe@gmail.com>
Date: Thu, 1 Feb 2024 10:58:20 -0600
Subject: [PATCH 03/11] zero-sum subgames

---
 experiments/tower_defense.jl | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/experiments/tower_defense.jl b/experiments/tower_defense.jl
index 381ad2c..4f2dec5 100644
--- a/experiments/tower_defense.jl
+++ b/experiments/tower_defense.jl
@@ -191,9 +191,9 @@ function project_onto_simplex(v; z=1.0)
 end 
 
 "Defender cost function"
-function J_1(u, v) 
-    norm_sqr(u - v)
-end
+function J_1(u, v, β)
+    -J_2(u, v, β) 
+end 
 
 """
 Attacker cost function
@@ -229,8 +229,8 @@ function build_stage_2(ps, βs)
     # Define Bayesian game player costs in Stage 2
     p_w_k_0(w_idx, θ) = (1 - θ[w_idx]) * ps[w_idx] / (1 - θ' * ps)
     fs = [
-        (x, θ) ->  sum([J_1(x[Block(1)], x[Block(w_idx + n + 1)]) * p_w_k_0(w_idx, θ) for w_idx in 1:n]), # u|s¹=0 IPI
-        [(x, θ) -> J_1(x[Block(w_idx + 1)], x[Block(w_idx + 2 * n + 1)]) for w_idx in 1:n]..., # u|s¹={1,2,3} PI
+        (x, θ) ->  sum([J_1(x[Block(1)], x[Block(w_idx + n + 1)], βs[w_idx]) * p_w_k_0(w_idx, θ) for w_idx in 1:n]), # u|s¹=0 IPI
+        [(x, θ) -> J_1(x[Block(w_idx + 1)], x[Block(w_idx + 2 * n + 1)], βs[w_idx]) for w_idx in 1:n]..., # u|s¹={1,2,3} PI
         [(x, θ) -> J_2(x[Block(1)], x[Block(w_idx + n + 1)], βs[w_idx]) for w_idx in 1:n]...,  # v|s¹=0 IPI
         [(x, θ) -> J_2(x[Block(w_idx + 1)], x[Block(w_idx + 2 * n + 1)], βs[w_idx]) for w_idx in 1:n]..., # v|s¹={1,2,3} PI
     ]
@@ -265,8 +265,8 @@ Compute objective at Stage 1
 """
 function compute_K(r, x, ps, βs)
     n = length(ps)
-    sum([(1 - r[j]) * ps[j] * J_1(x[Block(1)], x[Block(j + n + 1)]) for j in 1:n]) + 
-    sum([r[j] * ps[j] * J_1(x[Block(j + 1)], x[Block(j + 2 * n + 1)]) for j in 1:n])
+    sum([(1 - r[j]) * ps[j] * J_1(x[Block(1)], x[Block(j + n + 1)], βs[j]) for j in 1:n]) + 
+    sum([r[j] * ps[j] * J_1(x[Block(j + 1)], x[Block(j + 2 * n + 1)], βs[j]) for j in 1:n])
 end
 
 """

From 66a027a1bba1135ef6758c958756e00059644e46 Mon Sep 17 00:00:00 2001
From: fernandopalafox <ferpalafoxe@gmail.com>
Date: Thu, 1 Feb 2024 10:58:34 -0600
Subject: [PATCH 04/11] added formatter

---
 .JuliaFormatter.toml | 6 ++++++
 1 file changed, 6 insertions(+)
 create mode 100644 .JuliaFormatter.toml

diff --git a/.JuliaFormatter.toml b/.JuliaFormatter.toml
new file mode 100644
index 0000000..71488b2
--- /dev/null
+++ b/.JuliaFormatter.toml
@@ -0,0 +1,6 @@
+margin = 100
+always_for_in = true
+import_to_using = true
+remove_extra_newlines = true
+whitespace_typedefs = false
+whitespace_ops_in_indices = true
\ No newline at end of file

From d5781d6353b07c99597b9d66bd18097bd4e2da3e Mon Sep 17 00:00:00 2001
From: fernandopalafox <ferpalafoxe@gmail.com>
Date: Thu, 1 Feb 2024 10:59:52 -0600
Subject: [PATCH 05/11] formatting

---
 experiments/tower_defense.jl | 158 ++++++++++++++++++++---------------
 1 file changed, 91 insertions(+), 67 deletions(-)

diff --git a/experiments/tower_defense.jl b/experiments/tower_defense.jl
index 4f2dec5..145fbef 100644
--- a/experiments/tower_defense.jl
+++ b/experiments/tower_defense.jl
@@ -1,7 +1,7 @@
 using GamesVoI
 using BlockArrays
 using LinearAlgebra: norm_sqr, norm
-using Zygote  
+using Zygote
 using GLMakie: Figure, Axis, Colorbar, heatmap!, text!, surface!, Axis3
 
 """ Nomenclature
@@ -20,7 +20,6 @@ using GLMakie: Figure, Axis, Colorbar, heatmap!, text!, surface!, Axis3
     J                         : Stage 1's objective function  
 """
 
-
 """
 Solve Stage 1 to find optimal scout allocation r.
 
@@ -31,7 +30,15 @@ Inputs:
 Outputs:
     r: optimal scout allocation
 """
-function solve_r(ps, βs; r_init = [1/3, 1/3, 1/3], iter_limit=50, target_error=.00001, α=1, return_states = false)
+function solve_r(
+    ps,
+    βs;
+    r_init = [1 / 3, 1 / 3, 1 / 3],
+    iter_limit = 50,
+    target_error = 0.00001,
+    α = 1,
+    return_states = false,
+)
     @assert sum(r_init) ≈ 1.0 "Initial guess r must be a probability distribution"
     cur_iter = 0
     n = length(ps)
@@ -41,22 +48,25 @@ function solve_r(ps, βs; r_init = [1/3, 1/3, 1/3], iter_limit=50, target_error=
         x_list = []
         r_list = []
     end
-    game, _ = build_stage_2(ps, βs) 
+    game, _ = build_stage_2(ps, βs)
     r = r_init
     println("0: r = $r")
     x = compute_stage_2(r, ps, βs, game)
     dKdr = zeros(Float64, n)
     while cur_iter < iter_limit # TODO: Break if change from last iteration is small
         dKdr = compute_dKdr(r, x, ps, βs, game)
-        r_temp = r - α .* dKdr 
-        r = project_onto_simplex(r_temp) 
+        r_temp = r - α .* dKdr
+        r = project_onto_simplex(r_temp)
         x = compute_stage_2(
-            r, ps, βs, game;
-            initial_guess=vcat(x, zeros(total_dim(game) - n_players * var_dim))
+            r,
+            ps,
+            βs,
+            game;
+            initial_guess = vcat(x, zeros(total_dim(game) - n_players * var_dim)),
         )
         if return_states
-            push!(x_list,x)
-            push!(r_list,r)
+            push!(x_list, x)
+            push!(r_list, r)
         end
         # compute stage 1 cost function for current r and x 
         K = compute_K(r, x, ps, βs)
@@ -67,7 +77,7 @@ function solve_r(ps, βs; r_init = [1/3, 1/3, 1/3], iter_limit=50, target_error=
     if return_states
         r_matrix = reduce(hcat, r_list)
         x_matrix = reduce(hcat, x_list)
-        out = Dict("r"=>r, "x"=>x, "r_matrix"=>r_matrix, "x_matrix"=>x_matrix)
+        out = Dict("r" => r, "x" => x, "r_matrix" => r_matrix, "x_matrix" => x_matrix)
         return out
     end
     return r
@@ -77,14 +87,13 @@ end
 Temp. script to calculate and plot heatmap of Stage 1 cost function 
 """
 function run_visualization(; type = "heatmap")
-    ps = [1/3, 1/3, 1/3]
+    ps = [1 / 3, 1 / 3, 1 / 3]
     βs = [[2, 1, 1], [1, 2, 1], [1, 1, 2]]
     Ks = calculate_stage_1_costs(ps, βs)
     fig = type == "heatmap" ? display_heatmap(ps, Ks) : display_surface(ps, Ks)
     fig
 end
 
-
 """
 Calculate Stage 1's objective function for all possible values of r.
 
@@ -95,18 +104,18 @@ Inputs:
 Outputs:
     Ks: 2D Matrix of stage 1's objective function values for each r in the simplex. Normalized, by default 
 """
-function calculate_stage_1_costs(ps, βs; dr=0.05, normalize=true)
+function calculate_stage_1_costs(ps, βs; dr = 0.05, normalize = true)
     @assert sum(ps) ≈ 1.0 "Prior distribution ps must be a probability distribution"
     game, _ = build_stage_2(ps, βs)
     rs = 0:dr:1
-    Ks = NaN*ones(Float64, Int(1/dr + 1), Int(1/dr + 1))
+    Ks = NaN * ones(Float64, Int(1 / dr + 1), Int(1 / dr + 1))
     for (i, r1) in enumerate(rs)
         for (j, r2) in enumerate(rs)
             if r1 + r2 > 1
                 continue
             end
-            r3 = 1 - r1 - r2    
-            r = [r1, r2, r3]            
+            r3 = 1 - r1 - r2
+            r = [r1, r2, r3]
             x = compute_stage_2(r, ps, βs, game)
             K = compute_K(r, x, ps, βs)
             Ks[i, j] = K
@@ -132,10 +141,15 @@ Output:
     fig: Figure with simplex heatmap
 """
 function display_heatmap(ps, Ks)
-    rs = 0:1/(size(Ks)[1] - 1):1
-    fig = Figure(size = (600, 400))  
-    ax = Axis(fig[1, 1]; xlabel="r₁", ylabel="r₂", 
-    title="Stage 1 cost as a function of r \n priors = $(round.(ps, digits=2))", aspect=1)
+    rs = 0:(1 / (size(Ks)[1] - 1)):1
+    fig = Figure(size = (600, 400))
+    ax = Axis(
+        fig[1, 1];
+        xlabel = "r₁",
+        ylabel = "r₂",
+        title = "Stage 1 cost as a function of r \n priors = $(round.(ps, digits=2))",
+        aspect = 1,
+    )
     hmap = heatmap!(ax, rs, rs, Ks)
     Colorbar(fig[1, 2], hmap; label = "K̃", width = 15, ticksize = 15, tickalign = 1)
     text!(ax, "$(round(ps[1], digits=2))", position = (0.9, 0.15), font = "Bold")
@@ -153,8 +167,8 @@ Output:
     fig: Figure with simplex heatmap
 """
 function display_surface(ps, Ks)
-    rs = 0:1/(size(Ks)[1] - 1):1
-    fig = Figure(size = (600, 400))  
+    rs = 0:(1 / (size(Ks)[1] - 1)):1
+    fig = Figure(size = (600, 400))
     ax = Axis3(
         fig[1, 1],
         aspect = (1, 1, 1),
@@ -178,36 +192,35 @@ function display_surface(ps, Ks)
     fig
 end
 
-
 """
 Project onto simplex using Fig. 1 Duchi 2008
 """
-function project_onto_simplex(v; z=1.0)
-    μ = sort(v, rev=true)
-    ρ = findfirst([μ[j] - 1/j * (sum(μ[1:j]) - z) <= 0 for j in eachindex(v)]) 
+function project_onto_simplex(v; z = 1.0)
+    μ = sort(v, rev = true)
+    ρ = findfirst([μ[j] - 1 / j * (sum(μ[1:j]) - z) <= 0 for j in eachindex(v)])
     ρ = isnothing(ρ) ? length(v) : ρ - 1
-    θ = 1/ρ * (sum(μ[1:ρ]) - z)
+    θ = 1 / ρ * (sum(μ[1:ρ]) - z)
     return [maximum([v[i] - θ, 0]) for i in eachindex(v)]
-end 
+end
 
 "Defender cost function"
 function J_1(u, v, β)
-    -J_2(u, v, β) 
-end 
+    -J_2(u, v, β)
+end
 
 """
 Attacker cost function
 β: vector containing P2's (attacker) preference parameters for each world.
 """
 function J_2(u, v, β)
-#    δ = [β[ii] * v[ii] - u[ii] for ii in eachindex(β)]
-#    -sum([activate(δ[j]) for j in eachindex(β)])
-   -sum([β[ii]^(v[ii]-u[ii]) for ii in eachindex(β)])
-end 
+    #    δ = [β[ii] * v[ii] - u[ii] for ii in eachindex(β)]
+    #    -sum([activate(δ[j]) for j in eachindex(β)])
+    -sum([β[ii]^(v[ii] - u[ii]) for ii in eachindex(β)])
+end
 
 "Approximate Heaviside step function"
-function activate(δ; k=1.0)
-    return 1/(1 + exp(-2 * δ * k))
+function activate(δ; k = 1.0)
+    return 1 / (1 + exp(-2 * δ * k))
 end
 
 """
@@ -222,17 +235,25 @@ Outputs:
 
 """
 function build_stage_2(ps, βs)
-
     n = length(ps) # assume n_signals = n_worlds + 1
     n_players = 1 + n^2
 
     # Define Bayesian game player costs in Stage 2
     p_w_k_0(w_idx, θ) = (1 - θ[w_idx]) * ps[w_idx] / (1 - θ' * ps)
     fs = [
-        (x, θ) ->  sum([J_1(x[Block(1)], x[Block(w_idx + n + 1)], βs[w_idx]) * p_w_k_0(w_idx, θ) for w_idx in 1:n]), # u|s¹=0 IPI
-        [(x, θ) -> J_1(x[Block(w_idx + 1)], x[Block(w_idx + 2 * n + 1)], βs[w_idx]) for w_idx in 1:n]..., # u|s¹={1,2,3} PI
+        (x, θ) -> sum([
+            J_1(x[Block(1)], x[Block(w_idx + n + 1)], βs[w_idx]) * p_w_k_0(w_idx, θ) for
+            w_idx in 1:n
+        ]), # u|s¹=0 IPI
+        [
+            (x, θ) -> J_1(x[Block(w_idx + 1)], x[Block(w_idx + 2 * n + 1)], βs[w_idx]) for
+            w_idx in 1:n
+        ]..., # u|s¹={1,2,3} PI
         [(x, θ) -> J_2(x[Block(1)], x[Block(w_idx + n + 1)], βs[w_idx]) for w_idx in 1:n]...,  # v|s¹=0 IPI
-        [(x, θ) -> J_2(x[Block(w_idx + 1)], x[Block(w_idx + 2 * n + 1)], βs[w_idx]) for w_idx in 1:n]..., # v|s¹={1,2,3} PI
+        [
+            (x, θ) -> J_2(x[Block(w_idx + 1)], x[Block(w_idx + 2 * n + 1)], βs[w_idx]) for
+            w_idx in 1:n
+        ]..., # v|s¹={1,2,3} PI
     ]
 
     # equality constraints   
@@ -246,18 +267,19 @@ function build_stage_2(ps, βs)
     h̃ = (x, θ) -> [0]
 
     ParametricGame(;
-        objectives=fs,
-        equality_constraints=gs,
-        inequality_constraints=hs,
-        shared_equality_constraint=g̃,
-        shared_inequality_constraint=h̃,
-        parameter_dimension=3,
-        primal_dimensions=[3 for _ in 1:n_players],
-        equality_dimensions=[1 for _ in 1:n_players],
-        inequality_dimensions=[3 for _ in 1:n_players],
-        shared_equality_dimension=1,
-        shared_inequality_dimension=1
-    ), fs
+        objectives = fs,
+        equality_constraints = gs,
+        inequality_constraints = hs,
+        shared_equality_constraint = g̃,
+        shared_inequality_constraint = h̃,
+        parameter_dimension = 3,
+        primal_dimensions = [3 for _ in 1:n_players],
+        equality_dimensions = [1 for _ in 1:n_players],
+        inequality_dimensions = [3 for _ in 1:n_players],
+        shared_equality_dimension = 1,
+        shared_inequality_dimension = 1,
+    ),
+    fs
 end
 
 """
@@ -265,15 +287,14 @@ Compute objective at Stage 1
 """
 function compute_K(r, x, ps, βs)
     n = length(ps)
-    sum([(1 - r[j]) * ps[j] * J_1(x[Block(1)], x[Block(j + n + 1)], βs[j]) for j in 1:n]) + 
-    sum([r[j] * ps[j] * J_1(x[Block(j + 1)], x[Block(j + 2 * n + 1)], βs[j]) for j in 1:n])
+    sum([(1 - r[j]) * ps[j] * J_1(x[Block(1)], x[Block(j + n + 1)], βs[j]) for j in 1:n]) + sum([r[j] * ps[j] * J_1(x[Block(j + 1)], x[Block(j + 2 * n + 1)], βs[j]) for j in 1:n])
 end
 
 """
 Compute derivative of Stage 1's objective function w.r.t. x
 """
 function compute_dKdx(r, x, ps, βs)
-    gradient(x -> compute_K(r, x, ps, βs), x)[1] 
+    gradient(x -> compute_K(r, x, ps, βs), x)[1]
 end
 
 """
@@ -308,19 +329,22 @@ Inputs:
 Outputs:
     dxdr: Blocked Jacobian of Stage 2's decision variables w.r.t. Stage 1's decision variable
 """
-function compute_dxdr(r, x, ps, βs, game; verbose=false)
+function compute_dxdr(r, x, ps, βs, game; verbose = false)
     n = length(ps)
     n_players = 1 + n^2
     var_dim = n # TODO: Change this to be more general
 
     # Return Jacobian
-    dxdr = jacobian(r -> solve(
+    dxdr = jacobian(
+        r -> solve(
             game,
             r;
-            initial_guess=vcat(x, zeros(total_dim(game) - n_players * var_dim)),
-            verbose=false,
-            return_primals=false
-        ).variables[1:n_players*var_dim], r)[1]
+            initial_guess = vcat(x, zeros(total_dim(game) - n_players * var_dim)),
+            verbose = false,
+            return_primals = false,
+        ).variables[1:(n_players * var_dim)],
+        r,
+    )[1]
 
     BlockArray(dxdr, [var_dim for _ in 1:n_players], [var_dim])
 end
@@ -335,7 +359,7 @@ Input:
 Output: 
     x: decision variables of Stage 2 given r. BlockedArray with a block per player
 """
-function compute_stage_2(r, ps, βs, game; initial_guess = nothing, verbose=false)
+function compute_stage_2(r, ps, βs, game; initial_guess = nothing, verbose = false)
     n = length(ps) # assume n_signals = n_worlds + 1
     n_players = 1 + n^2
     var_dim = n # TODO: Change this to be more general
@@ -343,10 +367,10 @@ function compute_stage_2(r, ps, βs, game; initial_guess = nothing, verbose=fals
     solution = solve(
         game,
         r;
-        initial_guess=isnothing(initial_guess) ? zeros(total_dim(game)) : initial_guess,
-        verbose=verbose,
-        return_primals=false
+        initial_guess = isnothing(initial_guess) ? zeros(total_dim(game)) : initial_guess,
+        verbose = verbose,
+        return_primals = false,
     )
 
-    BlockArray(solution.variables[1:n_players * var_dim], [n for _ in 1:n_players])
+    BlockArray(solution.variables[1:(n_players * var_dim)], [n for _ in 1:n_players])
 end
\ No newline at end of file

From 6c2028c5460dea04bdcdbebbdf1cfc5ca66c5fa7 Mon Sep 17 00:00:00 2001
From: fernandopalafox <ferpalafoxe@gmail.com>
Date: Sat, 3 Feb 2024 13:20:06 -0600
Subject: [PATCH 06/11] added utility to plot cost of misidentification

---
 experiments/tower_defense.jl | 94 ++++++++++++++++++++++++++++++++++++
 1 file changed, 94 insertions(+)

diff --git a/experiments/tower_defense.jl b/experiments/tower_defense.jl
index 145fbef..daa6b92 100644
--- a/experiments/tower_defense.jl
+++ b/experiments/tower_defense.jl
@@ -94,6 +94,100 @@ function run_visualization(; type = "heatmap")
     fig
 end
 
+"""
+Temp. script to calculate and plot heatmap for cost of misidentifying a world 
+"""
+function run_misid_vis()
+    dr = 0.05
+    ps = [1/3, 1/3, 1/3]
+    βs = [[2, 1, 1], [1, 2, 1], [1, 1, 2]] 
+    world_1_misid_costs = calculate_misid_costs(ps, βs, 1; dr)
+    world_2_misid_costs = calculate_misid_costs(ps, βs, 2; dr)
+    world_3_misid_costs = calculate_misid_costs(ps, βs, 3; dr)
+
+    display_misid_costs([world_1_misid_costs, world_2_misid_costs, world_3_misid_costs], ps)
+end
+
+function calculate_misid_costs(ps, βs, world_idx; dr = 0.05, normalize = true)
+    @assert sum(ps) ≈ 1.0 "Prior distribution ps must be a probability distribution"
+    game, _ = build_stage_2(ps, βs)
+    rs = 0:dr:1
+    num_worlds = length(ps)
+    misid_costs = NaN * ones(Float64, Int(1 / dr + 1), Int(1 / dr + 1))
+    for (i, r1) in enumerate(rs)
+        for (j, r2) in enumerate(rs)
+            if r1 + r2 > 1
+                continue
+            end
+            r3 = 1 - r1 - r2
+            r = [r1, r2, r3]
+            x = compute_stage_2(r, ps, βs, game)
+            defender_signal_0 = x[Block(1)]
+            attacker_signal_0_world_idx = x[Block(world_idx + num_worlds + 1)]
+            misid_cost = J_1(defender_signal_0, attacker_signal_0_world_idx, βs[world_idx])
+            misid_costs[i, j] = misid_cost
+        end
+    end
+
+    if !normalize
+        return misid_costs
+    end
+
+    max_value = maximum(filter(!isnan, misid_costs))
+    misid_costs = [isnan(c) ? NaN : c / max_value for c in misid_costs]
+
+    return misid_costs
+end
+
+"""
+Display surface of Stage 1's objective function. Assumes number of worlds is 3.
+
+Input: 
+    Ks: 2D Matrix of stage 1's objective function values for each r in the simplex
+Output: 
+    fig: Figure with simplex heatmap
+"""
+function display_misid_costs(costs, ps)
+    rs = 0:(1 / (size(costs[1])[1] - 1)):1
+    num_worlds = length(costs)
+    fig = Figure(size = (1500, 500), title = "test")
+    axs = [
+        Axis3(
+            fig[1, world_idx],
+            aspect = (1, 1, 1),
+            perspectiveness = 0.5,
+            elevation = pi / 4,
+            azimuth = -π * (1 / 2 + 1 / 4),
+            zgridcolor = :grey,
+            ygridcolor = :grey,
+            xgridcolor = :grey;
+            xlabel = "r₁",
+            ylabel = "r₂",
+            zlabel = "Misid. cost",
+            title = "World $world_idx",
+        ) for world_idx in 1:num_worlds
+    ]
+    for world_idx in 1:num_worlds
+        cost_min = minimum(filter(!isnan, costs[world_idx]))
+        hmap = surface!(axs[world_idx], rs, rs, costs[world_idx], colormap = :viridis)
+        # text!(axs[world_idx], "$(round(ps[1], digits=2))", position = (0.9, 0.4, cost_min), font = "Bold")
+        # text!(axs[world_idx], "$(round(ps[2], digits=2))", position = (0.1, 0.95, cost_min), font = "Bold")
+        # text!(axs[world_idx], "$(round(ps[3], digits=2))", position = (0.2, 0.1, cost_min), font = "Bold")
+
+        if world_idx == num_worlds
+            Colorbar(
+                fig[1, world_idx + 1],
+                hmap;
+                label = "Misid. cost",
+                width = 15,
+                ticksize = 15,
+                tickalign = 1,
+            )
+        end
+    end
+    fig
+end
+
 """
 Calculate Stage 1's objective function for all possible values of r.
 

From fb9607f69e247059fd95eb9df2bb83956fae39ad Mon Sep 17 00:00:00 2001
From: fernandopalafox <ferpalafoxe@gmail.com>
Date: Sat, 3 Feb 2024 13:20:32 -0600
Subject: [PATCH 07/11] removed heatmap

---
 experiments/tower_defense.jl | 26 --------------------------
 1 file changed, 26 deletions(-)

diff --git a/experiments/tower_defense.jl b/experiments/tower_defense.jl
index daa6b92..a45e3cd 100644
--- a/experiments/tower_defense.jl
+++ b/experiments/tower_defense.jl
@@ -226,32 +226,6 @@ function calculate_stage_1_costs(ps, βs; dr = 0.05, normalize = true)
     return Ks
 end
 
-"""
-Display heatmap of Stage 1's objective function. Assumes number of worlds is 3.
-
-Input: 
-    Ks: 2D Matrix of stage 1's objective function values for each r in the simplex
-Output: 
-    fig: Figure with simplex heatmap
-"""
-function display_heatmap(ps, Ks)
-    rs = 0:(1 / (size(Ks)[1] - 1)):1
-    fig = Figure(size = (600, 400))
-    ax = Axis(
-        fig[1, 1];
-        xlabel = "r₁",
-        ylabel = "r₂",
-        title = "Stage 1 cost as a function of r \n priors = $(round.(ps, digits=2))",
-        aspect = 1,
-    )
-    hmap = heatmap!(ax, rs, rs, Ks)
-    Colorbar(fig[1, 2], hmap; label = "K̃", width = 15, ticksize = 15, tickalign = 1)
-    text!(ax, "$(round(ps[1], digits=2))", position = (0.9, 0.15), font = "Bold")
-    text!(ax, "$(round(ps[2], digits=2))", position = (0.1, 0.95), font = "Bold")
-    text!(ax, "$(round(ps[3], digits=2))", position = (0.1, 0.1), font = "Bold")
-    fig
-end
-
 """
 Display surface of Stage 1's objective function. Assumes number of worlds is 3.
 

From 968980ec19e59f04a0a1ad0017f3f688cdf85be9 Mon Sep 17 00:00:00 2001
From: fernandopalafox <ferpalafoxe@gmail.com>
Date: Sat, 3 Feb 2024 13:20:44 -0600
Subject: [PATCH 08/11] removed sigmoid cost

---
 experiments/tower_defense.jl | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/experiments/tower_defense.jl b/experiments/tower_defense.jl
index a45e3cd..ecbd453 100644
--- a/experiments/tower_defense.jl
+++ b/experiments/tower_defense.jl
@@ -281,16 +281,9 @@ Attacker cost function
 β: vector containing P2's (attacker) preference parameters for each world.
 """
 function J_2(u, v, β)
-    #    δ = [β[ii] * v[ii] - u[ii] for ii in eachindex(β)]
-    #    -sum([activate(δ[j]) for j in eachindex(β)])
     -sum([β[ii]^(v[ii] - u[ii]) for ii in eachindex(β)])
 end
 
-"Approximate Heaviside step function"
-function activate(δ; k = 1.0)
-    return 1 / (1 + exp(-2 * δ * k))
-end
-
 """
 Build parametric game for Stage 2.
 

From df0bc5758f84a1c3ad70aaea35458d7d50fa6709 Mon Sep 17 00:00:00 2001
From: fernandopalafox <ferpalafoxe@gmail.com>
Date: Sat, 3 Feb 2024 14:44:27 -0600
Subject: [PATCH 09/11] normalize according to max across ALL worlds

---
 experiments/tower_defense.jl | 29 ++++++++++++++++++-----------
 1 file changed, 18 insertions(+), 11 deletions(-)

diff --git a/experiments/tower_defense.jl b/experiments/tower_defense.jl
index ecbd453..853ddc4 100644
--- a/experiments/tower_defense.jl
+++ b/experiments/tower_defense.jl
@@ -105,10 +105,17 @@ function run_misid_vis()
     world_2_misid_costs = calculate_misid_costs(ps, βs, 2; dr)
     world_3_misid_costs = calculate_misid_costs(ps, βs, 3; dr)
 
+    # Normalize using maximum value across all worlds 
+    max_value =
+        maximum(filter(!isnan, vcat(world_1_misid_costs, world_2_misid_costs, world_3_misid_costs)))
+    world_1_misid_costs = [isnan(c) ? NaN : c / max_value for c in world_1_misid_costs]
+    world_2_misid_costs = [isnan(c) ? NaN : c / max_value for c in world_2_misid_costs]
+    world_3_misid_costs = [isnan(c) ? NaN : c / max_value for c in world_3_misid_costs]
+
     display_misid_costs([world_1_misid_costs, world_2_misid_costs, world_3_misid_costs], ps)
 end
 
-function calculate_misid_costs(ps, βs, world_idx; dr = 0.05, normalize = true)
+function calculate_misid_costs(ps, βs, world_idx; dr = 0.05)
     @assert sum(ps) ≈ 1.0 "Prior distribution ps must be a probability distribution"
     game, _ = build_stage_2(ps, βs)
     rs = 0:dr:1
@@ -129,13 +136,6 @@ function calculate_misid_costs(ps, βs, world_idx; dr = 0.05, normalize = true)
         end
     end
 
-    if !normalize
-        return misid_costs
-    end
-
-    max_value = maximum(filter(!isnan, misid_costs))
-    misid_costs = [isnan(c) ? NaN : c / max_value for c in misid_costs]
-
     return misid_costs
 end
 
@@ -156,7 +156,7 @@ function display_misid_costs(costs, ps)
             fig[1, world_idx],
             aspect = (1, 1, 1),
             perspectiveness = 0.5,
-            elevation = pi / 4,
+            elevation = pi / 5,
             azimuth = -π * (1 / 2 + 1 / 4),
             zgridcolor = :grey,
             ygridcolor = :grey,
@@ -165,11 +165,18 @@ function display_misid_costs(costs, ps)
             ylabel = "r₂",
             zlabel = "Misid. cost",
             title = "World $world_idx",
+            limits = (nothing, nothing, (0.01, 1)),
         ) for world_idx in 1:num_worlds
     ]
     for world_idx in 1:num_worlds
-        cost_min = minimum(filter(!isnan, costs[world_idx]))
-        hmap = surface!(axs[world_idx], rs, rs, costs[world_idx], colormap = :viridis)
+        hmap = surface!(
+            axs[world_idx],
+            rs,
+            rs,
+            costs[world_idx],
+            colormap = :viridis,
+            colorrange = (0, 1),
+        )
         # text!(axs[world_idx], "$(round(ps[1], digits=2))", position = (0.9, 0.4, cost_min), font = "Bold")
         # text!(axs[world_idx], "$(round(ps[2], digits=2))", position = (0.1, 0.95, cost_min), font = "Bold")
         # text!(axs[world_idx], "$(round(ps[3], digits=2))", position = (0.2, 0.1, cost_min), font = "Bold")

From 6a534ac44eed878c9dc7a4021dc8b4f363ca461f Mon Sep 17 00:00:00 2001
From: fernandopalafox <ferpalafoxe@gmail.com>
Date: Mon, 5 Feb 2024 14:29:08 -0600
Subject: [PATCH 10/11] cleand up old heatmap call

---
 experiments/tower_defense.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/experiments/tower_defense.jl b/experiments/tower_defense.jl
index 853ddc4..39c79d1 100644
--- a/experiments/tower_defense.jl
+++ b/experiments/tower_defense.jl
@@ -86,11 +86,11 @@ end
 """
 Temp. script to calculate and plot heatmap of Stage 1 cost function 
 """
-function run_visualization(; type = "heatmap")
+function run_visualization()
     ps = [1 / 3, 1 / 3, 1 / 3]
     βs = [[2, 1, 1], [1, 2, 1], [1, 1, 2]]
     Ks = calculate_stage_1_costs(ps, βs)
-    fig = type == "heatmap" ? display_heatmap(ps, Ks) : display_surface(ps, Ks)
+    fig = display_surface(ps, Ks)
     fig
 end
 

From 237081bc01f74849500f8847327798aa62a6e829 Mon Sep 17 00:00:00 2001
From: fernandopalafox <ferpalafoxe@gmail.com>
Date: Mon, 5 Feb 2024 14:29:31 -0600
Subject: [PATCH 11/11] changed axis limits for better perspective

---
 experiments/tower_defense.jl | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/experiments/tower_defense.jl b/experiments/tower_defense.jl
index 39c79d1..d5b1052 100644
--- a/experiments/tower_defense.jl
+++ b/experiments/tower_defense.jl
@@ -257,13 +257,14 @@ function display_surface(ps, Ks)
         ylabel = "r₂",
         zlabel = "K",
         title = "Stage 1 cost as a function of r \n priors = $(round.(ps, digits=2))",
+        limits = (nothing, nothing, (0.01, 1)),
     )
     Ks_min = minimum(filter(!isnan, Ks))
-    hmap = surface!(ax, rs, rs, Ks)
+    hmap = surface!(ax, rs, rs, Ks, colorrange = (0, 1))
     Colorbar(fig[1, 2], hmap; label = "K", width = 15, ticksize = 15, tickalign = 1)
-    text!(ax, "$(round(ps[1], digits=2))", position = (0.9, 0.4, Ks_min), font = "Bold")
-    text!(ax, "$(round(ps[2], digits=2))", position = (0.1, 0.95, Ks_min), font = "Bold")
-    text!(ax, "$(round(ps[3], digits=2))", position = (0.2, 0.1, Ks_min), font = "Bold")
+    text!(ax, "$(round(ps[1], digits=2))", position = (0.9, 0.4, 0.01), font = "Bold")
+    text!(ax, "$(round(ps[2], digits=2))", position = (0.1, 0.95, 0.01), font = "Bold")
+    text!(ax, "$(round(ps[3], digits=2))", position = (0.2, 0.1, 0.01), font = "Bold")
     fig
 end