トップページに戻る

SymPyー代数演算(2)応用:二重振り子

応用例

SymPyとNumPyの組み合わせが強力であることを示す例として、二重振り子を数値計算で解くということを試みます。
もちろん普通に解いてもしょうがないので、座標設定からラグランジアン、そして運動方程式をシンボリックに求め、最後に数値計算することにします。

モデル設定

水平にx軸、鉛直上向きにy軸があり、y軸負方向に重力がかかっています。
また、原点から長さ$r_1$の棒の先に1つ目の質点がついていて、そこから長さ$r_2$の棒で2つ目の質点がつながっているとします。
それぞれの棒がy軸とのなす角をそれぞれ$\theta_1, \theta_2$とします。

In [1]:
import sympy as sy
sy.init_printing()
In [2]:
import numpy as np
import matplotlib.pyplot as plt
%matplotlib qt5
In [3]:
# 変数と関数を定義
t = sy.Symbol("t", real=True)
r1, r2, m1, m2, g = sy.symbols(r"r_1 r_2 m_1 m_2 g", positive=True)
# theta1, theta2 = sy.symbols(r"\theta_1 \theta_2", real=True)
theta1 = sy.Function(r"\theta_1")(t)
theta2 = sy.Function(r"\theta_2")(t)
theta1_dot, theta2_dot = sy.symbols(r"\dot{\theta_1} \dot{\theta_2}")
In [4]:
# 質点1の位置と速度
x1 = sy.Matrix([r1 * sy.sin(theta1), - r1 * sy.cos(theta1)])
v1 = x1.diff(t)
x1, v1
Out[4]:
$$\left ( \left[\begin{matrix}r_{1} \sin{\left (\theta_{1}{\left (t \right )} \right )}\\- r_{1} \cos{\left (\theta_{1}{\left (t \right )} \right )}\end{matrix}\right], \quad \left[\begin{matrix}r_{1} \cos{\left (\theta_{1}{\left (t \right )} \right )} \frac{d}{d t} \theta_{1}{\left (t \right )}\\r_{1} \sin{\left (\theta_{1}{\left (t \right )} \right )} \frac{d}{d t} \theta_{1}{\left (t \right )}\end{matrix}\right]\right )$$
In [5]:
# 質点2の位置と速度
x2 = x1 + sy.Matrix([r2 * sy.sin(theta2), - r2 * sy.cos(theta2)])
v2 = x2.diff(t)
x2, v2
Out[5]:
$$\left ( \left[\begin{matrix}r_{1} \sin{\left (\theta_{1}{\left (t \right )} \right )} + r_{2} \sin{\left (\theta_{2}{\left (t \right )} \right )}\\- r_{1} \cos{\left (\theta_{1}{\left (t \right )} \right )} - r_{2} \cos{\left (\theta_{2}{\left (t \right )} \right )}\end{matrix}\right], \quad \left[\begin{matrix}r_{1} \cos{\left (\theta_{1}{\left (t \right )} \right )} \frac{d}{d t} \theta_{1}{\left (t \right )} + r_{2} \cos{\left (\theta_{2}{\left (t \right )} \right )} \frac{d}{d t} \theta_{2}{\left (t \right )}\\r_{1} \sin{\left (\theta_{1}{\left (t \right )} \right )} \frac{d}{d t} \theta_{1}{\left (t \right )} + r_{2} \sin{\left (\theta_{2}{\left (t \right )} \right )} \frac{d}{d t} \theta_{2}{\left (t \right )}\end{matrix}\right]\right )$$
In [6]:
# 運動エネルギー
T1 = (m1 * (v1[0]**2 + v1[1]**2) / 2).simplify()
T2 = (m2 * (v2[0]**2 + v2[1]**2) / 2).simplify()
T1, T2
Out[6]:
$$\left ( \frac{m_{1} r_{1}^{2}}{2} \left(\frac{d}{d t} \theta_{1}{\left (t \right )}\right)^{2}, \quad \frac{m_{2}}{2} \left(r_{1}^{2} \left(\frac{d}{d t} \theta_{1}{\left (t \right )}\right)^{2} + 2 r_{1} r_{2} \cos{\left (\theta_{1}{\left (t \right )} - \theta_{2}{\left (t \right )} \right )} \frac{d}{d t} \theta_{1}{\left (t \right )} \frac{d}{d t} \theta_{2}{\left (t \right )} + r_{2}^{2} \left(\frac{d}{d t} \theta_{2}{\left (t \right )}\right)^{2}\right)\right )$$
In [7]:
# ポテンシャル
V1 = m1 * g * x1[1]
V2 = m2 * g * x2[1]
V1, V2
Out[7]:
$$\left ( - g m_{1} r_{1} \cos{\left (\theta_{1}{\left (t \right )} \right )}, \quad g m_{2} \left(- r_{1} \cos{\left (\theta_{1}{\left (t \right )} \right )} - r_{2} \cos{\left (\theta_{2}{\left (t \right )} \right )}\right)\right )$$

運動エネルギーとポテンシャルが求まったので、ラグランジアン $L(\theta_1, \theta_2, \dot{\theta_1}, \dot{\theta_2})$ を得ることができます。

In [8]:
L = T1 + T2 - V1 - V2
L
Out[8]:
$$g m_{1} r_{1} \cos{\left (\theta_{1}{\left (t \right )} \right )} - g m_{2} \left(- r_{1} \cos{\left (\theta_{1}{\left (t \right )} \right )} - r_{2} \cos{\left (\theta_{2}{\left (t \right )} \right )}\right) + \frac{m_{1} r_{1}^{2}}{2} \left(\frac{d}{d t} \theta_{1}{\left (t \right )}\right)^{2} + \frac{m_{2}}{2} \left(r_{1}^{2} \left(\frac{d}{d t} \theta_{1}{\left (t \right )}\right)^{2} + 2 r_{1} r_{2} \cos{\left (\theta_{1}{\left (t \right )} - \theta_{2}{\left (t \right )} \right )} \frac{d}{d t} \theta_{1}{\left (t \right )} \frac{d}{d t} \theta_{2}{\left (t \right )} + r_{2}^{2} \left(\frac{d}{d t} \theta_{2}{\left (t \right )}\right)^{2}\right)$$

だんだん結果は複雑になってきましたが、人間側は相当楽をしています。
ラグランジアンから運動方程式はそれぞれ、 $$\frac{d}{dt} \frac{\partial L}{\partial \dot{\theta_1}}=\frac{\partial L}{\partial \theta_1}$$ $$\frac{d}{dt} \frac{\partial L}{\partial \dot{\theta_2}}=\frac{\partial L}{\partial \theta_2}$$ となるので、

In [9]:
# 質点1の運動方程式
EOM_1 = sy.Eq(L.diff(theta1.diff(t)).diff(t), L.diff(theta1)) 
EOM_1
Out[9]:
$$m_{1} r_{1}^{2} \frac{d^{2}}{d t^{2}} \theta_{1}{\left (t \right )} + \frac{m_{2}}{2} \left(2 r_{1}^{2} \frac{d^{2}}{d t^{2}} \theta_{1}{\left (t \right )} - 2 r_{1} r_{2} \left(\frac{d}{d t} \theta_{1}{\left (t \right )} - \frac{d}{d t} \theta_{2}{\left (t \right )}\right) \sin{\left (\theta_{1}{\left (t \right )} - \theta_{2}{\left (t \right )} \right )} \frac{d}{d t} \theta_{2}{\left (t \right )} + 2 r_{1} r_{2} \cos{\left (\theta_{1}{\left (t \right )} - \theta_{2}{\left (t \right )} \right )} \frac{d^{2}}{d t^{2}} \theta_{2}{\left (t \right )}\right) = - g m_{1} r_{1} \sin{\left (\theta_{1}{\left (t \right )} \right )} - g m_{2} r_{1} \sin{\left (\theta_{1}{\left (t \right )} \right )} - m_{2} r_{1} r_{2} \sin{\left (\theta_{1}{\left (t \right )} - \theta_{2}{\left (t \right )} \right )} \frac{d}{d t} \theta_{1}{\left (t \right )} \frac{d}{d t} \theta_{2}{\left (t \right )}$$
In [10]:
# 質点2の運動方程式
EOM_2 = sy.Eq(L.diff(theta2.diff(t)).diff(t), L.diff(theta2)) 
EOM_2
Out[10]:
$$\frac{m_{2}}{2} \left(- 2 r_{1} r_{2} \left(\frac{d}{d t} \theta_{1}{\left (t \right )} - \frac{d}{d t} \theta_{2}{\left (t \right )}\right) \sin{\left (\theta_{1}{\left (t \right )} - \theta_{2}{\left (t \right )} \right )} \frac{d}{d t} \theta_{1}{\left (t \right )} + 2 r_{1} r_{2} \cos{\left (\theta_{1}{\left (t \right )} - \theta_{2}{\left (t \right )} \right )} \frac{d^{2}}{d t^{2}} \theta_{1}{\left (t \right )} + 2 r_{2}^{2} \frac{d^{2}}{d t^{2}} \theta_{2}{\left (t \right )}\right) = - g m_{2} r_{2} \sin{\left (\theta_{2}{\left (t \right )} \right )} + m_{2} r_{1} r_{2} \sin{\left (\theta_{1}{\left (t \right )} - \theta_{2}{\left (t \right )} \right )} \frac{d}{d t} \theta_{1}{\left (t \right )} \frac{d}{d t} \theta_{2}{\left (t \right )}$$

$\ddot{\theta_1}, \ddot{\theta_2}$について解くと、

In [11]:
EOMs = sy.solve([EOM_1, EOM_2], (theta1.diff(t, 2), theta2.diff(t, 2)))
EOMs
Out[11]:
$$\left \{ \frac{d^{2}}{d t^{2}} \theta_{1}{\left (t \right )} : \frac{1}{r_{1} \left(- m_{1} + m_{2} \cos^{2}{\left (\theta_{1}{\left (t \right )} - \theta_{2}{\left (t \right )} \right )} - m_{2}\right)} \left(g m_{1} \sin{\left (\theta_{1}{\left (t \right )} \right )} + \frac{g m_{2}}{2} \sin{\left (\theta_{1}{\left (t \right )} - 2 \theta_{2}{\left (t \right )} \right )} + \frac{g m_{2}}{2} \sin{\left (\theta_{1}{\left (t \right )} \right )} + \frac{m_{2} r_{1}}{2} \sin{\left (2 \theta_{1}{\left (t \right )} - 2 \theta_{2}{\left (t \right )} \right )} \left(\frac{d}{d t} \theta_{1}{\left (t \right )}\right)^{2} + m_{2} r_{2} \sin{\left (\theta_{1}{\left (t \right )} - \theta_{2}{\left (t \right )} \right )} \left(\frac{d}{d t} \theta_{2}{\left (t \right )}\right)^{2}\right), \quad \frac{d^{2}}{d t^{2}} \theta_{2}{\left (t \right )} : \frac{1}{r_{2} \left(- m_{1} + m_{2} \cos^{2}{\left (\theta_{1}{\left (t \right )} - \theta_{2}{\left (t \right )} \right )} - m_{2}\right)} \left(\left(m_{1} + m_{2}\right) \left(g \sin{\left (\theta_{2}{\left (t \right )} \right )} - r_{1} \sin{\left (\theta_{1}{\left (t \right )} - \theta_{2}{\left (t \right )} \right )} \left(\frac{d}{d t} \theta_{1}{\left (t \right )}\right)^{2}\right) - \left(g m_{1} \sin{\left (\theta_{1}{\left (t \right )} \right )} + g m_{2} \sin{\left (\theta_{1}{\left (t \right )} \right )} + m_{2} r_{2} \sin{\left (\theta_{1}{\left (t \right )} - \theta_{2}{\left (t \right )} \right )} \left(\frac{d}{d t} \theta_{2}{\left (t \right )}\right)^{2}\right) \cos{\left (\theta_{1}{\left (t \right )} - \theta_{2}{\left (t \right )} \right )}\right)\right \}$$
In [12]:
# 関数を変数に置き換える
t1, t2 = sy.symbols(r"\theta_1 \theta_2")
diff2dot = [(theta1.diff(t), theta1_dot),
            (theta2.diff(t), theta2_dot)]
theta1_ddot = EOMs[theta1.diff(t, 2)].subs(diff2dot).subs([(theta1, t1), (theta2, t2)])
theta2_ddot = EOMs[theta2.diff(t, 2)].subs(diff2dot).subs([(theta1, t1), (theta2, t2)])
In [13]:
# NumPyの関数に変換する
args = (t1, t2, theta1_dot, theta2_dot, m1, m2, r1, r2, g)
func1 = sy.lambdify(args, theta1_ddot, "numpy")
func2 = sy.lambdify(args, theta2_ddot, "numpy")
In [15]:
# 数値計算
from scipy.integrate import ode

def time_evolve(t, y, params):
    theta1, theta2, theta1_dot, theta2_dot = y
    return [theta1_dot, theta2_dot,
            func1(*y, *params),
            func2(*y, *params)]

y0 = [np.pi * 1.01, np.pi * 1.01, 0, 0]
params = {"m1": 1, "m2": 1, "r1": 0.1, "r2": 0.1, "g": 9.8}
solver = ode(time_evolve).set_initial_value(y0, 0).set_f_params(params.values())

dt = 0.01
tmax = 100
results = []
while solver.t < tmax:
    y = solver.integrate(solver.t + dt)
    results.append([solver.t, *y])
In [16]:
import pandas as pd
results = pd.DataFrame(results, columns=["t", "theta1", "theta2", "theta1_dot", "theta2_dot"])
results["x1"] = params["r1"] * np.sin(results.theta1)
results["x2"] = results.x1 + params["r2"] * np.sin(results.theta2)
results["y1"] = -params["r1"] * np.cos(results.theta1)
results["y2"] = results.y1 - params["r2"] * np.cos(results.theta2)
In [17]:
import matplotlib.animation as animation
import matplotlib.patches as patches

def gen():
    for i, vals in results[["t", "x1", "x2", "y1", "y2"]].iterrows():
        yield vals.values

def plot_double_pendulum(data):
    t, x1, x2, y1, y2 = data
    ax.cla()
    R = params["r1"] + params["r2"]
    ax.set_xlim(-R, R)
    ax.set_ylim(-R, R)
    ax.scatter([x1, x2], [y1, y2])
    ax.add_patch(patches.Arrow(0, 0, x1, y1, width=0.01))
    ax.add_patch(patches.Arrow(x1, y1, (x2-x1), (y2-y1), width=0.01))
    ax.set_aspect("equal")
fig, ax = plt.subplots()
ani = animation.FuncAnimation(fig, plot_double_pendulum, gen, interval=30, save_count=500)
# ani.save("double_pendulum.gif", writer="imagemagick", dpi=100)
plt.show()

まとめると

一つのセルにまとめると以下のように書けます。

In [18]:
# 変数と関数を定義
t = sy.Symbol("t", real=True)
r1, r2, m1, m2, g = sy.symbols(r"r_1 r_2 m_1 m_2 g", positive=True)
# theta1, theta2 = sy.symbols(r"\theta_1 \theta_2", real=True)
theta1 = sy.Function(r"\theta_1")(t)
theta2 = sy.Function(r"\theta_2")(t)
theta1_dot, theta2_dot = sy.symbols(r"\dot{\theta_1} \dot{\theta_2}")

# 質点1の位置と速度
x1 = sy.Matrix([r1 * sy.sin(theta1), - r1 * sy.cos(theta1)])
v1 = x1.diff(t)

# 質点2の位置と速度
x2 = x1 + sy.Matrix([r2 * sy.sin(theta2), - r2 * sy.cos(theta2)])
v2 = x2.diff(t)

# 運動エネルギー
T1 = (m1 * (v1[0]**2 + v1[1]**2) / 2).simplify()
T2 = (m2 * (v2[0]**2 + v2[1]**2) / 2).simplify()

# ポテンシャル
V1 = m1 * g * x1[1]
V2 = m2 * g * x2[1]

# ラグランジアン
L = T1 + T2 - V1 - V2

# 質点1、2の運動方程式
EOM_1 = sy.Eq(L.diff(theta1.diff(t)).diff(t), L.diff(theta1)) 
EOM_2 = sy.Eq(L.diff(theta2.diff(t)).diff(t), L.diff(theta2)) 

# 加速度について解く
EOMs = sy.solve([EOM_1, EOM_2], (theta1.diff(t, 2), theta2.diff(t, 2)))

# 関数を変数に変換
t1, t2 = sy.symbols(r"\theta_1 \theta_2")
diff2dot = [(theta1.diff(t), theta1_dot),
            (theta2.diff(t), theta2_dot)]
theta1_ddot = EOMs[theta1.diff(t, 2)].subs(diff2dot).subs([(theta1, t1), (theta2, t2)])
theta2_ddot = EOMs[theta2.diff(t, 2)].subs(diff2dot).subs([(theta1, t1), (theta2, t2)])

# NumPy関数化
args = (t1, t2, theta1_dot, theta2_dot, m1, m2, r1, r2, g)
func1 = sy.lambdify(args, theta1_ddot, "numpy")
func2 = sy.lambdify(args, theta2_ddot, "numpy")

# 数値計算
from scipy.integrate import ode

def time_evolve(t, y, params):
    theta1, theta2, theta1_dot, theta2_dot = y
    return [theta1_dot, theta2_dot,
            func1(*y, *params),
            func2(*y, *params)]

# 初期値
y0 = [np.pi * 1.01, np.pi * 1.01, 0, 0]
params = {"m1": 1, "m2": 1, "r1": 0.1, "r2": 0.1, "g": 9.8}
solver = ode(time_evolve).set_initial_value(y0, 0).set_f_params(params.values())

dt = 0.1
tmax = 50
results = []
while solver.t < tmax:
    y = solver.integrate(solver.t + dt)
    results.append([solver.t, *y])

import pandas as pd
results = pd.DataFrame(results, columns=["t", "theta1", "theta2", "theta1_dot", "theta2_dot"])
results["x1"] = params["r1"] * np.sin(results.theta1)
results["x2"] = results.x1 + params["r2"] * np.sin(results.theta2)
results["y1"] = -params["r1"] * np.cos(results.theta1)
results["y2"] = results.y1 - params["r2"] * np.cos(results.theta2)

誤字やおかしい点などがあったら @zawawahoge (Twitter) にお気軽にご連絡ください。

トップページに戻る