from sympy import *
from sympy.solvers.solveset import linsolve
from sympy import Matrix, S
from sympy import symbols
from scipy.integrate import odeint
import numpy as np
init_printing()
n = 3
# arrays for all the constants in the system
l = []
m = []
k = []
g=9.81
# fill in the constant values
for i in range(n):
l.append(1)
m.append(1)
k.append(0.1)
# arrays for all the symbols in the system
phi = [] #angles
dphi = [] #first derivatives
ddphi = [] #second derivatives
diff2 = [] #list of tuples that map expressions phi.diff(t).diff(t) to symbols ddphi
diff1 = [] #see above for first derivatives
x = [] #x positions of pendulums masses
y = [] #y positions
Fr=[] #Friction terms
t = symbols("t")
# Lagrangian, an empty expression
L = Add(0, 0)
for i in range(n):
temp_phi = symbols("phi_" + str(i))
temp_phi = temp_phi(t)
phi.append(temp_phi)
temp_dphi = symbols("dphi_" + str(i))
dphi.append(temp_dphi)
temp_ddphi = symbols("ddphi_" + str(i))
ddphi.append(temp_ddphi)
diff1.append((temp_phi.diff(t), temp_dphi))
diff2.append((temp_phi.diff(t).diff(t), temp_ddphi))
temp_x = symbols("x_" + str(i))
temp_y = symbols("y_" + str(i))
temp_x = l[i] * sin(phi[-1])
temp_y = -l[i] * cos(phi[-1])
if i > 0:
temp_x = temp_x + x[-1]
temp_y = temp_y + y[-1]
x.append(temp_x)
y.append(temp_y)
temp_v = temp_x.diff(t) ** 2 + temp_y.diff(t) ** 2
Fr.append(-temp_dphi*k[i])
# adding the kinetic energy of this pendulum to L
L += 1 / 2 * temp_v * m[i]
# adding the potential energy of this pendulum to L
L -= m[i] * g * temp_y
eqs = []
for i in range(n):
a = L.diff(phi[i].diff(t))
da = a.diff(t)
b = L.diff(phi[i]) +Fr[i]
eq = da - b
eq=eq.subs(diff2).subs(diff1).simplify()
eqs.append(eq)
eqs
# this can be solved for the second order derivatives
import time
start=time.time()
res = linsolve(eqs, ddphi)
end=time.time()
print(end-start)
res = [res for res in res]
res = res[0]
lambdified_solns = [lambdify(phi + dphi, soln) for soln in res]
# state
# phi_1
# phi_1'
# phi_2
# phi_2'
# phi
# dphi
def derivs(state, t):
d=np.empty(state.shape)
d[:-n] = state[n:]
for i in range(n):
d[i+n]=lambdified_solns[i](*state)
return d
# solve ode
start_phi=[]
for i in range(n):
start_phi.append(3)
start_dphi=[]
for i in range(n):
start_dphi.append(0.5)
steps = 5000
t_end = 300
res_ode = odeint(derivs, start_phi + start_dphi, np.linspace(0, t_end, steps))