From ac009713d67031bc977d418310c1e215f13e936a Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Wed, 11 Sep 2024 14:13:35 +0200
Subject: [PATCH 01/32] Make graph largely independent of the parameters
 structure.

---
 src/main/create_graph.jl    | 43 ++++++++---------
 src/main/init_parameters.jl | 93 ++++++++++++++++++-------------------
 2 files changed, 65 insertions(+), 71 deletions(-)

diff --git a/src/main/create_graph.jl b/src/main/create_graph.jl
index bb3ef38..192ca5b 100644
--- a/src/main/create_graph.jl
+++ b/src/main/create_graph.jl
@@ -2,38 +2,35 @@
 # using LinearAlgebra
 
 """
-    create_graph(param, w = 11, h = 11; type = "map", kwargs...) -> Dict, Dict
+    create_graph(param, w = 11, h = 11; type = "map", N = 1, kwargs...) -> Dict
 
 Initialize the underlying graph, population and productivity parameters.
 
 # Arguments
-- `param::Dict`: Dict that contains the model parameters
-- `w::Int64=11`: Number of nodes along the width of the underlying graph (integer)  
-- `h::Int64=11`: Number of nodes along the height of the underlying graph (integer, odd if triangle)
+- `param::Dict`: Dict that contains the model parameters (only needed for checks)
+- `w::Int64=11`: Number of nodes along the width of the underlying graph if type != "custom" (integer)  
+- `h::Int64=11`: Number of nodes along the height of the underlying graph if type != "custom" (integer, odd if triangle)
 
 # Keyword Arguments
 - `type::String="map"`: Either "map", "square", "triangle", or "custom" 
-- `omega::Vector{Float64}`: Vector of Pareto weights for each node or region in partial mobility case (default ones(J or nregions))
-- `Zjn::Matrix{Float64}`: J x N matrix of producties per node (j = 1:J) and good (n = 1:N) (default ones(J, N))
 - `adjacency::BitMatrix`: J x J Adjacency matrix (only used for custom network)
 - `x::Vector{Float64}`: x coordinate (longitude) of each node (only used for custom network)
 - `y::Vector{Float64}`: y coordinate (latitude) of each node (only used for custom network)
+- `omega::Vector{Float64}`: Vector of Pareto weights for each node or region in partial mobility case (default ones(J or nregions))
+- `m::Vector{Float64}=ones(N)`: Vector of weights Nx1 in the cross congestion cost function
+- `Zjn::Matrix{Float64}`: J x N matrix of producties per node (j = 1:J) and good (n = 1:N) (default ones(J, N))
 - `Lj::Vector{Float64}`: Vector of populations in each node (j = 1:J) (only for fixed labour case)
 - `Hj::Vector{Float64}`: Vector of immobile good in each node (j = 1:J) (e.g. housing, default ones(J))
 - `Lr::Vector{Float64}`: Vector of populations in each region (r = 1:nregions) (only for partial mobility)
 - `region::Vector{Int64}`: Vector indicating region of each location (only for partial mobility)
 
-# Notes
-- `create_graph()` will overwrite any parameters `Zjn`, `Lj`, `Hj`, `omega`, `Lr`, already set in the `param` Dict with the default values. So these values should be set inside `create_graph()`, or after `create_graph()` has been called.
-
 # Examples
 ```julia
-param, graph = create_graph(init_parameters())
+graph = create_graph()
 ```
 """
 function create_graph(param, w = 11, h = 11; type = "map", kwargs...)
 
-    param = namedtuple_to_dict(param)
     options = retrieve_options_create_graph(param, w, h, type; kwargs...)
 
     if type == "map"
@@ -46,26 +43,26 @@ function create_graph(param, w = 11, h = 11; type = "map", kwargs...)
         graph = create_custom(options[:adjacency], options[:x], options[:y])
     end
 
-    param[:J] = graph[:J]
-    param[:Zjn] = haskey(options, :Zjn) ? options[:Zjn] : ones(graph[:J], param[:N])
-    param[:Hj] = haskey(options, :Hj) ? options[:Hj] : ones(graph[:J])
+    graph[:N] = param[:N] # Ensure graph has all information
+    graph[:Zjn] = haskey(options, :Zjn) ? options[:Zjn] : ones(graph[:J], graph[:N])
+    graph[:Hj] = haskey(options, :Hj) ? options[:Hj] : ones(graph[:J])
 
     if param[:mobility] == false
-        param[:Lj] = haskey(options, :Lj) ? options[:Lj] : ones(graph[:J]) / graph[:J]
-        param[:hj] = param[:Hj] ./ param[:Lj]
-        param[:hj][param[:Lj] .== 0] .= 1
-        param[:omegaj] = options[:omega]
+        graph[:Lj] = haskey(options, :Lj) ? options[:Lj] : ones(graph[:J]) / graph[:J]
+        graph[:hj] = graph[:Hj] ./ graph[:Lj]
+        graph[:hj][graph[:Lj] .== 0] .= 1
+        graph[:omegaj] = options[:omega]
     elseif param[:mobility] == 0.5
         graph[:region] = options[:region]
-        param[:omegar] = options[:omega]
-        param[:Lr] = options[:Lr]
-        param[:nregions] = options[:nregions]
+        graph[:omegar] = options[:omega]
+        graph[:Lr] = options[:Lr]
+        graph[:nregions] = options[:nregions]
     end
 
     # Running checks on population / productivity / housing parameters
-    check_graph_param(param)
+    check_graph_param(graph, param)
 
-    return param, graph
+    return graph
 end
 
 function isadjacency(M)
diff --git a/src/main/init_parameters.jl b/src/main/init_parameters.jl
index 5dbd2b2..cd593ab 100644
--- a/src/main/init_parameters.jl
+++ b/src/main/init_parameters.jl
@@ -1,35 +1,30 @@
 
-# ==============================================================
-#  OPTIMAL TRANSPORT NETWORKS IN SPATIAL EQUILIBRIUM
-#  by P. Fajgelbaum, E. Schaal, D. Henricot, C. Mantovani 2017-19
-# ================================================= version 1.0.4
 
 """
     init_parameters(; kwargs...) -> Dict
 
-Returns a `param` dict with the model parameters.
+Returns a `param` dict with the model parameters. These are independent of the graph structure and dimensions.
 
 # Keyword Arguments
 - `alpha::Float64=0.5`: Cobb-Douglas coefficient on final good c^alpha * h^(1-alpha)
 - `beta::Float64=1`: Parameter governing congestion in transport cost
 - `gamma::Float64=1`: Elasticity of transport cost relative to infrastructure
-- `K::Float64=1`: Amount of concrete/asphalt
 - `sigma::Float64=5`: Elasticity of substitution across goods (CES)
 - `rho::Float64=2`: Curvature in utility (c^alpha * h^(1-alpha))^(1-rho)/(1-rho)
 - `a::Float64=0.8`: Curvature of the production function L^alpha
-- `m::Vector{Float64}=ones(N)`: Vector of weights Nx1 in the cross congestion cost function
-- `N::Int64=1`: Number of goods
 - `nu::Float64=1`: Elasticity of substitution b/w goods in transport costs if cross-good congestion
+- `N::Int64=1`: Number of goods traded in the economy (used for checks in `create_graph()`)
+- `K::Float64=1`: Amount of concrete/asphalt
 - `labor_mobility::Any=false`: Switch for labor mobility (true/false or 'partial')
 - `cross_good_congestion::Bool=false`: Switch for cross-good congestion
 - `annealing::Bool=true`: Switch for the use of annealing at the end of iterations (only if gamma > beta)
 - `verbose::Bool=true`: Switch to turn on/off text output (from Ipopt or other optimizers)
 - `duality::Bool=false`: Switch to turn on/off duality whenever available
 - `warm_start::Bool=true`: Use the previous solution as a warm start for the next iteration
-- `kappa_min::Float64=1e-5`: Minimum value for road capacities κ
+- `kappa_min::Float64=1e-5`: Minimum value for road capacities K
 - `min_iter::Int64=20`: Minimum number of iterations
 - `max_iter::Int64=200`: Maximum number of iterations
-- `tol::Float64=1e-5`: Tolerance for convergence of road capacities κ
+- `tol::Float64=1e-5`: Tolerance for convergence of road capacities K
 
 # Optional Parameters
 - `optimizer = Ipopt.Optimizer`: Optimizer to be used
@@ -43,7 +38,7 @@ Returns a `param` dict with the model parameters.
 param = init_parameters(K = 10, labor_mobility = true)
 ```
 """
-function init_parameters(; alpha = 0.5, beta = 1, gamma = 1, K = 1, sigma = 5, rho = 2, a = 0.8, N = 1, m = ones(N), nu = 1, 
+function init_parameters(; alpha = 0.5, beta = 1, gamma = 1, sigma = 5, rho = 2, a = 0.8, nu = 1, K = 1,
                          labor_mobility = false, cross_good_congestion=false, annealing=true, 
                          verbose = true, duality = false, warm_start = true, 
                          kappa_min = 1e-5, min_iter = 20, max_iter = 200, tol = 1e-5, kwargs...)
@@ -65,8 +60,6 @@ function init_parameters(; alpha = 0.5, beta = 1, gamma = 1, K = 1, sigma = 5, r
     param[:sigma] = sigma
     param[:rho] = rho
     param[:a] = a
-    param[:m] = m
-    param[:N] = N
     param[:nu] = nu
     if labor_mobility === "partial" || labor_mobility === 0.5
         param[:mobility] = 0.5
@@ -98,57 +91,61 @@ function init_parameters(; alpha = 0.5, beta = 1, gamma = 1, K = 1, sigma = 5, r
     param[:F] = (L, a) -> L^a
     param[:Fprime] = (L, a) -> a * L^(a - 1)
 
-    # CHECK CONSISTENCY WITH ENDOWMENTS/PRODUCTIVITY (if applicable)
-    check_graph_param(param)
+    # # CHECK CONSISTENCY WITH ENDOWMENTS/PRODUCTIVITY (if applicable)
+    # check_graph_param(param)
 
     return param
 end
 
-function check_graph_param(param)
-
-    ## These are graph parameters that should not be in the parameters dict
-
-    # if haskey(param, :x) && length(param[:x]) != param[:J]
-    #     @warn "x does not have the right length J = $(param[:J])."
-    # end
-
-    # if haskey(param, :y) && length(param[:y]) != param[:J]
-    #     @warn "y does not have the right length J = $(param[:J])."
-    # end
-
-    # if haskey(param, :adjacency) && size(param[:adjacency]) != (param[:J], param[:J])
-    #     @warn "adjacency matrix does not have the right dimensions J x J."
-    # end
-
-    # if haskey(param, :region) && length(param[:region]) != param[:J]
-    #     @warn "region does not have the right length J = $(param[:J])."
-    # end
+# Check if the parameters are consistent with the graph structure
+function check_graph_param(graph, param)
 
-    if haskey(param, :omegaj) && length(param[:omegaj]) != param[:J]
-        @warn "omegaj does not have the right length J = $(param[:J])."
+    if haskey(param, :omegaj) && !haskey(graph, :omegaj)
+        graph[:omegaj] = param[:omegaj]
+    end
+    if haskey(graph, :omegaj) && length(graph[:omegaj]) != graph[:J]
+        @warn "omegaj does not have the right length J = $(graph[:J])."
     end
 
-    if haskey(param, :omegar) && length(param[:omegar]) != param[:nregions]
-        @warn "omegar does not have the right length nregions = $(param[:nregions])."
+    if haskey(param, :omegar) && !haskey(graph, :omegar)
+        graph[:omegar] = param[:omegar]
+    end
+    if haskey(graph, :omegar) && length(graph[:omegar]) != graph[:nregions]
+        @warn "omegar does not have the right length nregions = $(graph[:nregions])."
     end
 
-    if haskey(param, :Lr) && length(param[:Lr]) != param[:nregions]
-        @warn "Lr does not have the right length nregions = $(param[:nregions])."
+    if haskey(param, :Lr) && !haskey(graph, :Lr)
+        graph[:Lr] = param[:Lr]
+    end
+    if haskey(graph, :Lr) && length(graph[:Lr]) != graph[:nregions]
+        @warn "Lr does not have the right length nregions = $(graph[:nregions])."
     end
 
-    if haskey(param, :Lj) && param[:mobility] == 0 && length(param[:Lj]) != param[:J]
-        @warn "Lj does not have the right length J = $(param[:J])."
+    if haskey(param, :Lj) && !haskey(graph, :Lj) && param[:mobility] == 0
+        graph[:Lj] = param[:Lj]
+    end
+    if haskey(graph, :Lj) && param[:mobility] == 0 && length(graph[:Lj]) != graph[:J]
+        @warn "Lj does not have the right length J = $(graph[:J])."
     end
     
-    if haskey(param, :Zjn) && size(param[:Zjn]) != (param[:J], param[:N])
-        @warn "Zjn does not have the right size J ($(param[:J])) x N ($(param[:N]))."
+    if haskey(param, :Zjn) && !haskey(graph, :Zjn)
+        graph[:Zjn] = param[:Zjn]
+    end
+    if haskey(graph, :Zjn) && size(graph[:Zjn]) != (graph[:J], graph[:N])
+        @warn "Zjn does not have the right size J ($(graph[:J])) x N ($(graph[:N]))."
     end
 
-    if haskey(param, :Hj) && length(param[:Hj]) != param[:J]
-        @warn "Hj does not have the right length J = $(param[:J])."
+    if haskey(param, :Hj) && !haskey(graph, :Hj)
+        graph[:Hj] = param[:Hj]
+    end
+    if haskey(graph, :Hj) && length(graph[:Hj]) != graph[:J]
+        @warn "Hj does not have the right length J = $(graph[:J])."
     end
 
-    if haskey(param, :hj) && length(param[:hj]) != param[:J]
-        @warn "hj does not have the right length J = $(param[:J])."
+    if haskey(param, :hj) && !haskey(graph, :hj)
+        graph[:hj] = param[:hj]
+    end
+    if haskey(graph, :hj) && length(graph[:hj]) != graph[:J]
+        @warn "hj does not have the right length J = $(graph[:J])."
     end
 end
\ No newline at end of file

From ee7f7207e124dca4341ea88a3722198a6de34b1c Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Wed, 11 Sep 2024 15:00:12 +0200
Subject: [PATCH 02/32] Update to reflect changes.

---
 src/main/nodes.jl | 218 ++++++++++++++++++++--------------------------
 1 file changed, 95 insertions(+), 123 deletions(-)

diff --git a/src/main/nodes.jl b/src/main/nodes.jl
index 4bf5882..22978b4 100644
--- a/src/main/nodes.jl
+++ b/src/main/nodes.jl
@@ -1,113 +1,103 @@
 
 """
-    add_node(param, graph, x, y, neighbors) -> Dict, Dict
+    add_node(graph, x, y, neighbors) -> Dict
 
-Add a node in position (x,y) and list of neighbors. The new node is given an index J+1.
+Add a node in position (x,y) and list of neighbors. The new node is given an index J+1. Returns an updated `graph` object.
 
 # Arguments
-- `param::Dict`: Dict that contains the model's parameters, or `nothing` to only update graph
 - `graph::Dict`: Dict that contains the underlying graph (created by create_graph())
 - `x::Float64`: x coordinate of the new node (any real number)
 - `y::Float64`: y coordinate of the new node (any real number)
 - `neighbors::Vector{Int64}`: Vector of nodes to which it is connected (1 x n list of node indices between 1 and J, where n is an arbitrary # of neighbors) 
 
-The cost matrices `delta_tau` and `delta_i` are parametrized as a function of Euclidean distance between nodes.
-
-Returns the updated `graph` and `param` objects (`param` is affected too because the variable `Zjn`, `Lj`, `Hj` and others are reset to a uniform dist.)
+# Notes
+The cost matrices `delta_tau` and `delta_i` are parametrized as a function of Euclidean distance between nodes. 
+The new node is given population 1e-6 and productivity equal to the minimum productivity in the graph.
 """
-function add_node(param, graph, x, y, neighbors)
+function add_node(graph, x, y, neighbors)
 
-    graph = dict_to_namedtuple(graph)
+    if graph isa Dict
+        graph_new = copy(dict)
+    else
+        graph_new = Dict(pairs(namedtuple))
+    end
 
     # Check validity of neighbors list
-    if any(neighbors .!= floor.(neighbors)) || any(neighbors .< 1) || any(neighbors .> graph.J)
-        error("neighbors should be a list of integers between 1 and $(graph.J).")
+    if any(neighbors .!= floor.(neighbors)) || any(neighbors .< 1) || any(neighbors .> graph[:J])
+        error("neighbors should be a list of integers between 1 and $(graph[:J]).")
     end
-    Jnew = graph.J + 1
+    Jnew = graph[:J] + 1
 
     # Add node
-    nodes = [graph.nodes; [neighbors]]
-    new_x = [graph.x; x]
-    new_y = [graph.y; y]
+    nodes = [graph[:nodes]; [neighbors]]
+    new_x = [graph[:x]; x]
+    new_y = [graph[:y]; y]
 
     # Add new node to neighbors' neighbors
     # And update adjacency and cost matrices
     adjacency = zeros(Jnew, Jnew)
-    adjacency[1:graph.J, 1:graph.J] = graph.adjacency
+    adjacency[1:graph[:J], 1:graph[:J]] = graph[:adjacency]
 
     delta_tau = zeros(Jnew, Jnew)
-    delta_tau[1:graph.J, 1:graph.J] = graph.delta_tau
+    delta_tau[1:graph[:J], 1:graph[:J]] = graph[:delta_tau]
 
     delta_i = zeros(Jnew, Jnew)
-    delta_i[1:graph.J, 1:graph.J] = graph.delta_i
+    delta_i[1:graph[:J], 1:graph[:J]] = graph[:delta_i]
 
     for i in neighbors
         nodes[i] = [nodes[i]; Jnew]
         distance = sqrt((new_x[i] - x)^2 + (new_y[i] - y)^2)
 
         # Adjacency
-        adjacency[i, Jnew] = 1
-        adjacency[Jnew, i] = 1
+        adjacency[i, Jnew] = adjacency[Jnew, i] = 1
 
         # Travel cost: delta_tau
-        delta_tau[i, Jnew] = distance
-        delta_tau[Jnew, i] = distance
+        delta_tau[i, Jnew] = delta_tau[Jnew, i] = distance
 
         # Building cost: delta_i
-        delta_i[i, Jnew] = distance
-        delta_i[Jnew, i] = distance
+        delta_i[i, Jnew] = delta_i[Jnew, i] = distance
     end
 
     # Update number of degrees of liberty for Ijk
     ndeg = sum(tril(adjacency)) # nb of degrees of liberty in adjacency matrix
 
-    # Return new graph structure
-    graph_new = Dict(
-        :J => Jnew,
-        :x => new_x,
-        :y => new_y,
-        :nodes => nodes,
-        :adjacency => adjacency,
-        :delta_i => delta_i,
-        :delta_tau => delta_tau,
-        :ndeg => ndeg
-    )
-
-    if param !== nothing
-        # Now, update the param structure
-        if param isa Dict
-            param_new = copy(param)
-        else
-            param_new = Dict(pairs(param))
-        end
-        if haskey(param, :J)
-            param_new[:J] = Jnew
-        end
-        if haskey(param, :Lj)
-            param_new[:Lj] = ones(Jnew) / Jnew
-        end
-        if haskey(param, :Hj)
-            param_new[:Hj] = ones(Jnew)
-        end
-        if haskey(param, :hj)
-            param_new[:hj] = param[:Hj] ./ param[:Lj]
-        end
-        if haskey(param, :omegaj)
-            param_new[:omegaj] = ones(Jnew)
-        end
-        if haskey(param, :Zjn)
-            param_new[:Zjn] = ones(Jnew, param[:N])
-        end
-
-        return param_new, graph_new
+    # Update graph structure
+    graph_new[:J] = Jnew
+    graph_new[:x] = new_x
+    graph_new[:y] = new_y
+    graph_new[:nodes] = nodes
+    graph_new[:adjacency] = adjacency
+    graph_new[:delta_i] = delta_i
+    graph_new[:delta_tau] = delta_tau
+    graph_new[:ndeg] = ndeg
+
+    # Update other parameters if they exist
+    if haskey(graph, :Lj)
+        graph_new[:Lj] = [graph[:Lj]; 1e-6] # ones(Jnew) / Jnew
+    end
+    if haskey(graph, :Hj)
+        graph_new[:Hj] = [graph[:Hj]; 1e-6] # ones(Jnew)
+    end
+    if haskey(graph, :hj)
+        graph_new[:hj] = graph_new[:Hj] ./ graph_new[:Lj]
+    end
+    if haskey(graph, :omegaj)
+        graph_new[:omegaj] = [graph_new[:omegaj]; 1e-6] # ones(Jnew)
+    end
+    if haskey(graph, :Zjn)
+        Zjn = fill(minimum(graph[:Zjn]), Jnew, size(graph[:Zjn], 2))
+        Zjn[1:graph[:J], :] = graph[:Zjn]
+        graph_new[:Zjn] = Zjn # ones(Jnew, size(graph[:Zjn], 2))
+    end
+    if haskey(graph, :region)
+        closest_node = find_node(graph, x, y)
+        closest_region = graph[:region][closest_node]
+        graph_new[:region] = [graph[:region]; closest_region]
     end
 
     return graph_new
 end
 
-
-# Please note that in Julia, we use `Dict` to represent structures as in Matlab. The keys of the `Dict` are strings that correspond to the field names in the Matlab structure. Also, the `tril` function is not built-in in Julia, you may need to use a package like `LinearAlgebra` to use it.
-
 """
     find_node(graph, x, y) -> Int64
 
@@ -124,97 +114,79 @@ function find_node(graph, x, y)
     return id
 end
 
-
-
 """
-    remove_node(param, graph, i) -> updated_param, updated_graph
+    remove_node(graph, i) -> Dict
 
-Removes node i from the graph.
+Removes node i from the graph, returning an updated `graph` object.
 
 # Arguments
-- `param::Dict`: Dict that contains the model's parameters, or `nothing` to only update graph
 - `graph::Dict`: Dict that contains the underlying graph (created by create_graph())
-- `i::Int64`: index of the mode to be removed (integer between 1 and graph.J)
-
-Returns the updated graph and param objects (param is affected too because the variable Zjn, Lj, Hj and others are changed).
+- `i::Int64`: index of the mode to be removed (integer between 1 and graph[:J])
 """
-function remove_node(param, graph, i)
+function remove_node(graph, i)
 
     if graph isa Dict
         graph_new = copy(dict)
     else
         graph_new = Dict(pairs(namedtuple))
     end
-    graph = dict_to_namedtuple(graph)
 
-    if i < 1 || i > graph.J || i != floor(i)
-        error("remove_node: node i should be an integer between 1 and $(graph.J).")
+    if i < 1 || i > graph[:J] || i != floor(i)
+        error("remove_node: node i should be an integer between 1 and $(graph[:J]).")
     end
 
-    Jnew = graph.J - 1
+    Jnew = graph[:J] - 1
 
     graph_new[:J] = Jnew
-    graph_new[:x] = deleteat!(copy(graph.x), i)
-    graph_new[:y] = deleteat!(copy(graph.y), i)
-    graph_new[:nodes] = deleteat!(copy(graph.nodes), i)
-
-    nodes = graph_new[:nodes]
+    graph_new[:x] = deleteat!(copy(graph[:x]), i)
+    graph_new[:y] = deleteat!(copy(graph[:y]), i)
+    nodes = deleteat!(copy(graph[:nodes]), i)
     for k in 1:Jnew
-        node_k = filter(x -> x != i, nodes[k]) # nodes[k][nodes[k] .!= i]
+        node_k = filter(x -> x != i, nodes[k])
         nodes[k] = ifelse.(node_k .> i, node_k .- 1, node_k) # reindex nodes k > i to k-1
     end
+    graph_new[:nodes] = nodes  # Assign the modified nodes to graph_new
 
     # Rebuild adjacency matrix, delta_i and delta_tau
-    graph_new[:adjacency] = [graph.adjacency[1:i-1, 1:i-1] graph.adjacency[1:i-1, i+1:end];
-                             graph.adjacency[i+1:end, 1:i-1] graph.adjacency[i+1:end, i+1:end]]
+    graph_new[:adjacency] = [graph[:adjacency][1:i-1, 1:i-1] graph[:adjacency][1:i-1, i+1:end];
+                             graph[:adjacency][i+1:end, 1:i-1] graph[:adjacency][i+1:end, i+1:end]]
 
-    graph_new[:delta_i] = [graph.delta_i[1:i-1, 1:i-1] graph.delta_i[1:i-1, i+1:end];
-                           graph.delta_i[i+1:end, 1:i-1] graph.delta_i[i+1:end, i+1:end]]
+    graph_new[:delta_i] = [graph[:delta_i][1:i-1, 1:i-1] graph[:delta_i][1:i-1, i+1:end];
+                           graph[:delta_i][i+1:end, 1:i-1] graph[:delta_i][i+1:end, i+1:end]]
 
-    graph_new[:delta_tau] = [graph.delta_tau[1:i-1, 1:i-1] graph.delta_tau[1:i-1, i+1:end];
-                             graph.delta_tau[i+1:end, 1:i-1] graph.delta_tau[i+1:end, i+1:end]]
+    graph_new[:delta_tau] = [graph[:delta_tau][1:i-1, 1:i-1] graph[:delta_tau][1:i-1, i+1:end];
+                             graph[:delta_tau][i+1:end, 1:i-1] graph[:delta_tau][i+1:end, i+1:end]]
 
     if haskey(graph, :across_obstacle)
-        graph_new[:across_obstacle] = [graph.across_obstacle[1:i-1, 1:i-1] graph.across_obstacle[1:i-1, i+1:end];
-                                       graph.across_obstacle[i+1:end, 1:i-1] graph.across_obstacle[i+1:end, i+1:end]]
+        graph_new[:across_obstacle] = [graph[:across_obstacle][1:i-1, 1:i-1] graph[:across_obstacle][1:i-1, i+1:end];
+                                       graph[:across_obstacle][i+1:end, 1:i-1] graph[:across_obstacle][i+1:end, i+1:end]]
     end
 
     if haskey(graph, :along_obstacle)
-        graph_new[:along_obstacle] = [graph.along_obstacle[1:i-1, 1:i-1] graph.along_obstacle[1:i-1, i+1:end];
-                                      graph.along_obstacle[i+1:end, 1:i-1] graph.along_obstacle[i+1:end, i+1:end]]
+        graph_new[:along_obstacle] = [graph[:along_obstacle][1:i-1, 1:i-1] graph[:along_obstacle][1:i-1, i+1:end];
+                                      graph[:along_obstacle][i+1:end, 1:i-1] graph[:along_obstacle][i+1:end, i+1:end]]
     end
 
     graph_new[:ndeg] = sum(tril(graph_new[:adjacency]))
 
-    if graph.region !== nothing
-        graph_new[:region] = deleteat!(copy(graph.region), i)
-    end
-
-    if param !== nothing 
-        # Now, update the param structure
-        if param isa Dict
-            param_new = copy(param)
-        else
-            param_new = Dict(pairs(param))
-        end
-        param_new[:J] = Jnew
-        if haskey(param, :Lj)
-            param_new[:Lj] = deleteat!(copy(param[:Lj]), i)
-        end
-        if haskey(param, :Hj)
-            param_new[:Hj] = deleteat!(copy(param[:Hj]), i)
-        end
-        if haskey(param, :hj)
-            param_new[:hj] = deleteat!(copy(param[:hj]), i)
-        end
-        if haskey(param, :omegaj)
-            param_new[:omegaj] = deleteat!(copy(param[:omegaj]), i)
-        end
-        if haskey(param, :Zjn)
-            param_new[:Zjn] = vcat(param[:Zjn][1:i-1, :], param[:Zjn][i+1:end, :])
-        end
-
-        return param_new, graph_new
+    # Update other parameters if they exist
+    if haskey(graph, :Lj)
+        graph_new[:Lj] = deleteat!(copy(graph[:Lj]), i)
+    end
+    if haskey(graph, :Hj)
+        graph_new[:Hj] = deleteat!(copy(graph[:Hj]), i)
+    end
+    if haskey(graph, :hj)
+        graph_new[:hj] = deleteat!(copy(graph[:hj]), i)
+    end
+    if haskey(graph, :omegaj)
+        graph_new[:omegaj] = deleteat!(copy(graph[:omegaj]), i)
+    end
+    if haskey(graph, :Zjn)
+        graph_new[:Zjn] = vcat(graph[:Zjn][1:i-1, :], graph[:Zjn][i+1:end, :])
+    end
+    if haskey(graph, :region)
+        graph_new[:region] = deleteat!(copy(graph[:region]), i)
     end
 
     return graph_new

From 9362c246ed8935d68ab2045c90aa8fdfb28c14d9 Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Wed, 11 Sep 2024 15:54:34 +0200
Subject: [PATCH 03/32] Update helper files.

---
 src/main/helper.jl | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/src/main/helper.jl b/src/main/helper.jl
index e63e40c..8cfd691 100644
--- a/src/main/helper.jl
+++ b/src/main/helper.jl
@@ -164,16 +164,16 @@ end
 Construct the appropriate model based on the parameters and auxiliary data.
 
 # Arguments
-- `param`: A named tuple containing the model parameters.
 - `auxdata`: Auxiliary data required for constructing the model.
 
 # Returns
 - `model`: The constructed model.
 - `recover_allocation`: A function to recover the allocation from the model solution.
 """
-function get_model(param, auxdata)
+function get_model(auxdata)
+     param = auxdata.param
+     graph = auxdata.graph
      optimizer = get(param, :optimizer, Ipopt.Optimizer)
-     param = dict_to_namedtuple(param)
 
     if haskey(param, :model)
         model = param.model(optimizer, auxdata)
@@ -182,7 +182,7 @@ function get_model(param, auxdata)
         end
         recover_allocation = param.recover_allocation 
     elseif param.mobility == 1 && param.cong
-        if all(sum(param.Zjn .> 0, dims = 2) .<= 1) # Armington case
+        if all(sum(graph.Zjn .> 0, dims = 2) .<= 1) # Armington case
             model = model_mobility_cgc_armington(optimizer, auxdata)
             recover_allocation = recover_allocation_mobility_cgc_armington
         else
@@ -190,7 +190,7 @@ function get_model(param, auxdata)
             recover_allocation = recover_allocation_mobility_cgc
         end
     elseif param.mobility == 0.5 && param.cong
-        if all(sum(param.Zjn .> 0, dims = 2) .<= 1) # Armington case
+        if all(sum(graph.Zjn .> 0, dims = 2) .<= 1) # Armington case
             model = model_partial_mobility_cgc_armington(optimizer, auxdata)
             recover_allocation = recover_allocation_partial_mobility_cgc_armington
         else
@@ -198,7 +198,7 @@ function get_model(param, auxdata)
             recover_allocation = recover_allocation_partial_mobility_cgc
         end
     elseif param.mobility == 0 && param.cong
-        if all(sum(param.Zjn .> 0, dims = 2) .<= 1) # Armington case
+        if all(sum(graph.Zjn .> 0, dims = 2) .<= 1) # Armington case
             model = model_fixed_cgc_armington(optimizer, auxdata)
             recover_allocation = recover_allocation_fixed_cgc_armington
         else
@@ -206,7 +206,7 @@ function get_model(param, auxdata)
             recover_allocation = recover_allocation_fixed_cgc
         end
     elseif param.mobility == 1 && !param.cong
-        if all(sum(param.Zjn .> 0, dims = 2) .<= 1) # Armington case
+        if all(sum(graph.Zjn .> 0, dims = 2) .<= 1) # Armington case
             model = model_mobility_armington(optimizer, auxdata)
             recover_allocation = recover_allocation_mobility_armington
         else
@@ -214,7 +214,7 @@ function get_model(param, auxdata)
             recover_allocation = recover_allocation_mobility
         end
     elseif param.mobility == 0.5 && !param.cong
-        if all(sum(param.Zjn .> 0, dims = 2) .<= 1) # Armington case
+        if all(sum(graph.Zjn .> 0, dims = 2) .<= 1) # Armington case
             model = model_partial_mobility_armington(optimizer, auxdata)
             recover_allocation = recover_allocation_partial_mobility_armington    
         else
@@ -222,7 +222,7 @@ function get_model(param, auxdata)
             recover_allocation = recover_allocation_partial_mobility    
         end
     elseif param.mobility == 0 && !param.cong
-        if all(sum(param.Zjn .> 0, dims = 2) .<= 1) # Armington case
+        if all(sum(graph.Zjn .> 0, dims = 2) .<= 1) # Armington case
             model = model_fixed_armington(optimizer, auxdata)
             recover_allocation = recover_allocation_fixed_armington
         elseif param.beta <= 1 && param.a < 1 && param.duality

From c0b67cc133c478a96ed36e02cfc5489ba199c25a Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Wed, 11 Sep 2024 16:09:04 +0200
Subject: [PATCH 04/32] Delete nodes if applicable.

---
 src/main/apply_geography.jl | 78 +++++++++++++++++++++++++++----------
 1 file changed, 57 insertions(+), 21 deletions(-)

diff --git a/src/main/apply_geography.jl b/src/main/apply_geography.jl
index d57496f..6f943fa 100644
--- a/src/main/apply_geography.jl
+++ b/src/main/apply_geography.jl
@@ -1,6 +1,6 @@
 
 """
-    apply_geography(graph, geography; kwargs...) -> updated_graph
+    apply_geography(graph, geography; kwargs...) -> Dict
 
 Update the network building costs of a graph based on geographical features and remove edges impeded by geographical barriers. 
 Aversion to altitude changes rescales building infrastructure costs `delta_i` by (see also user manual to MATLAB toolbox):\n
@@ -90,6 +90,7 @@ function apply_geography(graph, geography; kwargs...)
     end
 
     # Remove edges where geographical barriers are (rivers)
+    remove_edge = false
     if obstacles !== nothing
 
         # Store initial delta matrics (avoid double counting)
@@ -101,6 +102,9 @@ function apply_geography(graph, geography; kwargs...)
         across_obstacle = falses(graph.J, graph.J)
         along_obstacle = falses(graph.J, graph.J)
         remove_edge = isinf(op.across_obstacle_delta_i) || isinf(op.across_obstacle_delta_tau) # remove the edge
+        if remove_edge
+            rm_nodes = falses(graph.J)
+        end
 
         for i in 1:graph.J
             neighbors = graph.nodes[i]
@@ -131,16 +135,16 @@ function apply_geography(graph, geography; kwargs...)
                                 rmi = findfirst(==(i), nodes_new[j])
                                 if rmi !== nothing
                                     deleteat!(nodes_new[j], rmi)
-                                    # if isempty(nodes_new[j])
-                                    #     deleteat!(nodes_new, j)
-                                    # end
+                                    if isempty(nodes_new[j])
+                                        rm_nodes[j] = true # deleteat!(nodes_new, j)
+                                    end
                                 end
                                 rmj = findfirst(==(j), nodes_new[i])
                                 if rmj !== nothing
                                     deleteat!(nodes_new[i], rmj)
-                                    # if isempty(nodes_new[i])
-                                    #     deleteat!(nodes_new, i)
-                                    # end
+                                    if isempty(nodes_new[i])
+                                        rm_nodes[i] = true # deleteat!(nodes_new, i)
+                                    end
                                 end
                                 
                                 adjacency_new[i, j] = false
@@ -165,20 +169,21 @@ function apply_geography(graph, geography; kwargs...)
         # This is for edges along obstable: allowing different delta (e.g., water transport on river)
         if isinf(op.along_obstacle_delta_i) || isinf(op.along_obstacle_delta_tau)
             # if infinite, remove edge
+            remove_edge = true
             for (io, jo) in zip(obstacles[:, 1], obstacles[:, 2])
                 rmio = findfirst(==(io), nodes_new[jo])
                 if rmio !== nothing
                     deleteat!(nodes_new[jo], rmio)
-                    # if isempty(nodes_new[jo])
-                    #     deleteat!(nodes_new, jo)
-                    # end
+                    if isempty(nodes_new[jo])
+                        rm_nodes[jo] = true # deleteat!(nodes_new, jo)
+                    end
                 end
                 rmjo = findfirst(==(jo), nodes_new[io])
                 if rmjo !== nothing
                     deleteat!(nodes_new[io], rmjo)
-                    # if isempty(nodes_new[io])
-                    #     deleteat!(nodes_new, io)
-                    # end
+                    if isempty(nodes_new[io])
+                        rm_nodes[io] = true # deleteat!(nodes_new, io)
+                    end
                 end
                 adjacency_new[io, jo] = false
                 adjacency_new[jo, io] = false
@@ -200,16 +205,47 @@ function apply_geography(graph, geography; kwargs...)
 
     # Creating new object
     graph_new = namedtuple_to_dict(graph)
-    graph_new[:delta_i] = delta_i_new
-    graph_new[:delta_tau] = delta_tau_new
 
-    if obstacles !== nothing
-        graph_new[:nodes] = nodes_new
-        graph_new[:adjacency] = adjacency_new
+    if remove_edge && any(rm_nodes)
+        keep = findall(.!rm_nodes)
+        graph_new[:delta_i] = delta_i_new[keep, keep]
+        graph_new[:delta_tau] = delta_tau_new[keep, keep]
+        graph_new[:nodes] = nodes_new[keep]
+        graph_new[:adjacency] = adjacency_new[keep, keep]
         # make sure that the degrees of freedom of the updated graph match the # of links
-        graph_new[:ndeg] = sum(tril(adjacency_new))
-        graph_new[:across_obstacle] = across_obstacle
-        graph_new[:along_obstacle] = along_obstacle
+        graph_new[:ndeg] = sum(tril(graph_new[:adjacency]))
+        graph_new[:across_obstacle] = across_obstacle[keep, keep]
+        graph_new[:along_obstacle] = along_obstacle[keep, keep]
+
+        if haskey(graph, :Lj)
+            graph_new[:Lj] = graph[:Lj][keep]
+        end
+        if haskey(graph, :Hj)
+            graph_new[:Hj] = graph[:Hj][keep]
+        end
+        if haskey(graph, :hj)
+            graph_new[:hj] = graph_new[:Hj] ./ graph_new[:Lj]
+        end
+        if haskey(graph, :omegaj)
+            graph_new[:omegaj] = graph[:omegaj][keep]
+        end
+        if haskey(graph, :Zjn)
+            graph_new[:Zjn] = graph[:Zjn][keep, :]
+        end
+        if haskey(graph, :region)
+            graph_new[:region] = graph[:region][keep]
+        end
+    else
+        graph_new[:delta_i] = delta_i_new
+        graph_new[:delta_tau] = delta_tau_new
+        if obstacles !== nothing
+            graph_new[:nodes] = nodes_new
+            graph_new[:adjacency] = adjacency_new
+            # make sure that the degrees of freedom of the updated graph match the # of links
+            graph_new[:ndeg] = sum(tril(adjacency_new))
+            graph_new[:across_obstacle] = across_obstacle
+            graph_new[:along_obstacle] = along_obstacle
+        end
     end
 
     return graph_new

From df86efbff672dc5049e6746c8367cc831d95d72a Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Wed, 11 Sep 2024 16:18:40 +0200
Subject: [PATCH 05/32] Also adjust for new structure.

---
 src/main/optimal_network.jl | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/src/main/optimal_network.jl b/src/main/optimal_network.jl
index bcef47a..ce547ab 100644
--- a/src/main/optimal_network.jl
+++ b/src/main/optimal_network.jl
@@ -22,8 +22,8 @@ Solve for the optimal network by solving the inner problem and the outer problem
 # Examples
 ```julia
 param = init_parameters(K = 10)
-param, graph = create_graph(param)
-param[:Zjn][61] = 10.0
+graph = create_graph(param)
+graph[:Zjn][61] = 10.0
 results = optimal_network(param, graph)
 plot_graph(graph, results[:Ijk])
 ```
@@ -32,8 +32,8 @@ function optimal_network(param, graph; I0=nothing, Il=nothing, Iu=nothing, verbo
 
     # I0=nothing; Il=nothing; Iu=nothing; verbose=false; return_model = false; return_model = 0;
     # Check param.Zjn and make matrix if vector:
-    if length(size(param[:Zjn])) == 1
-        param[:Zjn] = reshape(param[:Zjn], param[:J], 1)
+    if length(size(graph[:Zjn])) == 1
+        graph[:Zjn] = reshape(graph[:Zjn], graph[:J], 1)
     end   
     graph = dict_to_namedtuple(graph)
     param = dict_to_namedtuple(param)
@@ -67,7 +67,7 @@ function optimal_network(param, graph; I0=nothing, Il=nothing, Iu=nothing, verbo
     # INITIALIZATION
 
     auxdata = create_auxdata(param, graph, edges, I0)
-    model, recover_allocation = get_model(param, auxdata)
+    model, recover_allocation = get_model(auxdata)
 
     if return_model == 1
         return model, recover_allocation
@@ -85,7 +85,7 @@ function optimal_network(param, graph; I0=nothing, Il=nothing, Iu=nothing, verbo
     weight_old = 0.5
     I1 = zeros(J, J)
     # K in average infrastructure units
-    K_infra = (param.K / mean(graph.delta_i[graph.adjacency.==1]))
+    K_infra = (param.K / mean(graph.delta_i[graph.adjacency .== 1]))
     distance = 0.0
     all_vars = all_variables_except_kappa_ex(model)
     start_values = start_value.(all_vars)
@@ -121,7 +121,7 @@ function optimal_network(param, graph; I0=nothing, Il=nothing, Iu=nothing, verbo
             else
                 error("Solver returned with error code $(termination_status(model)).")
             end
-        elseif param[:warm_start] # Set next run starting values to previous run solution.
+        elseif param.warm_start # Set next run starting values to previous run solution.
             vars_solution = value.(all_vars)
             set_start_value.(all_vars, vars_solution)
             used_warm_start = true
@@ -133,7 +133,7 @@ function optimal_network(param, graph; I0=nothing, Il=nothing, Iu=nothing, verbo
             PQ = dropdims(sum(PQ + permutedims(PQ, [2, 1, 3]), dims=3), dims = 3)
         else
             PQ = repeat(results[:PCj], 1, J)
-            matm = permutedims(repeat(param[:m], 1, J, J), [3, 2, 1])
+            matm = permutedims(repeat(graph.m, 1, J, J), [3, 2, 1])
             cost = dropdims(sum(matm .* results[:Qjkn] .^ param.nu, dims=3), dims = 3) .^ ((param.beta + 1) / param.nu)
             PQ .*= cost
             PQ += PQ'

From 838bc316aded98f5a2a1f468690dda93aa2c55fd Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Wed, 11 Sep 2024 16:37:34 +0200
Subject: [PATCH 06/32] Minor reorganization + bring back 'm'.

---
 src/main/init_parameters.jl | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/src/main/init_parameters.jl b/src/main/init_parameters.jl
index cd593ab..c270bc0 100644
--- a/src/main/init_parameters.jl
+++ b/src/main/init_parameters.jl
@@ -9,14 +9,15 @@ Returns a `param` dict with the model parameters. These are independent of the g
 - `alpha::Float64=0.5`: Cobb-Douglas coefficient on final good c^alpha * h^(1-alpha)
 - `beta::Float64=1`: Parameter governing congestion in transport cost
 - `gamma::Float64=1`: Elasticity of transport cost relative to infrastructure
+- `K::Float64=1`: Amount of concrete/asphalt (infrastructure budget)
 - `sigma::Float64=5`: Elasticity of substitution across goods (CES)
 - `rho::Float64=2`: Curvature in utility (c^alpha * h^(1-alpha))^(1-rho)/(1-rho)
 - `a::Float64=0.8`: Curvature of the production function L^alpha
-- `nu::Float64=1`: Elasticity of substitution b/w goods in transport costs if cross-good congestion
 - `N::Int64=1`: Number of goods traded in the economy (used for checks in `create_graph()`)
-- `K::Float64=1`: Amount of concrete/asphalt
 - `labor_mobility::Any=false`: Switch for labor mobility (true/false or 'partial')
 - `cross_good_congestion::Bool=false`: Switch for cross-good congestion
+- `nu::Float64=1`: Elasticity of substitution b/w goods in transport costs if cross-good congestion
+- `m::Vector{Float64}=ones(N)`: Vector of weights Nx1 in the cross congestion cost function
 - `annealing::Bool=true`: Switch for the use of annealing at the end of iterations (only if gamma > beta)
 - `verbose::Bool=true`: Switch to turn on/off text output (from Ipopt or other optimizers)
 - `duality::Bool=false`: Switch to turn on/off duality whenever available
@@ -38,8 +39,9 @@ Returns a `param` dict with the model parameters. These are independent of the g
 param = init_parameters(K = 10, labor_mobility = true)
 ```
 """
-function init_parameters(; alpha = 0.5, beta = 1, gamma = 1, sigma = 5, rho = 2, a = 0.8, nu = 1, K = 1,
-                         labor_mobility = false, cross_good_congestion=false, annealing=true, 
+function init_parameters(; alpha = 0.5, beta = 1, gamma = 1, K = 1, sigma = 5, rho = 2, a = 0.8, N = 1, 
+                         labor_mobility = false, cross_good_congestion=false, nu = 1, m = ones(N), 
+                         annealing=true, 
                          verbose = true, duality = false, warm_start = true, 
                          kappa_min = 1e-5, min_iter = 20, max_iter = 200, tol = 1e-5, kwargs...)
     param = Dict()
@@ -61,6 +63,11 @@ function init_parameters(; alpha = 0.5, beta = 1, gamma = 1, sigma = 5, rho = 2,
     param[:rho] = rho
     param[:a] = a
     param[:nu] = nu
+    if length(m) != N
+        error("m must have length N = $N, but has length $(length(m))")
+    end
+    param[:N] = N
+    param[:m] = m
     if labor_mobility === "partial" || labor_mobility === 0.5
         param[:mobility] = 0.5
     else
@@ -100,6 +107,10 @@ end
 # Check if the parameters are consistent with the graph structure
 function check_graph_param(graph, param)
 
+    if haskey(param, :m) && length(param[:m]) != graph[:N]
+        @warn "m does not have the right length N = $(graph[:N])."
+    end
+
     if haskey(param, :omegaj) && !haskey(graph, :omegaj)
         graph[:omegaj] = param[:omegaj]
     end

From 0f6fdfdc3d63457a865c5eeef5e10fa900940ef9 Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Wed, 11 Sep 2024 16:38:57 +0200
Subject: [PATCH 07/32] Remove 'm' here.

---
 src/main/create_graph.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/main/create_graph.jl b/src/main/create_graph.jl
index 192ca5b..779ea24 100644
--- a/src/main/create_graph.jl
+++ b/src/main/create_graph.jl
@@ -17,7 +17,6 @@ Initialize the underlying graph, population and productivity parameters.
 - `x::Vector{Float64}`: x coordinate (longitude) of each node (only used for custom network)
 - `y::Vector{Float64}`: y coordinate (latitude) of each node (only used for custom network)
 - `omega::Vector{Float64}`: Vector of Pareto weights for each node or region in partial mobility case (default ones(J or nregions))
-- `m::Vector{Float64}=ones(N)`: Vector of weights Nx1 in the cross congestion cost function
 - `Zjn::Matrix{Float64}`: J x N matrix of producties per node (j = 1:J) and good (n = 1:N) (default ones(J, N))
 - `Lj::Vector{Float64}`: Vector of populations in each node (j = 1:J) (only for fixed labour case)
 - `Hj::Vector{Float64}`: Vector of immobile good in each node (j = 1:J) (e.g. housing, default ones(J))

From 6fff0bebfba639d1fb3e04d992bb4edee1ade12d Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Wed, 11 Sep 2024 16:41:59 +0200
Subject: [PATCH 08/32] 'm' is part of param again.

---
 src/main/optimal_network.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/main/optimal_network.jl b/src/main/optimal_network.jl
index ce547ab..9e240af 100644
--- a/src/main/optimal_network.jl
+++ b/src/main/optimal_network.jl
@@ -133,7 +133,7 @@ function optimal_network(param, graph; I0=nothing, Il=nothing, Iu=nothing, verbo
             PQ = dropdims(sum(PQ + permutedims(PQ, [2, 1, 3]), dims=3), dims = 3)
         else
             PQ = repeat(results[:PCj], 1, J)
-            matm = permutedims(repeat(graph.m, 1, J, J), [3, 2, 1])
+            matm = permutedims(repeat(param.m, 1, J, J), [3, 2, 1])
             cost = dropdims(sum(matm .* results[:Qjkn] .^ param.nu, dims=3), dims = 3) .^ ((param.beta + 1) / param.nu)
             PQ .*= cost
             PQ += PQ'

From 3ff93964ac75bbf283b15d7349e7a78e4124aa59 Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Wed, 11 Sep 2024 16:45:28 +0200
Subject: [PATCH 09/32] Minors reflecting changes.

---
 src/main/annealing.jl | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/main/annealing.jl b/src/main/annealing.jl
index 9d10c19..103969e 100644
--- a/src/main/annealing.jl
+++ b/src/main/annealing.jl
@@ -37,9 +37,9 @@ Runs the simulated annealing method starting from network `I0`. Only sensible if
 # Examples
 ```julia
 # Nonconvex case, disabling automatic annealing
-param = init_parameters(K = 10, annealing = false, gamma = 2)
-param, graph = create_graph(param)
-param[:Zjn][61] = 10.0
+param = init_parameters(K = 10, gamma = 2, annealing = false)
+graph = create_graph(param)
+graph[:Zjn][61] = 10.0
 results = optimal_network(param, graph)
 
 # Run annealing
@@ -103,7 +103,7 @@ function annealing(param, graph, I0; kwargs...)
         # TODO: set IO (kappa_ex)?
     else
         auxdata = create_auxdata(param, graph, edges, I0)
-        model, recover_allocation = get_model(param, auxdata)
+        model, recover_allocation = get_model(auxdata)
     end
 
     if !options.verbose

From 37461ec72fba445498487ad6317ca6a0e0c40f1c Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Wed, 11 Sep 2024 16:47:30 +0200
Subject: [PATCH 10/32] Minors documentation.

---
 src/main/helper.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/main/helper.jl b/src/main/helper.jl
index 8cfd691..89464f3 100644
--- a/src/main/helper.jl
+++ b/src/main/helper.jl
@@ -159,7 +159,7 @@ function create_auxdata(param, graph, edges, I)
 end
 
 """
-    get_model(param, auxdata)
+    get_model(auxdata)
 
 Construct the appropriate model based on the parameters and auxiliary data.
 

From 4867402181604c9fcc2472314d927e7d7922bf0f Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Wed, 11 Sep 2024 16:55:50 +0200
Subject: [PATCH 11/32] Let N only be part of parameters dict.

---
 src/main/create_graph.jl    | 3 +--
 src/main/init_parameters.jl | 8 ++++----
 src/main/optimal_network.jl | 2 +-
 3 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/src/main/create_graph.jl b/src/main/create_graph.jl
index 779ea24..4e10d13 100644
--- a/src/main/create_graph.jl
+++ b/src/main/create_graph.jl
@@ -42,8 +42,7 @@ function create_graph(param, w = 11, h = 11; type = "map", kwargs...)
         graph = create_custom(options[:adjacency], options[:x], options[:y])
     end
 
-    graph[:N] = param[:N] # Ensure graph has all information
-    graph[:Zjn] = haskey(options, :Zjn) ? options[:Zjn] : ones(graph[:J], graph[:N])
+    graph[:Zjn] = haskey(options, :Zjn) ? options[:Zjn] : ones(graph[:J], param[:N])
     graph[:Hj] = haskey(options, :Hj) ? options[:Hj] : ones(graph[:J])
 
     if param[:mobility] == false
diff --git a/src/main/init_parameters.jl b/src/main/init_parameters.jl
index c270bc0..216b3db 100644
--- a/src/main/init_parameters.jl
+++ b/src/main/init_parameters.jl
@@ -107,8 +107,8 @@ end
 # Check if the parameters are consistent with the graph structure
 function check_graph_param(graph, param)
 
-    if haskey(param, :m) && length(param[:m]) != graph[:N]
-        @warn "m does not have the right length N = $(graph[:N])."
+    if haskey(param, :m) && length(param[:m]) != param[:N]
+        @warn "m does not have the right length N = $(param[:N])."
     end
 
     if haskey(param, :omegaj) && !haskey(graph, :omegaj)
@@ -142,8 +142,8 @@ function check_graph_param(graph, param)
     if haskey(param, :Zjn) && !haskey(graph, :Zjn)
         graph[:Zjn] = param[:Zjn]
     end
-    if haskey(graph, :Zjn) && size(graph[:Zjn]) != (graph[:J], graph[:N])
-        @warn "Zjn does not have the right size J ($(graph[:J])) x N ($(graph[:N]))."
+    if haskey(graph, :Zjn) && size(graph[:Zjn]) != (graph[:J], param[:N])
+        @warn "Zjn does not have the right size J ($(graph[:J])) x N ($(param[:N]))."
     end
 
     if haskey(param, :Hj) && !haskey(graph, :Hj)
diff --git a/src/main/optimal_network.jl b/src/main/optimal_network.jl
index 9e240af..21ea17b 100644
--- a/src/main/optimal_network.jl
+++ b/src/main/optimal_network.jl
@@ -31,7 +31,7 @@ plot_graph(graph, results[:Ijk])
 function optimal_network(param, graph; I0=nothing, Il=nothing, Iu=nothing, verbose=false, return_model=0, solve_allocation = false)
 
     # I0=nothing; Il=nothing; Iu=nothing; verbose=false; return_model = false; return_model = 0;
-    # Check param.Zjn and make matrix if vector:
+    # Check graph.Zjn and make matrix if vector:
     if length(size(graph[:Zjn])) == 1
         graph[:Zjn] = reshape(graph[:Zjn], graph[:J], 1)
     end   

From 8128b29c76547dacef6b11be3a371808f3c5cf00 Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Wed, 11 Sep 2024 16:56:51 +0200
Subject: [PATCH 12/32] Update models to reflect new param and graph division.

---
 src/models/model_fixed.jl                     | 14 +++++-----
 src/models/model_fixed_armington.jl           | 18 ++++++-------
 src/models/model_fixed_cgc.jl                 | 16 ++++++------
 src/models/model_fixed_cgc_armington.jl       | 20 +++++++-------
 src/models/model_fixed_duality.jl             | 16 ++++++------
 src/models/model_mobility.jl                  |  8 +++---
 src/models/model_mobility_armington.jl        |  8 +++---
 src/models/model_mobility_cgc.jl              |  6 ++---
 src/models/model_mobility_cgc_armington.jl    |  8 +++---
 src/models/model_partial_mobility.jl          | 26 +++++++++----------
 .../model_partial_mobility_armington.jl       | 26 +++++++++----------
 src/models/model_partial_mobility_cgc.jl      | 24 ++++++++---------
 .../model_partial_mobility_cgc_armington.jl   | 26 +++++++++----------
 13 files changed, 108 insertions(+), 108 deletions(-)

diff --git a/src/models/model_fixed.jl b/src/models/model_fixed.jl
index ce47e1b..cad2c1f 100644
--- a/src/models/model_fixed.jl
+++ b/src/models/model_fixed.jl
@@ -10,7 +10,7 @@ function model_fixed(optimizer, auxdata)
     Apos = auxdata.edges.Apos
     Aneg = auxdata.edges.Aneg
     psigma = (param.sigma - 1) / param.sigma
-    Lj = param.Lj
+    Lj = graph.Lj
 
     # Model
     model = Model(optimizer)
@@ -27,13 +27,13 @@ function model_fixed(optimizer, auxdata)
 
     # Defining Utility Funcion: from Cjn + parameters (by operator overloading)
     @expression(model, Cj, sum(Cjn .^ psigma, dims=2) .^ (1 / psigma))
-    @expression(model, cj, ifelse.(param.Lj .== 0, 0.0, Cj ./ param.Lj))
-    @expression(model, uj, ((cj / param.alpha) .^ param.alpha .* (param.hj / (1-param.alpha)) .^ (1-param.alpha)) .^ (1-param.rho) / (1-param.rho))
-    @expression(model, U, sum(param.omegaj .* param.Lj .* uj))
+    @expression(model, cj, ifelse.(graph.Lj .== 0, 0.0, Cj ./ graph.Lj))
+    @expression(model, uj, ((cj / param.alpha) .^ param.alpha .* (graph.hj / (1-param.alpha)) .^ (1-param.alpha)) .^ (1-param.rho) / (1-param.rho))
+    @expression(model, U, sum(graph.omegaj .* graph.Lj .* uj))
     @objective(model, Max, U)
 
     # Define Yjn (production) as expression
-    @expression(model, Yjn[j=1:graph.J, n=1:param.N], param.Zjn[j, n] * Ljn[j, n]^param.a)
+    @expression(model, Yjn[j=1:graph.J, n=1:param.N], graph.Zjn[j, n] * Ljn[j, n]^param.a)
     # Balanced flow constraints
     @constraint(model, Pjn[j in 1:param.J, n in 1:param.N],
         Cjn[j, n] + sum(A[j, i] * Qin[i, n] for i in 1:graph.ndeg) -
@@ -62,9 +62,9 @@ function recover_allocation_fixed(model, auxdata)
     results[:Cjn] = value.(model_dict[:Cjn])
     results[:Cj] = dropdims(value.(model_dict[:Cj]), dims = 2)
     results[:Ljn] = value.(model_dict[:Ljn])
-    results[:Lj] = param.Lj
+    results[:Lj] = graph.Lj
     results[:cj] = dropdims(value.(model_dict[:cj]), dims = 2)
-    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, param.Hj ./ results[:Lj])
+    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, graph.Hj ./ results[:Lj])
     results[:uj] = dropdims(value.(model_dict[:uj]), dims = 2)
     # Prices
     results[:Pjn] = shadow_price.(model_dict[:Pjn])
diff --git a/src/models/model_fixed_armington.jl b/src/models/model_fixed_armington.jl
index a39a4d1..64b9133 100644
--- a/src/models/model_fixed_armington.jl
+++ b/src/models/model_fixed_armington.jl
@@ -10,9 +10,9 @@ function model_fixed_armington(optimizer, auxdata)
     Apos = auxdata.edges.Apos
     Aneg = auxdata.edges.Aneg
     psigma = (param.sigma - 1) / param.sigma
-    Lj = param.Lj
+    Lj = graph.Lj
     # Production: Fixed because just one good is produced by each location
-    Yjn = param.Zjn .* Lj .^ param.a
+    Yjn = graph.Zjn .* Lj .^ param.a
 
     # Model
     model = Model(optimizer)
@@ -29,8 +29,8 @@ function model_fixed_armington(optimizer, auxdata)
     # Defining Utility Funcion: from Cjn + parameters (by operator overloading)
     @expression(model, Cj, sum(Cjn .^ psigma, dims=2) .^ (1 / psigma))
     @expression(model, cj, ifelse.(Lj .== 0, 0.0, Cj ./ Lj))
-    @expression(model, uj, ((cj / param.alpha) .^ param.alpha .* (param.hj / (1-param.alpha)) .^ (1-param.alpha)) .^ (1-param.rho) / (1-param.rho))
-    @expression(model, U, sum(param.omegaj .* Lj .* uj))
+    @expression(model, uj, ((cj / param.alpha) .^ param.alpha .* (graph.hj / (1-param.alpha)) .^ (1-param.alpha)) .^ (1-param.rho) / (1-param.rho))
+    @expression(model, U, sum(graph.omegaj .* Lj .* uj))
     @objective(model, Max, U)
 
     # Balanced flow constraints
@@ -51,18 +51,18 @@ function recover_allocation_fixed_armington(model, auxdata)
     model_dict = model.obj_dict
     results = Dict()
     # Production: Fixed because just one good is produced by each location
-    Yj = dropdims(sum(param.Zjn, dims = 2) .* param.Lj .^ param.a, dims = 2)
-    WY = param.Zjn .> 0
+    Yj = dropdims(sum(graph.Zjn, dims = 2) .* graph.Lj .^ param.a, dims = 2)
+    WY = graph.Zjn .> 0
 
     results[:welfare] = value(model_dict[:U])
     results[:Yjn] = WY .* Yj
     results[:Yj] = Yj
     results[:Cjn] = value.(model_dict[:Cjn])
     results[:Cj] = dropdims(value.(model_dict[:Cj]), dims = 2)
-    results[:Ljn] = WY .* param.Lj
-    results[:Lj] = param.Lj
+    results[:Ljn] = WY .* graph.Lj
+    results[:Lj] = graph.Lj
     results[:cj] = dropdims(value.(model_dict[:cj]), dims = 2)
-    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, param.Hj ./ results[:Lj])
+    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, graph.Hj ./ results[:Lj])
     results[:uj] = dropdims(value.(model_dict[:uj]), dims = 2)
     # Prices
     results[:Pjn] = shadow_price.(model_dict[:Pjn])
diff --git a/src/models/model_fixed_cgc.jl b/src/models/model_fixed_cgc.jl
index dce2f6b..ce818e8 100644
--- a/src/models/model_fixed_cgc.jl
+++ b/src/models/model_fixed_cgc.jl
@@ -9,7 +9,7 @@ function model_fixed_cgc(optimizer, auxdata)
     A = auxdata.edges.A
     Apos = auxdata.edges.Apos
     Aneg = auxdata.edges.Aneg
-    Lj = param.Lj
+    Lj = graph.Lj
     m = param.m # 1:param.N: Vector of weights on each goods flow for aggregate congestion term
     psigma = (param.sigma - 1) / param.sigma
     beta_nu = (param.beta + 1) / param.nu
@@ -30,8 +30,8 @@ function model_fixed_cgc(optimizer, auxdata)
     set_parameter_value.(kappa_ex, kappa_ex_init)
 
     # Defining Utility Funcion: from cj + parameters (by operator overloading)
-    @expression(model, uj, ((cj / param.alpha) .^ param.alpha .* (param.hj / (1-param.alpha)) .^ (1-param.alpha)) .^ (1-param.rho) / (1-param.rho))
-    @expression(model, U, sum(param.omegaj .* param.Lj .* uj))      # Overall Utility
+    @expression(model, uj, ((cj / param.alpha) .^ param.alpha .* (graph.hj / (1-param.alpha)) .^ (1-param.alpha)) .^ (1-param.rho) / (1-param.rho))
+    @expression(model, U, sum(graph.omegaj .* graph.Lj .* uj))      # Overall Utility
     @objective(model, Max, U)
 
     # Create the matrix B_direct (resp. B_indirect) of transport cost along the direction of the edge (resp. in edge opposite direction)
@@ -39,10 +39,10 @@ function model_fixed_cgc(optimizer, auxdata)
     B_indirect = @expression(model, ((Qin_indirect .^ param.nu) * m) .^ beta_nu ./ kappa_ex)
     # Final good constraints
     @expression(model, Dj, sum(Djn .^ psigma, dims=2) .^ (1 / psigma))
-    @constraint(model, cj .* param.Lj + Apos * B_direct + Aneg * B_indirect - Dj .<= -1e-8)
+    @constraint(model, cj .* graph.Lj + Apos * B_direct + Aneg * B_indirect - Dj .<= -1e-8)
 
     # Balanced flow constraints
-    @expression(model, Yjn, param.Zjn .* (Ljn .^ param.a))
+    @expression(model, Yjn, graph.Zjn .* (Ljn .^ param.a))
     @constraint(model, Pjn, Djn + A * Qin_direct - A * Qin_indirect - Yjn .<= -1e-8)
 
     # Local labor availability constraints ( sum Ljn <= Lj )
@@ -61,12 +61,12 @@ function recover_allocation_fixed_cgc(model, auxdata)
     results[:Yjn] = value.(model_dict[:Yjn])
     results[:Yj] = dropdims(sum(results[:Yjn], dims=2), dims = 2) 
     results[:Ljn] = value.(model_dict[:Ljn])
-    results[:Lj] = param.Lj
+    results[:Lj] = graph.Lj
     results[:Djn] = value.(model_dict[:Djn]) # Consumption per good pre-transport cost
     results[:Dj] = dropdims(value.(model_dict[:Dj]), dims = 2)
     results[:cj] = value.(model_dict[:cj])
-    results[:Cj] = results[:cj] .* param.Lj
-    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, param.Hj ./ results[:Lj])
+    results[:Cj] = results[:cj] .* graph.Lj
+    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, graph.Hj ./ results[:Lj])
     results[:uj] = value.(model_dict[:uj])
     # Prices
     results[:Pjn] = shadow_price.(model_dict[:Pjn])
diff --git a/src/models/model_fixed_cgc_armington.jl b/src/models/model_fixed_cgc_armington.jl
index 1903c4b..a1e41a5 100644
--- a/src/models/model_fixed_cgc_armington.jl
+++ b/src/models/model_fixed_cgc_armington.jl
@@ -9,7 +9,7 @@ function model_fixed_cgc_armington(optimizer, auxdata)
     A = auxdata.edges.A
     Apos = auxdata.edges.Apos
     Aneg = auxdata.edges.Aneg
-    Lj = param.Lj
+    Lj = graph.Lj
     m = param.m # 1:param.N: Vector of weights on each goods flow for aggregate congestion term
     psigma = (param.sigma - 1) / param.sigma
     beta_nu = (param.beta + 1) / param.nu
@@ -29,8 +29,8 @@ function model_fixed_cgc_armington(optimizer, auxdata)
     set_parameter_value.(kappa_ex, kappa_ex_init)
 
     # Defining Utility Funcion: from cj + parameters (by operator overloading)
-    @expression(model, uj, ((cj / param.alpha) .^ param.alpha .* (param.hj / (1-param.alpha)) .^ (1-param.alpha)) .^ (1-param.rho) / (1-param.rho))
-    @expression(model, U, sum(param.omegaj .* Lj .* uj))      # Overall Utility
+    @expression(model, uj, ((cj / param.alpha) .^ param.alpha .* (graph.hj / (1-param.alpha)) .^ (1-param.alpha)) .^ (1-param.rho) / (1-param.rho))
+    @expression(model, U, sum(graph.omegaj .* Lj .* uj))      # Overall Utility
     @objective(model, Max, U)
 
     # Create the matrix B_direct (resp. B_indirect) of transport cost along the direction of the edge (resp. in edge opposite direction)
@@ -41,7 +41,7 @@ function model_fixed_cgc_armington(optimizer, auxdata)
     @constraint(model, cj .* Lj + Apos * B_direct + Aneg * B_indirect - Dj .<= -1e-8)
 
     # Balanced flow constraints
-    Yjn = param.Zjn .* Lj .^ param.a
+    Yjn = graph.Zjn .* Lj .^ param.a
     @constraint(model, Pjn, Djn + A * Qin_direct - A * Qin_indirect - Yjn .<= -1e-8)
 
     return model
@@ -53,19 +53,19 @@ function recover_allocation_fixed_cgc_armington(model, auxdata)
     model_dict = model.obj_dict
     results = Dict()
     # Production: Fixed because just one good is produced by each location
-    Yj = dropdims(sum(param.Zjn, dims = 2) .* param.Lj .^ param.a, dims = 2)
-    WY = param.Zjn .> 0
+    Yj = dropdims(sum(graph.Zjn, dims = 2) .* graph.Lj .^ param.a, dims = 2)
+    WY = graph.Zjn .> 0
 
     results[:welfare] = value(model_dict[:U])
     results[:Yjn] = WY .* Yj
     results[:Yj] = Yj
-    results[:Ljn] = WY .* param.Lj
-    results[:Lj] = param.Lj
+    results[:Ljn] = WY .* graph.Lj
+    results[:Lj] = graph.Lj
     results[:Djn] = value.(model_dict[:Djn]) # Consumption per good pre-transport cost
     results[:Dj] = dropdims(value.(model_dict[:Dj]), dims = 2)
     results[:cj] = value.(model_dict[:cj])
-    results[:Cj] = results[:cj] .* param.Lj
-    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, param.Hj ./ results[:Lj])
+    results[:Cj] = results[:cj] .* graph.Lj
+    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, graph.Hj ./ results[:Lj])
     results[:uj] = value.(model_dict[:uj])
     # Prices
     results[:Pjn] = shadow_price.(model_dict[:Pjn])
diff --git a/src/models/model_fixed_duality.jl b/src/models/model_fixed_duality.jl
index 2f3195e..b8b0a15 100644
--- a/src/models/model_fixed_duality.jl
+++ b/src/models/model_fixed_duality.jl
@@ -10,14 +10,14 @@ function model_fixed_duality(optimizer, auxdata)
     Aneg = auxdata.edges.Aneg
     es = auxdata.edges.edge_start
     ee = auxdata.edges.edge_end
-    omegaj = param.omegaj
-    Lj = param.Lj
+    omegaj = graph.omegaj
+    Lj = graph.Lj
     alpha = param.alpha
     sigma = param.sigma
     rho = param.rho
-    hj = param.hj
+    hj = graph.hj
     hj1malpha = (hj / (1-alpha)) .^ (1-alpha)
-    Zjn = param.Zjn
+    Zjn = graph.Zjn
     beta = param.beta
     a = param.a
 
@@ -92,12 +92,12 @@ function recover_allocation_fixed_duality(model, auxdata)
     results[:welfare] = value(model_dict[:U])
     results[:Yjn] = value.(model_dict[:Yjn])
     results[:Yj] = dropdims(sum(results[:Yjn], dims=2), dims = 2)
-    results[:Cjn] = value.(model_dict[:cjn]) .* repeat(param.Lj, 1, param.N)
+    results[:Cjn] = value.(model_dict[:cjn]) .* repeat(graph.Lj, 1, param.N)
     results[:cj] = value.(model_dict[:cj])
-    results[:Cj] = results[:cj] .* param.Lj
+    results[:Cj] = results[:cj] .* graph.Lj
     results[:Ljn] = value.(model_dict[:Ljn])
-    results[:Lj] = param.Lj
-    results[:hj] = param.hj
+    results[:Lj] = graph.Lj
+    results[:hj] = graph.hj
     results[:uj] = value.(model_dict[:uj]) # param.u.(results[:cj], results[:hj])
     # Prices
     results[:Pjn] = value.(model_dict[:Pjn])
diff --git a/src/models/model_mobility.jl b/src/models/model_mobility.jl
index 40c96c7..879d354 100644
--- a/src/models/model_mobility.jl
+++ b/src/models/model_mobility.jl
@@ -10,7 +10,7 @@ function model_mobility(optimizer, auxdata)
     Apos = auxdata.edges.Apos
     Aneg = auxdata.edges.Aneg
     psigma = (param.sigma - 1) / param.sigma
-    Hj = param.Hj
+    Hj = graph.Hj
 
     # Model
     model = Model(optimizer)
@@ -37,8 +37,8 @@ function model_mobility(optimizer, auxdata)
     end
 
     # balanced flow constraints
-    # Yjn = @expression(model, param.Zjn .* Ljn .^ param.a) # Same thing
-    @expression(model, Yjn[j=1:graph.J, n=1:param.N], param.Zjn[j, n] * Ljn[j, n]^param.a)
+    # Yjn = @expression(model, graph.Zjn .* Ljn .^ param.a) # Same thing
+    @expression(model, Yjn[j=1:graph.J, n=1:param.N], graph.Zjn[j, n] * Ljn[j, n]^param.a)
     @constraint(model, Pjn[j in 1:param.J, n in 1:param.N],
         Cjn[j, n] + sum(A[j, i] * Qin[i, n] for i in 1:graph.ndeg) -
         Yjn[j, n] + sum(
@@ -71,7 +71,7 @@ function recover_allocation_mobility(model, auxdata)
     results[:Ljn] = value.(model_dict[:Ljn])
     results[:Lj] = value.(model_dict[:Lj])
     results[:cj] = ifelse.(results[:Lj] .== 0, 0.0, results[:Cj] ./ results[:Lj])
-    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, param.Hj ./ results[:Lj])
+    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, graph.Hj ./ results[:Lj])
     results[:uj] = param.u.(results[:cj], results[:hj])
     # Prices
     results[:Pjn] = shadow_price.(model_dict[:Pjn])
diff --git a/src/models/model_mobility_armington.jl b/src/models/model_mobility_armington.jl
index b1b49ad..5537eef 100644
--- a/src/models/model_mobility_armington.jl
+++ b/src/models/model_mobility_armington.jl
@@ -10,7 +10,7 @@ function model_mobility_armington(optimizer, auxdata)
     Apos = auxdata.edges.Apos
     Aneg = auxdata.edges.Aneg
     psigma = (param.sigma - 1) / param.sigma
-    Hj = param.Hj
+    Hj = graph.Hj
 
     # Model
     model = Model(optimizer)
@@ -36,7 +36,7 @@ function model_mobility_armington(optimizer, auxdata)
     end
 
     # balanced flow constraints
-     @expression(model, Yjn[j=1:graph.J, n=1:param.N], param.Zjn[j, n] * Lj[j]^param.a)
+     @expression(model, Yjn[j=1:graph.J, n=1:param.N], graph.Zjn[j, n] * Lj[j]^param.a)
     @constraint(model, Pjn[j in 1:param.J, n in 1:param.N],
         Cjn[j, n] + sum(A[j, i] * Qin[i, n] for i in 1:graph.ndeg) -
         Yjn[j, n] + sum(
@@ -64,9 +64,9 @@ function recover_allocation_mobility_armington(model, auxdata)
     results[:Cjn] = value.(model_dict[:Cjn])
     results[:Cj] = dropdims(sum(results[:Cjn] .^ ((param.sigma-1)/param.sigma), dims=2), dims = 2) .^ (param.sigma/(param.sigma-1))
     results[:Lj] = value.(model_dict[:Lj])
-    results[:Ljn] = (param.Zjn .> 0) .* results[:Lj]
+    results[:Ljn] = (graph.Zjn .> 0) .* results[:Lj]
     results[:cj] = ifelse.(results[:Lj] .== 0, 0.0, results[:Cj] ./ results[:Lj])
-    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, param.Hj ./ results[:Lj])
+    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, graph.Hj ./ results[:Lj])
     results[:uj] = param.u.(results[:cj], results[:hj])
     # Prices
     results[:Pjn] = shadow_price.(model_dict[:Pjn])
diff --git a/src/models/model_mobility_cgc.jl b/src/models/model_mobility_cgc.jl
index 08386a1..f18b93b 100644
--- a/src/models/model_mobility_cgc.jl
+++ b/src/models/model_mobility_cgc.jl
@@ -34,7 +34,7 @@ function model_mobility_cgc(optimizer, auxdata)
     @objective(model, Max, u)
 
     # Utility constraint (Lj*u <= ... )
-    @constraint(model, Lj .* u - (cj .* Lj / param.alpha) .^ param.alpha .* (param.Hj / (1 - param.alpha)) .^ (1 - param.alpha) .<= -1e-8)
+    @constraint(model, Lj .* u - (cj .* Lj / param.alpha) .^ param.alpha .* (graph.Hj / (1 - param.alpha)) .^ (1 - param.alpha) .<= -1e-8)
 
     # Create the matrix B_direct (resp. B_indirect) of transport cost along the direction of the edge (resp. in edge opposite direction)
     B_direct = @expression(model, ((Qin_direct .^ param.nu) * m) .^ beta_nu ./ kappa_ex)
@@ -44,7 +44,7 @@ function model_mobility_cgc(optimizer, auxdata)
     @constraint(model, cj .* Lj + Apos * B_direct + Aneg * B_indirect - Dj .<= -1e-8)
 
     # Balanced flow constraints
-    @expression(model, Yjn, param.Zjn .* (Ljn .^ param.a))
+    @expression(model, Yjn, graph.Zjn .* (Ljn .^ param.a))
     @constraint(model, Pjn, Djn + A * Qin_direct - A * Qin_indirect - Yjn .<= -1e-8)
 
     # Labor resource constraint
@@ -71,7 +71,7 @@ function recover_allocation_mobility_cgc(model, auxdata)
     results[:Dj] = dropdims(value.(model_dict[:Dj]), dims = 2)
     results[:cj] = value.(model_dict[:cj])
     results[:Cj] = results[:cj] .* results[:Lj]
-    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, param.Hj ./ results[:Lj])
+    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, graph.Hj ./ results[:Lj])
     results[:uj] = param.u.(results[:cj], results[:hj])
     # Prices
     results[:Pjn] = shadow_price.(model_dict[:Pjn])
diff --git a/src/models/model_mobility_cgc_armington.jl b/src/models/model_mobility_cgc_armington.jl
index b6edae4..4f63102 100644
--- a/src/models/model_mobility_cgc_armington.jl
+++ b/src/models/model_mobility_cgc_armington.jl
@@ -33,7 +33,7 @@ function model_mobility_cgc_armington(optimizer, auxdata)
     @objective(model, Max, u)
 
     # Utility constraint (Lj*u <= ... )
-    @constraint(model, Lj .* u - (cj .* Lj / param.alpha) .^ param.alpha .* (param.Hj / (1 - param.alpha)) .^ (1 - param.alpha) .<= -1e-8)
+    @constraint(model, Lj .* u - (cj .* Lj / param.alpha) .^ param.alpha .* (graph.Hj / (1 - param.alpha)) .^ (1 - param.alpha) .<= -1e-8)
 
     # Create the matrix B_direct (resp. B_indirect) of transport cost along the direction of the edge (resp. in edge opposite direction)
     B_direct = @expression(model, ((Qin_direct .^ param.nu) * m) .^ beta_nu ./ kappa_ex)
@@ -43,7 +43,7 @@ function model_mobility_cgc_armington(optimizer, auxdata)
     @constraint(model, cj .* Lj + Apos * B_direct + Aneg * B_indirect - Dj .<= -1e-8)
 
     # Balanced flow constraints
-    @expression(model, Yjn, param.Zjn .* (Lj .^ param.a))
+    @expression(model, Yjn, graph.Zjn .* (Lj .^ param.a))
     @constraint(model, Pjn, Djn + A * Qin_direct - A * Qin_indirect - Yjn .<= -1e-8)
 
     # Labor resource constraint
@@ -62,12 +62,12 @@ function recover_allocation_mobility_cgc_armington(model, auxdata)
     results[:Yjn] = value.(model_dict[:Yjn])
     results[:Yj] = dropdims(sum(results[:Yjn], dims=2), dims = 2) 
     results[:Lj] = value.(model_dict[:Lj])
-    results[:Ljn] = (param.Zjn .> 0) .* results[:Lj]
+    results[:Ljn] = (graph.Zjn .> 0) .* results[:Lj]
     results[:Djn] = value.(model_dict[:Djn]) # Consumption per good pre-transport cost
     results[:Dj] = dropdims(value.(model_dict[:Dj]), dims = 2)
     results[:cj] = value.(model_dict[:cj])
     results[:Cj] = results[:cj] .* results[:Lj]
-    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, param.Hj ./ results[:Lj])
+    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, graph.Hj ./ results[:Lj])
     results[:uj] = param.u.(results[:cj], results[:hj])
     # Prices
     results[:Pjn] = shadow_price.(model_dict[:Pjn])
diff --git a/src/models/model_partial_mobility.jl b/src/models/model_partial_mobility.jl
index 27da94b..5179949 100644
--- a/src/models/model_partial_mobility.jl
+++ b/src/models/model_partial_mobility.jl
@@ -14,13 +14,13 @@ function model_partial_mobility(optimizer, auxdata)
     Apos = auxdata.edges.Apos
     Aneg = auxdata.edges.Aneg
     psigma = (param.sigma - 1) / param.sigma
-    Hj = param.Hj
-    Lr = param.Lr
-    if length(param.omegar) != param.nregions
-        error("length(param.omegar) = $(length(param.omegar)) does not match number of regions = $(param.nregions)")
+    Hj = graph.Hj
+    Lr = graph.Lr
+    if length(graph.omegar) != graph.nregions
+        error("length(graph.omegar) = $(length(graph.omegar)) does not match number of regions = $(graph.nregions)")
     end
-    if length(Lr) != param.nregions
-        error("Populations Lr = $(length(Lr)) does not match number of regions = $(param.nregions)")
+    if length(Lr) != graph.nregions
+        error("Populations Lr = $(length(Lr)) does not match number of regions = $(graph.nregions)")
     end
 
     # Model
@@ -28,14 +28,14 @@ function model_partial_mobility(optimizer, auxdata)
     set_string_names_on_creation(model, false)
 
     # Variables + bounds
-    @variable(model, ur[1:param.nregions], container=Array, start = 0.0)               # Utility per capita in each region
+    @variable(model, ur[1:graph.nregions], container=Array, start = 0.0)               # Utility per capita in each region
     @variable(model, Cjn[1:graph.J, 1:param.N] >= 1e-8, container=Array, start = 1e-6) # Good specific consumption
     @variable(model, Qin[1:graph.ndeg, 1:param.N], container=Array, start = 0.0)       # Good specific flow
     # NOTE: Fajgelbaum et al (2019) only optimize Lj and distribute it equally for goods with positive productivity
     @variable(model, Lj[1:graph.J] >= 1e-8, container=Array)                           # Total labour
     @variable(model, Ljn[1:graph.J, 1:param.N] >= 1e-8, container=Array)               # Good specific labour
     # Calculate start values for Lj and Ljn
-    pop_start = (Lr ./ gsum(ones(graph.J), param.nregions, region))[region]
+    pop_start = (Lr ./ gsum(ones(graph.J), graph.nregions, region))[region]
     set_start_value.(Lj, pop_start)
     pop_start_goods = repeat(pop_start / param.N, 1, param.N)
     set_start_value.(Ljn, pop_start_goods)
@@ -45,7 +45,7 @@ function model_partial_mobility(optimizer, auxdata)
     set_parameter_value.(kappa_ex, kappa_ex_init)
 
     # Objective
-    @expression(model, U, sum(param.omegar .* Lr .* ur))   # Overall utility
+    @expression(model, U, sum(graph.omegar .* Lr .* ur))   # Overall utility
     @objective(model, Max, U)
 
     # Utility constraint (Lj * ur <= ... )
@@ -55,8 +55,8 @@ function model_partial_mobility(optimizer, auxdata)
     end
 
     # Balanced flow constraints: same as with unrestricted mobility (no restrictions on goods)
-    # Yjn = @expression(model, param.Zjn .* Ljn .^ param.a) # Same thing
-    @expression(model, Yjn[j=1:graph.J, n=1:param.N], param.Zjn[j, n] * Ljn[j, n]^param.a)
+    # Yjn = @expression(model, graph.Zjn .* Ljn .^ param.a) # Same thing
+    @expression(model, Yjn[j=1:graph.J, n=1:param.N], graph.Zjn[j, n] * Ljn[j, n]^param.a)
     @constraint(model, Pjn[j in 1:param.J, n in 1:param.N],
         Cjn[j, n] + sum(A[j, i] * Qin[i, n] for i in 1:graph.ndeg) -
         Yjn[j, n] + sum(
@@ -67,7 +67,7 @@ function model_partial_mobility(optimizer, auxdata)
     )
 
     # Labor resource constraints (within each region)
-    @constraint(model, -1e-8 .<= gsum(Lj, param.nregions, region) .- Lr .<= 1e-8)
+    @constraint(model, -1e-8 .<= gsum(Lj, graph.nregions, region) .- Lr .<= 1e-8)
 
     # Local labor availability constraints ( sum Ljn <= Lj )
     @constraint(model, -1e-8 .<= sum(Ljn, dims=2) .- Lj .<= 1e-8)
@@ -90,7 +90,7 @@ function recover_allocation_partial_mobility(model, auxdata)
     results[:Ljn] = value.(model_dict[:Ljn])
     results[:Lj] = value.(model_dict[:Lj])
     results[:cj] = ifelse.(results[:Lj] .== 0, 0.0, results[:Cj] ./ results[:Lj])
-    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, param.Hj ./ results[:Lj])
+    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, graph.Hj ./ results[:Lj])
     results[:uj] = param.u.(results[:cj], results[:hj])
     # Prices
     results[:Pjn] = shadow_price.(model_dict[:Pjn])
diff --git a/src/models/model_partial_mobility_armington.jl b/src/models/model_partial_mobility_armington.jl
index ca66562..f3fab26 100644
--- a/src/models/model_partial_mobility_armington.jl
+++ b/src/models/model_partial_mobility_armington.jl
@@ -14,13 +14,13 @@ function model_partial_mobility_armington(optimizer, auxdata)
     Apos = auxdata.edges.Apos
     Aneg = auxdata.edges.Aneg
     psigma = (param.sigma - 1) / param.sigma
-    Hj = param.Hj
-    Lr = param.Lr
-    if length(param.omegar) != param.nregions
-        error("length(param.omegar) = $(length(param.omegar)) does not match number of regions = $(param.nregions)")
+    Hj = graph.Hj
+    Lr = graph.Lr
+    if length(graph.omegar) != graph.nregions
+        error("length(graph.omegar) = $(length(graph.omegar)) does not match number of regions = $(graph.nregions)")
     end
-    if length(Lr) != param.nregions
-        error("Populations Lr = $(length(Lr)) does not match number of regions = $(param.nregions)")
+    if length(Lr) != graph.nregions
+        error("Populations Lr = $(length(Lr)) does not match number of regions = $(graph.nregions)")
     end
 
     # Model
@@ -28,13 +28,13 @@ function model_partial_mobility_armington(optimizer, auxdata)
     set_string_names_on_creation(model, false)
 
     # Variables + bounds
-    @variable(model, ur[1:param.nregions], container=Array, start = 0.0)               # Utility per capita in each region
+    @variable(model, ur[1:graph.nregions], container=Array, start = 0.0)               # Utility per capita in each region
     @variable(model, Cjn[1:graph.J, 1:param.N] >= 1e-8, container=Array, start = 1e-6) # Good specific consumption
     @variable(model, Qin[1:graph.ndeg, 1:param.N], container=Array, start = 0.0)       # Good specific flow
     # NOTE: Fajgelbaum et al (2019) only optimize Lj and distribute it equally for goods with positive productivity
     @variable(model, Lj[1:graph.J] >= 1e-8, container=Array)                           # Total labour
     # Calculate start values for Lj
-    pop_start = (Lr ./ gsum(ones(graph.J), param.nregions, region))[region]
+    pop_start = (Lr ./ gsum(ones(graph.J), graph.nregions, region))[region]
     set_start_value.(Lj, pop_start)
 
     # Parameters: to be updated between solves
@@ -42,7 +42,7 @@ function model_partial_mobility_armington(optimizer, auxdata)
     set_parameter_value.(kappa_ex, kappa_ex_init)
 
     # Objective
-    @expression(model, U, sum(param.omegar .* Lr .* ur))   # Overall utility
+    @expression(model, U, sum(graph.omegar .* Lr .* ur))   # Overall utility
     @objective(model, Max, U)
 
     # Utility constraint (Lj * ur <= ... )
@@ -52,7 +52,7 @@ function model_partial_mobility_armington(optimizer, auxdata)
     end
 
     # Balanced flow constraints: same as with unrestricted mobility (no restrictions on goods)
-    @expression(model, Yjn[j=1:graph.J, n=1:param.N], param.Zjn[j, n] * Lj[j]^param.a)
+    @expression(model, Yjn[j=1:graph.J, n=1:param.N], graph.Zjn[j, n] * Lj[j]^param.a)
     @constraint(model, Pjn[j in 1:param.J, n in 1:param.N],
         Cjn[j, n] + sum(A[j, i] * Qin[i, n] for i in 1:graph.ndeg) -
         Yjn[j, n] + sum(
@@ -63,7 +63,7 @@ function model_partial_mobility_armington(optimizer, auxdata)
     )
 
     # Labor resource constraints (within each region)
-    @constraint(model, -1e-8 .<= gsum(Lj, param.nregions, region) .- Lr .<= 1e-8)
+    @constraint(model, -1e-8 .<= gsum(Lj, graph.nregions, region) .- Lr .<= 1e-8)
 
     return model
 end
@@ -81,9 +81,9 @@ function recover_allocation_partial_mobility_armington(model, auxdata)
     results[:Cjn] = value.(model_dict[:Cjn])
     results[:Cj] = dropdims(sum(results[:Cjn] .^ ((param.sigma-1)/param.sigma), dims=2), dims = 2) .^ (param.sigma/(param.sigma-1))
     results[:Lj] = value.(model_dict[:Lj])
-    results[:Ljn] = (param.Zjn .> 0).* results[:Lj]
+    results[:Ljn] = (graph.Zjn .> 0).* results[:Lj]
     results[:cj] = ifelse.(results[:Lj] .== 0, 0.0, results[:Cj] ./ results[:Lj])
-    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, param.Hj ./ results[:Lj])
+    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, graph.Hj ./ results[:Lj])
     results[:uj] = param.u.(results[:cj], results[:hj])
     # Prices
     results[:Pjn] = shadow_price.(model_dict[:Pjn])
diff --git a/src/models/model_partial_mobility_cgc.jl b/src/models/model_partial_mobility_cgc.jl
index 8a82554..ff620ad 100644
--- a/src/models/model_partial_mobility_cgc.jl
+++ b/src/models/model_partial_mobility_cgc.jl
@@ -16,12 +16,12 @@ function model_partial_mobility_cgc(optimizer, auxdata)
     m = param.m # Vector of weights on each goods flow for aggregate congestion term
     psigma = (param.sigma - 1) / param.sigma
     beta_nu = (param.beta + 1) / param.nu
-    Lr = param.Lr
-    if length(param.omegar) != param.nregions
-        error("length(param.omegar) = $(length(param.omegar)) does not match number of regions = $(param.nregions)")
+    Lr = graph.Lr
+    if length(graph.omegar) != graph.nregions
+        error("length(graph.omegar) = $(length(graph.omegar)) does not match number of regions = $(graph.nregions)")
     end
-    if length(Lr) != param.nregions
-        error("Populations Lr = $(length(Lr)) does not match number of regions = $(param.nregions)")
+    if length(Lr) != graph.nregions
+        error("Populations Lr = $(length(Lr)) does not match number of regions = $(graph.nregions)")
     end
 
     # Model
@@ -29,7 +29,7 @@ function model_partial_mobility_cgc(optimizer, auxdata)
     set_string_names_on_creation(model, false)
 
     # Variable declarations
-    @variable(model, ur[1:param.nregions], container=Array, start = 0.0)                          # Utility per capita in each region
+    @variable(model, ur[1:graph.nregions], container=Array, start = 0.0)                          # Utility per capita in each region
     @variable(model, Djn[1:graph.J, 1:param.N] >= 1e-8, container=Array, start = 1e-6)            # Consumption per good pre-transport cost (Dj)
     @variable(model, cj[1:graph.J] >= 1e-8, container=Array, start = 1e-6)                        # Overall consumption bundle, including transport costs
     @variable(model, Qin_direct[1:graph.ndeg, 1:param.N] >= 1e-8, container=Array, start = 0.0)   # Direct aggregate flow
@@ -38,7 +38,7 @@ function model_partial_mobility_cgc(optimizer, auxdata)
     @variable(model, Ljn[1:graph.J, 1:param.N] >= 1e-8, container=Array)            # Good specific labour
     @variable(model, Lj[1:graph.J] >= 1e-8, container=Array)                        # Overall labour
     # Calculate start values for Lj and Ljn
-    pop_start = (Lr ./ gsum(ones(graph.J), param.nregions, region))[region]
+    pop_start = (Lr ./ gsum(ones(graph.J), graph.nregions, region))[region]
     set_start_value.(Lj, pop_start)
     pop_start_goods = repeat(pop_start / param.N, 1, param.N)
     set_start_value.(Ljn, pop_start_goods)
@@ -48,11 +48,11 @@ function model_partial_mobility_cgc(optimizer, auxdata)
     set_parameter_value.(kappa_ex, kappa_ex_init)
 
     # Objective
-    @expression(model, U, sum(param.omegar .* Lr .* ur)) # Overall Utility
+    @expression(model, U, sum(graph.omegar .* Lr .* ur)) # Overall Utility
     @objective(model, Max, U)
 
     # Utility constraint (Lj * ur <= ... )
-    @constraint(model, Lj .* ur[region] - (cj .* Lj ./ param.alpha) .^ param.alpha .* (param.Hj ./ (1 - param.alpha)) .^ (1 - param.alpha) .<= -1e-8)
+    @constraint(model, Lj .* ur[region] - (cj .* Lj ./ param.alpha) .^ param.alpha .* (graph.Hj ./ (1 - param.alpha)) .^ (1 - param.alpha) .<= -1e-8)
 
     # Create the matrix B_direct (resp. B_indirect) of transport cost along the direction of the edge (resp. in edge opposite direction)
     B_direct = @expression(model, ((Qin_direct .^ param.nu) * m) .^ beta_nu ./ kappa_ex)
@@ -62,11 +62,11 @@ function model_partial_mobility_cgc(optimizer, auxdata)
     @constraint(model, cj .* Lj + Apos * B_direct + Aneg * B_indirect - Dj .<= -1e-8)
 
     # Balanced flow constraints: same as full mobility
-    @expression(model, Yjn, param.Zjn .* (Ljn .^ param.a))
+    @expression(model, Yjn, graph.Zjn .* (Ljn .^ param.a))
     @constraint(model, Pjn, Djn + A * Qin_direct - A * Qin_indirect - Yjn .<= -1e-8)
 
     # Labor resource constraints (within each region)
-    @constraint(model, -1e-8 .<= gsum(Lj, param.nregions, region) .- Lr .<= 1e-8)
+    @constraint(model, -1e-8 .<= gsum(Lj, graph.nregions, region) .- Lr .<= 1e-8)
 
     # Local labor availability constraints ( sum Ljn <= Lj )
     @constraint(model, -1e-8 .<= sum(Ljn, dims=2) .- Lj .<= 1e-8)
@@ -90,7 +90,7 @@ function recover_allocation_partial_mobility_cgc(model, auxdata)
     results[:Dj] = dropdims(value.(model_dict[:Dj]), dims=2)
     results[:cj] = value.(model_dict[:cj])
     results[:Cj] = results[:cj] .* results[:Lj]
-    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, param.Hj ./ results[:Lj])
+    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, graph.Hj ./ results[:Lj])
     results[:uj] = param.u.(results[:cj], results[:hj])
     # Prices
     results[:Pjn] = shadow_price.(model_dict[:Pjn])
diff --git a/src/models/model_partial_mobility_cgc_armington.jl b/src/models/model_partial_mobility_cgc_armington.jl
index 9e7eb58..d03bfb9 100644
--- a/src/models/model_partial_mobility_cgc_armington.jl
+++ b/src/models/model_partial_mobility_cgc_armington.jl
@@ -16,12 +16,12 @@ function model_partial_mobility_cgc_armington(optimizer, auxdata)
     m = param.m # Vector of weights on each goods flow for aggregate congestion term
     psigma = (param.sigma - 1) / param.sigma
     beta_nu = (param.beta + 1) / param.nu
-    Lr = param.Lr
-    if length(param.omegar) != param.nregions
-        error("length(param.omegar) = $(length(param.omegar)) does not match number of regions = $(param.nregions)")
+    Lr = graph.Lr
+    if length(graph.omegar) != graph.nregions
+        error("length(graph.omegar) = $(length(graph.omegar)) does not match number of regions = $(graph.nregions)")
     end
-    if length(Lr) != param.nregions
-        error("Populations Lr = $(length(Lr)) does not match number of regions = $(param.nregions)")
+    if length(Lr) != graph.nregions
+        error("Populations Lr = $(length(Lr)) does not match number of regions = $(graph.nregions)")
     end
 
     # Model
@@ -29,14 +29,14 @@ function model_partial_mobility_cgc_armington(optimizer, auxdata)
     set_string_names_on_creation(model, false)
 
     # Variable declarations
-    @variable(model, ur[1:param.nregions], container=Array, start = 0.0)                          # Utility per capita in each region
+    @variable(model, ur[1:graph.nregions], container=Array, start = 0.0)                          # Utility per capita in each region
     @variable(model, Djn[1:graph.J, 1:param.N] >= 1e-8, container=Array, start = 1e-6)            # Consumption per good pre-transport cost (Dj)
     @variable(model, cj[1:graph.J] >= 1e-8, container=Array, start = 1e-6)                        # Overall consumption bundle, including transport costs
     @variable(model, Qin_direct[1:graph.ndeg, 1:param.N] >= 1e-8, container=Array, start = 0.0)   # Direct aggregate flow
     @variable(model, Qin_indirect[1:graph.ndeg, 1:param.N] >= 1e-8, container=Array, start = 0.0) # Indirect aggregate flow
     @variable(model, Lj[1:graph.J] >= 1e-8, container=Array)                        # Overall labour
     # Calculate start values for Lj
-    pop_start = (Lr ./ gsum(ones(graph.J), param.nregions, region))[region]
+    pop_start = (Lr ./ gsum(ones(graph.J), graph.nregions, region))[region]
     set_start_value.(Lj, pop_start)
 
     # Parameters: to be updated between solves
@@ -44,11 +44,11 @@ function model_partial_mobility_cgc_armington(optimizer, auxdata)
     set_parameter_value.(kappa_ex, kappa_ex_init)
 
     # Objective
-    @expression(model, U, sum(param.omegar .* Lr .* ur)) # Overall Utility
+    @expression(model, U, sum(graph.omegar .* Lr .* ur)) # Overall Utility
     @objective(model, Max, U)
 
     # Utility constraint (Lj * ur <= ... )
-    @constraint(model, Lj .* ur[region] - (cj .* Lj ./ param.alpha) .^ param.alpha .* (param.Hj ./ (1 - param.alpha)) .^ (1 - param.alpha) .<= -1e-8)
+    @constraint(model, Lj .* ur[region] - (cj .* Lj ./ param.alpha) .^ param.alpha .* (graph.Hj ./ (1 - param.alpha)) .^ (1 - param.alpha) .<= -1e-8)
 
     # Create the matrix B_direct (resp. B_indirect) of transport cost along the direction of the edge (resp. in edge opposite direction)
     B_direct = @expression(model, ((Qin_direct .^ param.nu) * m) .^ beta_nu ./ kappa_ex)
@@ -58,11 +58,11 @@ function model_partial_mobility_cgc_armington(optimizer, auxdata)
     @constraint(model, cj .* Lj + Apos * B_direct + Aneg * B_indirect - Dj .<= -1e-8)
 
     # Balanced flow constraints: same as full mobility
-    @expression(model, Yjn, param.Zjn .* (Lj .^ param.a))
+    @expression(model, Yjn, graph.Zjn .* (Lj .^ param.a))
     @constraint(model, Pjn, Djn + A * Qin_direct - A * Qin_indirect - Yjn .<= -1e-8)
 
     # Labor resource constraints (within each region)
-    @constraint(model, -1e-8 .<= gsum(Lj, param.nregions, region) .- Lr .<= 1e-8)
+    @constraint(model, -1e-8 .<= gsum(Lj, graph.nregions, region) .- Lr .<= 1e-8)
     
     return model
 end
@@ -78,12 +78,12 @@ function recover_allocation_partial_mobility_cgc_armington(model, auxdata)
     results[:Yjn] = value.(model_dict[:Yjn])
     results[:Yj] = dropdims(sum(results[:Yjn], dims=2), dims = 2) 
     results[:Lj] = value.(model_dict[:Lj])
-    results[:Ljn] = (param.Zjn .> 0) .* results[:Lj]
+    results[:Ljn] = (graph.Zjn .> 0) .* results[:Lj]
     results[:Djn] = value.(model_dict[:Djn]) # Consumption per good pre-transport cost
     results[:Dj] = dropdims(value.(model_dict[:Dj]), dims=2)
     results[:cj] = value.(model_dict[:cj])
     results[:Cj] = results[:cj] .* results[:Lj]
-    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, param.Hj ./ results[:Lj])
+    results[:hj] = ifelse.(results[:Lj] .== 0, 0.0, graph.Hj ./ results[:Lj])
     results[:uj] = param.u.(results[:cj], results[:hj])
     # Prices
     results[:Pjn] = shadow_price.(model_dict[:Pjn])

From 582fa2d08450e42bfd3b4f2ae0c0c7a040d82246 Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Wed, 11 Sep 2024 17:04:24 +0200
Subject: [PATCH 13/32] Update documentation examples.

---
 src/OptimalTransportNetworks.jl | 10 +++++-----
 src/main/plot_graph.jl          |  4 ++--
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/OptimalTransportNetworks.jl b/src/OptimalTransportNetworks.jl
index e7df504..e39447f 100644
--- a/src/OptimalTransportNetworks.jl
+++ b/src/OptimalTransportNetworks.jl
@@ -51,15 +51,15 @@ a simplified interface and only exports key functions, while retaining full flex
 ```julia
 # Convex case
 param = init_parameters(K = 10)
-param, graph = create_graph(param)
-param[:Zjn][61] = 10.0
+graph = create_graph(param)
+graph[:Zjn][61] = 10.0
 results = optimal_network(param, graph)
 plot_graph(graph, results[:Ijk])
 
 # Nonconvex case, disabling automatic annealing
-param = init_parameters(K = 10, annealing = false, gamma = 2)
-param, graph = create_graph(param)
-param[:Zjn][61] = 10.0
+param = init_parameters(K = 10, gamma = 2, annealing = false)
+graph = create_graph(param)
+graph[:Zjn][61] = 10.0
 results = optimal_network(param, graph)
 
 # Run annealing
diff --git a/src/main/plot_graph.jl b/src/main/plot_graph.jl
index bd5fa02..2802307 100644
--- a/src/main/plot_graph.jl
+++ b/src/main/plot_graph.jl
@@ -50,8 +50,8 @@ Plot a graph visualization with various styling options.
 # Examples
 ```julia
 param = init_parameters(K = 10)
-param, graph = create_graph(param)
-param[:Zjn][51] = 10.0
+graph = create_graph(param)
+graph[:Zjn][51] = 10.0
 results = optimal_network(param, graph)
 plot_graph(graph, results[:Ijk])
 ```

From 5dfae36b62f9554106d3c11ef42184b0bdf5361f Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Wed, 11 Sep 2024 17:06:58 +0200
Subject: [PATCH 14/32] J is part of graph, not param.

---
 src/models/model_fixed.jl                      | 2 +-
 src/models/model_fixed_armington.jl            | 2 +-
 src/models/model_mobility.jl                   | 2 +-
 src/models/model_mobility_armington.jl         | 2 +-
 src/models/model_partial_mobility.jl           | 2 +-
 src/models/model_partial_mobility_armington.jl | 2 +-
 6 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/models/model_fixed.jl b/src/models/model_fixed.jl
index cad2c1f..06b3a4f 100644
--- a/src/models/model_fixed.jl
+++ b/src/models/model_fixed.jl
@@ -35,7 +35,7 @@ function model_fixed(optimizer, auxdata)
     # Define Yjn (production) as expression
     @expression(model, Yjn[j=1:graph.J, n=1:param.N], graph.Zjn[j, n] * Ljn[j, n]^param.a)
     # Balanced flow constraints
-    @constraint(model, Pjn[j in 1:param.J, n in 1:param.N],
+    @constraint(model, Pjn[j in 1:graph.J, n in 1:param.N],
         Cjn[j, n] + sum(A[j, i] * Qin[i, n] for i in 1:graph.ndeg) -
         Yjn[j, n] + sum(
             ifelse(Qin[i, n] > 0, Apos[j, i], Aneg[j, i]) *
diff --git a/src/models/model_fixed_armington.jl b/src/models/model_fixed_armington.jl
index 64b9133..e471707 100644
--- a/src/models/model_fixed_armington.jl
+++ b/src/models/model_fixed_armington.jl
@@ -34,7 +34,7 @@ function model_fixed_armington(optimizer, auxdata)
     @objective(model, Max, U)
 
     # Balanced flow constraints
-    @constraint(model, Pjn[j in 1:param.J, n in 1:param.N],
+    @constraint(model, Pjn[j in 1:graph.J, n in 1:param.N],
         Cjn[j, n] + sum(A[j, i] * Qin[i, n] for i in 1:graph.ndeg) -
         Yjn[j, n] + sum(
             ifelse(Qin[i, n] > 0, Apos[j, i], Aneg[j, i]) *
diff --git a/src/models/model_mobility.jl b/src/models/model_mobility.jl
index 879d354..05512d7 100644
--- a/src/models/model_mobility.jl
+++ b/src/models/model_mobility.jl
@@ -39,7 +39,7 @@ function model_mobility(optimizer, auxdata)
     # balanced flow constraints
     # Yjn = @expression(model, graph.Zjn .* Ljn .^ param.a) # Same thing
     @expression(model, Yjn[j=1:graph.J, n=1:param.N], graph.Zjn[j, n] * Ljn[j, n]^param.a)
-    @constraint(model, Pjn[j in 1:param.J, n in 1:param.N],
+    @constraint(model, Pjn[j in 1:graph.J, n in 1:param.N],
         Cjn[j, n] + sum(A[j, i] * Qin[i, n] for i in 1:graph.ndeg) -
         Yjn[j, n] + sum(
             ifelse(Qin[i, n] > 0, Apos[j, i], Aneg[j, i]) *
diff --git a/src/models/model_mobility_armington.jl b/src/models/model_mobility_armington.jl
index 5537eef..772b981 100644
--- a/src/models/model_mobility_armington.jl
+++ b/src/models/model_mobility_armington.jl
@@ -37,7 +37,7 @@ function model_mobility_armington(optimizer, auxdata)
 
     # balanced flow constraints
      @expression(model, Yjn[j=1:graph.J, n=1:param.N], graph.Zjn[j, n] * Lj[j]^param.a)
-    @constraint(model, Pjn[j in 1:param.J, n in 1:param.N],
+    @constraint(model, Pjn[j in 1:graph.J, n in 1:param.N],
         Cjn[j, n] + sum(A[j, i] * Qin[i, n] for i in 1:graph.ndeg) -
         Yjn[j, n] + sum(
             ifelse(Qin[i, n] > 0, Apos[j, i], Aneg[j, i]) *
diff --git a/src/models/model_partial_mobility.jl b/src/models/model_partial_mobility.jl
index 5179949..c69d7c0 100644
--- a/src/models/model_partial_mobility.jl
+++ b/src/models/model_partial_mobility.jl
@@ -57,7 +57,7 @@ function model_partial_mobility(optimizer, auxdata)
     # Balanced flow constraints: same as with unrestricted mobility (no restrictions on goods)
     # Yjn = @expression(model, graph.Zjn .* Ljn .^ param.a) # Same thing
     @expression(model, Yjn[j=1:graph.J, n=1:param.N], graph.Zjn[j, n] * Ljn[j, n]^param.a)
-    @constraint(model, Pjn[j in 1:param.J, n in 1:param.N],
+    @constraint(model, Pjn[j in 1:graph.J, n in 1:param.N],
         Cjn[j, n] + sum(A[j, i] * Qin[i, n] for i in 1:graph.ndeg) -
         Yjn[j, n] + sum(
             ifelse(Qin[i, n] > 0, Apos[j, i], Aneg[j, i]) *
diff --git a/src/models/model_partial_mobility_armington.jl b/src/models/model_partial_mobility_armington.jl
index f3fab26..197bed7 100644
--- a/src/models/model_partial_mobility_armington.jl
+++ b/src/models/model_partial_mobility_armington.jl
@@ -53,7 +53,7 @@ function model_partial_mobility_armington(optimizer, auxdata)
 
     # Balanced flow constraints: same as with unrestricted mobility (no restrictions on goods)
     @expression(model, Yjn[j=1:graph.J, n=1:param.N], graph.Zjn[j, n] * Lj[j]^param.a)
-    @constraint(model, Pjn[j in 1:param.J, n in 1:param.N],
+    @constraint(model, Pjn[j in 1:graph.J, n in 1:param.N],
         Cjn[j, n] + sum(A[j, i] * Qin[i, n] for i in 1:graph.ndeg) -
         Yjn[j, n] + sum(
             ifelse(Qin[i, n] > 0, Apos[j, i], Aneg[j, i]) *

From 6fda83e27c00248f3b1a50d7b98f5bbcfde66020 Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Wed, 11 Sep 2024 17:23:23 +0200
Subject: [PATCH 15/32] Update examples.

---
 examples/example01.jl       |  6 +++---
 examples/example02.jl       | 18 +++++++++---------
 examples/example03.jl       | 12 ++++++------
 examples/example04.jl       | 26 ++++++++++++-------------
 examples/paper_example01.jl | 22 ++++++++++-----------
 examples/paper_example02.jl | 38 ++++++++++++++++++-------------------
 examples/paper_example03.jl | 30 ++++++++++++++---------------
 examples/paper_example04.jl | 24 +++++++++++------------
 8 files changed, 88 insertions(+), 88 deletions(-)

diff --git a/examples/example01.jl b/examples/example01.jl
index bff0c1f..8ccc0cd 100644
--- a/examples/example01.jl
+++ b/examples/example01.jl
@@ -17,13 +17,13 @@ param = init_parameters(labor_mobility = true, K = 10, gamma = 1, beta = 1, verb
 # ------------
 # Init network
 
-param, graph = create_graph(param, 11, 11, type = "map") # create a map network of 11x11 nodes located in [0,10]x[0,10]
+graph = create_graph(param, 11, 11, type = "map") # create a map network of 11x11 nodes located in [0,10]x[0,10]
 # note: by default, productivity and population are equalized everywhere
 
 # Customize graph
-param[:Zjn] = fill(0.1, param[:J], param[:N]) # set most places to low productivity
+graph[:Zjn] = fill(0.1, graph[:J], param[:N]) # set most places to low productivity
 Ni = find_node(graph, 6, 6) # Find index of the central node at (6,6)
-param[:Zjn][Ni, :] .= 1 # central node more productive
+graph[:Zjn][Ni, :] .= 1 # central node more productive
 
 # ==========
 # RESOLUTION
diff --git a/examples/example02.jl b/examples/example02.jl
index 0d6ff1d..ca0eed0 100644
--- a/examples/example02.jl
+++ b/examples/example02.jl
@@ -13,15 +13,15 @@ param = init_parameters(labor_mobility = false, K = 100, gamma = 1, beta = 1, N
 # Init network
 
 w, h = 13, 13
-param, graph = create_graph(param, w, h, type = "triangle") # create a triangular network of 21x21
+graph = create_graph(param, w, h, type = "triangle") # create a triangular network of 21x21
 
 ncities = 20 # number of random cities
-param[:Lj] .= 0 # set population to minimum everywhere
-param[:Zjn] .= 0.01 # set low productivity everywhere
+graph[:Lj] .= 0 # set population to minimum everywhere
+graph[:Zjn] .= 0.01 # set low productivity everywhere
 
 Ni = find_node(graph, ceil(Int, w/2), ceil(Int, h/2)) # Find the central node
-param[:Zjn][Ni] = 1 # central node is more productive
-param[:Lj][Ni] = 1 # cities are equally populated
+graph[:Zjn][Ni] = 1 # central node is more productive
+graph[:Lj][Ni] = 1 # cities are equally populated
 
 # Draw the rest of the cities randomly
 Random.seed!(5) # reinit random number generator
@@ -29,15 +29,15 @@ for i in 1:ncities
     newdraw = false
     while !newdraw
         j = round(Int, 1 + rand() * (graph[:J] - 1))
-        if param[:Lj][j] != 1 / (ncities + 1) # make sure node j is unpopulated           
+        if graph[:Lj][j] != 1 / (ncities + 1) # make sure node j is unpopulated           
             newdraw = true
-            param[:Lj][j] = 1
+            graph[:Lj][j] = 1
         end
     end
 end
 
 # For Ipopt: population cannot be zero!
-param[:Lj][param[:Lj] .== 0] .= 1e-6
+graph[:Lj][graph[:Lj] .== 0] .= 1e-6
 
 # ==========
 # RESOLUTION
@@ -50,7 +50,7 @@ param[:Lj][param[:Lj] .== 0] .= 1e-6
     # second, in the nonconvex case (gamma>beta)
     param[:gamma] = 2 # change only gamma, keep other parameters
     res[2] = optimal_network(param, graph) # optimize by iterating on FOCs
-    res[3] = annealing(param, graph, res[2][:Ijk]) # improve with annealing, starting from previous result
+    res[3] = annealing(graph, res[2][:Ijk]) # improve with annealing, starting from previous result
 end
 
 # ==========
diff --git a/examples/example03.jl b/examples/example03.jl
index 64c0fa9..8201488 100644
--- a/examples/example03.jl
+++ b/examples/example03.jl
@@ -15,10 +15,10 @@ param = init_parameters(labor_mobility = true, K = 100, gamma = 2, beta = 1, N =
 # ------------
 # Init network
 
-param, graph = create_graph(param, 7, 7, type = "triangle") # create a triangular network of 21x21
+graph = create_graph(param, 7, 7, type = "triangle") # create a triangular network of 21x21
 
-param[:Zjn][:, 1:param[:N]-1] .= 0 # default locations cannot produce goods 1-10
-param[:Zjn][:, param[:N]] .= 1 # but they can all produce good 11 (agricultural)
+graph[:Zjn][:, 1:param[:N]-1] .= 0 # default locations cannot produce goods 1-10
+graph[:Zjn][:, param[:N]] .= 1 # but they can all produce good 11 (agricultural)
 
 # Draw the cities randomly
 Random.seed!(5) # reinit random number generator
@@ -26,10 +26,10 @@ for i in 1:param[:N]-1
     newdraw = false
     while newdraw == false
         j = round(Int, 1 + rand() * (graph[:J] - 1))
-        if any(param[:Zjn][j, 1:param[:N]-1] .> 0) == false # make sure node j does not produce any differentiated good
+        if any(graph[:Zjn][j, 1:param[:N]-1] .> 0) == false # make sure node j does not produce any differentiated good
             newdraw = true
-            param[:Zjn][j, 1:param[:N]] .= 0
-            param[:Zjn][j, i] = 1
+            graph[:Zjn][j, 1:param[:N]] .= 0
+            graph[:Zjn][j, i] = 1
         end
     end
 end
diff --git a/examples/example04.jl b/examples/example04.jl
index 4fbc356..fa28077 100644
--- a/examples/example04.jl
+++ b/examples/example04.jl
@@ -11,33 +11,33 @@ import Random
 param = init_parameters(K = 100, labor_mobility = false,
                         N = 1, gamma = 1, beta = 1, duality = false)
 w = 13; h = 13
-param, graph = create_graph(param, w, h, type = "map")
+graph = create_graph(param, w, h, type = "map")
 
 # ----------------
 # Draw populations
 
-param[:Zjn] = ones(param[:J], 1) .* 1e-6  # matrix of productivity (not 0 to avoid numerical glitches)
-param[:Lj] = ones(param[:J]) .* 1e-6  # matrix of population
+graph[:Zjn] = ones(graph[:J], 1) .* 1e-6  # matrix of productivity (not 0 to avoid numerical glitches)
+graph[:Lj] = ones(graph[:J]) .* 1e-6  # matrix of population
 
 Ni = find_node(graph, ceil(w/2), ceil(h/2))  # center
-param[:Zjn][Ni] = 1  # more productive node
-param[:Lj][Ni] = 1  # more productive node
+graph[:Zjn][Ni] = 1  # more productive node
+graph[:Lj][Ni] = 1  # more productive node
 
 Random.seed!(5)
 ncities = 20  # draw a number of random cities in space
 for i in 1:ncities-1
     newdraw = false
     while newdraw == false
-        j = round(Int, 1 + rand() * (param[:J] - 1))
-        if param[:Lj][j] <= 1e-6
+        j = round(Int, 1 + rand() * (graph[:J] - 1))
+        if graph[:Lj][j] <= 1e-6
             newdraw = true
-            param[:Lj][j] = 1
+            graph[:Lj][j] = 1
         end
     end
 end
 
-param[:hj] = param[:Hj] ./ param[:Lj]
-param[:hj][param[:Lj] .== 1e-6] .= 1  # catch errors in places with infinite housing per capita
+graph[:hj] = graph[:Hj] ./ graph[:Lj]
+graph[:hj][graph[:Lj] .== 1e-6] .= 1  # catch errors in places with infinite housing per capita
 
 
 # --------------
@@ -72,7 +72,7 @@ graph = apply_geography(graph, geography, alpha_up_i = 10, alpha_down_i = 10)
 plot_graph(graph, aspect_ratio = 3/4, 
             geography = geography, obstacles = true,
             mesh = true, mesh_transparency = 0.2, 
-            node_sizes = param[:Lj], node_shades = param[:Zjn], 
+            node_sizes = graph[:Lj], node_shades = graph[:Zjn], 
             node_color = :seismic,
             edge_min_thickness = 1, edge_max_thickness = 4)
 
@@ -86,8 +86,8 @@ plot_graph(graph, aspect_ratio = 3/4,
 # PLOT RESULTS
 # ============
 
-sizes = 2 .* results[:cj] .* (param[:Lj] .> 1e-6) / maximum(results[:cj])
-shades = results[:cj] .* (param[:Lj] .> 1e-6) / maximum(results[:cj])
+sizes = 2 .* results[:cj] .* (graph[:Lj] .> 1e-6) / maximum(results[:cj])
+shades = results[:cj] .* (graph[:Lj] .> 1e-6) / maximum(results[:cj])
 
 plot_graph(graph, results[:Ijk], aspect_ratio = 3/4,
            geography = geography, obstacles = true,
diff --git a/examples/paper_example01.jl b/examples/paper_example01.jl
index aa105fe..067ffd0 100644
--- a/examples/paper_example01.jl
+++ b/examples/paper_example01.jl
@@ -10,24 +10,24 @@ using Plots
 
 K = [1, 100]
 param = init_parameters()
-param, g = create_graph(param, 9, 9, type = "map")
+graph = create_graph(param, 9, 9, type = "map")
 
 #  Define fundamentals
-param[:Zjn] *= 0.1; # matrix of productivity
-Ni = find_node(g, 5, 5); # center
-param[:Zjn][Ni, :] .= 1; # more productive node
+graph[:Zjn] *= 0.1; # matrix of productivity
+Ni = find_node(graph, 5, 5); # center
+graph[:Zjn][Ni, :] .= 1; # more productive node
 
 # Plot the mesh with population 
-plot_graph(g, edges = false, mesh = true, node_sizes = param[:Lj])
+plot_graph(graph, edges = false, mesh = true, node_sizes = graph[:Lj])
 
 # Plot the mesh with productivity
-plot_graph(g, edges = false, mesh = true, node_sizes = param[:Zjn])
+plot_graph(graph, edges = false, mesh = true, node_sizes = graph[:Zjn])
 
 # Compute networks
 results = [] 
 for i in 1:length(K) 
     param[:K] = K[i]
-    push!(results, optimal_network(param, g))
+    push!(results, optimal_network(param, graph))
 end
 
 # Plot networks
@@ -35,16 +35,16 @@ end
 plots = [] # Initialize an empty array to hold the subplots
 for i in 1:length(K) 
     shades = sizes = results[i][:Cj] / maximum(results[i][:Cj])
-    p = plot_graph(g, results[i][:Ijk], node_sizes = sizes, node_shades = shades, node_sizes_scale = 40)
+    p = plot_graph(graph, results[i][:Ijk], node_sizes = sizes, node_shades = shades, node_sizes_scale = 40)
     title!(p, "(a) Transport Network (I_{jk})")
     push!(plots, p) 
-    p = plot_graph(g, results[1][:Qjkn][:, :, 1], edge_color = :brown, arrows = true, node_sizes = sizes, node_shades = shades, node_sizes_scale = 40)
+    p = plot_graph(graph, results[1][:Qjkn][:, :, 1], edge_color = :brown, arrows = true, node_sizes = sizes, node_shades = shades, node_sizes_scale = 40)
     title!(p, "(b) Shipping (Q_{jk})")
     push!(plots, p) 
-    p = plot_graph(g, results[i][:Ijk], map = results[i][:Pjn] / maximum(results[i][:Pjn]), node_sizes = sizes, node_shades = shades, node_sizes_scale = 40)
+    p = plot_graph(graph, results[i][:Ijk], map = results[i][:Pjn] / maximum(results[i][:Pjn]), node_sizes = sizes, node_shades = shades, node_sizes_scale = 40)
     title!(p, "(c) Prices (P_{j})")
     push!(plots, p) 
-    p = plot_graph(g, results[i][:Ijk], map = results[i][:cj] / maximum(results[i][:cj]), node_sizes = sizes, node_shades = shades, node_sizes_scale = 40)
+    p = plot_graph(graph, results[i][:Ijk], map = results[i][:cj] / maximum(results[i][:cj]), node_sizes = sizes, node_shades = shades, node_sizes_scale = 40)
     title!(p, "(d) Consumption (c_{j})")
     push!(plots, p) 
 end
diff --git a/examples/paper_example02.jl b/examples/paper_example02.jl
index e0f6ccb..5414a9f 100644
--- a/examples/paper_example02.jl
+++ b/examples/paper_example02.jl
@@ -11,42 +11,42 @@ height = 9
 nb_cities = 20
 
 param = init_parameters(K = 100, annealing = false)
-param, g = create_graph(param, width, height, type = "triangle") # case with random cities, one good
+graph = create_graph(param, width, height, type = "triangle") # case with random cities, one good
 
 # set fundamentals
 
 Random.seed!(5)
-minpop = minimum(param[:Lj])
-param[:Zjn] = fill(minpop, g[:J], 1)  # matrix of productivity
+minpop = minimum(graph[:Lj])
+graph[:Zjn] = fill(minpop, graph[:J], 1)  # matrix of productivity
 
-Ni = find_node(g, ceil(width/2), ceil(height/2)) # find center
-param[:Zjn][Ni] = 1 # more productive node
-param[:Lj][Ni] = 1  # more productive node
+Ni = find_node(graph, ceil(width/2), ceil(height/2)) # find center
+graph[:Zjn][Ni] = 1 # more productive node
+graph[:Lj][Ni] = 1  # more productive node
 
 for i in 1:nb_cities-1
     newdraw = false
     while newdraw == false
-        j = round(Int, 1 + rand() * (g[:J] - 1))
-        if param[:Lj][j] <= minpop
+        j = round(Int, 1 + rand() * (graph[:J] - 1))
+        if graph[:Lj][j] <= minpop
             newdraw = true
-            param[:Lj][j] = 1
+            graph[:Lj][j] = 1
         end
     end
 end
 
-param[:hj] = param[:Hj] ./ param[:Lj]
-param[:hj][param[:Lj] .== minpop] .= 1
+graph[:hj] = graph[:Hj] ./ graph[:Lj]
+graph[:hj][graph[:Lj] .== minpop] .= 1
 
 # Convex case
 results = Vector{Any}(undef, 3)
-results[1] = optimal_network(param, g)
+results[1] = optimal_network(param, graph)
 
 # Nonconvex - no annealing
 param[:gamma] = 2
-results[2] = optimal_network(param, g)
+results[2] = optimal_network(param, graph)
 
 # Nonconvex - annealing
-results[3] = annealing(param, g, results[2][:Ijk], perturbation_method = "rebranching")
+results[3] = annealing(graph, results[2][:Ijk], perturbation_method = "rebranching")
 
 welfare_increase = (results[3][:welfare] / results[2][:welfare]) ^ (1 / (param[:alpha] * (1 - param[:rho]))) # compute welfare increase in consumption equivalent
 
@@ -55,18 +55,18 @@ welfare_increase = (results[3][:welfare] / results[2][:welfare]) ^ (1 / (param[:
 
 plots = Vector{Any}(undef, 6) # Initialize an empty array to hold the subplots
 
-plots[1] = plot_graph(g, results[1][:Ijk], node_sizes = results[1][:Lj], node_shades = results[1][:Cj], node_sizes_scale = 40)
+plots[1] = plot_graph(graph, results[1][:Ijk], node_sizes = results[1][:Lj], node_shades = results[1][:Cj], node_sizes_scale = 40)
 title!(plots[1], "Convex Network (I_{jk})")
-plots[2] = plot_graph(g, results[1][:Qjkn][:, :, 1], edge_color = :brown, arrows = true, node_sizes = results[1][:Lj], node_shades = results[1][:Cj], node_sizes_scale = 40)
+plots[2] = plot_graph(graph, results[1][:Qjkn][:, :, 1], edge_color = :brown, arrows = true, node_sizes = results[1][:Lj], node_shades = results[1][:Cj], node_sizes_scale = 40)
 title!(plots[2], "Convex Shipping (Q_{jk})")
 
-plots[3] = plot_graph(g, results[2][:Ijk], node_sizes = results[2][:Lj], node_shades = results[2][:Cj], node_sizes_scale = 40)
+plots[3] = plot_graph(graph, results[2][:Ijk], node_sizes = results[2][:Lj], node_shades = results[2][:Cj], node_sizes_scale = 40)
 title!(plots[3], "Nonconvex Network (I_{jk})")
-plots[4] = plot_graph(g, results[2][:Qjkn][:, :, 1], edge_color = :brown, arrows = true, node_sizes = results[2][:Lj], node_shades = results[2][:Cj], node_sizes_scale = 40)
+plots[4] = plot_graph(graph, results[2][:Qjkn][:, :, 1], edge_color = :brown, arrows = true, node_sizes = results[2][:Lj], node_shades = results[2][:Cj], node_sizes_scale = 40)
 title!(plots[4], "Nonconvex Shipping (Q_{jk})")
 
 plots[5] = plots[3]
-plots[6] = plot_graph(g, results[3][:Ijk], node_sizes = results[3][:Lj], node_shades = results[3][:Cj], node_sizes_scale = 40)
+plots[6] = plot_graph(graph, results[3][:Ijk], node_sizes = results[3][:Lj], node_shades = results[3][:Cj], node_sizes_scale = 40)
 title!(plots[6], "Nonconvex Network (I_{jk})")
 annotate!(plots[6], [(0.5, 1.04, text("With Annealing. Welfare increase: $(round((welfare_increase-1)*100, digits = 2))%", :black, :center, 10))])
 
diff --git a/examples/paper_example03.jl b/examples/paper_example03.jl
index 4cd6b47..59e3c75 100644
--- a/examples/paper_example03.jl
+++ b/examples/paper_example03.jl
@@ -14,29 +14,29 @@ height = 9
 # Init and Solve network
 
 param = init_parameters(K = 10000, labor_mobility = true, annealing = false)
-param, g = create_graph(param, width, height, type = "triangle")
+graph = create_graph(param, width, height, type = "triangle")
 
 # set fundamentals
 
 param[:N] = 1 + ngoods
-param[:Zjn] = zeros(g[:J], param[:N]) # matrix of productivity
+graph[:Zjn] = zeros(graph[:J], param[:N]) # matrix of productivity
 
-param[:Zjn][:,1] .= 1 # the entire countryside produces the homogenenous good 1
+graph[:Zjn][:,1] .= 1 # the entire countryside produces the homogenenous good 1
 
 if ngoods > 0
-    Ni = find_node(g, 5, 5) # center
-    param[:Zjn][Ni, 2] = 1 # central node
-    param[:Zjn][Ni, 1] = 0 # central node
+    Ni = find_node(graph, 5, 5) # center
+    graph[:Zjn][Ni, 2] = 1 # central node
+    graph[:Zjn][Ni, 1] = 0 # central node
     
     Random.seed!(5)
     for i in 2:ngoods
         newdraw = false
         while newdraw == false
-            j = round(Int, 1 + rand() * (g[:J] - 1))
-            if param[:Zjn][j, 1] > 0
+            j = round(Int, 1 + rand() * (graph[:J] - 1))
+            if graph[:Zjn][j, 1] > 0
                 newdraw = true
-                param[:Zjn][j, i+1] = 1
-                param[:Zjn][j, 1] = 0
+                graph[:Zjn][j, i+1] = 1
+                graph[:Zjn][j, 1] = 0
             end
         end
     end
@@ -44,11 +44,11 @@ end
 
 # Convex case
 results = Array{Any}(undef, 2)
-results[1] = optimal_network(param, g)
+results[1] = optimal_network(param, graph)
 
 # Nonconvex
 param[:gamma] = 2
-results[2] = optimal_network(param, g, I0 = results[1][:Ijk])
+results[2] = optimal_network(graph, I0 = results[1][:Ijk])
 
 # Plot results
 
@@ -59,13 +59,13 @@ for j in 1:2
     # Initialize an empty array to hold the subplots
     plots = Vector{Any}(undef, (1 + param[:N])) 
     # Plot network
-    plots[1] = plot_graph(g, results[j][:Ijk], node_shades = results[j][:Lj], node_sizes = results[j][:Lj], node_sizes_scale = 40)
+    plots[1] = plot_graph(graph, results[j][:Ijk], node_shades = results[j][:Lj], node_sizes = results[j][:Lj], node_sizes_scale = 40)
     title!(plots[1], "(a) Transport Network")
     # Plot goods flows
     for i in 1:param[:N]
-        plots[i+1] = plot_graph(g, results[j][:Qjkn][:, :, i], edge_color = :brown, arrows = true, arrow_style = "thin", 
+        plots[i+1] = plot_graph(graph, results[j][:Qjkn][:, :, i], edge_color = :brown, arrows = true, arrow_style = "thin", 
                                 node_sizes = results[j][:Yjn][:, i], node_sizes_scale = 40,
-                                node_shades = param[:Zjn][:, i])
+                                node_shades = graph[:Zjn][:, i])
         title!(plots[i+1], string('(', Char(96 + i + 1), ')', " Flows Good ", i))
     end
     # Combine plots
diff --git a/examples/paper_example04.jl b/examples/paper_example04.jl
index 8685a77..877b042 100644
--- a/examples/paper_example04.jl
+++ b/examples/paper_example04.jl
@@ -10,31 +10,31 @@ width, height = 13, 13
 param, g0 = create_graph(param, width, height, type = "map")
 
 # Set fundamentals
-param[:Zjn] = ones(param[:J], 1) .* 1e-7  # matrix of productivity (not 0 to avoid numerical glitches)
-param[:Lj] = ones(param[:J]) .* 1e-7  # matrix of population
+g0[:Zjn] = ones(g0[:J], 1) .* 1e-7  # matrix of productivity (not 0 to avoid numerical glitches)
+g0[:Lj] = ones(g0[:J]) .* 1e-7  # matrix of population
 
 Ni = find_node(g0, ceil(width/2), ceil(height/2)) # center
-param[:Zjn][Ni] = 1 # more productive node
-param[:Lj][Ni] = 1 # more productive node
+g0[:Zjn][Ni] = 1 # more productive node
+g0[:Lj][Ni] = 1 # more productive node
 
 Random.seed!(10) # use 5, 8, 10 or 11
 nb_cities = 20 # draw a number of random cities in space
 for i in 1:nb_cities-1
     newdraw = false
     while newdraw == false
-        j = round(Int, 1 + rand() * (param[:J] - 1))
-        if param[:Lj][j] <= 1e-7
+        j = round(Int, 1 + rand() * (g0[:J] - 1))
+        if g0[:Lj][j] <= 1e-7
             newdraw = true
-            param[:Lj][j] = 1
+            g0[:Lj][j] = 1
         end
     end
 end
 
-param[:hj] = param[:Hj] ./ param[:Lj]
-param[:hj][param[:Lj] .== 1e-7] .= 1  # catch errors in places with infinite housing per capita
+g0[:hj] = g0[:Hj] ./ g0[:Lj]
+g0[:hj][g0[:Lj] .== 1e-7] .= 1  # catch errors in places with infinite housing per capita
 
 # Draw geography
-z = zeros(param[:J]) # altitude of each node
+z = zeros(g0[:J]) # altitude of each node
 geographies = Dict()
 geographies[:blank] = (z = z, obstacles = nothing)
 
@@ -116,8 +116,8 @@ for s in simulation
     plots[i] = plot_graph(g0, results[Symbol(s)][:Ijk], 
                           geography = geographies[Symbol(s)], obstacles = obstacles[i] == "on",
                           mesh = true, mesh_transparency = 0.2,
-                          node_sizes = results[Symbol(s)][:cj] .* (param[:Lj] .> 1e-7), 
-                          node_shades = param[:Zjn], node_color = :seismic,
+                          node_sizes = results[Symbol(s)][:cj], 
+                          node_shades = g0[:Zjn], node_color = :seismic,
                           edge_min_thickness = 1, edge_max_thickness = 4)
     title!(plots[i], "Geography $(s)")
 end

From 60bbe282fd72213c5ebf721b04e52670cb4326d2 Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Sat, 14 Sep 2024 12:44:28 +0200
Subject: [PATCH 16/32] Minor fixes.

---
 examples/example02.jl       | 2 +-
 examples/paper_example02.jl | 2 +-
 examples/paper_example03.jl | 5 +++--
 examples/paper_example04.jl | 4 ++--
 4 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/examples/example02.jl b/examples/example02.jl
index ca0eed0..7ff1a14 100644
--- a/examples/example02.jl
+++ b/examples/example02.jl
@@ -50,7 +50,7 @@ graph[:Lj][graph[:Lj] .== 0] .= 1e-6
     # second, in the nonconvex case (gamma>beta)
     param[:gamma] = 2 # change only gamma, keep other parameters
     res[2] = optimal_network(param, graph) # optimize by iterating on FOCs
-    res[3] = annealing(graph, res[2][:Ijk]) # improve with annealing, starting from previous result
+    res[3] = annealing(param, graph, res[2][:Ijk]) # improve with annealing, starting from previous result
 end
 
 # ==========
diff --git a/examples/paper_example02.jl b/examples/paper_example02.jl
index 5414a9f..4e135d3 100644
--- a/examples/paper_example02.jl
+++ b/examples/paper_example02.jl
@@ -46,7 +46,7 @@ param[:gamma] = 2
 results[2] = optimal_network(param, graph)
 
 # Nonconvex - annealing
-results[3] = annealing(graph, results[2][:Ijk], perturbation_method = "rebranching")
+results[3] = annealing(param, graph, results[2][:Ijk], perturbation_method = "rebranching")
 
 welfare_increase = (results[3][:welfare] / results[2][:welfare]) ^ (1 / (param[:alpha] * (1 - param[:rho]))) # compute welfare increase in consumption equivalent
 
diff --git a/examples/paper_example03.jl b/examples/paper_example03.jl
index 59e3c75..23a323f 100644
--- a/examples/paper_example03.jl
+++ b/examples/paper_example03.jl
@@ -48,7 +48,7 @@ results[1] = optimal_network(param, graph)
 
 # Nonconvex
 param[:gamma] = 2
-results[2] = optimal_network(graph, I0 = results[1][:Ijk])
+results[2] = optimal_network(param, graph, I0 = results[1][:Ijk])
 
 # Plot results
 
@@ -59,7 +59,8 @@ for j in 1:2
     # Initialize an empty array to hold the subplots
     plots = Vector{Any}(undef, (1 + param[:N])) 
     # Plot network
-    plots[1] = plot_graph(graph, results[j][:Ijk], node_shades = results[j][:Lj], node_sizes = results[j][:Lj], node_sizes_scale = 40)
+    plots[1] = plot_graph(graph, results[j][:Ijk], node_shades = results[j][:Lj], 
+                          node_sizes = results[j][:Lj], node_sizes_scale = 40)
     title!(plots[1], "(a) Transport Network")
     # Plot goods flows
     for i in 1:param[:N]
diff --git a/examples/paper_example04.jl b/examples/paper_example04.jl
index 877b042..a597139 100644
--- a/examples/paper_example04.jl
+++ b/examples/paper_example04.jl
@@ -3,11 +3,11 @@ using Random
 using Plots
 
 # Initialize parameters
-param = init_parameters(K = 100, labor_mobility = false)
+param = init_parameters(K = 100, rho = 0, labor_mobility = false)
 width, height = 13, 13
 
 # Create graph
-param, g0 = create_graph(param, width, height, type = "map")
+g0 = create_graph(param, width, height, type = "map")
 
 # Set fundamentals
 g0[:Zjn] = ones(g0[:J], 1) .* 1e-7  # matrix of productivity (not 0 to avoid numerical glitches)

From 96be68715a41d67ca9e04c266c1d524c6d4621a5 Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Sat, 14 Sep 2024 12:45:48 +0200
Subject: [PATCH 17/32] Revert node deletion.

---
 src/main/apply_geography.jl | 123 ++++++++++++++++++++----------------
 src/main/plot_graph.jl      |  35 +++++++---
 2 files changed, 95 insertions(+), 63 deletions(-)

diff --git a/src/main/apply_geography.jl b/src/main/apply_geography.jl
index 6f943fa..8a418b0 100644
--- a/src/main/apply_geography.jl
+++ b/src/main/apply_geography.jl
@@ -90,7 +90,7 @@ function apply_geography(graph, geography; kwargs...)
     end
 
     # Remove edges where geographical barriers are (rivers)
-    remove_edge = false
+    # remove_edge = false
     if obstacles !== nothing
 
         # Store initial delta matrics (avoid double counting)
@@ -102,9 +102,9 @@ function apply_geography(graph, geography; kwargs...)
         across_obstacle = falses(graph.J, graph.J)
         along_obstacle = falses(graph.J, graph.J)
         remove_edge = isinf(op.across_obstacle_delta_i) || isinf(op.across_obstacle_delta_tau) # remove the edge
-        if remove_edge
-            rm_nodes = falses(graph.J)
-        end
+        # if remove_edge
+        #     rm_nodes = falses(graph.J)
+        # end
 
         for i in 1:graph.J
             neighbors = graph.nodes[i]
@@ -135,16 +135,16 @@ function apply_geography(graph, geography; kwargs...)
                                 rmi = findfirst(==(i), nodes_new[j])
                                 if rmi !== nothing
                                     deleteat!(nodes_new[j], rmi)
-                                    if isempty(nodes_new[j])
-                                        rm_nodes[j] = true # deleteat!(nodes_new, j)
-                                    end
+                                    # if isempty(nodes_new[j])
+                                    #     rm_nodes[j] = true # deleteat!(nodes_new, j)
+                                    # end
                                 end
                                 rmj = findfirst(==(j), nodes_new[i])
                                 if rmj !== nothing
                                     deleteat!(nodes_new[i], rmj)
-                                    if isempty(nodes_new[i])
-                                        rm_nodes[i] = true # deleteat!(nodes_new, i)
-                                    end
+                                    # if isempty(nodes_new[i])
+                                    #     rm_nodes[i] = true # deleteat!(nodes_new, i)
+                                    # end
                                 end
                                 
                                 adjacency_new[i, j] = false
@@ -174,16 +174,16 @@ function apply_geography(graph, geography; kwargs...)
                 rmio = findfirst(==(io), nodes_new[jo])
                 if rmio !== nothing
                     deleteat!(nodes_new[jo], rmio)
-                    if isempty(nodes_new[jo])
-                        rm_nodes[jo] = true # deleteat!(nodes_new, jo)
-                    end
+                    # if isempty(nodes_new[jo])
+                    #     rm_nodes[jo] = true # deleteat!(nodes_new, jo)
+                    # end
                 end
                 rmjo = findfirst(==(jo), nodes_new[io])
                 if rmjo !== nothing
                     deleteat!(nodes_new[io], rmjo)
-                    if isempty(nodes_new[io])
-                        rm_nodes[io] = true # deleteat!(nodes_new, io)
-                    end
+                    # if isempty(nodes_new[io])
+                    #     rm_nodes[io] = true # deleteat!(nodes_new, io)
+                    # end
                 end
                 adjacency_new[io, jo] = false
                 adjacency_new[jo, io] = false
@@ -206,47 +206,60 @@ function apply_geography(graph, geography; kwargs...)
     # Creating new object
     graph_new = namedtuple_to_dict(graph)
 
-    if remove_edge && any(rm_nodes)
-        keep = findall(.!rm_nodes)
-        graph_new[:delta_i] = delta_i_new[keep, keep]
-        graph_new[:delta_tau] = delta_tau_new[keep, keep]
-        graph_new[:nodes] = nodes_new[keep]
-        graph_new[:adjacency] = adjacency_new[keep, keep]
-        # make sure that the degrees of freedom of the updated graph match the # of links
-        graph_new[:ndeg] = sum(tril(graph_new[:adjacency]))
-        graph_new[:across_obstacle] = across_obstacle[keep, keep]
-        graph_new[:along_obstacle] = along_obstacle[keep, keep]
+    # if remove_edge && any(rm_nodes)
+    #     keep = findall(.!rm_nodes)
+    #     graph_new[:x_orig] = graph[:x]
+    #     graph_new[:x] = graph[:x][keep]
+    #     graph_new[:y_orig] = graph[:y]
+    #     graph_new[:y] = graph[:y][keep]
+    #     graph_new[:J] = length(graph_new[:x])
+    #     graph_new[:rm_nodes] = rm_nodes
+    #     graph_new[:delta_i] = delta_i_new[keep, keep]
+    #     graph_new[:delta_tau] = delta_tau_new[keep, keep]
+    #     nodes_new = nodes_new[keep]
+    #     for i in keep
+    #         for k in 1:length(nodes_new)
+    #             node_k = filter(x -> x != i, nodes_new[k])
+    #             nodes_new[k] = ifelse.(node_k .> i, node_k .- 1, node_k) # reindex nodes k > i to k-1
+    #         end
+    #     end
+    #     graph_new[:nodes] = nodes_new
+    #     graph_new[:adjacency] = adjacency_new[keep, keep]
+    #     # make sure that the degrees of freedom of the updated graph match the # of links
+    #     graph_new[:ndeg] = sum(tril(graph_new[:adjacency]))
+    #     graph_new[:across_obstacle] = across_obstacle[keep, keep]
+    #     graph_new[:along_obstacle] = along_obstacle[keep, keep]
 
-        if haskey(graph, :Lj)
-            graph_new[:Lj] = graph[:Lj][keep]
-        end
-        if haskey(graph, :Hj)
-            graph_new[:Hj] = graph[:Hj][keep]
-        end
-        if haskey(graph, :hj)
-            graph_new[:hj] = graph_new[:Hj] ./ graph_new[:Lj]
-        end
-        if haskey(graph, :omegaj)
-            graph_new[:omegaj] = graph[:omegaj][keep]
-        end
-        if haskey(graph, :Zjn)
-            graph_new[:Zjn] = graph[:Zjn][keep, :]
-        end
-        if haskey(graph, :region)
-            graph_new[:region] = graph[:region][keep]
-        end
-    else
-        graph_new[:delta_i] = delta_i_new
-        graph_new[:delta_tau] = delta_tau_new
-        if obstacles !== nothing
-            graph_new[:nodes] = nodes_new
-            graph_new[:adjacency] = adjacency_new
-            # make sure that the degrees of freedom of the updated graph match the # of links
-            graph_new[:ndeg] = sum(tril(adjacency_new))
-            graph_new[:across_obstacle] = across_obstacle
-            graph_new[:along_obstacle] = along_obstacle
-        end
+    #     if haskey(graph, :Lj)
+    #         graph_new[:Lj] = graph[:Lj][keep]
+    #     end
+    #     if haskey(graph, :Hj)
+    #         graph_new[:Hj] = graph[:Hj][keep]
+    #     end
+    #     if haskey(graph, :hj)
+    #         graph_new[:hj] = graph_new[:Hj] ./ graph_new[:Lj]
+    #     end
+    #     if haskey(graph, :omegaj)
+    #         graph_new[:omegaj] = graph[:omegaj][keep]
+    #     end
+    #     if haskey(graph, :Zjn)
+    #         graph_new[:Zjn] = graph[:Zjn][keep, :]
+    #     end
+    #     if haskey(graph, :region)
+    #         graph_new[:region] = graph[:region][keep]
+    #     end
+    # else
+    graph_new[:delta_i] = delta_i_new
+    graph_new[:delta_tau] = delta_tau_new
+    if obstacles !== nothing
+        graph_new[:nodes] = nodes_new
+        graph_new[:adjacency] = adjacency_new
+        # make sure that the degrees of freedom of the updated graph match the # of links
+        graph_new[:ndeg] = sum(tril(adjacency_new))
+        graph_new[:across_obstacle] = across_obstacle
+        graph_new[:along_obstacle] = along_obstacle
     end
+    # end
 
     return graph_new
 end
diff --git a/src/main/plot_graph.jl b/src/main/plot_graph.jl
index 2802307..2152306 100644
--- a/src/main/plot_graph.jl
+++ b/src/main/plot_graph.jl
@@ -88,11 +88,12 @@ function plot_graph(graph, edges = nothing; kwargs...)
 
     # PLOT COLORMAP
     if op.map !== nothing || op.geography !== nothing
-        if op.geography !== nothing
-            vec_map = op.geography[:z]
-        else
-            vec_map = vec(op.map)
-        end
+        vec_map = op.geography !== nothing ? op.geography[:z] : vec(op.map)
+        # if length(vec_map) != length(vec_x) && haskey(graph, :rm_nodes)
+        #     if length(vec_map) == length(graph.rm_nodes)
+        #         vec_map = vec_map[.!graph.rm_nodes]
+        #     end
+        # end
         # Interpolate map onto grid
         xmap = range(minimum(vec_x), stop=maximum(vec_x), length=2*length(vec_x))
         ymap = range(minimum(vec_y), stop=maximum(vec_y), length=2*length(vec_y))
@@ -120,6 +121,23 @@ function plot_graph(graph, edges = nothing; kwargs...)
     # PLOT OBSTACLES
     if op.obstacles && length(op.geography[:obstacles]) > 0
         obstacles = op.geography[:obstacles]
+
+        # if haskey(graph, :x_orig)
+        #     vec_x_orig = graph.x_orig
+        #     vec_y_orig = graph.y_orig
+
+        #     for i in 1:size(obstacles, 1)
+        #         x1 = vec_x_orig[obstacles[i, 1]]
+        #         y1 = vec_y_orig[obstacles[i, 1]]
+        #         x2 = vec_x_orig[obstacles[i, 2]]
+        #         y2 = vec_y_orig[obstacles[i, 2]]
+    
+        #         plot!(pl, [x1, x2], [y1, y2], 
+        #               linecolor = op.obstacle_color, 
+        #               linewidth = op.obstacle_thickness, 
+        #               linealpha = 1, label = nothing)
+        #     end
+        # else
         for i in 1:size(obstacles, 1)
             x1 = vec_x[obstacles[i, 1]]
             y1 = vec_y[obstacles[i, 1]]
@@ -127,10 +145,11 @@ function plot_graph(graph, edges = nothing; kwargs...)
             y2 = vec_y[obstacles[i, 2]]
 
             plot!(pl, [x1, x2], [y1, y2], 
-                  linecolor = op.obstacle_color, 
-                  linewidth = op.obstacle_thickness, 
-                  linealpha = 1, label = nothing)
+                linecolor = op.obstacle_color, 
+                linewidth = op.obstacle_thickness, 
+                linealpha = 1, label = nothing)
         end
+        # end
     end
 
     # PLOT MESH

From 539ac9b2d5e7d3e6c021a5f07695d764b7578745 Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Sat, 14 Sep 2024 12:53:02 +0200
Subject: [PATCH 18/32] Update examples.

---
 src/main/apply_geography.jl | 2 +-
 src/main/create_graph.jl    | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/main/apply_geography.jl b/src/main/apply_geography.jl
index 8a418b0..28d52a0 100644
--- a/src/main/apply_geography.jl
+++ b/src/main/apply_geography.jl
@@ -32,7 +32,7 @@ and similarly for graph traversal costs `delta_tau`.
 
 # Examples
 ```julia
-param, graph = create_graph(init_parameters())
+graph = create_graph(init_parameters())
 geography = (z = 10*(rand(graph[:J]) .> 0.95), obstacles = [1 15; 70 72])
 updated_graph = apply_geography(graph, geography)
 
diff --git a/src/main/create_graph.jl b/src/main/create_graph.jl
index 4e10d13..fbd9975 100644
--- a/src/main/create_graph.jl
+++ b/src/main/create_graph.jl
@@ -25,7 +25,7 @@ Initialize the underlying graph, population and productivity parameters.
 
 # Examples
 ```julia
-graph = create_graph()
+graph = create_graph(init_parameters())
 ```
 """
 function create_graph(param, w = 11, h = 11; type = "map", kwargs...)

From 9c7b06a468be39308df49391b2442ed95bed07d7 Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Sat, 14 Sep 2024 13:53:25 +0200
Subject: [PATCH 19/32] Add/remove docstrings.

---
 src/main/helper.jl | 73 +++++++++++++++++++++++++++++++++++++---------
 1 file changed, 59 insertions(+), 14 deletions(-)

diff --git a/src/main/helper.jl b/src/main/helper.jl
index 89464f3..45c9254 100644
--- a/src/main/helper.jl
+++ b/src/main/helper.jl
@@ -1,4 +1,18 @@
 
+"""
+    dict_to_namedtuple(dict)
+
+Convert a dictionary to a NamedTuple.
+
+If the input is already a NamedTuple, it is returned unchanged.
+Otherwise, it creates a new NamedTuple from the dictionary's keys and values.
+
+# Arguments
+- `dict`: A dictionary or NamedTuple to be converted.
+
+# Returns
+- A NamedTuple equivalent to the input dictionary.
+"""
 function dict_to_namedtuple(dict)
     if dict isa NamedTuple
         return dict
@@ -7,6 +21,20 @@ function dict_to_namedtuple(dict)
     end
 end
 
+"""
+    namedtuple_to_dict(namedtuple)
+
+Convert a NamedTuple to a dictionary.
+
+If the input is already a dictionary, it is returned unchanged.
+Otherwise, it creates a new dictionary from the NamedTuple's pairs.
+
+# Arguments
+- `namedtuple`: A NamedTuple or dictionary to be converted.
+
+# Returns
+- A dictionary equivalent to the input NamedTuple.
+"""
 function namedtuple_to_dict(namedtuple)
     if namedtuple isa Dict
         return namedtuple
@@ -31,11 +59,7 @@ function gen_network_flows(Qin, graph, N)
     return Qjkn
 end
 
-"""
-    kappa_extract(graph, kappa)
-
-Description: auxiliary function that converts kappa_jk into kappa_i
-"""
+# Description: auxiliary function that converts kappa_jk into kappa_i
 function kappa_extract(graph, kappa)
     kappa_ex = zeros(graph.ndeg)
     nodes = graph.nodes
@@ -92,6 +116,22 @@ function rowmult(A, v)
     return r
 end 
 
+"""
+    represent_edges(graph)
+
+Creates a NamedTuple providing detailed representation of the graph edges. 
+
+# Arguments
+- `graph`: NamedTuple that contains the underlying graph (created by `dict_to_namedtuple(create_graph())`)
+
+# Returns
+- A NamedTuple with the following fields:
+  - `A`: J x ndeg matrix where each column represents an edge. The value is 1 if the edge starts at node J, -1 if it ends at node J, and 0 otherwise.
+  - `Apos`: J x ndeg matrix where each column represents an edge. The value is the positive part of the edge flow.
+  - `Aneg`: J x ndeg matrix where each column represents an edge. The value is the negative part of the edge flow.
+  - `edge_start`: J x ndeg matrix where each column represents an edge. The value is the starting node of the edge.
+  - `edge_end`: J x ndeg matrix where each column represents an edge. The value is the ending node of the edge.
+"""
 function represent_edges(graph)
     # Create matrix A
     # Note: matrix A is of dimension J*ndeg and takes value 1 when node J is
@@ -127,13 +167,18 @@ end
 Creates the auxdata structure that contains all the auxiliary parameters for estimation
 
 # Arguments
-- `param`: structure that contains the model's parameters
-- `graph`: structure that contains the underlying graph (created by create_graph function)
-- `edges`: structure that contains the edges of the graph (created by represent_edges function)
-- `I`: provides the current JxJ symmetric matrix of infrastructure investment
+- `param`: NamedTuple that contains the model's parameters (created by `dict_to_namedtuple(init_parameters())`)
+- `graph`: NamedTuple that contains the underlying graph (created by `dict_to_namedtuple(create_graph())`)
+- `edges`: NamedTuple that contains the edges of the graph (created by `represent_edges()`)
+- `I`: J x J symmetric matrix of current infrastructure (investments)
 
-# Output
-- `auxdata`: structure auxdata to be used by IPOPT bundle.
+# Returns
+- A NamedTuple with the following fields:
+  - `param`: The input parameter NamedTuple.
+  - `graph`: The input graph NamedTuple.
+  - `edges`: The edges of the graph.
+  - `kappa`: The kappa matrix: I^gamma / delta_tau
+  - `kappa_ex`: The extracted kappa values (ndeg x 1)
 """
 function create_auxdata(param, graph, edges, I)
     # Make named tuples
@@ -161,13 +206,13 @@ end
 """
     get_model(auxdata)
 
-Construct the appropriate model based on the parameters and auxiliary data.
+Construct the appropriate JuMP model based on the parameters and auxiliary data.
 
 # Arguments
-- `auxdata`: Auxiliary data required for constructing the model.
+- `auxdata`: Auxiliary data required for constructing the model (created by `create_auxdata()`).
 
 # Returns
-- `model`: The constructed model.
+- `model`: The constructed JuMP model.
 - `recover_allocation`: A function to recover the allocation from the model solution.
 """
 function get_model(auxdata)

From 732045a260a89920d4da9e0c7c0dcc616ee7988d Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Sat, 14 Sep 2024 13:53:54 +0200
Subject: [PATCH 20/32] Remove some docstrings.

---
 src/main/annealing.jl | 15 +++++----------
 1 file changed, 5 insertions(+), 10 deletions(-)

diff --git a/src/main/annealing.jl b/src/main/annealing.jl
index 103969e..ae9949c 100644
--- a/src/main/annealing.jl
+++ b/src/main/annealing.jl
@@ -514,12 +514,10 @@ function shake_network(param, graph, I0, results, options)
 end
 
 
-"""
-    rebranch_network(param, graph, I0, results, options)
 
-This function implements the rebranching algorithm described in the paper.
-Links are reshuffled everywhere so that each node is better connected to its best neighbor
-(those with the lowest price index for traded goods, i.e., more central places in the trading network).
+# This function implements the rebranching algorithm described in the paper.
+# Links are reshuffled everywhere so that each node is better connected to its best neighbor
+# (those with the lowest price index for traded goods, i.e., more central places in the trading network).
 """
 function rebranch_network(param, graph, I0, results, options)
     J = graph.J
@@ -558,12 +556,9 @@ function rebranch_network(param, graph, I0, results, options)
 end
 
 
-"""
-    random_rebranch_network(param, graph, I0, results, options)
 
-This function does the same as `rebranch_network` except that only a few nodes
-(#num_random_perturbations) are selected for rebranching at random.
-"""
+# This function does the same as `rebranch_network` except that only a few nodes
+# (#num_random_perturbations) are selected for rebranching at random.
 function random_rebranch_network(param, graph, I0, results, options)
     J = graph.J
 

From 987c132cd29004009047e16b9f252c5976fd35c7 Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Sat, 14 Sep 2024 14:00:58 +0200
Subject: [PATCH 21/32] Specify order of functions.

---
 docs/src/api.md | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/docs/src/api.md b/docs/src/api.md
index 77b6bda..a422025 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -22,6 +22,18 @@ Private = false
 Order = [:function]
 ```
 
+```@docs
+init_parameters
+create_graph
+apply_geography
+optimal_network
+annealing
+plot_graph
+find_node
+add_node
+remove_node
+```
+
 # Documented Internal Functions
 
 ```@autodocs

From c85d018df56352ee2c625b3179a318d90edafc56 Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Sat, 14 Sep 2024 14:13:15 +0200
Subject: [PATCH 22/32] Add NEWS.md.

---
 NEWS.md | 13 +++++++++++++
 1 file changed, 13 insertions(+)
 create mode 100644 NEWS.md

diff --git a/NEWS.md b/NEWS.md
new file mode 100644
index 0000000..4ef8d56
--- /dev/null
+++ b/NEWS.md
@@ -0,0 +1,13 @@
+# 0.1.6
+## Breaking Changes
+* All spatial parameters, including `Lj`, `Lr`, `Hj`, `hj`, `Zjn`, `omegaj`, `omegar`, and `region` are now stored in the `graph` structure created by `create_graph()`. `create_graph()` therefore only returns the `graph` structure, instead of both the (updated) parameters and the graph. Converesely, `init_parameters()` only contains parameters that are independent of the particular geography defined by `create_graph()`.
+
+## Improvements
+* Minor improvements to Simulated Annealing.
+* Better spline options for plotting frictions surface (geography).
+* More faithful translation of `apply_geography()`. 
+
+
+# 0.1.5
+
+* Removed the MATLAB toolbox and corresponding documentation (PDF files) from the repo to decrease size. A new repo was created for the MATLAB toolbox at [SebKrantz/OptimalTransportNetworkToolbox](https://github.com/SebKrantz/OptimalTransportNetworkToolbox). This repo and especially the `docs` folder continue to be very useful for Julia users, but are no longer part of the Julia library. 
\ No newline at end of file

From 3c2692a5d1d3062f82302804785f6c6e519848d8 Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Sat, 14 Sep 2024 14:26:31 +0200
Subject: [PATCH 23/32] Docs minors.

---
 docs/src/api.md | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index a422025..5b2b1cb 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -16,12 +16,6 @@ Order   = [:function]
 
 # Exported Functions
 
-```@autodocs
-Modules = [OptimalTransportNetworks]
-Private = false
-Order = [:function]
-```
-
 ```@docs
 init_parameters
 create_graph

From 67a7314dd7dd1a868f066afdf9580f03b524312c Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Sat, 14 Sep 2024 14:42:18 +0200
Subject: [PATCH 24/32] Fixes.

---
 src/main/annealing.jl | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/main/annealing.jl b/src/main/annealing.jl
index ae9949c..eb4068d 100644
--- a/src/main/annealing.jl
+++ b/src/main/annealing.jl
@@ -518,7 +518,6 @@ end
 # This function implements the rebranching algorithm described in the paper.
 # Links are reshuffled everywhere so that each node is better connected to its best neighbor
 # (those with the lowest price index for traded goods, i.e., more central places in the trading network).
-"""
 function rebranch_network(param, graph, I0, results, options)
     J = graph.J
 

From 9371210478e03ccb1a7f7ad347c5d8df84a7c6f5 Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Sat, 14 Sep 2024 14:43:43 +0200
Subject: [PATCH 25/32] Documentation minors.

---
 src/main/create_graph.jl | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/src/main/create_graph.jl b/src/main/create_graph.jl
index fbd9975..ffaa9e9 100644
--- a/src/main/create_graph.jl
+++ b/src/main/create_graph.jl
@@ -253,10 +253,8 @@ Creates a triangular graph structure with width `w` and height `h`
 horizontal and along the two diagonals)
 
 # Arguments
-- `w`: Width of graph (i.e. the max number of nodes along horizontal dimension), 
-must be an integer
-- `h`: Height of graph (i.e. the max number of nodes along vertical dimension),  
-must be an odd integer
+- `w`: Width of graph (i.e. the max number of nodes along horizontal dimension), must be an integer
+- `h`: Height of graph (i.e. the max number of nodes along vertical dimension), must be an odd integer
 """
 function create_triangle(w, h)
     if h % 2 == 0
@@ -361,10 +359,8 @@ with width w and height h (nodes have 4 neighbors in total, along
 horizontal and vertical dimensions, NOT diagonals)
 
 # Arguments
-- `w`: width of graph (ie. the number of nodes along horizontal
-dimension), must be an integer 
-- `h`: height of graph (ie. the number of nodes along vertical
-dimension), must be an integer
+- `w`: width of graph (ie. the number of nodes along horizontal dimension), must be an integer 
+- `h`: height of graph (ie. the number of nodes along vertical dimension), must be an integer
 """
 function create_square(w, h)
     J = w * h

From 868a6ae105366d135c42b381e285d1c32893c2a7 Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Sat, 14 Sep 2024 14:45:36 +0200
Subject: [PATCH 26/32] Also fixing order.

---
 docs/src/api.md | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index 5b2b1cb..33bcb40 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -30,9 +30,22 @@ remove_node
 
 # Documented Internal Functions
 
-```@autodocs
+```@docs
+dict_to_namedtuple
+namedtuple_to_dict
+create_map
+create_square
+create_triangle
+create_custom
+represent_edges
+create_auxdata
+get_model
+```
+
+<!-- ```@autodocs
+<!-- ```@autodocs
 Modules = [OptimalTransportNetworks]
 Public = false
 Private = true
 Order = [:function]
-```
\ No newline at end of file
+``` -->
\ No newline at end of file

From eab8875a0a0981f613eb9e1abd9d43572ef1b5cd Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Sat, 14 Sep 2024 14:58:44 +0200
Subject: [PATCH 27/32] Add qualified names.

---
 docs/src/api.md | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index 33bcb40..e7c8d80 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -31,15 +31,15 @@ remove_node
 # Documented Internal Functions
 
 ```@docs
-dict_to_namedtuple
-namedtuple_to_dict
-create_map
-create_square
-create_triangle
-create_custom
-represent_edges
-create_auxdata
-get_model
+OptimalTransportNetworks.dict_to_namedtuple
+OptimalTransportNetworks.namedtuple_to_dict
+OptimalTransportNetworks.create_map
+OptimalTransportNetworks.create_square
+OptimalTransportNetworks.create_triangle
+OptimalTransportNetworks.create_custom
+OptimalTransportNetworks.represent_edges
+OptimalTransportNetworks.create_auxdata
+OptimalTransportNetworks.get_model
 ```
 
 <!-- ```@autodocs

From 044b30a11ccb1b81b690727489c88eb1fbf8e2ed Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Sat, 14 Sep 2024 15:23:19 +0200
Subject: [PATCH 28/32] Minors.

---
 docs/src/api.md | 1 -
 1 file changed, 1 deletion(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index e7c8d80..ddf0cc3 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -42,7 +42,6 @@ OptimalTransportNetworks.create_auxdata
 OptimalTransportNetworks.get_model
 ```
 
-<!-- ```@autodocs
 <!-- ```@autodocs
 Modules = [OptimalTransportNetworks]
 Public = false

From fb53949b682e922693961da7dcc8838c8a78f106 Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Sat, 14 Sep 2024 15:42:33 +0200
Subject: [PATCH 29/32] Removing comments.

---
 docs/src/api.md | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/docs/src/api.md b/docs/src/api.md
index ddf0cc3..322a57b 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -41,10 +41,3 @@ OptimalTransportNetworks.represent_edges
 OptimalTransportNetworks.create_auxdata
 OptimalTransportNetworks.get_model
 ```
-
-<!-- ```@autodocs
-Modules = [OptimalTransportNetworks]
-Public = false
-Private = true
-Order = [:function]
-``` -->
\ No newline at end of file

From 57c77d665b19c9f2f030d5fa962522624036053e Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Sat, 14 Sep 2024 15:57:51 +0200
Subject: [PATCH 30/32] Clarify API changes.

---
 src/OptimalTransportNetworks.jl | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/OptimalTransportNetworks.jl b/src/OptimalTransportNetworks.jl
index e39447f..3439c0c 100644
--- a/src/OptimalTransportNetworks.jl
+++ b/src/OptimalTransportNetworks.jl
@@ -22,7 +22,8 @@ Fajgelbaum, P. D., & Schaal, E. (2020). Optimal transport networks in spatial eq
 
 The library is based on the JuMP modeling framework for mathematical optimization in Julia, and roughly follows the 
 MATLAB OptimalTransportNetworkToolbox (v1.0.4b) provided by the authors. Compared to MATLAB, the Julia library presents
-a simplified interface and only exports key functions, while retaining full flexibility. 
+a simplified interface. Notably, the `graph` object contains both the graph structure and all data parameterizing 
+the network, whereas the `param` object only contains (non-spatial) model parameters.
 
 # Exported Functions
 

From e9b02eb98b64fda02a48fb0d65ba089c5496887e Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Sat, 14 Sep 2024 15:59:36 +0200
Subject: [PATCH 31/32] Better indentation.

---
 README.md         | 4 ++--
 docs/src/index.md | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index c693e86..7268c60 100644
--- a/README.md
+++ b/README.md
@@ -42,7 +42,7 @@ import MathOptInterface as MOI
 import MathOptSymbolicAD
 
 param[:model_attr] = Dict(:backend => (MOI.AutomaticDifferentiationBackend(), 
-                                        MathOptSymbolicAD.DefaultBackend())) 
+                                       MathOptSymbolicAD.DefaultBackend())) 
                                 # Or:  MathOptSymbolicAD.ThreadedBackend()
 ```
 
@@ -50,7 +50,7 @@ param[:model_attr] = Dict(:backend => (MOI.AutomaticDifferentiationBackend(),
 
 ```julia
 param[:optimizer_attr] = Dict(:hsllib => "/usr/local/lib/libhsl.dylib", # Adjust path
-                                :linear_solver => "ma57") # Use ma57, ma86 or ma97
+                              :linear_solver => "ma57") # Use ma57, ma86 or ma97
 ```
 
 The [Ipopt.jl README](https://github.com/jump-dev/Ipopt.jl?tab=readme-ov-file#linear-solvers) suggests to use the larger LibHSL package for which there exists a Julia module and proceed similarly. In addition, users may try an [optimized BLAS](https://github.com/jump-dev/Ipopt.jl?tab=readme-ov-file#blas-and-lapack) and see if it yields significant performance gains (and let me know if it does). 
\ No newline at end of file
diff --git a/docs/src/index.md b/docs/src/index.md
index 673a43f..faa9fd5 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -37,7 +37,7 @@ import MathOptInterface as MOI
 import MathOptSymbolicAD
 
 param[:model_attr] = Dict(:backend => (MOI.AutomaticDifferentiationBackend(), 
-                                        MathOptSymbolicAD.DefaultBackend())) 
+                                       MathOptSymbolicAD.DefaultBackend())) 
                                 # Or:  MathOptSymbolicAD.ThreadedBackend()
 ```
 
@@ -45,7 +45,7 @@ param[:model_attr] = Dict(:backend => (MOI.AutomaticDifferentiationBackend(),
 
 ```julia
 param[:optimizer_attr] = Dict(:hsllib => "/usr/local/lib/libhsl.dylib", # Adjust path
-                                :linear_solver => "ma57") # Use ma57, ma86 or ma97
+                              :linear_solver => "ma57") # Use ma57, ma86 or ma97
 ```
 
 The [Ipopt.jl README](https://github.com/jump-dev/Ipopt.jl?tab=readme-ov-file#linear-solvers) suggests to use the larger LibHSL package for which there exists a Julia module and proceed similarly. In addition, users may try an [optimized BLAS](https://github.com/jump-dev/Ipopt.jl?tab=readme-ov-file#blas-and-lapack) and see if it yields significant performance gains (and let me know if it does). 
\ No newline at end of file

From d30f21388198038cab2a7d4baf3a5ac808520f16 Mon Sep 17 00:00:00 2001
From: Sebastian Krantz <sebastian.krantz@graduateinstitute.ch>
Date: Sat, 14 Sep 2024 16:11:44 +0200
Subject: [PATCH 32/32] Minors.

---
 examples/paper_example04.jl | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/examples/paper_example04.jl b/examples/paper_example04.jl
index a597139..5b70a01 100644
--- a/examples/paper_example04.jl
+++ b/examples/paper_example04.jl
@@ -3,7 +3,7 @@ using Random
 using Plots
 
 # Initialize parameters
-param = init_parameters(K = 100, rho = 0, labor_mobility = false)
+param = init_parameters(K = 100, sigma = 2, rho = 0, labor_mobility = false)
 width, height = 13, 13
 
 # Create graph
@@ -17,7 +17,7 @@ Ni = find_node(g0, ceil(width/2), ceil(height/2)) # center
 g0[:Zjn][Ni] = 1 # more productive node
 g0[:Lj][Ni] = 1 # more productive node
 
-Random.seed!(10) # use 5, 8, 10 or 11
+Random.seed!(11) # use 5, 8, 10 or 11
 nb_cities = 20 # draw a number of random cities in space
 for i in 1:nb_cities-1
     newdraw = false
@@ -113,11 +113,13 @@ plots = Vector{Any}(undef, length(simulation))
 i = 0
 for s in simulation
     i += 1
+    sizes = 2 .* results[Symbol(s)][:cj] .* (g0[:Lj] .> 1e-6) / maximum(results[Symbol(s)][:cj])
+    shades = results[Symbol(s)][:cj] .* (g0[:Lj] .> 1e-6) / maximum(results[Symbol(s)][:cj])
     plots[i] = plot_graph(g0, results[Symbol(s)][:Ijk], 
                           geography = geographies[Symbol(s)], obstacles = obstacles[i] == "on",
                           mesh = true, mesh_transparency = 0.2,
-                          node_sizes = results[Symbol(s)][:cj], 
-                          node_shades = g0[:Zjn], node_color = :seismic,
+                          node_sizes = sizes, node_shades = shades, 
+                          node_color = :seismic,
                           edge_min_thickness = 1, edge_max_thickness = 4)
     title!(plots[i], "Geography $(s)")
 end