Skip to content

Commit

Permalink
Deduplicate test dynamics (#410)
Browse files Browse the repository at this point in the history
  • Loading branch information
@calcmogul
calcmogul authored Feb 26, 2024
1 parent ea80030 commit db99092
Show file tree
Hide file tree
Showing 4 changed files with 162 additions and 273 deletions.
157 changes: 60 additions & 97 deletions jormungandr/test/optimization/optimization_problem_cart_pole_test.py
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import math

import jormungandr.autodiff as autodiff
from jormungandr.autodiff import ExpressionType, VariableMatrix
from jormungandr.autodiff import ExpressionType, VariableMatrix, Variable
from jormungandr.optimization import OptimizationProblem, SolverExitCondition
import numpy as np
import pytest
Expand All @@ -18,69 +18,63 @@ def rk4(f, x, u, dt):
return x + h / 6.0 * (k1 + 2.0 * k2 + 2.0 * k3 + k4)


def cart_pole_dynamics_double(x, u):
# https://underactuated.mit.edu/acrobot.html#cart_pole
#
# θ is CCW+ measured from negative y-axis.
#
# q = [x, θ]ᵀ
# q̇ = [ẋ, θ̇]ᵀ
# u = f_x
#
# M(q)q̈ + C(q, q̇)q̇ = τ_g(q) + Bu
# M(q)q̈ = τ_g(q) − C(q, q̇)q̇ + Bu
# q̈ = M⁻¹(q)(τ_g(q) − C(q, q̇)q̇ + Bu)
#
# [ m_c + m_p m_p l cosθ]
# M(q) = [m_p l cosθ m_p l² ]
#
# [0 −m_p lθ̇ sinθ]
# C(q, q̇) = [0 0 ]
#
# [ 0 ]
# τ_g(q) = [-m_p gl sinθ]
#
# [1]
# B = [0]
m_c = 5.0 # Cart mass (kg)
m_p = 0.5 # Pole mass (kg)
l = 0.5 # Pole length (m)
g = 9.806 # Acceleration due to gravity (m/s²)
# https://underactuated.mit.edu/acrobot.html#cart_pole
#
# θ is CCW+ measured from negative y-axis.
#
# q = [x, θ]ᵀ
# q̇ = [ẋ, θ̇]ᵀ
# u = f_x
#
# M(q)q̈ + C(q, q̇)q̇ = τ_g(q) + Bu
# M(q)q̈ = τ_g(q) − C(q, q̇)q̇ + Bu
# q̈ = M⁻¹(q)(τ_g(q) − C(q, q̇)q̇ + Bu)
#
# [ m_c + m_p m_p l cosθ]
# M(q) = [m_p l cosθ m_p l² ]
#
# [0 −m_p lθ̇ sinθ]
# C(q, q̇) = [0 0 ]
#
# [ 0 ]
# τ_g(q) = [-m_p gl sinθ]
#
# [1]
# B = [0]

m_c = 5.0 # Cart mass (kg)
m_p = 0.5 # Pole mass (kg)
l = 0.5 # Pole length (m)
g = 9.806 # Acceleration due to gravity (m/s²)


def cart_pole_dynamics_double(x, u):
q = x[:2, :]
qdot = x[2:, :]
theta = q[1, 0]
thetadot = qdot[1, 0]

# [ m_c + m_p m_p l cosθ]
# M(q) = [m_p l cosθ m_p l² ]
M = np.empty((2, 2))
M[0, 0] = m_c + m_p
M[0, 1] = m_p * l * math.cos(theta)
M[1, 0] = m_p * l * math.cos(theta)
M[1, 1] = m_p * l**2

Minv = np.empty((2, 2))
Minv[0, 0] = M[1, 1]
Minv[0, 1] = -M[0, 1]
Minv[1, 0] = -M[1, 0]
Minv[1, 1] = M[0, 0]
M = np.array(
[
[m_c + m_p, m_p * l * math.cos(theta)],
[m_p * l * math.cos(theta), m_p * l**2],
]
)

detM = M[0, 0] * M[1, 1] - M[0, 1] * M[1, 0]
Minv /= detM
Minv = np.array(
[[M[1, 1] / detM, -M[0, 1] / detM], [-M[1, 0] / detM, M[0, 0] / detM]]
)

# [0 −m_p lθ̇ sinθ]
# C(q, q̇) = [0 0 ]
C = np.empty((2, 2))
C[0, 0] = 0.0
C[0, 1] = -m_p * l * thetadot * math.sin(theta)
C[1, 0] = 0.0
C[1, 1] = 0.0
C = np.array([[0.0, -m_p * l * thetadot * math.sin(theta)], [0.0, 0.0]])

# [ 0 ]
# τ_g(q) = [-m_p gl sinθ]
tau_g = np.empty((2, 1))
tau_g[0, :] = 0.0
tau_g[1, :] = -m_p * g * l * math.sin(theta)
tau_g = np.array([[0.0], [-m_p * g * l * math.sin(theta)]])

# [1]
# B = [0]
Expand All @@ -94,68 +88,37 @@ def cart_pole_dynamics_double(x, u):


def cart_pole_dynamics(x, u):
# https://underactuated.mit.edu/acrobot.html#cart_pole
#
# θ is CCW+ measured from negative y-axis.
#
# q = [x, θ]ᵀ
# q̇ = [ẋ, θ̇]ᵀ
# u = f_x
#
# M(q)q̈ + C(q, q̇)q̇ = τ_g(q) + Bu
# M(q)q̈ = τ_g(q) − C(q, q̇)q̇ + Bu
# q̈ = M⁻¹(q)(τ_g(q) − C(q, q̇)q̇ + Bu)
#
# [ m_c + m_p m_p l cosθ]
# M(q) = [m_p l cosθ m_p l² ]
#
# [0 −m_p lθ̇ sinθ]
# C(q, q̇) = [0 0 ]
#
# [ 0 ]
# τ_g(q) = [-m_p gl sinθ]
#
# [1]
# B = [0]
m_c = 5.0 # Cart mass (kg)
m_p = 0.5 # Pole mass (kg)
l = 0.5 # Pole length (m)
g = 9.806 # Acceleration due to gravity (m/s²)

q = x[:2, :]
qdot = x[2:, :]
theta = q[1, 0]
thetadot = qdot[1, 0]

# [ m_c + m_p m_p l cosθ]
# M(q) = [m_p l cosθ m_p l² ]
M = VariableMatrix(2, 2)
M[0, 0] = m_c + m_p
M[0, 1] = m_p * l * autodiff.cos(theta)
M[1, 0] = m_p * l * autodiff.cos(theta)
M[1, 1] = m_p * l**2

Minv = VariableMatrix(2, 2)
Minv[0, 0] = M[1, 1]
Minv[0, 1] = -M[0, 1]
Minv[1, 0] = -M[1, 0]
Minv[1, 1] = M[0, 0]
M = VariableMatrix(
[
[Variable(m_c + m_p), m_p * l * autodiff.cos(theta)],
[m_p * l * autodiff.cos(theta), Variable(m_p * l**2)],
]
)

detM = M[0, 0] * M[1, 1] - M[0, 1] * M[1, 0]
Minv /= detM
Minv = VariableMatrix(
[[M[1, 1] / detM, -M[0, 1] / detM], [-M[1, 0] / detM, M[0, 0] / detM]]
)

# [0 −m_p lθ̇ sinθ]
# C(q, q̇) = [0 0 ]
C = VariableMatrix(2, 2)
C[0, 0] = 0.0
C[0, 1] = -m_p * l * thetadot * autodiff.sin(theta)
C[1, 0] = 0.0
C[1, 1] = 0.0
C = VariableMatrix(
[
[Variable(0.0), -m_p * l * thetadot * autodiff.sin(theta)],
[Variable(0.0), Variable(0.0)],
]
)

# [ 0 ]
# τ_g(q) = [-m_p gl sinθ]
tau_g = VariableMatrix(2, 1)
tau_g[0, :] = 0.0
tau_g[1, :] = -m_p * g * l * autodiff.sin(theta)
tau_g = VariableMatrix([[Variable(0.0)], [-m_p * g * l * autodiff.sin(theta)]])

# [1]
# B = [0]
Expand Down
52 changes: 20 additions & 32 deletions jormungandr/test/optimization/optimization_problem_differential_drive_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,55 +21,43 @@ def rk4(f, x, u, dt):
return x + h / 6.0 * (k1 + 2.0 * k2 + 2.0 * k3 + k4)


def differential_drive_dynamics_double(x, u):
# x = [x, y, heading, left velocity, right velocity]ᵀ
# u = [left voltage, right voltage]ᵀ
trackwidth = 0.699 # m
Kv_linear = 3.02 # V/(m/s)
Ka_linear = 0.642 # V/(m/s²)
Kv_angular = 1.382 # V/(m/s)
Ka_angular = 0.08495 # V/(m/s²)

v = (x[3, 0] + x[4, 0]) / 2.0
# x = [x, y, heading, left velocity, right velocity]ᵀ
# u = [left voltage, right voltage]ᵀ
trackwidth = 0.699 # m
Kv_linear = 3.02 # V/(m/s)
Ka_linear = 0.642 # V/(m/s²)
Kv_angular = 1.382 # V/(m/s)
Ka_angular = 0.08495 # V/(m/s²)

A1 = -(Kv_linear / Ka_linear + Kv_angular / Ka_angular) / 2.0
A2 = -(Kv_linear / Ka_linear - Kv_angular / Ka_angular) / 2.0
B1 = 0.5 / Ka_linear + 0.5 / Ka_angular
B2 = 0.5 / Ka_linear - 0.5 / Ka_angular
A = np.array([[A1, A2], [A2, A1]])
B = np.array([[B1, B2], [B2, B1]])

A1 = -(Kv_linear / Ka_linear + Kv_angular / Ka_angular) / 2.0
A2 = -(Kv_linear / Ka_linear - Kv_angular / Ka_angular) / 2.0
B1 = 0.5 / Ka_linear + 0.5 / Ka_angular
B2 = 0.5 / Ka_linear - 0.5 / Ka_angular
A = np.array([[A1, A2], [A2, A1]])
B = np.array([[B1, B2], [B2, B1]])

def differential_drive_dynamics_double(x, u):
xdot = np.empty((5, 1))

v = (x[3, 0] + x[4, 0]) / 2.0
xdot[0, :] = v * math.cos(x[2, 0])
xdot[1, :] = v * math.sin(x[2, 0])
xdot[2, :] = (x[4, 0] - x[3, 0]) / trackwidth
xdot[3:5, :] = A @ x[3:5, :] + B @ u

return xdot


def differential_drive_dynamics(x, u):
# x = [x, y, heading, left velocity, right velocity]ᵀ
# u = [left voltage, right voltage]ᵀ
trackwidth = 0.699 # m
Kv_linear = 3.02 # V/(m/s)
Ka_linear = 0.642 # V/(m/s²)
Kv_angular = 1.382 # V/(m/s)
Ka_angular = 0.08495 # V/(m/s²)
xdot = VariableMatrix(5, 1)

v = (x[3, 0] + x[4, 0]) / 2.0

A1 = -(Kv_linear / Ka_linear + Kv_angular / Ka_angular) / 2.0
A2 = -(Kv_linear / Ka_linear - Kv_angular / Ka_angular) / 2.0
B1 = 0.5 / Ka_linear + 0.5 / Ka_angular
B2 = 0.5 / Ka_linear - 0.5 / Ka_angular
A = np.array([[A1, A2], [A2, A1]])
B = np.array([[B1, B2], [B2, B1]])

xdot = VariableMatrix(5, 1)
xdot[0, :] = v * autodiff.cos(x[2, 0])
xdot[1, :] = v * autodiff.sin(x[2, 0])
xdot[2, :] = (x[4, 0] - x[3, 0]) / trackwidth
xdot[3:5, :] = A @ x[3:5, :] + B @ u

return xdot


Expand Down
Loading

0 comments on commit db99092

Please sign in to comment.