How to define and integrate a HamiltonianProblem?

[cdawg]

Dear friends. I have a hamiltonian for the cart and pole system which is a standard problem in RL. Lagrangian version shown here, I just legendre transformed it.

I don’t need a full solution. I just want to step it forward in time in a symplectic, energy preserving way. Here’s what I got so far

using DiffEqPhysics
using DifferentialEquations

function cart_hamiltonian(q, p; g = 9.8, m1 = 2.0 , m2 = 1.0, L = 0.5) 
    (x, θ) = q
    (p_x, p_θ) = p
    h =  g * L * m2 * cos(θ) + 
        (L^2 * m2 * p_x^2 + (m1 + m2)* p_θ^2 - 
        2 *L * m2 * p_x * p_θ * cos(θ) ) / 
        ( 2 * L^2*m2*(m1 + m2 * sin(θ)^2))
    return h
end

a = HamiltonianProblem(cart_hamiltonian,[0,0.3],[0,.1],[0,.2]) # gives me ODE problem?
sol = init(a, Tsit5())

which yields

MethodError: no method matching cart_hamiltonian(::Vector{Float64}, ::Vector{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqPhysics.PhysicsTag, Float64}, Float64, 2}}, ::Nothing)

when cart_hamiltonian sees dual numbers.

ultimately I just want a function

function step(p,q,δt)
return (p_new, q_new) #new p q at slightly later time

I could write an euler integrator, but I feel like I’m close. I see step! in the documentation. Part of the problem with Julia is that it’s so easy to write blazing fast integrators I never needed to learn the details. But autodiffed hamiltonians would be so freaking awesome. Thanks

[contradict]

I think you have p and q reversed in the definition of H. The documentation is a little unclear on this, but H actually takes 3 arguments H(p, q, params). The params argument defaults to nothing, but can be useful if you want to run multiple integrations of the same equations with different masses for example.

I think you need

function cart_hamiltonian(p, q, params; g = 9.8, m1 = 2.0 , m2 = 1.0, L = 0.5) 
    (x, θ) = q
    (p_x, p_θ) = p
    h =  g * L * m2 * cos(θ) + 
        (L^2 * m2 * p_x^2 + (m1 + m2)* p_θ^2 - 
        2 *L * m2 * p_x * p_θ * cos(θ) ) / 
        ( 2 * L^2*m2*(m1 + m2 * sin(θ)^2))
    return h
end

function initialize()
    a = HamiltonianProblem(cart_hamiltonian,[0,0.3],[0,.1],[0,.2])
    integrator = init(a, Tsit5())
end

Then you can step like this

integrator = initialize()
step!(integrator, dt, true)

[genkuroki]

Working Example: https://github.com/genkuroki/public/blob/main/0030/DiffEqPhysics%20HamiltonianProblem.ipynb

using DiffEqPhysics
using DifferentialEquations
using Plots
using StaticArrays

params = (g = 9.8, m1 = 2.0 , m2 = 1.0, L = 0.5)

function cart_hamiltonian(p, q, params, t)
    (; g, m1, m2, L) = params
    (x, θ) = q
    (p_x, p_θ) = p
    h =  g * L * m2 * cos(θ) + 
        (L^2 * m2 * p_x^2 + (m1 + m2)* p_θ^2 - 
        2 *L * m2 * p_x * p_θ * cos(θ) ) / 
        ( 2 * L^2*m2*(m1 + m2 * sin(θ)^2))
    return h
end

p0 = SVector(0.0, 0.3)
q0 = SVector(0.0, 0.1)
tspan = (0.0, 0.2)
a = HamiltonianProblem(cart_hamiltonian, p0, q0, tspan, params)
sol = solve(a, Tsit5())
plot(sol; legend=:outertopright)

Output:

[cdawg]

Thanks for the great replies. I think this is fantastic and I’m marking as solved. But now I think I should have disclosed my full use case:

I need really this function with the reinforcement learning form: update(state, action) → state

Now I see this “param” field, and it seems like a great place to put in the action (i.e. the control choice as an added linear potential.) I have

 
function cart_hamiltonian(p, q, control; g = 9.8, m1 = 2.0 , m2 = 1.0, L = 0.5) 
    (x, θ) = q
    (p_x, p_θ) = p
    h =  g * L * m2 * cos(θ) + 
        (L^2 * m2 * p_x^2 + (m1 + m2)* p_θ^2 - 
        2 *L * m2 * p_x * p_θ * cos(θ) ) / 
        ( 2 * L^2*m2*(m1 + m2 * sin(θ)^2)) 
        - control * x # this is the control, an added linear potential that kicks the system
    return h
end

So the following code works for my purposes, but I feel like it could be made more elegant than just rebuilding an integrator each time step. Can I change the parameter “on the fly” somehow with the same integrator?

function initialize(p,q,c)
    a = HamiltonianProblem(cart_hamiltonian,p,q,[0.0,1.0],c)
    return init(a, Tsit5())
end

function update(s,c; dt = .2)
   (t,p,q) = s
   integrator = initialize(p,q,c)
   step!(integrator, dt, true)
   (p1,q1) = integrator.u
   return (t+dt,p1,q1)
end

update((1.0,[0.1,0.2],[.3,.4]),.2) # the propagator in general control form

Thanks for the great help!

[genkuroki]

Second simple example: https://github.com/genkuroki/public/blob/main/0030/example%20of%20init%20and%20step!%20of%20integrator.ipynb

using DiffEqPhysics
using DifferentialEquations
using Plots
default(fmt = :png, size = (450, 250))
using StaticArrays

module My
mutable struct ParamSinglePendulum{Tg, Tl}
    g::Tg
    l::Tl
end
end

function singlependulum(p, q, param, t = nothing)
    (; g, l) = param
    dθ = p[1]
    θ = q[1]
    (1/2)*(dθ/l)^2 - g*l*cos(θ)
end

plot_singlependulum(sol) = plot(sol; label=["dθ/dt" "θ"], legend=:outertopright)

param = My.ParamSinglePendulum(9.8, 1.0) # on the earth
q0 = SVector(-0.9π)
p0 = SVector(0.0)
tspan = (0.0, 10.0)
prob = HamiltonianProblem(singlependulum, p0, q0, tspan, param)

First, a single pendulum is shaken on the earth.

integrator = init(prob, Tsit5())
param.g = 9.8 # on the earth
dt = 0.1
for _ in 1:40
    step!(integrator, dt, true)
end
plot_singlependulum(integrator.sol)

Then, suddenly, the gravitational acceleration becomes equal to the gravitational acceleration of the moon.

param.g = 1.6 # on the moon
dt = 0.1
for _ in 1:210
    step!(integrator, dt, true)
end
plot_singlependulum(integrator.sol)

The gravitational acceleration returns to its original level.

param.g = 9.8 # on the earth
dt = 0.1
for _ in 1:200
    step!(integrator, dt, true)
end
plot_singlependulum(integrator.sol)

[cdawg]

very cool. @genkuroki. just read your post about the Problem-Algorithm-Solver pattern As someone who never uses globals, or modules outside of the parent source file in a package, I am amazed at your awareness and control of Julia scope.

However, I wonder if there is an irremovable disadvantage to your pattern in the case where you want to have multiple actors acting in parallel, each with different state, each changing this potential parameter . I think I need this “update(state, action) → state” mapping as a closed, stateless function.

Sadly for my use case, the transparency of a hard coded Euler solver may ultimately be unbeatable here. But I still learned a lot and hope this post will be helpful for HamiltonianProblem users in the future. Thanks very much!