Add Scale predictor (#89)

docs/src/api.md

@@ -86,6 +86,11 @@ ReLUQuadratic
 ReLUSOS1
 ```
 
+## `Scale`
+```@docs
+Scale
+```
+
 ## `Sigmoid`
 ```@docs
 Sigmoid

docs/src/manual/predictors.md

@@ -20,6 +20,7 @@ The following predictors are supported. See their docstrings for details:
 | [`ReLUBigM`](@ref) |  $f(x) = \max.(0, x)$                  | $M \rightarrow M$   |
 | [`ReLUQuadratic`](@ref) |  $f(x) = \max.(0, x)$             | $M \rightarrow M$   |
 | [`ReLUSOS1`](@ref) |  $f(x) = \max.(0, x)$                  | $M \rightarrow M$   |
+| [`Scale`](@ref)    |  $f(x) = scale .* x .+ bias$           | $M \rightarrow M$   |
 | [`Sigmoid`](@ref)  |  $f(x) = \frac{1}{1 + e^{-x}}$         | $M \rightarrow M$   |
 | [`SoftMax`](@ref)  |  $f(x) = \frac{e^{x_i}}{\sum e^{x_i}}$ | $M \rightarrow 1$   |
 | [`SoftPlus`](@ref) |  $f(x) = \log(1 + e^x)$                | $M \rightarrow M$   |

ext/MathOptAIFluxExt.jl

@@ -24,6 +24,7 @@ Add a trained neural network from Flux.jl to `model`.
 ## Supported layers
 
  * `Flux.Dense`
+ * `Flux.Scale`
  * `Flux.softmax`
 
 ## Supported activation functions
@@ -86,6 +87,7 @@ Convert a trained neural network from Flux.jl to a [`Pipeline`](@ref).
 ## Supported layers
 
  * `Flux.Dense`
+ * `Flux.Scale`
  * `Flux.softmax`
 
 ## Supported activation functions
@@ -173,4 +175,14 @@ function _add_predictor(
     return
 end
 
+function _add_predictor(
+    predictor::MathOptAI.Pipeline,
+    layer::Flux.Scale,
+    config::Dict,
+)
+    push!(predictor.layers, MathOptAI.Scale(layer.scale, layer.bias))
+    _add_predictor(predictor, layer.σ, config)
+    return
+end
+
 end  # module

ext/MathOptAILuxExt.jl

@@ -24,6 +24,7 @@ Add a trained neural network from Lux.jl to `model`.
 ## Supported layers
 
  * `Lux.Dense`
+ * `Lux.Scale`
 
 ## Supported activation functions
 
@@ -96,6 +97,7 @@ Convert a trained neural network from Lux.jl to a [`Pipeline`](@ref).
 ## Supported layers
 
  * `Lux.Dense`
+ * `Lux.Scale`
 
 ## Supported activation functions
 
@@ -188,4 +190,15 @@ function _add_predictor(
     return
 end
 
+function _add_predictor(
+    predictor::MathOptAI.Pipeline,
+    layer::Lux.Scale,
+    p,
+    config::Dict,
+)
+    push!(predictor.layers, MathOptAI.Scale(p.weight, p.bias))
+    _add_predictor(predictor, layer.activation, config)
+    return
+end
+
 end  # module

src/predictors/Scale.jl

@@ -0,0 +1,82 @@
+# Copyright (c) 2024: Oscar Dowson and contributors
+# Copyright (c) 2024: Triad National Security, LLC
+#
+# Use of this source code is governed by a BSD-style license that can be found
+# in the LICENSE.md file.
+
+"""
+    Scale(
+        scale::Vector{T},
+        bias::Vector{T},
+    ) where {T} <: AbstractPredictor
+
+An [`AbstractPredictor`](@ref) that represents the affine relationship:
+```math
+f(x) = scale .* x .+ bias
+```
+
+## Example
+
+```jldoctest
+julia> using JuMP, MathOptAI
+
+julia> model = Model();
+
+julia> @variable(model, x[1:2]);
+
+julia> f = MathOptAI.Scale([2.0, 3.0], [4.0, 5.0])
+Scale(scale, bias)
+
+julia> y = MathOptAI.add_predictor(model, f, x)
+1-element Vector{VariableRef}:
+ moai_Scale[1]
+ moai_Scale[2]
+
+julia> print(model)
+Feasibility
+Subject to
+ 2 x[1] + moai_Scale[1] = -4
+ 3 x[2] + moai_Scale[2] = -5
+
+julia> y = MathOptAI.add_predictor(model, MathOptAI.ReducedSpace(f), x)
+2-element Vector{AffExpr}:
+ 2 x[1] + 4
+ 3 x[2] + 5
+```
+"""
+struct Scale{T} <: AbstractPredictor
+    scale::Vector{T}
+    bias::Vector{T}
+end
+
+function Base.show(io::IO, ::Scale)
+    return print(io, "Scale(scale, bias)")
+end
+
+function add_predictor(
+    model::JuMP.AbstractModel,
+    predictor::Scale,
+    x::Vector,
+)
+    m = length(predictor.scale)
+    y = JuMP.@variable(model, [1:m], base_name = "moai_Scale")
+    bounds = _get_variable_bounds.(x)
+    for (i, scale) in enumerate(predictor.scale)
+        y_lb = y_ub = predictor.bias[i]
+        lb, ub = bounds[i]
+        y_ub += scale * ifelse(scale >= 0, ub, lb)
+        y_lb += scale * ifelse(scale >= 0, lb, ub)
+        _set_bounds_if_finite(y[i], y_lb, y_ub)
+    end
+    JuMP.@constraint(model, predictor.scale .* x .+ predictor.bias .== y)
+    return y
+end
+
+function add_predictor(
+    model::JuMP.AbstractModel,
+    predictor::ReducedSpace{<:Scale},
+    x::Vector,
+)
+    scale, bias = predictor.predictor.scale, predictor.predictor.bias
+    return JuMP.@expression(model, scale .* x .+ bias)
+end

test/test_Flux.jl

@@ -41,6 +41,31 @@ function _train_lux_model(model)
     return model
 end
 
+function test_end_to_end_with_scale()
+    chain = _train_lux_model(
+        Flux.Chain(
+            Flux.Scale(1),
+            Flux.Dense(1 => 16, Flux.relu),
+            Flux.Dense(16 => 1),
+        ),
+    )
+    model = Model(HiGHS.Optimizer)
+    set_silent(model)
+    @variable(model, x)
+    y = MathOptAI.add_predictor(
+        model,
+        chain,
+        [x];
+        config = Dict(Flux.relu => MathOptAI.ReLUBigM(100.0)),
+    )
+    @constraint(model, only(y) <= 4)
+    @objective(model, Min, x)
+    optimize!(model)
+    @test is_solved_and_feasible(model)
+    @test isapprox(value(x), -1.24; atol = 1e-1)
+    return
+end
+
 function test_end_to_end_ReLUBigM()
     chain = _train_lux_model(
         Flux.Chain(Flux.Dense(1 => 16, Flux.relu), Flux.Dense(16 => 1)),

test/test_Lux.jl

@@ -49,6 +49,31 @@ function _train_lux_model(model)
     return (model, parameters, state)
 end
 
+function test_end_to_end_with_scale()
+    state = _train_lux_model(
+        Lux.Chain(
+            Lux.Scale(1),
+            Lux.Dense(1 => 16, Lux.relu),
+            Lux.Dense(16 => 1),
+        ),
+    )
+    model = Model(HiGHS.Optimizer)
+    set_silent(model)
+    @variable(model, x)
+    y = MathOptAI.add_predictor(
+        model,
+        state,
+        [x];
+        config = Dict(Lux.relu => MathOptAI.ReLUBigM(100.0)),
+    )
+    @constraint(model, only(y) <= 4)
+    @objective(model, Min, x)
+    optimize!(model)
+    @test is_solved_and_feasible(model)
+    @test isapprox(value(x), -1.24; atol = 1e-2)
+    return
+end
+
 function test_end_to_end_ReLUBigM()
     state = _train_lux_model(
         Lux.Chain(Lux.Dense(1 => 16, Lux.relu), Lux.Dense(16 => 1)),

test/test_predictors.jl

@@ -337,6 +337,26 @@ function test_ReducedSpace_ReducedSpace()
     return
 end
 
+function test_Scale()
+    model = Model()
+    @variable(model, x[1:2])
+    f = MathOptAI.Scale([2.0, 3.0], [4.0, 5.0])
+    @test sprint(show, f) == "Scale(scale, bias)"
+    y = MathOptAI.add_predictor(model, f, x)
+    cons = all_constraints(model; include_variable_in_set_constraints = false)
+    @test length(cons) == 2
+    objs = constraint_object.(cons)
+    @test objs[1].set == MOI.EqualTo(-4.0)
+    @test objs[2].set == MOI.EqualTo(-5.0)
+    @test isequal_canonical(objs[1].func, 2.0 * x[1] - y[1])
+    @test isequal_canonical(objs[2].func, 3.0 * x[2] - y[2])
+    y = MathOptAI.add_predictor(model, MathOptAI.ReducedSpace(f), x)
+    cons = all_constraints(model; include_variable_in_set_constraints = false)
+    @test length(cons) == 2
+    @test isequal_canonical(y, [2.0 * x[1] + 4.0, 3.0 * x[2] + 5.0])
+    return
+end
+
 end  # module
 
 TestPredictors.runtests()