Physics Informed Machine Learning with TensorLayer
In this tutorial, we show how to use the DiffEqFlux
TensorLayer to solve problems in Physics Informed Machine Learning.
Let's consider the anharmonic oscillator described by the ODE
\[ẍ = - kx - αx³ - βẋ -γẋ³.\]
To obtain the training data, we solve the equation of motion using one of the solvers in DifferentialEquations
:
using DiffEqFlux, DifferentialEquations, LinearAlgebra
k, α, β, γ = 1, 0.1, 0.2, 0.3
tspan = (0.0,10.0)
function dxdt_train(du,u,p,t)
du[1] = u[2]
du[2] = -k*u[1] - α*u[1]^3 - β*u[2] - γ*u[2]^3
end
u0 = [1.0,0.0]
ts = collect(0.0:0.1:tspan[2])
prob_train = ODEProblem{true}(dxdt_train,u0,tspan,p=nothing)
data_train = Array(solve(prob_train,Tsit5(),saveat=ts))
Now, we create a TensorLayer that will be able to perform 10th order expansions in a Legendre Basis:
A = [LegendreBasis(10), LegendreBasis(10)]
nn = TensorLayer(A, 1)
and we also instantiate the model we are trying to learn, "informing" the neural about the ∝x
and ∝v
dependencies in the equation of motion:
f = x -> min(30one(x),x)
function dxdt_pred(du,u,p,t)
du[1] = u[2]
du[2] = -p[1]*u[1] - p[2]*u[2] + f(nn(u,p[3:end])[1])
end
α = zeros(102)
prob_pred = ODEProblem{true}(dxdt_pred,u0,tspan,p=nothing)
Note that we introduced a "cap" in the neural network term to avoid instabilities in the solution of the ODE. We also initialized the vector of parameters to zero in order to obtain a faster convergence for this particular example.
Finally, we introduce the corresponding loss function:
function predict_adjoint(θ)
x = Array(solve(prob_pred,Tsit5(),p=θ,saveat=ts))
end
function loss_adjoint(θ)
x = predict_adjoint(θ)
loss = sum(norm.(x - data_train))
return loss
end
function cb(θ,l)
@show θ, l
return false
end
and we train the network using two rounds of ADAM
:
res1 = DiffEqFlux.sciml_train(loss_adjoint, α, ADAM(0.05), cb = cb, maxiters = 150)
res2 = DiffEqFlux.sciml_train(loss_adjoint, res1.u, ADAM(0.001), cb = cb,maxiters = 150)
opt = res2.u
We plot the results and we obtain a fairly accurate learned model:
using Plots
data_pred = predict_adjoint(opt)
plot(ts, data_train[1,:], label = "X (ODE)")
plot!(ts, data_train[2,:], label = "V (ODE)")
plot!(ts, data_pred[1,:], label = "X (NN)")
plot!(ts, data_pred[2,:],label = "V (NN)")