feat: add SymbolicNeuralNetwork as a callable parameter interface

Author: SebastianM-C

Project.toml

@@ -5,6 +5,7 @@ version = "1.6.1"
 
 [deps]
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
@@ -17,6 +18,7 @@ Aqua = "0.8"
 ComponentArrays = "0.15.11"
 DifferentiationInterface = "0.6"
 ForwardDiff = "0.10.36"
+IntervalSets = "0.7.10"
 JET = "0.8, 0.9"
 Lux = "1"
 LuxCore = "1"

src/ModelingToolkitNeuralNets.jl

@@ -1,14 +1,15 @@
 module ModelingToolkitNeuralNets
 
 using ModelingToolkit: @parameters, @named, ODESystem, t_nounits
+using IntervalSets: (..)
 using ModelingToolkitStandardLibrary.Blocks: RealInputArray, RealOutputArray
 using Symbolics: Symbolics, @register_array_symbolic, @wrapped
 using LuxCore: stateless_apply
 using Lux: Lux
 using Random: Xoshiro
 using ComponentArrays: ComponentArray
 
-export NeuralNetworkBlock, multi_layer_feed_forward
+export NeuralNetworkBlock, SymbolicNeuralNetwork, multi_layer_feed_forward, get_network
 
 include("utils.jl")
 
@@ -56,4 +57,65 @@ function lazyconvert(T, x::Symbolics.Arr)
     Symbolics.array_term(convert, T, x, size = size(x))
 end
 
+"""
+    SymbolicNeuralNetwork(; n_input = 1, n_output = 1,
+        chain = multi_layer_feed_forward(n_input, n_output),
+        rng = Xoshiro(0),
+        init_params = Lux.initialparameters(rng, chain),
+        eltype = Float64)
+
+Create symbolic parameter for a neural network and one for its parameters.
+Example:
+```
+chain = multi_layer_feed_forward(2, 2)
+NN, p = SymbolicNeuralNetwork(; chain, n_input=2, n_output=2, rng = StableRNG(42))
+```
+
+The NN and p are symbolic parameters that can be used lates as part of a system.
+To get the predictions of the neural network, use
+```
+pred ~ NN(input, p)
+```
+where `pred` and `input` are a symbolic vector variable with the lengths `n_output` and `n_input`.
+
+To use this outside of an equation, you can get the default values for the symbols and make a similar call
+```
+defaults(sys)[sys.NN](input, nn_p)
+```
+
+where `sys` is a system (e.g. `ODESystem`) that contains `NN`, `input` is a vector of `n_input` length and
+`nn_p` is a vector representing parameter values for the neural network.
+
+To get the underlying Lux model you can use `get_network(defaults(sys)[sys.NN])` or
+"""
+function SymbolicNeuralNetwork(; n_input = 1, n_output = 1,
+        chain = multi_layer_feed_forward(n_input, n_output),
+        rng = Xoshiro(0),
+        init_params = Lux.initialparameters(rng, chain),
+        eltype = Float64)
+    ca = ComponentArray{eltype}(init_params)
+    wrapper = StatelessApplyWrapper(chain, ca)
+
+    @parameters p[1:length(ca)] = Vector(ca)
+    @parameters (NN::typeof(wrapper))(..)[1:n_output] = wrapper
+
+    return NN, p
+end
+
+struct StatelessApplyWrapper{NN}
+    lux_model::NN
+    T::DataType
+end
+
+function (wrapper::StatelessApplyWrapper)(input::AbstractArray, nn_p::AbstractVector)
+    stateless_apply(get_network(wrapper), input, convert(wrapper.T, nn_p))
+end
+
+function Base.show(io::IO, m::MIME"text/plain", wrapper::StatelessApplyWrapper)
+    printstyled(io, "LuxCore.stateless_apply wrapper for:\n", color=:gray)
+    show(io, m, get_network(wrapper))
+end
+
+get_network(wrapper::StatelessApplyWrapper) = wrapper.lux_model
+
 end