improve docstrings and add docs

Author: Giovanni3A

docs/Project.toml

@@ -0,0 +1,2 @@
+[deps]
+Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"

docs/make.jl

@@ -0,0 +1,16 @@
+using Documenter
+
+push!(LOAD_PATH, "../src")
+using ApplicationDrivenLearning
+
+makedocs(;
+    modules=[ApplicationDrivenLearning],
+    doctest=false,
+    clean=true,
+    sitename="ApplicationDrivenLearning.jl",
+    authors="Giovanni Amorim, Joaquim Garcia",
+    pages=[
+        "Home" => "index.md",
+        "API Reference" => "reference.md"
+    ]
+)

docs/src/index.md

@@ -0,0 +1,78 @@
+# ApplicationDrivenLearning.jl Documentation
+
+## Introduction
+
+ApplicationDrivenLearning.jl is a Julia package for training time series models using the application driven learning framework, that connects the optimization problem final cost with predictive model parameters in order to achieve the best model for a given application.
+
+## Usage
+
+```julia
+import Pkg
+
+Pkg.add("ApplicationDrivenLearning")  # not working yet! clone the repo instead
+
+using ApplicationDrivenLearning
+
+## Single power plan problem
+
+# data
+X = reshape([1 1], (2, 1))
+Y = reshape([0 2], (2, 1))
+
+# main model and policy / forecast variables
+model = ApplicationDrivenLearning.Model()
+@variables(model, begin
+    z, ApplicationDrivenLearning.Policy
+    θ, ApplicationDrivenLearning.Forecast
+end)
+
+# plan model
+@variables(ApplicationDrivenLearning.Plan(model), begin
+    c1 ≥ 0
+    c2 ≥ 0
+end)
+@constraints(ApplicationDrivenLearning.Plan(model), begin
+    c1 ≥ 100 * (θ.plan-z.plan)
+    c2 ≥ 20 * (z.plan-θ.plan)
+end)
+@objective(ApplicationDrivenLearning.Plan(model), Min, 10*z.plan + c1 + c2)
+
+# assess model
+@variables(ApplicationDrivenLearning.Assess(model), begin
+    c1 ≥ 0
+    c2 ≥ 0
+end)
+@constraints(ApplicationDrivenLearning.Assess(model), begin
+    c1 ≥ 100 * (θ.assess-z.assess)
+    c2 ≥ 20 * (z.assess-θ.assess)
+end)
+@objective(ApplicationDrivenLearning.Assess(model), Min, 10*z.assess + c1 + c2)
+
+# basic setting
+set_optimizer(model, HiGHS.Optimizer)
+set_silent(model)
+
+# forecast model
+nn = Chain(Dense(1 => 1; bias=false))
+ApplicationDrivenLearning.set_forecast_model(model, nn)
+
+# training and getting solution
+solution = ApplicationDrivenLearning.train!(
+    model,
+    X,
+    Y,
+    ApplicationDrivenLearning.Options(
+        ApplicationDrivenLearning.NelderMeadMode
+    )
+)
+print(solution.params)
+```
+
+## Installation
+
+This package is **not yet** registered so if you want to use or test the code clone this repo and include source code from `src` directory.
+
+## Contributing
+
+* PRs such as adding new models and fixing bugs are very welcome!
+* For nontrivial changes, you'll probably want to first discuss the changes via issue.

docs/src/reference.md

@@ -0,0 +1,69 @@
+# [API](@id API)
+
+This section documents the ApplicationDrivenLearning API.
+
+## Constructors
+
+```@docs
+Model
+PredictiveModel
+Plan
+Assess
+```
+
+## JuMP variable types
+
+```@docs
+Policy
+Forecast
+```
+
+## Structs
+
+```@docs
+ApplicationDrivenLearning.Options
+ApplicationDrivenLearning.Solution
+```
+
+## Modes
+
+```@docs
+ApplicationDrivenLearning.NelderMeadMode
+ApplicationDrivenLearning.GradientMode
+ApplicationDrivenLearning.NelderMeadMPIMode
+ApplicationDrivenLearning.GradientMPIMode
+ApplicationDrivenLearning.BilevelMode
+```
+
+## Attributes getters and setters
+
+```@docs
+ApplicationDrivenLearning.plan_policy_vars
+ApplicationDrivenLearning.assess_policy_vars
+ApplicationDrivenLearning.plan_forecast_vars
+ApplicationDrivenLearning.assess_forecast_vars
+ApplicationDrivenLearning.set_forecast_model
+ApplicationDrivenLearning.extract_params
+ApplicationDrivenLearning.apply_params
+```
+
+### Flux attributes getters and setters
+
+```@docs
+ApplicationDrivenLearning.extract_flux_params
+ApplicationDrivenLearning.fix_flux_params_single_model
+ApplicationDrivenLearning.fix_flux_params_multi_model
+ApplicationDrivenLearning.has_params
+ApplicationDrivenLearning.apply_gradient!
+```
+
+## Other functions
+
+```@docs
+forecast
+compute_cost
+train!
+ApplicationDrivenLearning.build_plan_model_forecast_params
+ApplicationDrivenLearning.build_assess_model_policy_constraint
+ApplicationDrivenLearning.build
+```

src/ApplicationDrivenLearning.jl

@@ -11,7 +11,11 @@ import Base.*, Base.+
 include("flux_utils.jl")
 include("predictive_model.jl")
 
-# special variable types
+"""
+    Policy{T}
+
+Policy variable type that holds plan and assess variables.
+"""
 struct Policy{T}
     plan::T
     assess::T
@@ -20,6 +24,11 @@ end
 +(p1::Policy, p2::Policy) = Policy(p1.plan + p2.plan, p1.assess + p2.assess)
 *(c::Number, p::Policy) = Policy(c*p.plan, c*p.assess)
 
+"""
+    Forecast{T}
+
+Forecast variable type that holds plan and assess variables.
+"""
 struct Forecast{T}
     plan::T
     assess::T
@@ -28,7 +37,12 @@ end
 +(p1::Forecast, p2::Forecast) = Forecast(p1.plan + p2.plan, p1.assess + p2.assess)
 *(c::Number, p::Forecast) = Forecast(c*p.plan, c*p.assess)
 
-# main model
+"""
+    Model <: JuMP.AbstractModel
+
+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
     assess::JuMP.Model
@@ -100,6 +114,11 @@ function set_forecast_model(
     model.forecast = forecast
 end
 
+"""
+    forecast(model, X)
+
+Return forecast model output for given input.
+"""
 function forecast(model::Model, X::AbstractMatrix)
     return model.forecast(X)
 end
@@ -159,7 +178,11 @@ include("optimizers/nelder_mead_mpi.jl")
 include("optimizers/gradient_mpi.jl")
 include("optimizers/bilevel.jl")
 
+"""
+    train!(model, X, y, options)
 
+Train model using given data and options.
+"""
 function train!(
     model::Model, 
     X::Matrix{<:Real}, 

src/flux_utils.jl

@@ -1,10 +1,20 @@
 using Flux
 
+"""
+    extract_flux_params(model)
+
+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)
     return reduce(vcat, [vec(p) for p in θ])
 end
 
+"""
+    fix_flux_params_single_model(model, θ)
+
+Return model after fixing the parameters from an adequate vector of parameters.
+"""
 function fix_flux_params_single_model(model::Union{Flux.Chain, Flux.Dense}, θ::Vector{<:Real})
     i = 1
     for p in Flux.params(model)
@@ -15,6 +25,11 @@ function fix_flux_params_single_model(model::Union{Flux.Chain, Flux.Dense}, θ::
     return model
 end
 
+"""
+    fix_flux_params_multi_model(models, θ)
+
+Return iterable of models after fixing the parameters from an adequate vector of parameters.
+"""
 function fix_flux_params_multi_model(
     models, 
     θ::Vector{<:Real}
@@ -30,6 +45,11 @@ function fix_flux_params_multi_model(
     return models
 end
 
+"""
+    has_params(layer)
+
+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

src/jump.jl

@@ -80,11 +80,20 @@ function JuMP.add_variable(
     return forecast
 end
 
-# plan and assess models
+"""
+    Upper(model::ApplicationDrivenLearning.Model)
+
+Create a reference to the plan model of an application driven learning model.
+"""
 function Plan(model::Model)
     return model.plan::JuMP.Model
 end
 
+"""
+    Assess(model::ApplicationDrivenLearning.Model)
+
+Create a reference to the assess model of an application driven learning model.
+"""
 function Assess(model::Model)
     return model.assess::JuMP.Model
 end

src/options.jl

@@ -1,11 +1,82 @@
 abstract type AbstractOptimizationMode end
 
+"""
+    BilevelMode <: AbstractOptimizationMode
+
+Used to solve the application driven learning training problem as a bilevel optimization problem
+by using the BilevelJuMP.jl package.
+
+...
+# Parameters
+- `optimizer::Function` is equivalent to `solver` in BilevelJuMP.BilevelModel.
+- `silent::Bool` is equivalent to `silent` in BilevelJuMP.BilevelModel.
+- `mode::Union{Nothing, BilevelJuMP.BilevelMode}` is equivalent to `mode` in BilevelJuMP.BilevelModel. 
+...
+"""
 struct BilevelMode <: AbstractOptimizationMode end
+
+"""
+    NelderMeadMode <: AbstractOptimizationMode
+
+Used to solve the application driven learning training problem using the Nelder-Mead optimization method
+implementation from Optim.jl package.
+
+...
+# Parameters
+- `initial_simplex` is the initial simplex of solutions to be applied.
+- `parameters` is the parameters to be applied to the Nelder-Mead optimization method.
+...
+"""
 struct NelderMeadMode <: AbstractOptimizationMode end
+
+"""
+    GradientMode <: AbstractOptimizationMode
+
+Used to solve the application driven learning training problem using the gradient optimization method
+
+...
+# Parameters
+- `rule` is the optimiser object to be used in the gradient optimization process.
+- 'epochs' is the number of epochs to be used in the gradient optimization process.
+- 'batch_size' is the batch size to be used in the gradient optimization process.
+- 'verbose' is the flag of whether to print the training process.
+- 'compute_cost_every' is the epoch frequency for computing the cost and evaluating best solution.
+- 'time_limit' is the time limit for the training process.
+...
+"""
 struct GradientMode <: AbstractOptimizationMode end
+
+"""
+    NelderMeadMPIMode <: AbstractOptimizationMode
+
+MPI implementation of NelderMeadMode.
+"""
 struct NelderMeadMPIMode <: AbstractOptimizationMode end
+
+"""
+    GradientMPIMode <: AbstractOptimizationMode
+
+MPI implementation of GradientMode.
+"""
 struct GradientMPIMode <: AbstractOptimizationMode end
 
+"""
+    Options(mode; params...)
+
+Options struct to hold optimization mode and mode parameters.
+
+...
+# Example
+```julia
+options = Options(
+    GradientMode; 
+    rule=Optim.RMSProp(0.01),
+    epochs=100,
+    batch_size=10
+)
+```
+...
+"""
 struct Options
     mode
     params::Dict{Symbol, Any}