Merge pull request #2 from EarthyScience/bounds

bounds

Author: lazarusA

src/HybridSymbolics.jl

@@ -9,6 +9,7 @@ using ADTypes
 using ForwardDiff
 using ProgressMeter
 using ProtoStructs
+using Random
 
 include("hybridModel.jl")
 include("macroHybrid.jl")

src/hybridModel.jl

@@ -1,4 +1,4 @@
-export Global, Varying, Fixed, PartitionedFunction, HybridModel
+export Global, Varying, Fixed, PartitionedFunction, HybridModel, setbounds, setup
 export HybridSymbolic, SymbolTypes
 
 abstract type SymbolTypes end
@@ -39,9 +39,15 @@ struct PartitionedFunction{F,O,A1,A2,A3,A4,V} <: HybridSymbolic
     end
 end
 
-struct HybridModel <: HybridSymbolic
+@proto struct HybridModel{T} <: HybridSymbolic
     nn::Lux.Chain
     func::PartitionedFunction
+    p_min::T 
+    p_max::T 
+end
+
+function HybridModel(nn::Lux.Chain, func::PartitionedFunction)
+    return HybridModel(nn, func, nothing, nothing)
 end
 # TODO: This needs to be more general. i.e. ŷ = NN(α * NN(x) + β).
 
@@ -60,4 +66,24 @@ function (m::HybridModel)(X::Vector{Float32}, params, st)
     out_NN = m.nn(X, ps, st)[1]
     out = m.func.opt_func(tuple([[out_NN[1]] for i = 1:n_varargs]...), globals)
     return out[1]
-end
+end
+
+# Assumes that the last layer has sigmoid activation function
+function setbounds(m::HybridModel, bounds::Dict{Symbol, Tuple{T,T}}) where {T}
+    n_args = length(m.func.varying_args)
+    p_min = zeros(Float32, n_args)
+    p_max = zeros(Float32, n_args)
+    for (i,arg) in enumerate(Symbol.(m.func.varying_args))
+        @assert arg in keys(bounds)
+        p_min[i] = bounds[arg][1]
+        p_max[i] = bounds[arg][2]
+    end
+    p_range = p_max .- p_min
+    wf = WrappedFunction((x) -> x .* (p_range) .+ p_min)
+    new_nn = Chain(m.nn, wf)
+    return HybridModel(new_nn, m.func, p_min, p_max)
+end
+
+function setup(rng::AbstractRNG, m::HybridModel)
+    return Lux.setup(rng, m.nn)
+end

test/core.jl

@@ -47,6 +47,38 @@ function test_hybridmodel()
     @test model(rand(Float32, 5,5), model_params, st) isa Vector{Float32}
 end
 
+function test_bounds()
+    local input_size = 5
+    local α, β, γ, δ
+    @syms α::Real β::Real γ::Real δ::Real
+    structured = @hybrid function testfunc(α::Varying, β::Varying, γ::Fixed=1.0, δ::Global)
+        return (exp.(α) .- β)./(γ .* δ)
+    end
+    NN = Chain(
+       Dense(input_size => 4, sigmoid_fast),
+       Dense(4 => 2, sigmoid_fast)
+    )
+    NN = f32(NN)
+    rng = MersenneTwister()
+    model = HybridModel(
+        NN,
+        structured
+    )
+    model = setbounds(model, Dict(:α => (-1.0f0, 1.0f0), :β => (-1.0f0, 1.0f0)))
+    @test model isa HybridModel
+    ps, st = setup(rng, model)
+    globals = [1.2f0]
+    model_params = (nn = ps, globals = globals)
+    @test model(rand(Float32, 5), model_params, st) isa Float32
+    @test model(rand(Float32, 5,5), model_params, st) isa Vector{Float32}
+    @test all(model.p_min .== -1.0f0)
+    @test all(model.p_max .== 1.0f0)
+    @testset "Testing Bounds: input $item" for item in [rand(Float32, 5) for _ in 1:10]
+        output_params = model.nn(item, ps, st)[1]
+        @test all(output_params .>= -1.0f0) && all(output_params .<= 1.0f0)
+    end
+end
+
 function test_gradcalc()
     local input_size = 5
     local α, β, γ, δ

test/runtests.jl

@@ -5,5 +5,6 @@ include("./core.jl")
 @testset "HybridSymbolics.jl" begin
     @testset test_structuredfunc()
     @testset test_hybridmodel()
+    @testset test_bounds()
     @testset test_gradcalc()
 end

test/sample.jl

@@ -20,13 +20,14 @@ NN = Chain(
 
 NN = f32(NN)
 rng = MersenneTwister()
-ps, st = Lux.setup(rng, NN)
 
 model = HybridModel(
     NN,
     structured
 )
+model = setbounds(model, Dict(:α => (-1.0f0, 1.0f0), :β => (-1.0f0, 1.0f0)))
 
+ps, st = setup(rng, model)
 globals = [1.2f0]
 model_params = (nn = ps, globals = globals)