How do I plot/check the trained neural network.

[TorkelE]

Considering the example in https://sciml.github.io/ModelingToolkitNeuralNets.jl/dev/friction/, it is shown how to plot the friction over time, and compared that to the tru value.

However, we also started with friction as a function of velocity:

function friction(v)
    sqrt(2 * MathConstants.e) * (Fbrk - Fc) * exp(-(v / vst)^2) * (v / vst) +
    Fc * tanh(v / vcol)
end

Would it be possible to evaluate the trained neural network (representing the friction) at various velocity values? Furthermore, could than the plot the real and trained frictions functions?

[SebastianM-C]

If you want to just look at the inputs and outputs of the neural network block, you can use

plot(res_sol, idxs=(sys.nn.input.u[1], sys.nn.output.u[1]))

which should show something like

which looks similar to the plot for sys.nn.output.u since the velocity is linear.

You can also leverage the plotting interface for comparing the friction with the results from the neural network

plot(res_sol, idxs=(sys.nn.input.u[1], sys.nn.output.u[1]))
plot!(res_sol, idxs=(sys.nn.input.u[1], friction(sys.nn.input.u[1])))

Bonus: now what if you want to evaluate the embedded neural network at arbitrary values?
You can actually recover the Lux network embedded in the equations

julia> arguments(arguments(equations(sys.nn)[1].rhs[1])[1])[1]
Chain(
    layer_1 = Dense(1 => 10, mish, use_bias=false),  # 10 parameters
    layer_2 = Dense(10 => 10, mish, use_bias=false),  # 100 parameters
    layer_3 = Dense(10 => 1, use_bias=false),  # 10 parameters
)         # Total: 120 parameters,
          #        plus 0 states.

and you can call this with whatever inputs are needed

julia> LuxCore.stateless_apply(arguments(arguments(equations(sys.nn)[1].rhs[1])[1])[1], [1.23], convert(res_sol.ps[sys.nn.T], res_sol.ps[sys.nn.p]))
1-element Vector{Float64}:
 18.04731050719053

but note that you also have to reconstruct the ComponentArray structure of the NN parameters.

We should make this easier by storing the network as a callable parameter, I don't consider the above "official recommendations" 😅 , but I wanted to point out that it's possible. Since it relies on implementation details of the NN block, this is not public API in any form, use at your own risk.

[TorkelE]

Thanks a lot! Yeah, something like this would be really useful. Especially when evaluating it, and there is a known function from which the data is fetched. Then I could check if this function is recovered (which is a lot more intuitive then just checking what it is evaluated for across the simulation.

Just checking so I get it right. Say that I want to evaluate the friction NN at velocity = 1.5, then I would do:

velocity = 1.5
friction = LuxCore.stateless_apply(arguments(arguments(equations(sys.nn)[1].rhs[1])[1])[1], [velocity], convert(res_sol.ps[sys.nn.T], res_sol.ps[sys.nn.p]))

?

[SebastianM-C]

yes, if you have access to the neural network from the "outside" of the model it's much easier and that's the recommended thing to do for now, but I wanted to point out that it should be technically possible even if you just have the model.

if you save the Lux model separately, you can just call it with LuxCore.stateless_apply(lux_model, nn_input, convert(res_sol.ps[sys.nn.T], res_sol.ps[sys.nn.p]). You can make it even easier if you have the neural network parameters on the "outside" too, as you can just convert res_sol.ps[sys.nn.p] to their type. Essentially lux models can use ComponentArrays to store the neural network parameters, so you just need to convert the vector of plain numbers to that type to restore the layer structure.