# 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 = u
du = u
du = -(19.62*sin(u)+sin(u)*cos(u)*u^2+F*cos(u))/(2-cos(u)^2)
du = -(sin(u)*u^2 + 9.81*sin(u)*cos(u)+F)/(2-cos(u)^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))
push!(polys, cos.(u))
push!(polys, sin.(u)^2)
push!(polys, cos.(u)^2)
push!(polys, sin.(u).*u[3:4]...)
push!(polys, sin.(u).*u[3:4].^2...)
push!(polys, sin.(u).*cos.(u)...)
push!(polys, sin.(u).*cos.(u).*u[3:4]...)
push!(polys, sin.(u).*cos.(u).*u[3:4].^2...)

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

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

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 ? Inf : 2*x-2*log(x)

res = solve(ddprob, basis, opt, du, maxiter = 1000, g = g, show_progress = true)
system = result(res)
Model ##Basis#529 with 2 equations
States : u u u u
Parameters : 10
Independent variable: t
Equations
du = (p₃*sin(u) + p₅*cos(u)*x + p₄*(u^2)*cos(u)*sin(u)) / (-p₁ - p₂*(cos(u)^2))
du = (p₁₀*x + p₈*(u^2)*sin(u) + p₉*cos(u)*sin(u)) / (-p₆ - p₇*(cos(u)^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 = u
du = u
du = -(19.62*sin(u)+sin(u)*cos(u)*u^2+F*cos(u))/(2-cos(u)^2)
du = -(sin(u)*u^2 + 9.81*sin(u)*cos(u)+F)/(2-cos(u)^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))
push!(polys, cos.(u))
push!(polys, sin.(u)^2)
push!(polys, cos.(u)^2)
push!(polys, sin.(u).*u[3:4]...)
push!(polys, sin.(u).*u[3:4].^2...)
push!(polys, sin.(u).*cos.(u)...)
push!(polys, sin.(u).*cos.(u).*u[3:4]...)
push!(polys, sin.(u).*cos.(u).*u[3:4].^2...)

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

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 ? Inf : 2*x-2*log(x)

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