In [1]:
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()

Dynamics


In [2]:
# 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()


Out[2]:
$$\left[\begin{matrix}- g l m \operatorname{sin}\left(\theta\right) + l^{2} m \ddot{\theta} - \tau\end{matrix}\right]$$

In [3]:
sympy.Eq(sympy.Symbol('L'), L)


Out[3]:
$$L = - g l m \operatorname{cos}\left(\theta\right) + \frac{l^{2} m}{2} \left(\dot{\theta}\right)^{2}$$

In [4]:
f = lm.rhs().subs(no_t_sub)
f


Out[4]:
$$\left[\begin{matrix}Q\\\frac{1}{l^{2} m} \left(g l m \operatorname{sin}\left(\theta\right) + \tau\right)\end{matrix}\right]$$

Measurements

Accelerometer


In [5]:
accel = point_cm.acc(frame_i)  - frame_i.z*g
accel


Out[5]:
$$- l \ddot{\theta}\mathbf{\hat{b}_x} + l \left(\dot{\theta}\right)^{2}\mathbf{\hat{b}_z} - g\mathbf{\hat{i}_z}$$

In [6]:
g_accel = accel.to_matrix(frame_b).subs(
    theta(t).diff(t,2), f[1]).subs(no_t_sub).simplify()
g_accel


Out[6]:
$$\left[\begin{matrix}- \frac{\tau}{l m}\\0\\Q^{2} l - g \operatorname{cos}\left(\theta\right)\end{matrix}\right]$$

Gyroscope


In [7]:
gyro = frame_b.ang_vel_in(frame_i)
gyro


Out[7]:
$$\dot{\theta}\mathbf{\hat{b}_y}$$

In [8]:
g_gyro = gyro.to_matrix(frame_b).subs(no_t_sub)
g_gyro


Out[8]:
$$\left[\begin{matrix}0\\Q\\0\end{matrix}\right]$$

Save


In [9]:
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'))

Generate Code


In [10]:
import sympy_utils

In [11]:
pendulum = sympy_utils.create_fortran_module_from_sympy_save(
    'pendulum', os.path.join('save','pendulum.sympy'))

In [12]:
%%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.


In [13]:
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)

In [14]:
%%time
do_sim();


CPU times: user 754 ms, sys: 19.8 ms, total: 774 ms
Wall time: 773 ms