# Implicit Nonlinear Dynamics : Cartpole

The following is another example on how to use the ImplicitOptimizer that is taken from the original paper. As always, we start by creating a corresponding dataset:

using DataDrivenDiffEq
using ModelingToolkit
using OrdinaryDiffEq
using LinearAlgebra
using Plots
gr()

function cart_pole(u, p, t)
du = similar(u)
F = -0.2 + 0.5*sin(6*t) # the input
du[1] = u[3]
du[2] = u[4]
du[3] = -(19.62*sin(u[1])+sin(u[1])*cos(u[1])*u[3]^2+F*cos(u[1]))/(2-cos(u[1])^2)
du[4] = -(sin(u[1])*u[3]^2 + 9.81*sin(u[1])*cos(u[1])+F)/(2-cos(u[1])^2)
return du
end

u0 = [0.3; 0; 1.0; 0]
tspan = (0.0, 5.0)
dt = 0.05
cart_pole_prob = ODEProblem(cart_pole, u0, tspan)
solution = solve(cart_pole_prob, Tsit5(), saveat = dt)

X = solution[:,:]
DX = similar(X)
for (i, xi) in enumerate(eachcol(X))
DX[:, i] = cart_pole(xi, [], solution.t[i])
end
t = solution.t

X , t, DX = DX[3:4, :], U = (u,p,t) -> [-0.2 + 0.5*sin(6*t)]
)

plot(ddprob)

Note that we just included the third and forth time derivative, assuming that we already know that the velocity x[3:4] is equal to the time derivative of the position x[1:2]. Next, we define a sufficient Basis. Again, we need to include implicits in the definition of our candidate functions and inform the Basis of it.

@parameters t
@variables u[1:4] du[1:2] x[1:1]
u, du, x = map(collect, [u, du, x])

polys = polynomial_basis(u, 2)
push!(polys, sin.(u[1]))
push!(polys, cos.(u[1]))
push!(polys, sin.(u[1])^2)
push!(polys, cos.(u[1])^2)
push!(polys, sin.(u[1]).*u[3:4]...)
push!(polys, sin.(u[1]).*u[3:4].^2...)
push!(polys, sin.(u[1]).*cos.(u[1])...)
push!(polys, sin.(u[1]).*cos.(u[1]).*u[3:4]...)
push!(polys, sin.(u[1]).*cos.(u[1]).*u[3:4].^2...)

implicits = [du;  du[1] .* u; du[2] .* u; du .* cos(u[1]);   du .* cos(u[1])^2; polys]
push!(implicits, x...)
push!(implicits, x[1]*cos(u[1]))
push!(implicits, x[1]*sin(u[1]))

basis= Basis(implicits, u, implicits = du, controls = x,  iv = t);
Model ##Basis#548 with 45 equations
States : u[1] u[2] u[3] u[4]
Independent variable: t
Equations
φ₁ = du[1]
φ₂ = du[2]
φ₃ = du[1]*u[1]
φ₄ = du[1]*u[2]
...
φ₄₅ = sin(u[1])*x[1]

We solve the problem by varying over a sufficient set of thresholds for the associated optimizer. Additionally we activate the scale_coefficients option for the ImplicitOptimizer, which helps to find sparse equations by normalizing the resulting coefficient matrix after each suboptimization. To evaluate the pareto optimal solution, we use the functions f and g which can be passed as keyword arguments into the solve function. f is a function with different signatures for different optimizers, but returns the $L_0$ norm of the coefficients and the $L_2$ error of the current model. g takes this vector and projects it down onto a scalar, using the $L_2$ norm per default. However, here we want to use the AIC of the output of f. A noteworthy exception is of course, that we want only results with two or more active coefficents. Hence, we modify g accordingly.

λ = [1e-4;5e-4;1e-3;2e-3;3e-3;4e-3;5e-3;6e-3;7e-3;8e-3;9e-3;1e-2;2e-2;3e-2;4e-2;5e-2]

opt = ImplicitOptimizer(λ)

g(x) = x[1] <= 1 ? Inf : 2*x[1]-2*log(x[2])

res = solve(ddprob, basis, opt, du, maxiter = 1000, g = g, show_progress = true)
system = result(res)
Model ##Basis#551 with 2 equations
States : u[1] u[2] u[3] u[4]
Parameters : 10
Independent variable: t
Equations
du[1] = (p₃*sin(u[1]) + p₅*cos(u[1])*x[1] + p₄*(u[3]^2)*cos(u[1])*sin(u[1])) / (-p₁ - p₂*(cos(u[1])^2))
du[2] = (p₁₀*x[1] + p₈*(u[3]^2)*sin(u[1]) + p₉*cos(u[1])*sin(u[1])) / (-p₆ - p₇*(cos(u[1])^2))

We have recovered the correct equations of motion! Another visual check using the problem and the result yields

plot(
plot(ddprob), plot(res), layout = (1,2)
)

## Copy-Pasteable Code

using DataDrivenDiffEq
using ModelingToolkit
using OrdinaryDiffEq
using LinearAlgebra

function cart_pole(u, p, t)
du = similar(u)
F = -0.2 + 0.5*sin(6*t) # the input
du[1] = u[3]
du[2] = u[4]
du[3] = -(19.62*sin(u[1])+sin(u[1])*cos(u[1])*u[3]^2+F*cos(u[1]))/(2-cos(u[1])^2)
du[4] = -(sin(u[1])*u[3]^2 + 9.81*sin(u[1])*cos(u[1])+F)/(2-cos(u[1])^2)
return du
end

u0 = [0.3; 0; 1.0; 0]
tspan = (0.0, 5.0)
dt = 0.05
cart_pole_prob = ODEProblem(cart_pole, u0, tspan)
solution = solve(cart_pole_prob, Tsit5(), saveat = dt)

X = solution[:,:]
DX = similar(X)
for (i, xi) in enumerate(eachcol(X))
DX[:, i] = cart_pole(xi, [], solution.t[i])
end
t = solution.t

X , t, DX = DX[3:4, :], U = (u,p,t) -> [-0.2 + 0.5*sin(6*t)]
)

@parameters t
@variables u[1:4] du[1:2] x[1:1]
u, du, x = map(collect, [u, du, x])

polys = polynomial_basis(u, 2)
push!(polys, sin.(u[1]))
push!(polys, cos.(u[1]))
push!(polys, sin.(u[1])^2)
push!(polys, cos.(u[1])^2)
push!(polys, sin.(u[1]).*u[3:4]...)
push!(polys, sin.(u[1]).*u[3:4].^2...)
push!(polys, sin.(u[1]).*cos.(u[1])...)
push!(polys, sin.(u[1]).*cos.(u[1]).*u[3:4]...)
push!(polys, sin.(u[1]).*cos.(u[1]).*u[3:4].^2...)

implicits = [du;  du[1] .* u; du[2] .* u; du .* cos(u[1]);   du .* cos(u[1])^2; polys]
push!(implicits, x...)
push!(implicits, x[1]*cos(u[1]))
push!(implicits, x[1]*sin(u[1]))

basis= Basis(implicits, u, implicits = du, controls = x,  iv = t);
println(basis) # hide

λ = [1e-4;5e-4;1e-3;2e-3;3e-3;4e-3;5e-3;6e-3;7e-3;8e-3;9e-3;1e-2;2e-2;3e-2;4e-2;5e-2]

opt = ImplicitOptimizer(λ)

g(x) = x[1] <= 1 ? Inf : 2*x[1]-2*log(x[2])

res = solve(ddprob, basis, opt, du, maxiter = 1000, g = g, show_progress = true)
system = result(res)
println(system) #hide

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl