add .JuliaFormatter.toml and apply custom setting format

Author: Giovanni3A

.JuliaFormatter.toml

@@ -0,0 +1,9 @@
+# Configuration file for JuliaFormatter.jl
+# For more information, see: https://domluna.github.io/JuliaFormatter.jl/stable/config/
+
+indent = 4
+always_use_return = true
+margin = 80
+remove_extra_newlines = true
+short_to_long_function_def = true
+format_docstrings = true

src/ApplicationDrivenLearning.jl

@@ -1,13 +1,11 @@
 module ApplicationDrivenLearning
 
-
 using Flux
 using JuMP
 using DiffOpt
 import ParametricOptInterface as POI
 import Base.*, Base.+
 
-
 include("flux_utils.jl")
 include("predictive_model.jl")
 
@@ -34,14 +32,16 @@ struct Forecast{T}
     assess::T
 end
 
-+(p1::Forecast, p2::Forecast) = Forecast(p1.plan + p2.plan, p1.assess + p2.assess)
+function +(p1::Forecast, p2::Forecast)
+    return Forecast(p1.plan + p2.plan, p1.assess + p2.assess)
+end
 *(c::Number, p::Forecast) = Forecast(c * p.plan, c * p.assess)
 
 """
     Model <: JuMP.AbstractModel
 
-Create an empty ApplicationDrivenLearning.Model with empty plan and assess models,
-missing forecast model and default settings.
+Create an empty ApplicationDrivenLearning.Model with empty plan and assess
+models, missing forecast model and default settings.
 """
 mutable struct Model <: JuMP.AbstractModel
     plan::JuMP.Model
@@ -116,7 +116,7 @@ function set_forecast_model(
         forecast = PredictiveModel(network)
     end
     @assert forecast.output_size == size(model.forecast_vars, 1)
-    model.forecast = forecast
+    return model.forecast = forecast
 end
 
 """
@@ -151,7 +151,11 @@ end
 Creates new constraint to assess model that fixes policy variables.
 """
 function build_assess_model_policy_constraint(model::Model)
-    @constraint(model.assess, assess_policy_fix, assess_policy_vars(model) .== 0)
+    @constraint(
+        model.assess,
+        assess_policy_fix,
+        assess_policy_vars(model) .== 0
+    )
 end
 
 """
@@ -166,7 +170,7 @@ function build(model::Model)
 
     # build plan model
     build_plan_model_forecast_params(model)
-    build_assess_model_policy_constraint(model)
+    return build_assess_model_policy_constraint(model)
 end
 
 include("jump.jl")
@@ -184,7 +188,12 @@ include("optimizers/bilevel.jl")
 
 Train model using given data and options.
 """
-function train!(model::Model, X::Matrix{<:Real}, y::Matrix{<:Real}, options::Options)
+function train!(
+    model::Model,
+    X::Matrix{<:Real},
+    y::Matrix{<:Real},
+    options::Options,
+)
     if options.mode == NelderMeadMode
         return train_with_nelder_mead!(model, X, y, options.params)
     elseif options.mode == GradientMode
@@ -194,8 +203,8 @@ function train!(model::Model, X::Matrix{<:Real}, y::Matrix{<:Real}, options::Opt
     elseif options.mode == GradientMPIMode
         return train_with_gradient_mpi!(model, X, y, options.params)
     elseif options.mode == BilevelMode
-        assert_msg = "BilevelMode not implemented for multiple forecasting models"
-        @assert length(model.forecast.networks) == 1 assert_msg
+        asr_msg = "BilevelMode not implemented for multiple forecasting models"
+        @assert length(model.forecast.networks) == 1 asr_msg
         return solve_bilevel(model, X, y, options.params)
     else
         # should never get here

src/flux_utils.jl

@@ -3,7 +3,8 @@ using Flux
 """
     extract_flux_params(model)
 
-Extract the parameters of a Flux model (Flux.Chain or Flux.Dense) into a single vector.
+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)
@@ -31,7 +32,8 @@ end
 """
     fix_flux_params_multi_model(models, θ)
 
-Return iterable of models after fixing the parameters from an adequate vector of parameters.
+Return iterable of models after fixing the parameters from an adequate vector
+of parameters.
 """
 function fix_flux_params_multi_model(models, θ::Vector{<:Real})
     i = 1
@@ -55,7 +57,7 @@ function has_params(layer)
         # Attempt to get parameters; if it works and isn't empty, return true
         return !isempty(Flux.params(layer))
     catch e
-        # If there is an error (e.g., method not matching), assume no parameters
+        # If there is an error (e.g. method not matching), assume no parameters
         return false
     end
 end

src/jump.jl

@@ -47,7 +47,11 @@ function JuMP.build_variable(
     return ForecastInfo(info, info, kwargs)
 end
 
-function JuMP.add_variable(model::Model, forecast_info::ForecastInfo, name::String)
+function JuMP.add_variable(
+    model::Model,
+    forecast_info::ForecastInfo,
+    name::String,
+)
     forecast = Forecast(
         JuMP.add_variable(
             model.plan,
@@ -84,11 +88,14 @@ end
 
 # jump functions
 function JuMP.objective_sense(model::Model)
-    @assert JuMP.objective_sense(model.plan) == JuMP.objective_sense(model.assess)
+    @assert JuMP.objective_sense(model.plan) ==
+            JuMP.objective_sense(model.assess)
     return JuMP.objective_sense(model.plan)
 end
 
-JuMP.num_variables(m::Model) = JuMP.num_variables(m.plan) + JuMP.num_variables(m.assess)
+function JuMP.num_variables(m::Model)
+    return JuMP.num_variables(m.plan) + JuMP.num_variables(m.assess)
+end
 
 function JuMP.show_constraints_summary(io::IO, model::Model)
     println("Plan Model:")
@@ -113,16 +120,17 @@ function JuMP.set_optimizer(model::Model, builder, evaluate_duals::Bool = true)
     new_diff_optimizer = DiffOpt.diff_optimizer(builder)
     JuMP.set_optimizer(
         model.plan,
-        () -> POI.Optimizer(new_diff_optimizer; evaluate_duals = evaluate_duals),
+        () ->
+            POI.Optimizer(new_diff_optimizer; evaluate_duals = evaluate_duals),
     )
 
     # basic setting for assess model
-    JuMP.set_optimizer(model.assess, builder)
+    return JuMP.set_optimizer(model.assess, builder)
 end
 
 function JuMP.set_silent(model::Model)
     MOI.set(model.plan, MOI.Silent(), true)
-    MOI.set(model.assess, MOI.Silent(), true)
+    return MOI.set(model.assess, MOI.Silent(), true)
 end
 
 function JuMP.num_variables(model::Model)

src/optimizers/bilevel.jl

@@ -2,7 +2,6 @@ using JuMP
 using Flux
 using BilevelJuMP
 
-
 function solve_bilevel(
     model::Model,
     X::Matrix{<:Real},
@@ -27,10 +26,12 @@ function solve_bilevel(
     T = size(Y, 1)
 
     # lower model variables
-    low_var_map = Dict{JuMP.VariableRef,Vector{BilevelJuMP.BilevelVariableRef}}()
+    low_var_map =
+        Dict{JuMP.VariableRef,Vector{BilevelJuMP.BilevelVariableRef}}()
     for pre_var in all_variables(model.plan)
         low_var_name = string(name(pre_var), "_low")
-        low_var_ref = @variable(Lower(bilevel_model), [1:T], base_name = low_var_name)
+        low_var_ref =
+            @variable(Lower(bilevel_model), [1:T], base_name = low_var_name)
         if has_lower_bound(pre_var)
             set_lower_bound.(low_var_ref, lower_bound(pre_var))
         end
@@ -45,7 +46,8 @@ function solve_bilevel(
     for post_var in all_variables(model.assess)
         if !(post_var in assess_policy_vars(model))
             up_var_name = string(name(post_var), "_up")
-            up_var_ref = @variable(Upper(bilevel_model), [1:T], base_name = up_var_name)
+            up_var_ref =
+                @variable(Upper(bilevel_model), [1:T], base_name = up_var_name)
             if has_lower_bound(post_var)
                 set_lower_bound.(up_var_ref, lower_bound(post_var))
             end
@@ -65,20 +67,30 @@ function solve_bilevel(
     end
 
     # lower model base constraints
-    for pre_con in
-        JuMP.all_constraints(model.plan, include_variable_in_set_constraints = false)
+    for pre_con in JuMP.all_constraints(
+        model.plan,
+        include_variable_in_set_constraints = false,
+    )
         pre_con_func = JuMP.constraint_object(pre_con).func
         lhs = [value(x -> low_var_map[x][t], pre_con_func) for t = 1:T]
-        @constraint(Lower(bilevel_model), lhs .∈ JuMP.constraint_object(pre_con).set)
+        @constraint(
+            Lower(bilevel_model),
+            lhs .∈ JuMP.constraint_object(pre_con).set
+        )
     end
 
     # upper model base constraints
-    for post_con in
-        JuMP.all_constraints(model.assess, include_variable_in_set_constraints = false)
+    for post_con in JuMP.all_constraints(
+        model.assess,
+        include_variable_in_set_constraints = false,
+    )
         if name(post_con) != "assess_policy_fix"
             post_con_func = JuMP.constraint_object(post_con).func
             lhs = [value(x -> up_var_map[x][t], post_con_func) for t = 1:T]
-            @constraint(Upper(bilevel_model), lhs .∈ JuMP.constraint_object(post_con).set)
+            @constraint(
+                Upper(bilevel_model),
+                lhs .∈ JuMP.constraint_object(post_con).set
+            )
         end
     end
 
@@ -97,7 +109,10 @@ function solve_bilevel(
     # fix upper model observations
     i_obs_var = 1
     for obs_var in assess_forecast_vars(model)
-        @constraint(Upper(bilevel_model), up_var_map[obs_var] - Y[1:T, i_obs_var] .== 0)
+        @constraint(
+            Upper(bilevel_model),
+            up_var_map[obs_var] - Y[1:T, i_obs_var] .== 0
+        )
         i_obs_var += 1
     end
 
@@ -114,7 +129,10 @@ function solve_bilevel(
         if has_params(layer)
             # get size and parameters W and b
             (layer_size_out, layer_size_in) = size(layer.weight)
-            W = @variable(Upper(bilevel_model), [1:layer_size_out, 1:layer_size_in])
+            W = @variable(
+                Upper(bilevel_model),
+                [1:layer_size_out, 1:layer_size_in]
+            )
             if layer.bias == false
                 b = zeros(layer_size_out)
             else
@@ -123,7 +141,8 @@ function solve_bilevel(
             predictive_model_vars[i_layer] = Dict(:W => W, :b => b)
             # build layer output as next layer input
             for output_idx in values(model.forecast.input_output_map[1])
-                layers_inpt[output_idx] = layer.σ(W * layers_inpt[output_idx]' .+ b)'
+                layers_inpt[output_idx] =
+                    layer.σ(W * layers_inpt[output_idx]' .+ b)'
             end
             # if activation function layer, just apply
         elseif supertype(typeof(layer)) == Function
@@ -144,7 +163,10 @@ function solve_bilevel(
     ipred_var_count = 1
     for pred_var in plan_forecast_vars(model)
         low_pred_var = low_var_map[pred_var]
-        @constraint(Lower(bilevel_model), low_pred_var .- y_hat[:, ipred_var_count] .== 0)
+        @constraint(
+            Lower(bilevel_model),
+            low_pred_var .- y_hat[:, ipred_var_count] .== 0
+        )
         ipred_var_count += 1
     end
 
@@ -165,5 +187,8 @@ function solve_bilevel(
         ilayer += 1
     end
 
-    return Solution(objective_value(bilevel_model), extract_params(model.forecast))
+    return Solution(
+        objective_value(bilevel_model),
+        extract_params(model.forecast),
+    )
 end

src/optimizers/gradient.jl

@@ -1,7 +1,8 @@
 using Flux
 
 """
-Compute assess cost and cost gradient (with respect to predicted values) based on incomplete batch of examples.
+Compute assess cost and cost gradient (with respect to predicted values) based
+on incomplete batch of examples.
 """
 function stochastic_compute(model, X, Y, batch, compute_full_cost::Bool)
     C, dC = compute_cost(model, X[batch, :], Y[batch, :], true)
@@ -12,7 +13,8 @@ function stochastic_compute(model, X, Y, batch, compute_full_cost::Bool)
 end
 
 """
-Compute assess cost and cost gradient (with respect to predicted values) based on complete batch of examples.
+Compute assess cost and cost gradient (with respect to predicted values) based
+on complete batch of examples.
 """
 function deterministic_compute(model, X, Y)
     C, dC = compute_cost(model, X, Y, true)
@@ -53,7 +55,13 @@ function train_with_gradient!(
 
         if stochastic
             epochx = X[batches[epoch, :], :]
-            C, dC = stochastic_compute(model, X, Y, batches[epoch, :], compute_full_cost)
+            C, dC = stochastic_compute(
+                model,
+                X,
+                Y,
+                batches[epoch, :],
+                compute_full_cost,
+            )
         else
             epochx = X
             C, dC = deterministic_compute(model, X, Y)
@@ -80,7 +88,6 @@ function train_with_gradient!(
 
         # take gradient step
         apply_gradient!(model.forecast, dC, epochx, rule)
-
     end
 
     # fix best model

src/optimizers/gradient_mpi.jl

@@ -3,7 +3,6 @@ import ParametricOptInterface as POI
 using MPI
 import JobQueueMPI as JQM
 
-
 function train_with_gradient_mpi!(
     model::Model,
     X::Matrix{<:Real},
@@ -52,7 +51,6 @@ function train_with_gradient_mpi!(
         end
 
         return step_cost, step_grad
-
     end
 
     # call optim as the controller
@@ -75,18 +73,21 @@ function train_with_gradient_mpi!(
                     (v) -> compute_cost_and_gradients(v[1], v[2], true),
                     [[curr_θ, i] for i in batches[epoch, :]],
                 )
-                dCdy = sum([r[2] for r in pmap_result_with_gradients]) ./ batch_size
+                dCdy =
+                    sum([r[2] for r in pmap_result_with_gradients]) ./ batch_size
 
                 if compute_full_cost
                     # broadcast `is_done = false` again
                     MPI.bcast(is_done, MPI.COMM_WORLD)
 
                     # compute full cost
                     pmap_result_without_gradients = JQM.pmap(
-                        (v) -> compute_cost_and_gradients(v[1], v[2], false),
+                        (v) ->
+                            compute_cost_and_gradients(v[1], v[2], false),
                         [[curr_θ, i] for i = 1:T],
                     )
-                    curr_C = sum([r[1] for r in pmap_result_without_gradients]) ./ T
+                    curr_C =
+                        sum([r[1] for r in pmap_result_without_gradients]) ./ T
                 end
 
             else