Incorporating the ode within my neural network

Good morning everyone, I am working on using neural ode for household consumption prediction using Rc models. Once I have created my solution and array. How can I make sure to incorporate the ode within the neural network? I clarify that my goal is not to incorporate the solution but the PROBLEM ODE within my chain. in diffeqflux the chain is used and then the neural ode to minimize the difference from the solution of the ode then incorporating the physics only in the loss calculation. my goal now is to incorporate the physics directly within the chain, is this possible?

Yeah, it’s just a differentiable function. You can do anything you want with it. What did you try?

Hi chris, I haven’t solved the problem yet. I’m working on using Neural Ode to predict an RC model. Do you have documentation or code examples to suggest for adding physical losses to the total loss? can it be done for neural odes? another approach could be to use my problem ode within the chain but I don’t really know how to do it


using DifferentialEquations, Plots, Flux,Optim, DiffEqFlux, DataInterpolations,Random, ComponentArrays, Lux
using Optimization, OptimizationOptimisers, OptimizationOptimJL,OptimizationNLopt
rng = Random.default_rng()
using CSV
using DataFrames
using Plots
using Flux
using Statistics: mean
using DiffEqFlux, Optimization, OptimizationOptimJL,Plots
using ComponentArrays, Lux, DiffEqFlux, Optimization, OptimizationPolyalgorithms, DifferentialEquations, Plots
using DiffEqFlux: group_ranges
# Load the data
df = CSV.read("/Users/marco/OneDrive/Desktop/Tesi/Julia_files/measured_data_ex_00.csv", DataFrame)
df=repeat(df, outer=10)

registered_gas_flow = df[:, :2]
Registered_Temperature = df[:, 3]
sec_Temp = [mean(Registered_Temperature[i:i+59]) for i in 1:60:length(Registered_Temperature)-59]
sec_gas = [mean(registered_gas_flow[i:i+59]) for i in 1:60:length(registered_gas_flow)-59]
#create a 3600 time vector
function RC!(du,u,p,t)
    Rv, Ci,Rv2,Ci2 = p
    Text = ext_Temp(t)
    P= ext_power(t)
    P2= ext_power2(t)
    
    du[1] = 1/(Rv*Ci) .* (Text .- u[1]) .+ P/Ci
    end

u0= [20.0]


tspan= (0.0f0, 3000.0f0)
datasize = 3000
tsteps= range(tspan[1], tspan[2], length = datasize)
p= [10000, 10, 10000, 10]

Ext_temperature = LinearInterpolation(sec_Temp,tsteps);
power = LinearInterpolation(sec_gas,tsteps);
power2= LinearInterpolation(sec_gas,tsteps);
function ext_Temp(tsteps)
return Ext_temperature(tsteps)
end
    
function ext_power(tsteps)
return power(tsteps)
end   

function ext_power2(tsteps)
return power2(tsteps)
end

prob= ODEProblem(RC!, u0, tspan, p)
sol= solve(prob, Tsit5(), saveat=tsteps)[1,:]

# Verify ODE solution
ode_data =Array(solve(prob, Tsit5(), saveat=tsteps))

plot(tsteps, ode_data[1,:], label = "data")


#resize ode data 
ode_data=ode_data[:,1:30:3000]
tsteps=tsteps[1:30:3000]

plot(tsteps, ode_data[1,:], label = "data")

ode_data_train=ode_data[:,1:50]
ode_data_test=ode_data[:,51:100]
tspan= (0.0f0, 50.0f0)
tsteps= range(tspan[1], tspan[2], length = 50)
datasize=50

plot(ode_data_test[1,:], label = "data")

anim = Plots.Animation()

# Define the Neural Network
nn = Lux.Chain(
                Lux.Dense(1, 84, tanh),
                Lux.Dense(84, 44, tanh),
                Lux.Dense(44,1))
p_init, st = Lux.setup(rng, nn)

neuralode = NeuralODE(nn, tspan, Tsit5(), saveat = tsteps)
prob_node = ODEProblem((u,p,t)->nn(u,p,st)[1], u0, tspan, ComponentArray(p_init))

function plot_multiple_shoot(plt, preds, group_size)
	step = group_size-1
	ranges = group_ranges(datasize, group_size)

	for (i, rg) in enumerate(ranges)
		plot!(plt, tsteps[rg], preds[i][1,:], markershape=:circle)
	end
end

# Animate training, cannot make animation on CI server
# anim = Plots.Animation()
iter = 0
callback = function (p, l, preds; doplot = true)
  display(l)
  global iter
  iter += 1
  if doplot && iter%1 == 0
    # plot the original data
    plt = scatter(tsteps, ode_data_train[1,:], label = "Data")

    # plot the different predictions for individual shoot
    plot_multiple_shoot(plt, preds, group_size)

    frame(anim)
    display(plot(plt))
  end
  return false
end

# Define parameters for Multiple Shooting
group_size = 4
continuity_term = 100

function loss_function(data, pred)
	return sum(abs2, data - pred)
end

function loss_multiple_shooting(p)
    return multiple_shoot(p, ode_data_train, tsteps, prob_node, loss_function, Tsit5(),
                          group_size; continuity_term)
end

adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x,p) -> loss_multiple_shooting(x), adtype)
optprob = Optimization.OptimizationProblem(optf, ComponentArray(p_init))
res_ms = Optimization.solve(optprob, PolyOpt(), callback = callback)


gif(anim, "multiple_shooting.gif", fps=15)

optimized_params= res_ms.u

u1 = Float32[ode_data_train[1,end]]
tspan1=(50.0f0, 100.0f0)
tsteps1=range(tspan1[1], tspan1[2], length = datasize)

prob_neuralode2 = NeuralODE(nn, tspan1, Tsit5(), saveat = tsteps1)

function predict_neuralode2(p)
  Array(prob_neuralode2(u1, p, st)[1])
end
prediction=predict_neuralode2(optimized_params)
tstepsfinal=range(0.0f0, 100.0f0, length = 100)

plot(tstepsfinal, ode_data[1,:], label = "data")
plot!(tsteps1,ode_data_test[1,:], label = "data to check prediction")
scatter!(tsteps1, prediction[1,:], label = "prediction to check prediction")