Update API ProblemIterator

.github/workflows/CI.yml

@@ -18,7 +18,6 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.6'
           - '1.9'
         os:
           - ubuntu-latest

Project.toml

@@ -1,7 +1,7 @@
 name = "L2O"
 uuid = "e1d8bfa7-c465-446a-84b9-451470f6e76c"
 authors = ["andrewrosemberg <[email protected]> and contributors"]
-version = "1.2.0-DEV"
+version = "1.0.0"
 
 [deps]
 Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45"
@@ -21,12 +21,12 @@ UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
-Arrow = "2"
-CSV = "0.10"
-JuMP = "1"
-ParametricOptInterface = "0.7"
+Arrow = "^2"
+CSV = "^0.10"
+JuMP = "^1"
+ParametricOptInterface = "^0.8"
 Zygote = "^0.6.68"
-julia = "1.6"
+julia = "^1.9"
 
 [extras]
 CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2"

src/FullyConnected.jl

@@ -66,14 +66,20 @@ end
 # @forward((ConvexRegressor, :model), MLJFlux.Regressor)
 
 # Define a container to hold any optimiser specific parameters (if any):
-struct ConvexRule <: Flux.Optimise.AbstractOptimiser
-    rule::Flux.Optimise.AbstractOptimiser
+struct ConvexRule <: Optimisers.AbstractRule
+    rule::Optimisers.AbstractRule
     tol::Real
 end
-function ConvexRule(rule::Flux.Optimise.AbstractOptimiser; tol=1e-6)
+function ConvexRule(rule::Optimisers.AbstractRule; tol=1e-6)
     return ConvexRule(rule, tol)
 end
 
+Optimisers.init(o::ConvexRule, x::AbstractArray) = Optimisers.init(o.rule, x)
+
+function Optimisers.apply!(o::ConvexRule, mvel, x::AbstractArray{T}, dx) where T
+    return Optimisers.apply!(o.rule, mvel, x, dx)
+end
+
 """
     function make_convex!(chain::PairwiseFusion; tol = 1e-6)
 
@@ -102,24 +108,48 @@ function make_convex!(model::Chain; tol=1e-6)
     end
 end
 
-function MLJFlux.train!(
-    model::MLJFlux.MLJFluxDeterministic, penalty, chain, optimiser::ConvexRule, X, y
-)
+function MLJFlux.train(
+    model,
+    chain,
+    optimiser::ConvexRule,
+    optimiser_state,
+    epochs,
+    verbosity,
+    X,
+    y,
+    )
+
     loss = model.loss
+
+    # intitialize and start progress meter:
+    meter = MLJFlux.Progress(epochs + 1, dt=0, desc="Optimising neural net:",
+        barglyphs=MLJFlux.BarGlyphs("[=> ]"), barlen=25, color=:yellow)
+    verbosity != 1 || MLJFlux.next!(meter)
+
+    # initiate history:
     n_batches = length(y)
-    training_loss = zero(Float32)
-    for i in 1:n_batches
-        parameters = Flux.params(chain)
-        gs = Flux.gradient(parameters) do
-            yhat = chain(X[i])
-            batch_loss = loss(yhat, y[i]) + penalty(parameters) / n_batches
-            training_loss += batch_loss
-            return batch_loss
-        end
-        Flux.update!(optimiser.rule, parameters, gs)
+
+    losses = (loss(chain(X[i]), y[i]) for i in 1:n_batches)
+    history = [mean(losses),]
+
+    for i in 1:epochs
+        chain, optimiser_state, current_loss = MLJFlux.train_epoch(
+            model,
+            chain,
+            optimiser,
+            optimiser_state,
+            X,
+            y,
+        )
         make_convex!(chain; tol=optimiser.tol)
+        verbosity < 2 ||
+            @info "Loss is $(round(current_loss; sigdigits=4))"
+        verbosity != 1 || next!(meter)
+        push!(history, current_loss)
     end
-    return training_loss / n_batches
+
+    return chain, optimiser_state, history
+
 end
 
 function train!(model, loss, opt_state, X, Y; _batchsize=32, shuffle=true)

src/datasetgen.jl

@@ -105,6 +105,14 @@ end
 
 abstract type AbstractProblemIterator end
 
+abstract type AbstractParameterType end
+
+abstract type POIParamaterType <: AbstractParameterType end
+
+abstract type JuMPNLPParameterType <: AbstractParameterType end
+
+abstract type JuMPParameterType <: AbstractParameterType end
+
 """
     ProblemIterator(ids::Vector{UUID}, pairs::Dict{VariableRef, Vector{Real}})
 
@@ -115,24 +123,30 @@ struct ProblemIterator{T<:Real} <: AbstractProblemIterator
     ids::Vector{UUID}
     pairs::Dict{VariableRef,Vector{T}}
     early_stop::Function
+    param_type::Type{<:AbstractParameterType}
+    pre_solve_hook::Function
     function ProblemIterator(
         ids::Vector{UUID},
         pairs::Dict{VariableRef,Vector{T}},
         early_stop::Function=(args...) -> false,
+        param_type::Type{<:AbstractParameterType}=POIParamaterType,
+        pre_solve_hook::Function=(args...) -> nothing
     ) where {T<:Real}
         model = JuMP.owner_model(first(keys(pairs)))
         for (p, val) in pairs
             @assert length(ids) == length(val)
         end
-        return new{T}(model, ids, pairs, early_stop)
+        return new{T}(model, ids, pairs, early_stop, param_type, pre_solve_hook)
     end
 end
 
 function ProblemIterator(
-    pairs::Dict{VariableRef,Vector{T}}; early_stop::Function=(args...) -> false
-) where {T<:Real}
+    pairs::Dict{VariableRef,Vector{T}}; early_stop::Function=(args...) -> false,
+    pre_solve_hook::Function=(args...) -> nothing,
+    param_type::Type{<:AbstractParameterType}=POIParamaterType,
     ids = [uuid1() for _ in 1:length(first(values(pairs)))]
-    return ProblemIterator(ids, pairs, early_stop)
+) where {T<:Real}
+    return ProblemIterator(ids, pairs, early_stop, param_type, pre_solve_hook)
 end
 
 """
@@ -174,7 +188,8 @@ end
 
 function load(model_file::AbstractString, input_file::AbstractString, ::Type{T}; 
     batch_size::Union{Nothing, Integer}=nothing,
-    ignore_ids::Vector{UUID}=UUID[]
+    ignore_ids::Vector{UUID}=UUID[],
+    param_type::Type{<:AbstractParameterType}=JuMPParameterType
 ) where {T<:FileType}
     # Load full set
     df = load(input_file, T)
@@ -191,31 +206,40 @@ function load(model_file::AbstractString, input_file::AbstractString, ::Type{T};
     # No batch
     if isnothing(batch_size)
         pairs = _dataframe_to_dict(df, model_file)
-        return ProblemIterator(ids, pairs)
+        return ProblemIterator(pairs; ids=ids, param_type=param_type)
     end
     # Batch
     num_batches = ceil(Int, length(ids) / batch_size)
     idx_range = (i) -> (i-1)*batch_size+1:min(i*batch_size, length(ids))
-    return (i) -> ProblemIterator(ids[idx_range(i)], _dataframe_to_dict(df[idx_range(i), :], model_file)), num_batches
+    return (i) -> ProblemIterator(_dataframe_to_dict(df[idx_range(i), :], model_file);
+        ids=ids[idx_range(i)], param_type=param_type), num_batches
 end
 
 """
     update_model!(model::JuMP.Model, p::VariableRef, val::Real)
 
 Update the value of a parameter in a JuMP model.
 """
-function update_model!(model::JuMP.Model, p::VariableRef, val)
+function update_model!(::Type{POIParamaterType}, model::JuMP.Model, p::VariableRef, val)
     return MOI.set(model, POI.ParameterValue(), p, val)
 end
 
+function update_model!(::Type{JuMPNLPParameterType}, model::JuMP.Model, p::VariableRef, val)
+    return set_parameter_value(p, val)
+end
+
+function update_model!(::Type{JuMPParameterType}, model::JuMP.Model, p::VariableRef, val)
+    return fix(p, val)
+end
+
 """
     update_model!(model::JuMP.Model, pairs::Dict, idx::Integer)
 
 Update the values of parameters in a JuMP model.
 """
-function update_model!(model::JuMP.Model, pairs::Dict, idx::Integer)
+function update_model!(model::JuMP.Model, pairs::Dict, idx::Integer, param_type::Type{<:AbstractParameterType})
     for (p, val) in pairs
-        update_model!(model, p, val[idx])
+        update_model!(param_type, model, p, val[idx])
     end
 end
 
@@ -228,7 +252,8 @@ function solve_and_record(
     problem_iterator::ProblemIterator, recorder::Recorder, idx::Integer
 )
     model = problem_iterator.model
-    update_model!(model, problem_iterator.pairs, idx)
+    problem_iterator.pre_solve_hook(model)
+    update_model!(model, problem_iterator.pairs, idx, problem_iterator.param_type)
     optimize!(model)
     status = recorder.filterfn(model)
     early_stop_bool = problem_iterator.early_stop(model, status, recorder)

src/samplers.jl

@@ -119,7 +119,7 @@ end
 Load the parameters from a JuMP model.
 """
 function load_parameters(model::JuMP.Model)
-    cons = constraint_object.(all_constraints(model, VariableRef, MOI.Parameter{Float64}))
+    cons = constraint_object.([all_constraints(model, VariableRef, MOI.Parameter{Float64}); all_constraints(model, VariableRef, MOI.EqualTo{Float64})])
     parameters = [cons[i].func for i in 1:length(cons)]
     vals = [cons[i].set.value for i in 1:length(cons)]
     return parameters, vals

test/datasetgen.jl

@@ -5,7 +5,7 @@
 Test dataset generation for different filetypes
 """
 function test_problem_iterator(path::AbstractString)
-    @testset "Dataset Generation Type: $filetype" for filetype in [CSVFile, ArrowFile]
+    @testset "Dataset Generation (POI) Type: $filetype" for filetype in [CSVFile, ArrowFile]
         # The problem to iterate over
         model = JuMP.Model(() -> POI.Optimizer(HiGHS.Optimizer()))
         @variable(model, x)
@@ -55,6 +55,17 @@ function test_problem_iterator(path::AbstractString)
             @test num_p * successfull_solves == 1
         end
 
+        @testset "pre_solve_hook" begin
+            file_dual_output = joinpath(path, "test_$(string(uuid1()))_output") # file path
+            recorder_dual = Recorder{filetype}(file_dual_output; dual_variables=[cons])
+            sum_p = 0
+            problem_iterator = ProblemIterator(
+                Dict(p => collect(1.0:num_p)); pre_solve_hook=(args...) -> sum_p += 1
+            )
+            successfull_solves = solve_batch(problem_iterator, recorder_dual)
+            @test sum_p == num_p
+        end
+
         @testset "solve_batch" begin
             successfull_solves = solve_batch(problem_iterator, recorder)
 
@@ -96,6 +107,58 @@ function test_problem_iterator(path::AbstractString)
             end
         end
     end
+    @testset "Dataset Generation JuMP" begin
+        model = JuMP.Model(HiGHS.Optimizer)
+        @variable(model, x)
+        p = @variable(model, _p)
+        @constraint(model, cons, x + _p >= 3)
+        @objective(model, Min, 2x)
+        num_p = 10
+        batch_id = string(uuid1())
+        problem_iterator = ProblemIterator(Dict(p => collect(1.0:num_p)); param_type=L2O.JuMPParameterType)
+        file_output = joinpath(path, "test_$(batch_id)_output") # file path
+        recorder = Recorder{ArrowFile}(
+            file_output; primal_variables=[x], dual_variables=[cons]
+        )
+        successfull_solves = solve_batch(problem_iterator, recorder)
+        iter_files = readdir(joinpath(path))
+        iter_files = filter(x -> occursin(string(ArrowFile), x), iter_files)
+        file_outs = [
+            joinpath(path, file) for
+            file in iter_files if occursin("$(batch_id)_output", file)
+        ]
+        # test output file
+        df = Arrow.Table(file_outs)
+        @test length(df) == 8
+        @test length(df[1]) == num_p * successfull_solves
+        rm.(file_outs)
+    end
+    @testset "Dataset Generation JuMPNLP" begin
+        model = JuMP.Model(Ipopt.Optimizer)
+        @variable(model, x)
+        p = @variable(model, _p in MOI.Parameter(1.0))
+        @constraint(model, cons, x + _p >= 3)
+        @objective(model, Min, 2x)
+        num_p = 10
+        batch_id = string(uuid1())
+        problem_iterator = ProblemIterator(Dict(p => collect(1.0:num_p)); param_type=L2O.JuMPNLPParameterType)
+        file_output = joinpath(path, "test_$(batch_id)_output") # file path
+        recorder = Recorder{ArrowFile}(
+            file_output; primal_variables=[x], dual_variables=[cons]
+        )
+        successfull_solves = solve_batch(problem_iterator, recorder)
+        iter_files = readdir(joinpath(path))
+        iter_files = filter(x -> occursin(string(ArrowFile), x), iter_files)
+        file_outs = [
+            joinpath(path, file) for
+            file in iter_files if occursin("$(batch_id)_output", file)
+        ]
+        # test output file
+        df = Arrow.Table(file_outs)
+        @test length(df) == 8
+        @test length(df[1]) == num_p * successfull_solves
+        rm.(file_outs)
+    end
 end
 
 function test_load(model_file::AbstractString, input_file::AbstractString, ::Type{T}, ids::Vector{UUID};

test/test_flux_forecaster.jl

@@ -21,7 +21,7 @@ function test_flux_forecaster(file_in::AbstractString, file_out::AbstractString)
             rng=123,
             epochs=20,
             optimiser=ConvexRule(
-                Flux.Optimise.Adam(0.001, (0.9, 0.999), 1.0e-8, IdDict{Any,Any}())
+                Optimisers.Adam()
             ),
         )