{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to SciPy\n",
    "Tutorial at EuroSciPy 2019, Bilbao\n",
    "## 3.1. Optimization – The hanging chain\n",
    "![chain link](images/chainlink.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy import optimize\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n",
    "class Chain:\n",
    "    def __init__(self, nlinks, hdist):\n",
    "        if nlinks < hdist:\n",
    "            raise ValueError('a hanging chain with {} links of '\n",
    "                             'length 1 cannot have endpoints '\n",
    "                             'separated by a distance {}'.format(\n",
    "                                    nlinks, hdist))\n",
    "        self.nlinks = nlinks\n",
    "        self.hdist = hdist\n",
    " \n",
    "    def equilibrium(self):\n",
    "        phi0 = np.arcsin(self.hdist/self.nlinks)\n",
    "        result = optimize.minimize(\n",
    "                     self.f_energy,\n",
    "                     np.linspace(-phi0, phi0, self.nlinks),\n",
    "                     method='SLSQP',\n",
    "                     constraints=[{'type': 'eq', 'fun': self.x_constraint},\n",
    "                                  {'type': 'eq', 'fun': self.y_constraint}])\n",
    "        return result.x\n",
    "    \n",
    "    def x_constraint(self, phi_vals):\n",
    "        \"\"\"ensure the correct horizontal distance\n",
    "        \n",
    "           phi_vals: angles of links with respect to the horizontal\n",
    "           \n",
    "        \"\"\"\n",
    "        return np.sum(np.cos(phi_vals))-self.hdist \n",
    "       \n",
    "    def y_constraint(self, phi_vals):\n",
    "        \"\"\"ensure that endpoints are at the same height\n",
    "        \n",
    "           phi_vals: angles of links with respect to the horizontal\n",
    "           \n",
    "        \"\"\"\n",
    "        return np.sum(np.sin(phi_vals))\n",
    "\n",
    "    def f_energy(self, phi_vals):\n",
    "        \"\"\"potential energy of all links\n",
    "        \n",
    "        \"\"\"\n",
    "        return np.sum(np.arange(self.nlinks, 0, -1)*np.sin(phi_vals))\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def angles_to_coords(phi_vals):\n",
    "    \"\"\"convert angles to coordinates of link endpoints\n",
    "        \n",
    "       phi_vals: angles of links with respect to the horizontal\n",
    "       x, y: coordinates of link endpoints\n",
    "           \n",
    "    \"\"\"\n",
    "    dim = phi_vals.shape[0]+1\n",
    "    x = np.zeros(dim)\n",
    "    x[1:] = np.cumsum(np.cos(phi_vals))\n",
    "    y = np.zeros(dim)\n",
    "    y[1:] = np.cumsum(np.sin(phi_vals))\n",
    "    return x, y\n",
    "\n",
    "nlinks = 20\n",
    "hdist = 16\n",
    "\n",
    "c = Chain(nlinks, hdist)\n",
    "plt.axes().set_aspect('equal')\n",
    "_ = plt.plot(*angles_to_coords(c.equilibrium()), 'o-')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.2. Solving a set of ordinary differential equations – The falling chain"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Literature    \n",
    "W. Tomaszewski, P. Pieranski, J.-C. Geminard    \n",
    "*The motion of a free falling chain tip*    \n",
    "[American Journal of Physics **74**, 776 (2006)](https://dx.doi.org/10.1119/1.2204074)\n",
    "\n",
    "#### Equations of motion\n",
    "\n",
    "$$\\sum_{j=1}^n m_{i,j}\\cos(\\varphi_i-\\varphi_j)\\ddot\\varphi_j\n",
    "  = -\\sum_{j=1}^n m_{i, j}\\sin(\\varphi_i-\\varphi_j)\\dot\\varphi_j^2\n",
    "    +\\frac{r}{m\\ell^2}(\\dot\\varphi_{i-1}-2\\dot\\varphi_i+\\dot\\varphi_{i+1})\n",
    "    -\\frac{g}{\\ell}a_i\\cos(\\varphi_i)$$\n",
    "with\n",
    "$$a_i = n-i+\\frac{1}{2}$$\n",
    "and\n",
    "$$m_{i,j} = \\begin{cases}\n",
    "            n-i+\\frac{1}{3} & i=j \\\\\n",
    "            n-\\max(i,j)+\\frac{1}{2} & i\\neq j\n",
    "            \\end{cases}$$\n",
    "            \n",
    "parameters:    \n",
    "$m$ mass of chain element    \n",
    "$\\ell$ length of chain element     \n",
    "$g$ acceleration due to gravity    \n",
    "$r$ damping constant\n",
    "\n",
    "#### How to deal with second-order differential equations\n",
    "\n",
    "Introduce a new variable\n",
    "$$p_\\varphi = \\dot\\varphi$$\n",
    "to replace the second-order differential equation\n",
    "$$\\ddot\\varphi = f(\\dot\\varphi, \\varphi)$$\n",
    "by two first-order differential equations\n",
    "$$\n",
    " \\begin{aligned}\n",
    "  \\dot\\varphi &= p_\\varphi\\\\\n",
    "  \\dot p_\\varphi &= f(p_\\varphi, \\varphi)\n",
    " \\end{aligned}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy.linalg as LA\n",
    "from scipy import integrate\n",
    "\n",
    "class FallingChain(Chain):\n",
    "    def __init__(self, nlinks, hdist, damping):\n",
    "        super(FallingChain, self).__init__(nlinks, hdist)\n",
    "        self.set_matrix_m()\n",
    "        self.set_vector_a()\n",
    "        self.set_matrix_damping(damping)\n",
    "        \n",
    "    def set_matrix_m(self):\n",
    "        m = np.fromfunction(lambda i, j: self.nlinks-np.maximum(i, j)-0.5,\n",
    "                            (self.nlinks, self.nlinks))\n",
    "        self.m = m-np.identity(self.nlinks)/6\n",
    "\n",
    "    def set_vector_a(self):\n",
    "        self.a = np.arange(self.nlinks, 0, -1)-0.5\n",
    "    \n",
    "    def set_matrix_damping(self, damping):\n",
    "        self.damping = (-2*np.identity(self.nlinks, dtype=np.float64)\n",
    "                        + np.eye(self.nlinks, k=1)\n",
    "                        + np.eye(self.nlinks, k=-1))\n",
    "        self.damping[0, 0] = -1\n",
    "        self.damping[self.nlinks-1, self.nlinks-1] = -1\n",
    "        self.damping = damping*self.damping\n",
    "\n",
    "    def solve_eq_of_motion(self, time_i, time_f, nt):\n",
    "        y0 = np.zeros(2*self.nlinks, dtype=np.float64)\n",
    "        y0[self.nlinks:] = self.equilibrium()\n",
    "        self.solution = integrate.solve_ivp(\n",
    "                            self.diff, (time_i, time_f), y0,\n",
    "                            method='RK45',\n",
    "                            t_eval=np.linspace(time_i, time_f, nt))\n",
    "\n",
    "    def diff(self, t, y):\n",
    "        momenta = y[:self.nlinks]\n",
    "        angles = y[self.nlinks:]\n",
    "        d_angles = momenta\n",
    "        ci = np.cos(angles)\n",
    "        cij = np.cos(angles[:, np.newaxis]-angles)\n",
    "        sij = np.sin(angles[:, np.newaxis]-angles)\n",
    "        mcinv = LA.inv(self.m*cij)\n",
    "        d_momenta = -np.dot(self.m*sij, momenta*momenta)\n",
    "        d_momenta = d_momenta + np.dot(self.damping, momenta)\n",
    "        d_momenta = d_momenta - self.a*ci\n",
    "        d_momenta = np.dot(mcinv, d_momenta)\n",
    "        d = np.empty_like(y)\n",
    "        d[:self.nlinks] = d_momenta\n",
    "        d[self.nlinks:] = d_angles\n",
    "        return d"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Set-up for animation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import rc\n",
    "import matplotlib.animation as animation\n",
    "rc('animation', html='jshtml')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following cell should always be executed before running the animation cell below.\n",
    "In this way, an undesired extra static image will be avoided."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture\n",
    "fig = plt.figure()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Animation of a falling chain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def animate(i):\n",
    "    x, y = angles_to_coords(c.solution.y[:, i][c.nlinks:])\n",
    "    line.set_data(x, y)\n",
    "    return line,\n",
    "\n",
    "def init():\n",
    "    line.set_data([], [])\n",
    "    return line,\n",
    "\n",
    "nlinks = 50\n",
    "hdist = 40\n",
    "damping = 1\n",
    "ti = 0\n",
    "tf = 200\n",
    "tsteps = 1000\n",
    "\n",
    "c = FallingChain(nlinks, hdist, damping)\n",
    "c.solve_eq_of_motion(ti, tf, tsteps)\n",
    "ax = fig.add_subplot(111, autoscale_on=False,\n",
    "                     xlim=(-nlinks, nlinks), ylim=(-nlinks, 0.3*nlinks))\n",
    "ax.set_aspect('equal')\n",
    "line, = ax.plot([], [])\n",
    "anim = animation.FuncAnimation(fig, animate, tsteps,\n",
    "                               interval=20, blit=True,\n",
    "                               init_func=init)\n",
    "anim"
   ]
  }
 ],
 "metadata": {
  "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}