import sympy
import sympy.physics.mechanics as mech
import os
import matplotlib.pyplot as plt
from sympy.printing.theanocode import theano_function
from sympy_utils import rhs_to_scipy_ode, \
save_sympy_expr, load_sympy_expr, gen_fortran_module
from numpy import f2py
import numpy as np
from sympy.utilities import codegen
%matplotlib inline
%load_ext autoreload
%autoreload 2
sympy.init_printing()
mech.init_vprinting()
# symbols
m, g, l, t, tau, phi, theta, Q, psi = \
sympy.symbols('m, g, l, t, tau, phi, theta, Q, psi')
# state space
x = (theta, Q)
u = (tau,)
no_t_sub = {theta(t).diff(t):Q}
no_t_sub.update({v(t): v for v in x+u})
# frames
frame_i = mech.ReferenceFrame('i') # inertial
frame_b = frame_i.orientnew( # body
newname='b', rot_type='Body',
amounts=(0, theta(t), 0), rot_order='321')
# point at center of rotation
point_o = mech.Point('o')
point_o.set_vel(frame_i, 0)
# point at center of mass
point_cm = point_o.locatenew('cm', -l*frame_b.z)
point_cm.set_vel(frame_b, 0)
point_cm.v1pt_theory(otherpoint=point_o,
outframe=frame_i, interframe=frame_b)
# particle
ball = mech.Particle(name='ball', point=point_cm, mass=m)
ball.set_potential_energy(m*g*l*sympy.cos(theta(t)))
# derive eoms
L = mech.Lagrangian(frame_i, ball)
lm = mech.LagrangesMethod(
Lagrangian=L,
qs=[theta(t)],
forcelist=[(frame_b, tau(t)*frame_b.y)], frame=frame_i)
lm.form_lagranges_equations()
sympy.Eq(sympy.Symbol('L'), L)
f = lm.rhs().subs(no_t_sub)
f
accel = point_cm.acc(frame_i) - frame_i.z*g
accel
g_accel = accel.to_matrix(frame_b).subs(
theta(t).diff(t,2), f[1]).subs(no_t_sub).simplify()
g_accel
gyro = frame_b.ang_vel_in(frame_i)
gyro
g_gyro = gyro.to_matrix(frame_b).subs(no_t_sub)
g_gyro
save_sympy_expr({
't': t, 'x': x, 'u': u, 'f': f,
'g_dict': {
'g_accel': g_accel,
'g_gyro': g_gyro},
'const': (m, g, l),
}, os.path.join('save','pendulum.sympy'))
import sympy_utils
pendulum = sympy_utils.create_fortran_module_from_sympy_save(
'pendulum', os.path.join('save','pendulum.sympy'))
%%timeit
pendulum.compute_a(t=0, x=[1,2], u=1, m=1, l=1, g=9.8)
pendulum.compute_b(t=0, x=[1,2], u=1, m=1, l=1, g=9.8)
pendulum.compute_f(t=0, x=[1,2], u=1, m=1, l=1, g=9.8)
pendulum.compute_g_accel(t=0, x=[0,0], u=0, m=1, l=1, g=9.8)
pendulum.compute_g_accel_h(t=0, x=[0,0], u=0, m=1, l=1, g=9.8)
pendulum.compute_g_gyro(t=0, x=[0,0], u=0, m=1, l=1, g=9.8)
pendulum.compute_g_gyro_h(t=0, x=[0,0], u=0, m=1, l=1, g=9.8)
100000 loops, best of 3: 12.7 µs per loop
The lambdify function which uses numpy is sufficiently fast.
def do_sim():
import scipy.integrate
ode = scipy.integrate.ode(pendulum.compute_f, pendulum.compute_a)
ode.set_initial_value([0.1,0])
dt = 0.001
x = []
t = []
u = 0
m = 1
g = 9.8
l = 1
while ode.t +dt < 100:
ode.set_f_params(u, m, g, l)
ode.set_jac_params(u, m, g, l)
x.append(ode.y)
t.append(ode.t)
ode.integrate(ode.t + dt)
plt.plot(t, x)
%%time
do_sim();
CPU times: user 754 ms, sys: 19.8 ms, total: 774 ms Wall time: 773 ms