Add metadata to neural network parameters, designating them as such

[TorkelE]

Not directly needed, but it could be useful for downstream packages if theta (and maybe NN as well) had some metadata tag to designate what kind of parameters they are. I think in optimisation, one might want to treat these a bit differently (i.e. on sys bio normal parameters are typically fitted on log scale, but I don't think this is desirable for neural network parameters). It might be worth considering, might be useful for downstream packages like PEtab, when they integrate UDE support.

[SebastianM-C]

What metadata are you thinking about? The tunable can be added, we already add it to the things that are strictly not tunable, but what other information would be needed?

[TorkelE]

I think just a flag might be useful, i.e. is_nn_parameter = true or something. I can double-check with Sebastian Person when he starts to work on UDEs for PEtab, but I think something like that would be enough.

[TorkelE]

I had a discussion yesterday with Sebastian (author of PEtab.jl and who works on UDE research) regarding this, and we do think some from of metadata would be very helpful for the development of downstream packages.

i.e., in this situation:

using ModelingToolkitBase, ModelingToolkitNeuralNets, Lux
chain = Lux.Chain(
    Lux.Dense(1 => 3, Lux.softplus, use_bias = false),
    Lux.Dense(3 => 3, Lux.softplus, use_bias = false),
    Lux.Dense(3 => 1, Lux.softplus, use_bias = false)
)
@parameters d
@SymbolicNeuralNetwork NN, θ = chain

it would be very useful to have the following utility functions:

isneuralnetwork
isneuralnetworkps

with these two potentially also being useful (but less so, and probably also requiring more consideration for implementing):

hasneuralnetworktag
getneuralnetworktag

(not sure what appropriate names would be)

Here

isneuralnetwork(d) # false
isneuralnetwork(NN) # true
isneuralnetwork(θ) # false
isneuralnetworkps(d) # false
isneuralnetworkps(NN) # false
isneuralnetworkps(θ) # true
hasneuralnetworktag(d) # false
hasneuralnetworktag(NN) # true
hasneuralnetworktag(θ) # true
getneuralnetworktag(d) # Error: ...
getneuralnetworktag(NN) # Returns some identifier
getneuralnetworktag(θ) # Returns some identifier

I.e. isneuralnetwork and isneuralnetworkps can be used to identify parameters relating to neural networks. Next, hasneuralnetworktag and getneuralnetworktag provides the functionality to find a tag coupling a neural network (NN) to its parameter set (θ). Maybe there is an alternative way to achieve this functionality, but something like this should work.

If ModelingToolkitNeuralNets provided functionality like this, it PEtab's new UDE functionality that Sebastian is developing could be directly applied to fit UDEs generated through ModelingToolkitNeuralNets. I.e. workflows like this could be made to work:

using ModelingToolkitBase, ModelingToolkitNeuralNets, Lux
using ModelingToolkitBase: t_nounits as t, D_nounits as D

# Create UDE.
@variables X(t) = 1.0 Y(t) = 0.2
@parameters d
chain = Lux.Chain(
    Lux.Dense(1 => 3, Lux.softplus, use_bias = false),
    Lux.Dense(3 => 3, Lux.softplus, use_bias = false),
    Lux.Dense(3 => 1, Lux.softplus, use_bias = false)
)
@SymbolicNeuralNetwork NN, θ = chain
eqs_ude = [
    D(X) ~ NN(Y, θ)[1] - d*X
    D(Y) ~ X - d*Y
]
@mtkcompile sys_ude = System(eqs_ude, t)

# Fit parameters through PEtab
using PEtab, OptimizationOptimisers
ps_est = [PEtabMLParameter(θ), PEtabParameter(d)]
observables = [PEtabObservable(:obs_X, X, 0.5)]
measurements = [] # Either load from file or generate a synthetic measurements DataFrame.
petab_model = PEtabModel(sys_ude, observables, measurements, ps_est)
petab_prob = PEtabODEProblem(model_ude)
fitted_ude = ms_res = calibrate_multistart(petab_prob, Adam(1e-3), 100; maxiters = 10000)

[TorkelE]

@SebastianM-C