{ "cells": [ { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "import sympy as smp\n", "from scipy.integrate import odeint\n", "import matplotlib.pyplot as plt\n", "from matplotlib import animation\n", "from mpl_toolkits.mplot3d import Axes3D\n", "from matplotlib.animation import PillowWriter" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "
\n", " \n", "
" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Define all appropriate symbols using sympy. " ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [], "source": [ "t, g = smp.symbols('t g')\n", "m1, m2 = smp.symbols('m1 m2')\n", "L1, L2 = smp.symbols('L1, L2')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "$\\theta_1$ and $\\theta_2$ are functions of time (which we will eventually solve for). We need to define them carefully." ] }, { "cell_type": "code", "execution_count": 10, "metadata": {}, "outputs": [], "source": [ "the1, the2 = smp.symbols(r'\\theta_1, \\theta_2', cls=smp.Function)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Explicitly write them as functions of time $t$:" ] }, { "cell_type": "code", "execution_count": 12, "metadata": {}, "outputs": [], "source": [ "the1 = the1(t)\n", "the2 = the2(t)" ] }, { "cell_type": "code", "execution_count": 13, "metadata": {}, "outputs": [ { "data": { "text/latex": [ "$\\displaystyle \\theta_{1}{\\left(t \\right)}$" ], "text/plain": [ "\\theta_1(t)" ] }, "execution_count": 13, "metadata": {}, "output_type": "execute_result" } ], "source": [ "the1" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Define derivatives and second derivatives" ] }, { "cell_type": "code", "execution_count": 14, "metadata": {}, "outputs": [], "source": [ "the1_d = smp.diff(the1, t)\n", "the2_d = smp.diff(the2, t)\n", "the1_dd = smp.diff(the1_d, t)\n", "the2_dd = smp.diff(the2_d, t)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Define $x_1$, $y_1$, $x_2$, and $y_2$ written in terms of the parameters above." ] }, { "cell_type": "code", "execution_count": 17, "metadata": {}, "outputs": [], "source": [ "x1 = L1*smp.sin(the1)\n", "y1 = -L1*smp.cos(the1)\n", "x2 = L1*smp.sin(the1)+L2*smp.sin(the2)\n", "y2 = -L1*smp.cos(the1)-L2*smp.cos(the2)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Then use these to define kinetic and potential energy for each mass. Obtain the Lagrangian" ] }, { "cell_type": "code", "execution_count": 23, "metadata": {}, "outputs": [], "source": [ "# Kinetic\n", "T1 = 1/2 * m1 * (smp.diff(x1, t)**2 + smp.diff(y1, t)**2)\n", "T2 = 1/2 * m2 * (smp.diff(x2, t)**2 + smp.diff(y2, t)**2)\n", "T = T1+T2\n", "# Potential\n", "V1 = m1*g*y1\n", "V2 = m2*g*y2\n", "V = V1 + V2\n", "# Lagrangian\n", "L = T-V" ] }, { "cell_type": "code", "execution_count": 24, "metadata": {}, "outputs": [ { "data": { "text/latex": [ "$\\displaystyle L_{1} g m_{1} \\cos{\\left(\\theta_{1}{\\left(t \\right)} \\right)} - g m_{2} \\left(- L_{1} \\cos{\\left(\\theta_{1}{\\left(t \\right)} \\right)} - L_{2} \\cos{\\left(\\theta_{2}{\\left(t \\right)} \\right)}\\right) + 0.5 m_{1} \\left(L_{1}^{2} \\sin^{2}{\\left(\\theta_{1}{\\left(t \\right)} \\right)} \\left(\\frac{d}{d t} \\theta_{1}{\\left(t \\right)}\\right)^{2} + L_{1}^{2} \\cos^{2}{\\left(\\theta_{1}{\\left(t \\right)} \\right)} \\left(\\frac{d}{d t} \\theta_{1}{\\left(t \\right)}\\right)^{2}\\right) + 0.5 m_{2} \\left(\\left(L_{1} \\sin{\\left(\\theta_{1}{\\left(t \\right)} \\right)} \\frac{d}{d t} \\theta_{1}{\\left(t \\right)} + L_{2} \\sin{\\left(\\theta_{2}{\\left(t \\right)} \\right)} \\frac{d}{d t} \\theta_{2}{\\left(t \\right)}\\right)^{2} + \\left(L_{1} \\cos{\\left(\\theta_{1}{\\left(t \\right)} \\right)} \\frac{d}{d t} \\theta_{1}{\\left(t \\right)} + L_{2} \\cos{\\left(\\theta_{2}{\\left(t \\right)} \\right)} \\frac{d}{d t} \\theta_{2}{\\left(t \\right)}\\right)^{2}\\right)$" ], "text/plain": [ "L1*g*m1*cos(\\theta_1(t)) - g*m2*(-L1*cos(\\theta_1(t)) - L2*cos(\\theta_2(t))) + 0.5*m1*(L1**2*sin(\\theta_1(t))**2*Derivative(\\theta_1(t), t)**2 + L1**2*cos(\\theta_1(t))**2*Derivative(\\theta_1(t), t)**2) + 0.5*m2*((L1*sin(\\theta_1(t))*Derivative(\\theta_1(t), t) + L2*sin(\\theta_2(t))*Derivative(\\theta_2(t), t))**2 + (L1*cos(\\theta_1(t))*Derivative(\\theta_1(t), t) + L2*cos(\\theta_2(t))*Derivative(\\theta_2(t), t))**2)" ] }, "execution_count": 24, "metadata": {}, "output_type": "execute_result" } ], "source": [ "L" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Get Lagrange's equations\n", "\n", "$$\\frac{\\partial L}{\\partial \\theta_1} - \\frac{d}{dt}\\frac{\\partial L}{\\partial \\dot{\\theta_1}} = 0$$\n", "$$\\frac{\\partial L}{\\partial \\theta_2} - \\frac{d}{dt}\\frac{\\partial L}{\\partial \\dot{\\theta_2}} = 0$$" ] }, { "cell_type": "code", "execution_count": 25, "metadata": {}, "outputs": [], "source": [ "LE1 = smp.diff(L, the1) - smp.diff(smp.diff(L, the1_d), t).simplify()\n", "LE2 = smp.diff(L, the2) - smp.diff(smp.diff(L, the2_d), t).simplify()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Solve Lagranges equations (this assumes that `LE1` and `LE2` are both equal to zero)" ] }, { "cell_type": "code", "execution_count": 28, "metadata": {}, "outputs": [], "source": [ "sols = smp.solve([LE1, LE2], (the1_dd, the2_dd),\n", " simplify=False, rational=False)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now we have \n", "\n", "* $\\frac{d^2 \\theta_1}{dt^2} = ...$\n", "* $\\frac{d^2 \\theta_2}{dt^2} = ...$\n", "\n", "These are two second order ODEs! In python we can only solve systems of first order ODEs. Any system of second order ODEs can be converted as follows:\n", "\n", "1. Define $z_1 = d\\theta_1/dt$ and $z_2=d\\theta_2/dt$\n", "2. Then $dz_1/dt = d^2\\theta_1/dt^2$ and $dz_2/dt = d^2\\theta_2/dt^2$\n", "\n", "Now we get a system of 4 first order ODEs (as opposed to 2 second order ones)\n", "\n", "* $d z_1/dt = ...$\n", "* $d\\theta_1/dt = z_1$\n", "* $d z_2/dt = ...$\n", "* $d\\theta_2/dt = z_1$\n", "\n", "We need to convert the **symbolic** expressions above to numerical functions so we can use them in a numerical python solver. For this we use `smp.lambdify`" ] }, { "cell_type": "code", "execution_count": 33, "metadata": {}, "outputs": [], "source": [ "dz1dt_f = smp.lambdify((t,g,m1,m2,L1,L2,the1,the2,the1_d,the2_d), sols[the1_dd])\n", "dz2dt_f = smp.lambdify((t,g,m1,m2,L1,L2,the1,the2,the1_d,the2_d), sols[the2_dd])\n", "dthe1dt_f = smp.lambdify(the1_d, the1_d)\n", "dthe2dt_f = smp.lambdify(the2_d, the2_d)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now define $\\vec{S} = (\\theta_1, z_1, \\theta_2, z_2)$. IF we're going to use an ODE solver in python, we need to write a function that takes in $\\vec{S}$ and $t$ and returns $d\\vec{S}/dt$. In other words, we need to define $d\\vec{S}/dt (\\vec{S}, t)$\n", "\n", "* Our system of ODEs can be fully specified using $d\\vec{S}/dt$ and depends only on $\\vec{S}$ and $t$" ] }, { "cell_type": "code", "execution_count": 35, "metadata": {}, "outputs": [], "source": [ "def dSdt(S, t, g, m1, m2, L1, L2):\n", " the1, z1, the2, z2 = S\n", " return [\n", " dthe1dt_f(z1),\n", " dz1dt_f(t, g, m1, m2, L1, L2, the1, the2, z1, z2),\n", " dthe2dt_f(z2),\n", " dz2dt_f(t, g, m1, m2, L1, L2, the1, the2, z1, z2),\n", " ]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Solve the system of ODEs using scipys `odeint` method" ] }, { "cell_type": "code", "execution_count": 36, "metadata": {}, "outputs": [], "source": [ "t = np.linspace(0, 40, 1001)\n", "g = 9.81\n", "m1=2\n", "m2=1\n", "L1 = 2\n", "L2 = 1\n", "ans = odeint(dSdt, y0=[1, -3, -1, 5], t=t, args=(g,m1,m2,L1,L2))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "25 times per second (number of data points). This will be important for animating later on." ] }, { "cell_type": "code", "execution_count": 41, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "array([[ 1. , 0.87413906, 0.73333306, ..., -0.75268561,\n", " -0.60481044, -0.45270209],\n", " [-3. , -3.31025173, -3.75998052, ..., 3.64397936,\n", " 3.7514106 , 3.84881428],\n", " [-1. , -0.78277685, -0.52139856, ..., 79.49932146,\n", " 79.50044235, 79.55942868],\n", " [ 5. , 5.91440903, 7.24155048, ..., -0.66192636,\n", " 0.73534285, 2.22745168]])" ] }, "execution_count": 41, "metadata": {}, "output_type": "execute_result" } ], "source": [ "ans.T" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Can obtain $\\theta_1(t)$ and $\\theta_2(t)$ from the answer" ] }, { "cell_type": "code", "execution_count": 42, "metadata": {}, "outputs": [], "source": [ "the1 = ans.T[0]\n", "the2 = ans.T[2]" ] }, { "cell_type": "code", "execution_count": 43, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "[