{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> This is one of the 100 recipes of the [IPython Cookbook](http://ipython-books.github.io/), the definitive guide to high-performance scientific computing and data science in Python.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 9.4. Finding the equilibrium state of a physical system by minimizing its potential energy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1. Let's import NumPy, SciPy and matplotlib."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.optimize as opt\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. We define a few constants in the International System of Units."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "g = 9.81  # gravity of Earth\n",
    "m = .1  # mass, in kg\n",
    "n = 20  # number of masses\n",
    "e = .1  # initial distance between the masses\n",
    "l = e  # relaxed length of the springs\n",
    "k = 10000  # spring stiffness"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. We define the initial positions of the masses. They are arranged on a two-dimensional grid with two lines and $n/2$ columns."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "P0 = np.zeros((n, 2))\n",
    "P0[:,0] = np.repeat(e*np.arange(n//2), 2)\n",
    "P0[:,1] = np.tile((0,-e), n//2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. Now, let's define the connectivity matrix between the masses. Coefficient $i,j$ is $1$ if masses $i$ and $j$ are connected by a spring."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "A = np.eye(n, n, 1) + np.eye(n, n, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5. We also specify the spring stiffness of each spring. It is $l$, except for *diagonal* springs where it is $l \\times \\sqrt{2}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "L = l * (np.eye(n, n, 1) + np.eye(n, n, 2))\n",
    "for i in range(n//2-1):\n",
    "    L[2*i+1,2*i+2] *= np.sqrt(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "6. We also need the indices of the spring connections."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "I, J = np.nonzero(A)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "7. The `dist` function computes the distance matrix (distance between any pair of masses)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "dist = lambda P: np.sqrt((P[:,0]-P[:,0][:, np.newaxis])**2 + \n",
    "                         (P[:,1]-P[:,1][:, np.newaxis])**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "7. We define a function that displays the system. The springs are colored according to their tension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def spring_color_map(c):\n",
    "    min_c, max_c = -0.00635369422326, 0.00836362559722\n",
    "    ratio = (max_c-c) / (max_c-min_c)\n",
    "    color = plt.cm.coolwarm(ratio)\n",
    "    shading = np.sqrt(abs(ratio-0.5)*2)\n",
    "    return (shading*color[0], shading*color[1], shading*color[2], color[3])\n",
    "\n",
    "def show_bar(P):\n",
    "    plt.figure(figsize=(5,4));\n",
    "    # Wall.\n",
    "    plt.axvline(0, color='k', lw=3);\n",
    "    # Distance matrix.\n",
    "    D = dist(P)\n",
    "    # We plot the springs.\n",
    "    for i, j in zip(I, J):\n",
    "        # The color depends on the spring tension, which\n",
    "        # is proportional to the spring elongation.\n",
    "        c = D[i,j] - L[i,j]\n",
    "        plt.plot(P[[i,j],0], P[[i,j],1], \n",
    "                 lw=2, color=spring_color_map(c));\n",
    "    # We plot the masses.\n",
    "    plt.plot(P[[I,J],0], P[[I,J],1], 'ok',);\n",
    "    # We configure the axes.\n",
    "    plt.axis('equal');\n",
    "    plt.xlim(P[:,0].min()-e/2, P[:,0].max()+e/2);\n",
    "    plt.ylim(P[:,1].min()-e/2, P[:,1].max()+e/2);\n",
    "    plt.xticks([]); plt.yticks([]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "8. Here is the system in its initial configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "show_bar(P0);\n",
    "plt.title(\"Initial configuration\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "9. To find the equilibrium state, we need to minimize the total potential energy of the system. The following function computes the energy of the system, given the positions of the masses. This function is explained in *How it works...*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def energy(P):\n",
    "    # The argument P is a vector (flattened matrix).\n",
    "    # We convert it to a matrix here.\n",
    "    P = P.reshape((-1, 2))\n",
    "    # We compute the distance matrix.\n",
    "    D = dist(P)\n",
    "    # The potential energy is the sum of the\n",
    "    # gravitational and elastic potential energies.\n",
    "    return (g * m * P[:,1].sum() + \n",
    "            .5 * (k * A * (D - L)**2).sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "10. Let's compute the potential energy of the initial configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "energy(P0.ravel())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "11. Now, let's minimize the potential energy with a function minimization method. We need a **constrained optimization algorithm**, because we make the assumption that the two first masses are fixed to the wall. Therefore, their positions cannot change. The **L-BFGS-B** algorithm, a variant of the BFGS algorithm, accepts bound constraints. Here, we force the first two points to stay at their initial positions, whereas there are no constraints on the other points. The `minimize` function accepts a `bounds` list containing, for each dimension, a pair of `[min, max]` values. (http://en.wikipedia.org/wiki/Limited-memory_BFGS#L-BFGS-B)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "bounds = np.c_[P0[:2,:].ravel(), P0[:2,:].ravel()].tolist() + \\\n",
    "         [[None, None]] * (2*(n-2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "P1 = opt.minimize(energy, P0.ravel(),\n",
    "                  method='L-BFGS-B',\n",
    "                  bounds=bounds).x.reshape((-1, 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "12. Let's display the stable configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "show_bar(P1);\n",
    "plt.title(\"Equilibrium configuration\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This configuration looks realistic. The tension appears to be maximal on the top springs near the wall."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> You'll find all the explanations, figures, references, and much more in the book (to be released later this summer).\n",
    "\n",
    "> [IPython Cookbook](http://ipython-books.github.io/), by [Cyrille Rossant](http://cyrille.rossant.net), Packt Publishing, 2014 (500 pages)."
   ]
  }
 ],
 "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.4.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}