format

Author: andrewrosemberg

docs/make.jl

@@ -1,23 +1,18 @@
 using L2OALM
 using Documenter
 
-DocMeta.setdocmeta!(L2OALM, :DocTestSetup, :(using L2OALM); recursive=true)
+DocMeta.setdocmeta!(L2OALM, :DocTestSetup, :(using L2OALM); recursive = true)
 
 makedocs(;
-    modules=[L2OALM],
-    authors="Andrew <[email protected]> and contributors",
-    sitename="L2OALM.jl",
-    format=Documenter.HTML(;
-        canonical="https://LearningToOptimize.github.io/L2OALM.jl",
-        edit_link="main",
-        assets=String[],
+    modules = [L2OALM],
+    authors = "Andrew <[email protected]> and contributors",
+    sitename = "L2OALM.jl",
+    format = Documenter.HTML(;
+        canonical = "https://LearningToOptimize.github.io/L2OALM.jl",
+        edit_link = "main",
+        assets = String[],
     ),
-    pages=[
-        "Home" => "index.md",
-    ],
+    pages = ["Home" => "index.md"],
 )
 
-deploydocs(;
-    repo="github.com/LearningToOptimize/L2OALM.jl",
-    devbranch="main",
-)
+deploydocs(; repo = "github.com/LearningToOptimize/L2OALM.jl", devbranch = "main")

src/L2OALM.jl

@@ -9,8 +9,8 @@ using Lux.Training
 using CUDA
 using Statistics
 
-export LagrangianDualLoss, LagrangianPrimalLoss, TrainingStepLoop,
-       L2OALM_epoch!, L2OALM_train!
+export LagrangianDualLoss,
+    LagrangianPrimalLoss, TrainingStepLoop, L2OALM_epoch!, L2OALM_train!
 
 """
     LagrangianDualLoss(;max_dual=1e6)
@@ -23,27 +23,27 @@ Target dual variables are clipped from zero to `max_dual`.
 Keywords:
     - `max_dual`: Maximum value for the target dual variables.
 """
-function LagrangianDualLoss(num_equal::Int; max_dual=1e6)
+function LagrangianDualLoss(num_equal::Int; max_dual = 1e6)
     return (dual_model, ps_dual, st_dual, data) -> begin
         x, hpm, dual_hat_k, gh = data
         ρ = hpm[:ρ]
         # Get current dual predictions
         dual_hat, st_dual_new = dual_model(x, ps_dual, st_dual)
-        
+
         # Separate bound and equality constraints
-        gh_bound = gh[1:end-num_equal,:]
-        gh_equal = gh[end-num_equal+1:end,:]
-        dual_hat_bound = dual_hat_k[1:end-num_equal,:]
-        dual_hat_equal = dual_hat_k[end-num_equal+1:end,:]
-        
+        gh_bound = gh[1:end-num_equal, :]
+        gh_equal = gh[end-num_equal+1:end, :]
+        dual_hat_bound = dual_hat_k[1:end-num_equal, :]
+        dual_hat_equal = dual_hat_k[end-num_equal+1:end, :]
+
         # Target for dual variables
         dual_target = vcat(
             min.(max.(dual_hat_bound + ρ .* gh_bound, 0), max_dual),
-            min.(dual_hat_equal + ρ .* gh_equal, max_dual)
+            min.(dual_hat_equal + ρ .* gh_equal, max_dual),
         )
-        
-        loss = mean((dual_hat .- dual_target).^2)
-        return loss, st_dual_new, (dual_loss=loss,)
+
+        loss = mean((dual_hat .- dual_target) .^ 2)
+        return loss, st_dual_new, (dual_loss = loss,)
     end
 end
 
@@ -56,31 +56,31 @@ from current dual predictions `dual_hat` for the batch model `bm` under paramete
 Arguments:
     - `bm`: A `BatchModel` instance that contains the model and batch configuration.
 """
-function LagrangianPrimalLoss(bm::BatchModel)    
+function LagrangianPrimalLoss(bm::BatchModel)
     return (model, ps, st, data) -> begin
         Θ, hpm, dual_hat = data
         ρ = hpm[:ρ]
         num_s = size(Θ, 2)
 
         # Forward pass for prediction
         X̂, st_new = model(Θ, ps, st)
-        
+
         # Calculate violations and objectives
         objs = BNK.objective!(bm, X̂, Θ)
         # gh = constraints!(bm, X̂, Θ)
         Vc, Vb = BNK.all_violations!(bm, X̂, Θ)
         V = vcat(Vb, Vc)
         total_loss = (
-            sum(abs.(dual_hat .* V)) / num_s +
-            ρ / 2 * sum((V).^2) / num_s +
-            mean(objs)
+            sum(abs.(dual_hat .* V)) / num_s + ρ / 2 * sum((V) .^ 2) / num_s + mean(objs)
         )
 
-        return total_loss, st_new, (
-            total_loss=total_loss,
-            mean_violations=mean(V),
-            new_max_violation=maximum(V),
-            mean_objs=mean(objs),
+        return total_loss,
+        st_new,
+        (
+            total_loss = total_loss,
+            mean_violations = mean(V),
+            new_max_violation = maximum(V),
+            mean_objs = mean(objs),
         )
     end
 end
@@ -100,7 +100,7 @@ Fields:
 mutable struct TrainingStepLoop
     loss_fn::Function
     stopping_criteria::Vector{Function}
-    hyperparameters::Dict{Symbol, Any}
+    hyperparameters::Dict{Symbol,Any}
     parameter_update_fns::Vector{Function}
     reconcile_state::Function
     pre_hook::Function
@@ -112,7 +112,14 @@ end
 Default pre-hook function for the primal model in the L2O-ALM algorithm.
 This function performs a forward pass through the dual model to obtain the dual predictions.
 """
-function _pre_hook_primal(θ, primal_model, train_state_primal, dual_model, train_state_dual, bm)
+function _pre_hook_primal(
+    θ,
+    primal_model,
+    train_state_primal,
+    dual_model,
+    train_state_dual,
+    bm,
+)
     # Forward pass for dual model
     dual_hat_k, _ = dual_model(θ, train_state_dual.parameters, train_state_dual.states)
 
@@ -125,15 +132,22 @@ end
 Default pre-hook function for the dual model in the L2O-ALM algorithm.
 This function performs a forward pass through the primal model to obtain the predicted state and constraints.
 """
-function _pre_hook_dual(θ, primal_model, train_state_primal, dual_model, train_state_dual, bm)
+function _pre_hook_dual(
+    θ,
+    primal_model,
+    train_state_primal,
+    dual_model,
+    train_state_dual,
+    bm,
+)
     # # Forward pass for primal model
     X̂, _ = primal_model(θ, train_state_primal.parameters, train_state_primal.states)
     gh = constraints!(bm, X̂, Θ)
-    
+
     # Forward pass for dual model
     dual_hat, _ = dual_model(θ, train_state_dual.parameters, train_state_dual.states)
 
-    return (dual_hat, gh,)
+    return (dual_hat, gh)
 end
 
 """
@@ -149,10 +163,10 @@ function _reconcile_alm_primal_state(batch_states::Vector{NamedTuple})
     mean_objs = mean([s.mean_objs for s in batch_states])
     mean_loss = mean([s.total_loss for s in batch_states])
     return (;
-        new_max_violation=max_violation,
-        mean_violations=mean_violations,
-        mean_objs=mean_objs,
-        total_loss=mean_loss,
+        new_max_violation = max_violation,
+        mean_violations = mean_violations,
+        mean_objs = mean_objs,
+        total_loss = mean_loss,
     )
 end
 
@@ -164,7 +178,7 @@ This function computes the mean dual loss from the batch states.
 """
 function _reconcile_alm_dual_state(batch_states::Vector{NamedTuple})
     dual_loss = mean([s.dual_loss for s in batch_states])
-    return (dual_loss=dual_loss,)
+    return (dual_loss = dual_loss,)
 end
 
 """
@@ -173,7 +187,11 @@ end
 Default function to update the hyperparameter ρ in the ALM algorithm.
 This function increases ρ by a factor of α if the new maximum violation exceeds τ times the previous maximum violation.
 """
-function _update_ALM_ρ!(hpm_primal::Dict{Symbol, Any}, hpm_dual::Dict{Symbol, Any}, current_state::NamedTuple)
+function _update_ALM_ρ!(
+    hpm_primal::Dict{Symbol,Any},
+    hpm_dual::Dict{Symbol,Any},
+    current_state::NamedTuple,
+)
     if current_state.new_max_violation > hpm_primal.τ * hpm_primal.max_violation
         hpm_primal[:ρ] = min(hpm_primal[:ρmax], hpm_primal[:ρ] * hpm_primal[:α])
         hpm_dual[:ρ] = hpm_primal[:ρ]  # Ensure dual model uses the same ρ
@@ -191,7 +209,7 @@ function _default_primal_loop(bm::BatchModel)
     return TrainingStepLoop(
         LagrangianPrimalLoss(bm),
         [(iter, current_state, hpm) -> iter >= 100 ? true : false],
-        Dict{Symbol, Any}(
+        Dict{Symbol,Any}(
             :ρ => 1.0,
             :ρmax => 1e6,
             :τ => 0.8,
@@ -200,7 +218,7 @@ function _default_primal_loop(bm::BatchModel)
         ),
         [_update_ALM_ρ!],
         _reconcile_alm_primal_state,
-        _pre_hook_primal
+        _pre_hook_primal,
     )
 end
 
@@ -213,13 +231,10 @@ function _default_dual_loop(num_equal::Int)
     return TrainingStepLoop(
         LagrangianDualLoss(num_equal),
         [(iter, current_state, hpm) -> iter >= 100 ? true : false],
-        Dict{Symbol, Any}(
-            :max_dual => 1e6,
-            :ρ => 1.0,
-        ),
+        Dict{Symbol,Any}(:max_dual => 1e6, :ρ => 1.0),
         [],
         _reconcile_alm_dual_state,
-        _pre_hook_dual
+        _pre_hook_dual,
     )
 end
 
@@ -249,7 +264,7 @@ function L2OALM_epoch!(
     train_state_dual::Lux.Training.TrainState,
     training_step_loop_primal::TrainingStepLoop,
     training_step_loop_dual::TrainingStepLoop,
-    data
+    data,
 )
     iter_primal = 1
     iter_dual = 1
@@ -258,36 +273,73 @@ function L2OALM_epoch!(
     current_state_dual = (;)
 
     # primal loop
-    while all(stopping_criterion(iter_primal, current_state_primal, training_step_loop_primal.hyperparameters) for stopping_criterion in training_step_loop_primal.stopping_criteria)
+    while all(
+        stopping_criterion(
+            iter_primal,
+            current_state_primal,
+            training_step_loop_primal.hyperparameters,
+        ) for stopping_criterion in training_step_loop_primal.stopping_criteria
+    )
         current_states_primal = Vector{NamedTuple}(undef, num_batches)
         iter_batch = 1
         for (θ) in data
             _, loss_val, stats, train_state_primal = Training.single_train_step!(
                 AutoZygote(),          # AD backend
                 training_step_loop_primal.loss_fn,  # Loss function
-                (θ, training_step_loop_primal.hyperparameters, training_step_loop_primal.pre_hook(θ, primal_model, train_state_primal, dual_model, train_state_dual)...), # Data
-                train_state_primal # Training state
+                (
+                    θ,
+                    training_step_loop_primal.hyperparameters,
+                    training_step_loop_primal.pre_hook(
+                        θ,
+                        primal_model,
+                        train_state_primal,
+                        dual_model,
+                        train_state_dual,
+                    )...,
+                ), # Data
+                train_state_primal, # Training state
             )
             current_states_primal[iter_batch] = stats
             iter_batch += 1
         end
-        current_state_primal = training_step_loop_primal.reconcile_state(current_states_primal)
+        current_state_primal =
+            training_step_loop_primal.reconcile_state(current_states_primal)
         iter_primal += 1
     end
     for fn in training_step_loop_primal.parameter_update_fns
-        fn(training_step_loop_primal.hyperparameters, training_step_loop_dual.hyperparameters, current_state_primal)
+        fn(
+            training_step_loop_primal.hyperparameters,
+            training_step_loop_dual.hyperparameters,
+            current_state_primal,
+        )
     end
 
     # dual loop
-    while all(stopping_criterion(iter_dual, current_state_dual, training_step_loop_dual.hyperparameters) for stopping_criterion in training_step_loop_dual.stopping_criteria)
+    while all(
+        stopping_criterion(
+            iter_dual,
+            current_state_dual,
+            training_step_loop_dual.hyperparameters,
+        ) for stopping_criterion in training_step_loop_dual.stopping_criteria
+    )
         current_states_dual = Vector{NamedTuple}(undef, num_batches)
         iter_batch = 1
         for (θ) in data
             _, loss_val, stats, train_state_dual = Training.single_train_step!(
                 AutoZygote(),          # AD backend
                 training_step_loop_dual.loss_fn,  # Loss function
-                (θ, training_step_loop_dual.hyperparameters, training_step_loop_dual.pre_hook(θ, primal_model, train_state_primal, dual_model, train_state_dual)...), # Data
-                train_state_dual # Training state
+                (
+                    θ,
+                    training_step_loop_dual.hyperparameters,
+                    training_step_loop_dual.pre_hook(
+                        θ,
+                        primal_model,
+                        train_state_primal,
+                        dual_model,
+                        train_state_dual,
+                    )...,
+                ), # Data
+                train_state_dual, # Training state
             )
             current_states_dual[iter_batch] = stats
             iter_batch += 1
@@ -296,7 +348,11 @@ function L2OALM_epoch!(
         iter_dual += 1
     end
     for fn in training_step_loop_dual.parameter_update_fns
-        fn(training_step_loop_primal.hyperparameters, training_step_loop_dual.hyperparameters, current_state_dual)
+        fn(
+            training_step_loop_primal.hyperparameters,
+            training_step_loop_dual.hyperparameters,
+            current_state_dual,
+        )
     end
     return
 end
@@ -334,20 +390,33 @@ function L2OALM_train!(
     train_state_primal::Lux.Training.TrainState,
     train_state_dual::Lux.Training.TrainState,
     data;
-    training_step_loop_primal::TrainingStepLoop=_default_primal_loop(bm),
-    training_step_loop_dual::TrainingStepLoop=_default_dual_loop(num_equal),
-    stopping_criteria::Vector{F}=[(iter, primal_model, dual_model, tr_st_primal, tr_st_dual, hpm_primal) -> iter >= 100 ? true : false],
-) where F<:Function
+    training_step_loop_primal::TrainingStepLoop = _default_primal_loop(bm),
+    training_step_loop_dual::TrainingStepLoop = _default_dual_loop(num_equal),
+    stopping_criteria::Vector{F} = [
+        (iter, primal_model, dual_model, tr_st_primal, tr_st_dual, hpm_primal) ->
+            iter >= 100 ? true : false,
+    ],
+) where {F<:Function}
     iter = 1
-    while all(stopping_criterion(iter, primal_model, dual_model, train_state_primal, train_state_dual, training_step_loop_primal.hyperparameters, training_step_loop_dual.hyperparameters) for stopping_criterion in stopping_criteria)
+    while all(
+        stopping_criterion(
+            iter,
+            primal_model,
+            dual_model,
+            train_state_primal,
+            train_state_dual,
+            training_step_loop_primal.hyperparameters,
+            training_step_loop_dual.hyperparameters,
+        ) for stopping_criterion in stopping_criteria
+    )
         L2OALM_epoch!(
             primal_model,
             train_state_primal,
             dual_model,
             train_state_dual,
             training_step_loop_primal,
             training_step_loop_dual,
-            data
+            data,
         )
         iter += 1
     end