biomech

import sympy
from sympy import symbols, init_printing
import sympy.physics.mechanics as me

init_printing()
import matplotlib.pyplot as plt
import numpy as np
from pydy.system import System
from numpy import linspace

# Create the variables
L, theta = me.dynamicsymbols('L theta')

# Create the velocities
L_dot, theta_dot = me.dynamicsymbols('L_dot theta_dot')

# Create the constants
m, k, g, t, L_0 = sympy.symbols('m k g t L_0')

# Create the world frame
A = me.ReferenceFrame('A')

# Create the pendulum frame
B = A.orientnew('B', 'axis', [theta, A.z])

# Set the rotation of the pendulum frame
B.set_ang_vel(A, theta_dot * A.z)

# Create the Origin
O = me.Point('O')

# Create the mass point
P = O.locatenew('P', L * B.x)

# Set origin velocity to zero
O.set_vel(A, 0)

# Create the velocity of the mass point
P.set_vel(B, L_dot * B.x)
P.v1pt_theory(O, A, B)

pendulum = me.Particle('pend', P, m)

# Create damping forces for spring and rotation
spring = k * (L_0 - abs(L)) * B.x
gravity = m * g * -1*A.y
forces = spring + gravity

kane = me.KanesMethod(A,
                      q_ind=[L, theta],
                      u_ind=[L_dot, theta_dot],
                      kd_eqs=[L_dot - L.diff(t),
                              theta_dot - theta.diff(t)])

fr, frstar = kane.kanes_equations([(P, forces)], [pendulum])

M = kane.mass_matrix_full
f = kane.forcing_full
M.inv() * f

M = kane.mass_matrix_full
f = kane.forcing_full
M.inv() * f

sys = System(kane)

sys.constants = {m: 80,
                 g: 9.81,
                 k: 10000.0,
                 L_0: 1.1}
sys.initial_conditions = {L: 1.1,
                          L_dot: -1.37,
                          theta: np.deg2rad(20),
                          theta_dot: -np.deg2rad(3.4)}

runtime = 0.4
# For 30fps I want 30 datapoints per second so runtime is multiplied by 30
sys.times = linspace(0.0, runtime, runtime * 30)
sys.generate_ode_function(generator='cython')
resp = sys.integrate()

sim_time = sys.times
fig = plt.figure(0, figsize=(12, 5))
fig.add_subplot(121)
plt.plot(sim_time, resp[:, 0], label='Unshaped')
plt.title(r'L', fontsize=18)
plt.xlabel('Time', fontsize=18)
plt.ylabel('Meters', fontsize=18)
fig.add_subplot(122)
plt.plot(sim_time, np.rad2deg(resp[:, 1]), label='Unshaped')
plt.title(r'$\theta$', fontsize=18)
plt.xlabel('Time', fontsize=18)
plt.ylabel('Degrees', fontsize=18)
plt.tight_layout()
plt.show()

fig = plt.figure(0, figsize=(12, 7))
plt.plot(resp[0, 0] * np.sin(resp[1, 1]),
         resp[0, 0] * np.cos(resp[1, 1]), 'g^')
plt.plot(resp[:, 0] * np.sin(resp[:, 1]),
         resp[:, 0] * np.cos(resp[:, 1]), label='Unshaped')
# plt.gca().invert_yaxis()
plt.axes().set_aspect('equal', 'datalim')
plt.xlabel('Horizontal Motion')
plt.ylabel('Vertical Motion')
plt.show()