{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "dependent-girlfriend",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sympy\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "taken-evolution",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Building Constraints in [SymPy](https://www.sympy.org)\n",
    "\n",
    "It can be difficult to build constraints for general dynamic solutions. Each constraint is realtively straight forward, but keeping track of Jacobian terms or acceleration constraints\n",
    "- $\\mathbf{C_q}$ \n",
    "- $\\mathbf{Q_d} = -(\\mathbf{C_q \\dot{q}})_\\mathbf{q}\\mathbf{\\dot{q}} - 2\\mathbf{C}_{\\mathbf{q}t}\\dot{\\mathbf{q}}- \\mathbf{C}_{tt}$ \n",
    "\n",
    "can be cumbersome and confusing. In this notebook, you will use [SymPy](https://www.sympy.org) to \n",
    "\n",
    "1. define $\\mathbf{q}~and~\\dot{\\mathbf{q}}$\n",
    "2. take a Jacobian and find the constraints on acceleration\n",
    "3. [`lambdify`](https://docs.sympy.org/latest/tutorial/basic_operations.html#lambdify) the constraints, Jacobian, and $\\mathbf{Q}_d$\n",
    "\n",
    "## 1. define $\\mathbf{q}~and~\\dot{\\mathbf{q}}$\n",
    "\n",
    "First, define the variables using `sympy.Matrix`. SymPy uses `var` to define variables, here you create\n",
    "\n",
    "$q = \\left[\\begin{matrix}q_{1}\\\\q_{2}\\\\q_{3}\\\\q_{4}\\\\q_{5}\\\\q_{6}\\end{matrix}\\right],~\n",
    "\\dot{q} = \\left[\\begin{matrix}dq_{1}\\\\dq_{2}\\\\dq_{3}\\\\dq_{4}\\\\dq_{5}\\\\dq_{6}\\end{matrix}\\right],~t,~and~L$\n",
    "\n",
    "as SymPy variables. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "living-corruption",
   "metadata": {},
   "outputs": [],
   "source": [
    "sympy.var('q1, q2, q3, q4, q5, q6')\n",
    "sympy.var('dq1, dq2, dq3, dq4, dq5, dq6')\n",
    "sympy.var('t, L2')\n",
    "q = sympy.Matrix([q1, q2, q3, q4, q5, q6])\n",
    "dq = sympy.Matrix([dq1, dq2, dq3, dq4, dq5, dq6])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "junior-imagination",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left[\\begin{matrix}q_{1}\\\\q_{2}\\\\q_{3}\\\\q_{4}\\\\q_{5}\\\\q_{6}\\end{matrix}\\right]$"
      ],
      "text/plain": [
       "Matrix([\n",
       "[q1],\n",
       "[q2],\n",
       "[q3],\n",
       "[q4],\n",
       "[q5],\n",
       "[q6]])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "q"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "distinct-capital",
   "metadata": {},
   "source": [
    "$\\mathbf{q} = \\left[\\begin{matrix}q_{1}\\\\q_{2}\\\\q_{3}\\\\q_{4}\\\\q_{5}\\\\q_{6}\\end{matrix}\\right]$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "derived-honey",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left[\\begin{matrix}dq_{1}\\\\dq_{2}\\\\dq_{3}\\\\dq_{4}\\\\dq_{5}\\\\dq_{6}\\end{matrix}\\right]$"
      ],
      "text/plain": [
       "Matrix([\n",
       "[dq1],\n",
       "[dq2],\n",
       "[dq3],\n",
       "[dq4],\n",
       "[dq5],\n",
       "[dq6]])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dq"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ultimate-deviation",
   "metadata": {},
   "source": [
    "## 2. take a Jacobian and find the constraints on acceleration\n",
    "\n",
    "![two bodies: sliding block and compound pendulum](../images/spring_compound-2_bodies.svg)\n",
    "\n",
    "In this example, you have two rigid bodies. Body one $[q_1,~q_2,~q_3]$ slides on the x-axis. Body 1 and body 2, $[q_4,~q_5,~q_6]$ are connected by a pin in the center of body 1, $\\mathbf{R}_{pin} = q_1\\hat{i}+q_2\\hat{j}$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "gothic-packet",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left[\\begin{matrix}\\frac{L_{2} \\cos{\\left(q_{6} \\right)}}{2} + q_{1} - q_{4}\\\\\\frac{L_{2} \\sin{\\left(q_{6} \\right)}}{2} + q_{2} - q_{5}\\\\q_{2}\\\\q_{3}\\end{matrix}\\right]$"
      ],
      "text/plain": [
       "Matrix([\n",
       "[L2*cos(q6)/2 + q1 - q4],\n",
       "[L2*sin(q6)/2 + q2 - q5],\n",
       "[                    q2],\n",
       "[                    q3]])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "C = sympy.Matrix([q1  - q4 + L2/2*sympy.cos(q6),\n",
    "                 q2 - q5 + L2/2*sympy.sin(q6),\n",
    "                 q2,\n",
    "                 q3])\n",
    "C"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "confused-marriage",
   "metadata": {},
   "source": [
    "Next, take the Jacobian, $\\mathbf{C_q}$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "timely-microwave",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cq dimensions: (4, 6)\n"
     ]
    },
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left[\\begin{matrix}1 & 0 & 0 & -1 & 0 & - \\frac{L_{2} \\sin{\\left(q_{6} \\right)}}{2}\\\\0 & 1 & 0 & 0 & -1 & \\frac{L_{2} \\cos{\\left(q_{6} \\right)}}{2}\\\\0 & 1 & 0 & 0 & 0 & 0\\\\0 & 0 & 1 & 0 & 0 & 0\\end{matrix}\\right]$"
      ],
      "text/plain": [
       "Matrix([\n",
       "[1, 0, 0, -1,  0, -L2*sin(q6)/2],\n",
       "[0, 1, 0,  0, -1,  L2*cos(q6)/2],\n",
       "[0, 1, 0,  0,  0,             0],\n",
       "[0, 0, 1,  0,  0,             0]])"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Cq = C.jacobian(q)\n",
    "print('Cq dimensions:', Cq.shape)\n",
    "Cq"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "subsequent-warning",
   "metadata": {},
   "source": [
    "Now, you can calculate \n",
    "\n",
    "$\\mathbf{Q_d} = -(\\mathbf{C_q \\dot{q}})_\\mathbf{q}\\mathbf{\\dot{q}} - 2\\mathbf{C}_{\\mathbf{q}t}\\dot{\\mathbf{q}}- \\mathbf{C}_{tt}$ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "native-candy",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\left[\\begin{matrix}\\frac{L_{2} dq_{6}^{2} \\cos{\\left(q_{6} \\right)}}{2}\\\\\\frac{L_{2} dq_{6}^{2} \\sin{\\left(q_{6} \\right)}}{2}\\\\0\\\\0\\end{matrix}\\right]$"
      ],
      "text/plain": [
       "Matrix([\n",
       "[L2*dq6**2*cos(q6)/2],\n",
       "[L2*dq6**2*sin(q6)/2],\n",
       "[                  0],\n",
       "[                  0]])"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "Qd = -(C.jacobian(q)@dq).jacobian(q)@dq\\\n",
    "    - 2*sympy.diff(Cq,t)@dq\\\n",
    "    - sympy.diff(C, t, 2)\n",
    "\n",
    "Qd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "available-responsibility",
   "metadata": {},
   "source": [
    "## 3. [`lambdify`](https://docs.sympy.org/latest/tutorial/basic_operations.html#lambdify) the constraints, Jacobian, and $\\mathbf{Q}_d$\n",
    "\n",
    "Finally, you can create functions that return NumPy arrays with `sympy.lambdify`. The inputs are\n",
    "\n",
    "```python\n",
    "sympy.lambdify( inputs, function, \"numpy\" )\n",
    "```\n",
    "\n",
    "where the `inputs` are $\\mathbf{q}$, $\\dot{\\mathbf{q}}$, and other parameters like time or dimensions, $L^2$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "restricted-lease",
   "metadata": {},
   "outputs": [],
   "source": [
    "Cq_sys = sympy.lambdify((q, t, L2), Cq, 'numpy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "recent-springfield",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 1. ,  0. ,  0. , -1. ,  0. , -0. ],\n",
       "       [ 0. ,  1. ,  0. ,  0. , -1. ,  0.5],\n",
       "       [ 0. ,  1. ,  0. ,  0. ,  0. ,  0. ],\n",
       "       [ 0. ,  0. ,  1. ,  0. ,  0. ,  0. ]])"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Cq_sys(np.zeros(6), 1, L2 = 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "geographic-highlight",
   "metadata": {},
   "outputs": [],
   "source": [
    "Q = sympy.lambdify((q, dq, t, L2), Qd, \"numpy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "ready-jordan",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0.27015115],\n",
       "       [0.42073549],\n",
       "       [0.        ],\n",
       "       [0.        ]])"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Q(np.ones(6), np.ones(6), 0, L2 = 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "mental-handy",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,md:myst"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}