{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dancing particles\n",
    "\n",
    "In this chapter we focus on developing the particle dynamics code. First we install required packages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import Pkg; Pkg.add([\"Plots\", \"Printf\", \"Zygote\"])\n",
    "using LinearAlgebra"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting in Julia\n",
    "\n",
    "In many day-to-day tasks there is not yet the one canonical way to do it in Julia. This also concerns plotting and visualisation of data. Right now quite a few different packages with their own APIs and feature sets exist. To give you some examples:\n",
    "- [PyPlot](https://github.com/JuliaPy/PyPlot.jl) (Plotting via python and matplotlib)\n",
    "- [PGFPlots](https://github.com/JuliaTeX/PGFPlots.jl) (Plotting via TeX, TikZ and pgf)\n",
    "- [Gladify](https://github.com/GiovineItalia/Gadfly.jl) (Pure Julia plotting implementation, widespread)\n",
    "- [PlotlyJS](https://github.com/sglyon/PlotlyJS.jl) (Ploting via Javascript and Plotly)\n",
    "- [GR](https://github.com/jheinen/GR.jl) (Plotting based on the GR framework, pretty fast)\n",
    "- [Makie](https://github.com/JuliaPlots/Makie.jl) (Extremely feature-rich; can plot on the GPU directly)\n",
    "\n",
    "What we will use in this course is [Plots](https://github.com/JuliaPlots/Plots.jl). I like it because it is a metapackage unifying some of the above under a common interface. Flipping the backend is just a single function call, so you can change between them as needed. Also it will probably not disappear any time soon ;). We load `Plots` first and explicitly switch to the default `GR` backend."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Plots\n",
    "Plots.gr()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plotting is then done with the `plot` function. Additional data series can be added using `plot!` and attributes can be changed via keywords arguments or standalone functions like `xaxis!`, `title!` and so on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Vho(r, a=0.5)  = (r - a)^2   # Shifted harmonic oscillator\n",
    "dVho(r, a=0.5) = 2r -2a      # Derivative\n",
    "\n",
    "r = collect(-5:0.1:5)        # Create series of x-values\n",
    "p = plot(r, Vho.(r), label=\"Vho a=0.5\")\n",
    "plot!(p, r, dVho.(r), label=\"∂Vho\")\n",
    "xaxis!(p, \"relative radial distance\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Of course scatter plots can be made as well ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scatter(r, Vho.(r), label=\"Vho\", marker=:cross)\n",
    "scatter!(r, dVho.(r), label=\"∂Vho\", marker=:circle)\n",
    "xaxis!(\"relative radial distance\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "... and in fact there are plenty of things you can do to style your plot further, see [attributes](http://docs.juliaplots.org/latest/attributes) and [layouts](http://docs.juliaplots.org/latest/layouts/) for some ideas if you are curious."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### More details\n",
    "- http://docs.juliaplots.org/\n",
    "- http://docs.juliaplots.org/latest/tutorial\n",
    "- https://github.com/sswatson/cheatsheets/blob/master/plotsjl-cheatsheet.pdf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Taking derivatives without headaches\n",
    "\n",
    "Recall that for our dynamics with a given potential, say\n",
    "$$ V\\left( \\begin{array}{c} x_1 \\\\ x_2 \\end{array}\\right) = x_1^2 + x_2^2 $$\n",
    "we needed the acceleration\n",
    "$$ \\vec{A}_V = - \\nabla V. $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the harmonic oscillator this is kind of ok, but for more complicated potentials $V$, taking the derivative is error prone and most importantly quite *boring*. The solution is the great package `Zygote`, which is capable of adjoint-mode automatic differentiation. Without going into details, this means that the Julia compiler takes the derivatives for us and the result is (usually almost) as fast as if we did it ourselves. Let's try this in 1D on `Vho` as defined above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Zygote\n",
    "\n",
    "r = collect(-5:0.1:5)\n",
    "p = plot(r, Vho.(r), label=\"Vho a=0.5\")\n",
    "xaxis!(p, \"relative radial distance\")\n",
    "plot!(p, r, Vho'.(r), label=\"∂Vho\")   # Notice the dash"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With ease we add the second derivative to the plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot!(p, r, Vho''.(r), label=\"∂∂Vho\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is extremely helpful for more complicated potentials, for example the **Morse potential**, which is a common model for a chemical bond:\n",
    "$$ V_\\text{Morse} = D (1 - exp(-\\alpha (x-x_0)))^2 - D$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "Vmorse(r; re=0.7, α=1.3, D=10) = D*(1 - exp(-α * (r - re)))^2 - D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r = collect(0:0.1:7)\n",
    "p = plot(r, Vmorse.(r), label=\"Vmorse\", ylim=(-25, 25))\n",
    "xaxis!(p, \"Radial distance\")\n",
    "plot!(p, r, Vmorse'.(r), label=\"Vmorse'\")\n",
    "plot!(p, r, Vmorse''.(r), label=\"Vmorse''\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By the way: This works for higher dimensions and more complicated expressions, too, we will use this in a sec."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inspecting Euler dynamics in 1D\n",
    "\n",
    "Now we return to our forward-euler implementation, which works for both 1D and higher dimensions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "euler(A, Δt, xn, vn) = xn + vn * Δt, vn + A(xn) * Δt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "where $A_V$ (or in the code $A$) was the negative gradient of a potential $V$.\n",
    "With plotting at hand, let us actually see the dynamics. We define a plot function to plot the potential and animate the particle over time:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Printf\n",
    "\n",
    "# The arguments after the ; are so-called keyword arguments, they can be omitted\n",
    "# and then the default value after the = is used\n",
    "function plot_dynamics_line(V, Δt, n_steps; x0=0, v0=randn(), integrator=euler,\n",
    "                            xlim=(-5, 5), ylim=(0, 10))\n",
    "    t, x, v = 0, x0, v0  # Initialise state\n",
    "    \n",
    "    # Compute potential values (for plotting)\n",
    "    xrange = xlim[1]:0.1:xlim[2]\n",
    "    Vrange = V.(xrange)\n",
    "    \n",
    "    # @gif is a macro to \"record\" an animation of the dynamics,\n",
    "    # each loop iteration gives a frame\n",
    "    @gif for i in 1:n_steps\n",
    "        # Propagate dynamics (notice the derivative)\n",
    "        x, v = integrator(x -> -V'(x), Δt, x, v)\n",
    "        t += Δt\n",
    "\n",
    "        # Plot potential\n",
    "        p = plot(xrange, Vrange, label=\"Potential\", xlim=xlim, ylim=ylim)\n",
    "        \n",
    "        # Plot the particle and its velocity (as an arrow)\n",
    "        timestr = @sprintf \"%.2f\" t   # Format time as a nice string\n",
    "        scatter!(p, [x], [V(x)], label=\"Particle at t=$timestr\")\n",
    "        plot!(p, [(x, V(x)), (x + 0.5v, V(x))], arrow=1.0,\n",
    "              label=\"particle velocity / 2\")\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now let's actually see it ....\n",
    "plot_dynamics_line(Vho, 0.1, 200)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Wow ... that's strange. \n",
    "Our plot points at the well-known problem that energy is not conserved for a simple forward Euler scheme.\n",
    "We can also diagnose this using a phase-space plot, where the time evolution of particle position $x$ and particle velocity $v$ is shown. This should be a closed curve if energy is conserved, so let's see ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function plot_phase_space(A, Δt, n_steps; x0=0, v0=randn(), integrator=euler,\n",
    "                          xlim=(-7, 7), ylim=(-7, 7))\n",
    "    x, v = x0, v0\n",
    "    p = plot([x], [v], xlim=xlim, ylim=ylim, label=\"\")  # Plot first position\n",
    "    xaxis!(p, \"position\")\n",
    "    yaxis!(p, \"velocity\")\n",
    "    @gif for N in 1:n_steps\n",
    "        x, v = integrator(A, Δt, x, v)\n",
    "        push!(p, x, v)  # Add new positions to the plot ...\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_phase_space(x -> -Vho'(x), 0.1, 200)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise:**\n",
    "\n",
    "A much more stable integrator than the `euler` we used so far is the verlocity Verlet:\n",
    "\n",
    "$$\\left\\{\\begin{array}{l}\n",
    "x^{(n+1)} = x^{(n)} + v^{(n)} \\Delta t + \\frac{A_V(x^{(n)})}{2} \\Delta t^2\\\\\n",
    "v^{(n+1)} = v^{(n)} + \\frac{A_V(x^{(n))} + A_V(x^{(n+1)})}{2} \\Delta t\\\\\n",
    "\\end{array}\\right. $$\n",
    "\n",
    "- Program this algorithm, taking care that it supports multi-dimensional positions and velocities as well. In practice we would like to avoid recomputing $A_V(x)$ as much as possible, since this is usually the expensive step \n",
    "  of the dynamics. For our purposes there is no need to keep an eye on that.\n",
    "- How does the previous dynamics look like in this example. Does this algorithm conserve energy (phase-space plot)?\n",
    "- Also look at the Morse potential"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Some example plots and parameters:\n",
    "# plot_dynamics_line(Vho, 0.1, 200, integrator=verlet, ylim=(0, 2.5), xlim=(-1, 3))\n",
    "# plot_phase_space(x -> -Vho'(x), 0.1, 200, integrator=verlet, xlim=(-2, 2), ylim=(-2, 2))\n",
    "# plot_dynamics_line(Vmorse, 0.03, 200, integrator=verlet, xlim=(-1, 4), ylim=(-10, 5), x0=2.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multiple particles"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we want to simulate multiple identical partices in 2D. For a system of $N$ particles in 2D, we define the particle positions as the matrix\n",
    "$$ \\textbf{x} = \\left(\n",
    "    \\begin{array}{cccc}\n",
    "        x_{1x} & x_{2x} & \\cdots & x_{Nx} \\\\\n",
    "        x_{1y} & x_{2y} & \\cdots & x_{Ny}\n",
    "    \\end{array}\n",
    "   \\right) = \\left( \\vec{x}_1 \\vec{x}_2 \\cdots \\vec{x}_N \\right). $$\n",
    "with the individual particle vectors as columns.\n",
    "We assume our particles interact via the idential pair potential $V_\\text{pair}(r)$\n",
    "depending only on particle distance $r$. The total potential is therefore\n",
    "$$ V_\\text{tot}(\\textbf{x}) = \\sum_{i=1}^N \\sum_{j=i+1}^N V_\\text{pair}(\\| \\vec{x}_i - \\vec{x}_j \\|). $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Exercise\n",
    "Program the total potential function for a matrix $\\textbf{x}$. A useful function is `norm` from the `LinearAlgebra` package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using LinearAlgebra\n",
    "# You're code here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we want to plot the dynamics in a plane. In the following function the acceleration is computed using automatically generated derivatives."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function plot_dynamics_plane(Vpair, Δt, n_steps; x0=randn(2, 2), v0=zero(x0), integrator=verlet,\n",
    "                             xlim=(-3, 3), ylim=(-3, 3))\n",
    "    # Total acceleration via automatic differentiation\n",
    "    V(x) = Vtot(Vpair, x)\n",
    "    Atot(x) = -V'(x)\n",
    "    \n",
    "    t, x, v = 0, x0, v0  # Initialise state\n",
    "    @gif for i in 1:n_steps\n",
    "        # Propagate dynamics\n",
    "        x, v = integrator(Atot, Δt, x, v)\n",
    "        t += Δt\n",
    "        timestr = @sprintf \"%.2f\" t   # Format time\n",
    "        \n",
    "        # Plot the particles and their velocities\n",
    "        p = scatter(x[1, :], x[2, :], xlim=xlim, ylim=ylim, label=\"Particles at t=$timestr\")\n",
    "        label = \"Velocity / 4\"\n",
    "        for (xi, vi) in zip(eachcol(x), eachcol(v))\n",
    "            plot!(p, [Tuple(xi), Tuple(xi + 0.25vi)], arrow=1.0, colour=:red, label=label)\n",
    "            label = \"\"\n",
    "        end\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Δt = 0.07\n",
    "n_steps = 300\n",
    "x0 = [[0.; 0.] [1.; 0.] [-1.; 0.0]]\n",
    "plot_dynamics_plane(Vmorse, Δt, n_steps; x0=x0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A few nice examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Δt = 0.07\n",
    "n_steps = 300\n",
    "x0 = [[0.; 0.] [1.; 0.] [-1.; 0.15]]\n",
    "plot_dynamics_plane(Vmorse, Δt, n_steps; x0=x0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Δt = 0.05\n",
    "n_steps = 300\n",
    "x0 = [[0.; 1.] [1.; 0.] [-1.; 0] [0; -1.2]]\n",
    "plot_dynamics_plane(Vmorse, Δt, n_steps; x0=x0, ylim=(-10, 10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Δt = 0.05\n",
    "n_steps = 300\n",
    "x0 = 4randn(2, 10)\n",
    "plot_dynamics_plane(Vmorse, Δt, n_steps; x0=x0, xlim=(-10, 10), ylim=(-10, 10))"
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": "0ee45cb602ba4de49f9d8b52e7e3ca3d",
   "lastKernelId": "15fb271e-7c73-4f9c-be84-edddd39fe586"
  },
  "kernelspec": {
   "display_name": "Julia 1.3.0",
   "language": "julia",
   "name": "julia-1.3"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.3.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}