Merge branch 'main' into compathelper/new_version/2025-03-12-01-17-18-775-00042309767

Author: joaquimg

Project.toml

@@ -1,7 +1,7 @@
 name = "ApplicationDrivenLearning"
 uuid = "0856f1c8-ef17-4e14-9230-2773e47a789e"
-authors = ["Giovanni Amorin"]
-version = "0.1.0"
+authors = ["Giovanni Amorim", "Joaquim Garcia"]
+version = "0.1.1"
 
 [deps]
 BilevelJuMP = "485130c0-026e-11ea-0f1a-6992cd14145c"
@@ -11,18 +11,20 @@ JobQueueMPI = "32d208e1-246e-420c-b6ff-18b71b410923"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 ParametricOptInterface = "0ce4ce61-57bf-432b-a095-efac525d185e"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 BilevelJuMP = "0.6.2"
 DiffOpt = "0.5.0"
-Flux = "0.14.25"
+Flux = "0.16.3"
 JobQueueMPI = "0.1.1"
 JuMP = "1.24"
 MPI = "0.20.22"
 Optim = "1.11"
-ParametricOptInterface = "0.9.0"
+Optimisers = "0.4.5"
+ParametricOptInterface = "0.9.0, 0.10"
 Zygote = "0.6.75, 0.7"
 julia = "1.10"

src/ApplicationDrivenLearning.jl

@@ -135,15 +135,10 @@ and new constraint fixing to original forecast variables.
 function build_plan_model_forecast_params(model::Model)
     # adds parametrized forecast variables using MOI.Parameter
     forecast_size = size(model.forecast_vars)[1]
-    model.plan_forecast_params = @variable(
-        model.plan,
-        _forecast[1:forecast_size] in MOI.Parameter.(zeros(forecast_size))
-    )
-    # fixes old and new prediction variables together
+    model.plan_forecast_params = plan_forecast_vars(model)
     @constraint(
         model.plan,
-        plan_forecast_fix,
-        model.plan_forecast_params .== plan_forecast_vars(model)
+        model.plan_forecast_params .∈ MOI.Parameter.(zeros(forecast_size))
     )
 end
 

src/flux_utils.jl

@@ -7,7 +7,7 @@ Extract the parameters of a Flux model (Flux.Chain or Flux.Dense) into a single
 vector.
 """
 function extract_flux_params(model::Union{Flux.Chain,Flux.Dense})
-    θ = Flux.params(model)
+    θ = Flux.trainables(model)
     return reduce(vcat, [vec(p) for p in θ])
 end
 
@@ -21,7 +21,7 @@ function fix_flux_params_single_model(
     θ::Vector{<:Real},
 )
     i = 1
-    for p in Flux.params(model)
+    for p in Flux.trainables(model)
         psize = prod(size(p))
         p .= reshape(θ[i:i+psize-1], size(p))
         i += psize
@@ -38,7 +38,7 @@ of parameters.
 function fix_flux_params_multi_model(models, θ::Vector{<:Real})
     i = 1
     for model in models
-        for p in Flux.params(model)
+        for p in Flux.trainables(model)
             psize = prod(size(p))
             p .= reshape(θ[i:i+psize-1], size(p))
             i += psize
@@ -54,8 +54,8 @@ Check if a Flux layer has parameters.
 """
 function has_params(layer)
     try
-        # Attempt to get parameters; if it works and isn't empty, return true
-        return !isempty(Flux.params(layer))
+        # Attempt to get trainable parameters; if it works and isn't empty, return true
+        return !isempty(Flux.trainable(layer))
     catch e
         # If there is an error (e.g. method not matching), assume no parameters
         return false

src/jump.jl

@@ -64,6 +64,20 @@ function JuMP.add_variable(
             name * "_assess",
         ),
     )
+
+    # forecast variables can't have bounds
+    if JuMP.has_lower_bound(forecast.plan)
+        @warn "Forecast variable lower bound will be removed."
+        JuMP.delete_lower_bound(forecast.plan)
+        JuMP.delete_lower_bound(forecast.assess)
+    end
+
+    if JuMP.has_upper_bound(forecast.plan)
+        @warn "Forecast variable upper bound will be removed."
+        JuMP.delete_upper_bound(forecast.plan)
+        JuMP.delete_upper_bound(forecast.assess)
+    end
+
     push!(model.forecast_vars, forecast)
     return forecast
 end

src/optimizers/bilevel.jl

@@ -177,10 +177,10 @@ function solve_bilevel(
     ilayer = 1
     for layer in model.forecast.networks[1]
         if has_params(layer)
-            for p in Flux.params(layer.weight)
+            for p in Flux.trainables(layer.weight)
                 p .= value.(predictive_model_vars[ilayer][:W])
             end
-            for p in Flux.params(layer.bias)
+            for p in Flux.trainables(layer.bias)
                 p .= value.(predictive_model_vars[ilayer][:b])
             end
         end

src/optimizers/gradient.jl

@@ -42,6 +42,7 @@ function train_with_gradient!(
     best_θ = extract_params(model.forecast)
     trace = Array{Float64}(undef, epochs)
     stochastic = batch_size > 0
+    opt_state = Flux.setup(rule, model.forecast)
 
     # precompute batches
     batches = repeat(1:T, outer = (1, epochs))'
@@ -87,7 +88,7 @@ function train_with_gradient!(
         end
 
         # take gradient step
-        apply_gradient!(model.forecast, dC, epochx, rule)
+        apply_gradient!(model.forecast, dC, epochx, opt_state)
     end
 
     # fix best model

src/optimizers/gradient_mpi.jl

@@ -32,6 +32,7 @@ function train_with_gradient_mpi!(
     T = size(X)[1]
     stochastic = batch_size > 0
     compute_full_cost = true
+    opt_state = Flux.setup(rule, model.forecast)
 
     # precompute batches
     batches = repeat(1:T, outer = (1, epochs))'
@@ -121,7 +122,7 @@ function train_with_gradient_mpi!(
             end
 
             # take gradient step (if not last epoch)
-            apply_gradient!(model.forecast, dCdy, epochx, rule)
+            apply_gradient!(model.forecast, dCdy, epochx, opt_state)
         end
 
         # release workers

src/predictive_model.jl

@@ -1,6 +1,7 @@
 using Flux
 using Statistics
 import Zygote
+import Optimisers
 
 """
     PredictiveModel(networks, input_output_map, input_size, output_size)
@@ -198,7 +199,7 @@ function apply_params(model::PredictiveModel, θ)
 end
 
 """
-    apply_gradient!(model, dCdy, X, optimizer)
+    apply_gradient!(model, dCdy, X, rule)
 
 Apply a gradient vector to the model parameters.
 
@@ -209,17 +210,16 @@ Apply a gradient vector to the model parameters.
   - `model::PredictiveModel`: model to be updated.
   - `dCdy::Vector{<:Real}`: gradient vector.
   - `X::Matrix{<:Real}`: input data.
-  - `optimizer`: Optimiser to be used.
+  - `rule`: Optimisation rule.
     ...
 """
 function apply_gradient!(
     model::PredictiveModel,
     dCdy::Vector{<:Real},
     X::Matrix{<:Real},
-    optimizer,
+    opt_state,
 )
-    ps = Flux.params(model.networks)
-    loss(x, y) = mean(dCdy'model(x))
-    train_data = [(X', 0.0)]
-    return Flux.train!(loss, ps, train_data, optimizer)
+    loss3(m, X) = mean(dCdy'm(X'))
+    grad = Zygote.gradient(loss3, model, X)[1]
+    return Optimisers.update!(opt_state, model, grad)
 end

src/simulation.jl

@@ -3,17 +3,33 @@ function compute_single_step_cost(
     y::Vector{<:Real},
     yhat::Vector{<:Real},
 )
+    # set forecast params as prediction output
     MOI.set.(model.plan, POI.ParameterValue(), model.plan_forecast_params, yhat)
+    # optimize plan model
     optimize!(model.plan)
-    @assert termination_status(model.plan) == MOI.OPTIMAL "Optimization failed for PLAN model"
+    # check for solution and fix assess policy vars
+    try
+        set_normalized_rhs.(
+            model.assess[:assess_policy_fix],
+            value.(plan_policy_vars(model)),
+        )
+    catch e
+        println("Optimization failed for PLAN model.")
+        throw(e)
+    end
+    # fix assess forecast vars on observer values
     fix.(assess_forecast_vars(model), y; force = true)
-    set_normalized_rhs.(
-        model.assess[:assess_policy_fix],
-        value.(plan_policy_vars(model)),
-    )
+    # optimize assess model
     optimize!(model.assess)
-    @assert termination_status(model.assess) == MOI.OPTIMAL "Optimization failed for ASSESS model"
-    return objective_value(model.assess)
+    # check for optimization
+    try
+        return objective_value(model.assess)
+    catch e
+        println("Optimization failed for ASSESS model")
+        throw(e)
+    end
+    # should never get here
+    return 0
 end
 
 """

test/Project.toml

@@ -8,6 +8,7 @@ JobQueueMPI = "32d208e1-246e-420c-b6ff-18b71b410923"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 ParametricOptInterface = "0ce4ce61-57bf-432b-a095-efac525d185e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"