NeuralODE depending on external disturbances

[mariaade26]

Hi everyone, I am trying to solve this neuralODE which describes a RC building thermal model. I start having problems when I include the dependence of the equation on Text and Phih, two external vectors.
Does anyone have any suggestion?
I would be very grateful!

This is the definition of the ODE

using DifferentialEquations, Plots, Flux,Optim, DiffEqFlux, DataInterpolations,Random, ComponentArrays, Lux
using Optimization, OptimizationOptimisers, OptimizationOptimJL

function RC!(du,u,p,t)
    Rv, Ci = p
    Text = ext_disturbance(t)
    phih= ext_disturbance2(t)
    
    du[1] = 1/(Rv*Ci) .* (Text .- u[1]) .+ phih/Ci

end 
 
u0= [5.0]
tspan= (0.0f0, 5.0f00)
t= range(tspan[1], tspan[2], length=10)
p= [0.05, 10.0]


Testerna = [1,2,3,4,5,6,7,8,9,10]
phih = [10, 20, 30, 40, 50, 60 ,70 ,80 ,90 ,100];

disturbance = LinearInterpolation(Testerna,t);
disturbance2 = LinearInterpolation(phih,t);

function ext_disturbance(t)
return disturbance(t)
end

function ext_disturbance2(t)
return disturbance2(t)
end 
    
#ODE

prob=ODEProblem(RC!, u0, tspan, p)
sol= solve(prob, Tsit5(), saveat=t)
plot(t,sol[1,:], label = "reale",xlabel = "tempo [s]", ylabel = "Temperatura [°C]"); `

and this is where I try to go on with the NeuralODE and its optimization

prob3 = ODEProblem(RC!, u0, tspan, p_random)

dudt= Lux.Chain( p-> solve(prob3, Tsit5(), saveat=t, p=p)[1,:],
vec,
Lux.Dense(10,50, tanh),
Lux.Dense(50,1))

n_ode= NeuralODE(dudt, tspan, Tsit5(), saveat=t)
p = [0.05, 10]
rng= Random.default_rng()
p, st = Lux.setup(rng,dudt)

function predict_neuralode(p)
    Array(n_ode(u0,p,st)[1])
end 

function loss_neuralode(p)
    pred= predict_neuralode(p)
    loss = sum(abs2,  sol .- pred)
    return loss, pred
end 

callback = function ()
display(loss())
    end 
    
pinit = ComponentArray(p)

adtype = Optimization.AutoZygote()

optf = Optimization.OptimizationFunction((x, p) -> loss_neuralode(x), adtype)
optprob = Optimization.OptimizationProblem(optf, pinit)

result_neuralode = Optimization.solve(optprob,
                                      Optimisers.Adam(0.05),
                                       callback = callback,
                                       maxiters = 400)

In particolar, I get this kind of error "

BoundsError: attempt to access 1-element Vector{Float64} at index [2]