Physics-Informed Machine Learning (PIML) 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, Optimization, OptimizationFlux, 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)
data_train = Array(solve(prob_train,Tsit5(),saveat=ts))

2×101 Matrix{Float64}:
 1.0   0.994543   0.97839    0.951986  …  -0.260537  -0.249301  -0.235816
 0.0  -0.108662  -0.213651  -0.313419      0.100489   0.12388    0.145526

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)

(::TensorLayer{Array{DiffEqFlux.TensorProductBasis}, Vector{Float64}, Int64}) (generic function with 2 methods)

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)

ODEProblem with uType Vector{Float64} and tType Float64. In-place: true
timespan: (0.0, 10.0)
u0: 2-element Vector{Float64}:
 1.0
 0.0

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,
                  sensealg=InterpolatingAdjoint(autojacvec=ReverseDiffVJP(true))))
end

function loss_adjoint(θ)
  x = predict_adjoint(θ)
  loss = sum(norm.(x - data_train))
  return loss
end

iter = 0
function callback(θ,l)
  global iter
  iter += 1
  if iter%10 == 0
    println(l)
  end
  return false
end

callback (generic function with 1 method)

and we train the network using two rounds of ADAM:

adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x,p) -> loss_adjoint(x), adtype)
optprob = Optimization.OptimizationProblem(optf, α)
res1 = Optimization.solve(optprob, ADAM(0.05), callback = callback, maxiters = 150)

optprob2 = Optimization.OptimizationProblem(optf, res1.u)
res2 = Optimization.solve(optprob2, ADAM(0.001), callback = callback,maxiters = 150)
opt = res2.u

102-element Vector{Float64}:
  0.47688995148237284
  0.22519780001081344
 -0.07238117917716692
 -0.22519780001081344
  0.10756214889496968
  0.09504474981925565
 -0.08203421686257009
 -0.11684810058218227
  0.0719837755471003
  0.2884832566786953
  ⋮
  0.3085277811910385
  0.29409589155385435
 -0.3224055496906616
 -0.2910325368895399
  0.33367887770993193
  0.2953125773536795
 -0.36361513534723583
 -0.29576398738443804
 -0.1489235000987413

We plot the results and we obtain a fairly accurate learned model:

using Plots
data_pred = predict_adjoint(res1.u)
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)")