notebooks/customizing_potentials_cookbook.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Custom Potentials.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/google/jax-md/blob/master/notebooks/customizing_potentials_cookbook.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vTSgycb_WCAx"
      },
      "source": [
        "# Customizing Potentials in JAX MD\n",
        "\n",
        "This cookbook was contributed by Carl Goodrich."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y9jcXj44Bvcj",
        "cellView": "form"
      },
      "source": [
        "#@title Imports & Utils\n",
        "!pip install -q git+https://www.github.com/google/jax-md\n",
        "\n",
        "import numpy as onp\n",
        "\n",
        "import jax.numpy as np\n",
        "from jax import random\n",
        "from jax import jit, grad, vmap, value_and_grad\n",
        "from jax import lax\n",
        "from jax import ops\n",
        "\n",
        "from jax.config import config\n",
        "config.update(\"jax_enable_x64\", True)\n",
        "\n",
        "from jax_md import space, smap, energy, minimize, quantity, simulate, partition\n",
        "\n",
        "from functools import partial\n",
        "import time\n",
        "\n",
        "f32 = np.float32\n",
        "f64 = np.float64\n",
        "\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "plt.rcParams.update({'font.size': 16})\n",
        "#import seaborn as sns \n",
        "#sns.set_style(style='white')\n",
        "\n",
        "def format_plot(x, y):  \n",
        "  plt.grid(True)\n",
        "  plt.xlabel(x, fontsize=20)\n",
        "  plt.ylabel(y, fontsize=20)\n",
        "  \n",
        "def finalize_plot(shape=(1, 0.7)):\n",
        "  plt.gcf().set_size_inches(\n",
        "    shape[0] * 1.5 * plt.gcf().get_size_inches()[1], \n",
        "    shape[1] * 1.5 * plt.gcf().get_size_inches()[1])\n",
        "  plt.tight_layout()\n",
        "\n",
        "def calculate_bond_data(displacement_or_metric, R, dr_cutoff, species=None):\n",
        "  if( not(species is None)):\n",
        "    assert(False)\n",
        "    \n",
        "  metric = space.map_product(space.canonicalize_displacement_or_metric(displacement))\n",
        "  dr = metric(R,R)\n",
        "\n",
        "  dr_include = np.triu(np.where(dr<dr_cutoff, 1, 0)) - np.eye(R.shape[0],dtype=np.int32)\n",
        "  index_list=np.dstack(np.meshgrid(np.arange(N), np.arange(N), indexing='ij'))\n",
        "\n",
        "  i_s = np.where(dr_include==1, index_list[:,:,0], -1).flatten()\n",
        "  j_s = np.where(dr_include==1, index_list[:,:,1], -1).flatten()\n",
        "  ij_s = np.transpose(np.array([i_s,j_s]))\n",
        "\n",
        "  bonds = ij_s[(ij_s!=np.array([-1,-1]))[:,1]]\n",
        "  lengths = dr.flatten()[(ij_s!=np.array([-1,-1]))[:,1]]\n",
        "\n",
        "  return bonds, lengths\n",
        "\n",
        "def plot_system(R,box_size,species=None,ms=20):\n",
        "  R_plt = onp.array(R)\n",
        "\n",
        "  if(species is None):\n",
        "    plt.plot(R_plt[:, 0], R_plt[:, 1], 'o', markersize=ms)\n",
        "  else:\n",
        "    for ii in range(np.amax(species)+1):\n",
        "      Rtemp = R_plt[species==ii]\n",
        "      plt.plot(Rtemp[:, 0], Rtemp[:, 1], 'o', markersize=ms)\n",
        "\n",
        "  plt.xlim([0, box_size])\n",
        "  plt.ylim([0, box_size])\n",
        "  plt.xticks([], [])\n",
        "  plt.yticks([], [])\n",
        "\n",
        "  finalize_plot((1,1))\n",
        "  \n",
        "key = random.PRNGKey(0)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lHhvM9_nZ9TT"
      },
      "source": [
        "##Prerequisites\n",
        "\n",
        "This cookbook assumes a working knowledge of Python and Numpy. The concept of broadcasting is particularly important both in this cookbook and in JAX MD. \n",
        "\n",
        "We also assume a basic knowlege of [JAX](https://github.com/google/jax/), which JAX MD is built on top of. Here we briefly review a few JAX basics that are important for us:\n",
        "\n",
        "\n",
        "*   ```jax.vmap``` allows for automatic vectorization of a function. What this means is that if you have a function that takes an input ```x``` and returns an output ```y```, i.e. ```y = f(x)```, then ```vmap``` will transform this function to act on an array of ```x```'s and return an array of ```y```'s, i.e. ```Y = vmap(f)(X)```, where ```X=np.array([x1,x2,...,xn])``` and ```Y=np.array([y1,y2,...,yn])```. \n",
        "\n",
        "*   ```jax.grad``` employs automatic differentiation to transform a function into a new function that calculates its gradient, for example: ```dydx = grad(f)(x)```. \n",
        "\n",
        "*   ```jax.lax.scan``` allows for efficient for-loops that can be compiled and differentiated over. See [here](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html#jax.lax.scan) for more details.\n",
        "\n",
        "*   [Random numbers are different in JAX.](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#%F0%9F%94%AA-Random-Numbers) The details aren't necessary for this cookbook, but if things look a bit different, this is why.\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uUb4pEymgfpW"
      },
      "source": [
        "##The basics of user-defined potentials"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AGQNsI7Q00_3"
      },
      "source": [
        "###Create a user defined potential function to use throughout this cookbook\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RLo12Gsf1G9k"
      },
      "source": [
        "Here we create a custom potential that has a short-ranged, non-diverging repulsive interaction and a medium-ranged Morse-like attractive interaction. It takes the following form:\n",
        "\\begin{equation}\n",
        "V(r) =\n",
        "\\begin{cases}\n",
        "    \\frac{1}{2} k (r-r_0)^2 - D_0,&  r < r_0\\\\\n",
        "    D_0\\left( e^{-2\\alpha (r-r_0)} -2 e^{-\\alpha(r-r_0)}\\right),              & r \\geq r_0\n",
        "\\end{cases}\n",
        "\\end{equation}\n",
        "and has 4 parameters: $D_0$, $\\alpha$, $r_0$, and $k$.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hZysrOXi0xdp"
      },
      "source": [
        "def harmonic_morse(dr, D0=5.0, alpha=5.0, r0=1.0, k=50.0, **kwargs):\n",
        "  U = np.where(dr < r0, \n",
        "               0.5 * k * (dr - r0)**2 - D0,\n",
        "               D0 * (np.exp(-2. * alpha * (dr - r0)) - 2. * np.exp(-alpha * (dr - r0)))\n",
        "               )\n",
        "  return np.array(U, dtype=dr.dtype)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3yVG0w4oBV1q"
      },
      "source": [
        "plot $V(r)$."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TgIbQa_D1kIM"
      },
      "source": [
        "drs = np.arange(0,3,0.01)\n",
        "U = harmonic_morse(drs)\n",
        "plt.plot(drs,U)\n",
        "format_plot(r'$r$', r'$V(r)$')\n",
        "finalize_plot()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C1lMC0LeCSe9"
      },
      "source": [
        "###Calculate the energy of a system of interacting particles"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kyuIoCz3maxs"
      },
      "source": [
        "We now want to calculate the energy of a system of $N$ spheres in $d$ dimensions, where each particle interacts with every other particle via our user-defined function $V(r)$. The total energy is\n",
        "\\begin{equation}\n",
        "E_\\text{total} = \\sum_{i<j}V(r_{ij}),\n",
        "\\end{equation}\n",
        "where $r_{ij}$ is the distance between particles $i$ and $j$. \n",
        "\n",
        "Our first task is to set up the system by specifying the $N$, $d$, and the size of the simulation box. We then use JAX's internal random number generator to pick positions for each particle. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q77Cepyy4M14"
      },
      "source": [
        "N = 50\n",
        "dimension = 2\n",
        "box_size = 6.8\n",
        "\n",
        "key, split = random.split(key)\n",
        "R = random.uniform(split, (N,dimension), minval=0.0, maxval=box_size, dtype=f64) \n",
        "\n",
        "plot_system(R,box_size)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lB5ISBfB4txR"
      },
      "source": [
        "At this point, we could manually loop over all particle pairs and calculate the energy, keeping track of boundary conditions, etc. Fortunately, JAX MD has machinery to automate this. \n",
        "\n",
        "First, we must define two functions, ```displacement``` and ```shift```, which contain all the information of the simulation box, boundary conditions, and underlying metric. ```displacement``` is used to calculate the vector displacement between particles, and ```shift``` is used to move particles. For most cases, it is recommended to use JAX MD's built in functions, which can be called using:\n",
        "*   ``` displacement, shift = space.free()```\n",
        "*   ``` displacement, shift = space.periodic(box_size)```\n",
        "*   ``` displacement, shift = space.periodic_general(T)```\n",
        "\n",
        "For demonstration purposes, we will define these manually for a square periodic box, though without proper error handling, etc. The following should have the same functionality as ```displacement, shift = space.periodic(box_size)```."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z93nLpSk4Krm"
      },
      "source": [
        "def setup_periodic_box(box_size):\n",
        "  def displacement_fn(Ra, Rb, **unused_kwargs):\n",
        "    dR = Ra - Rb\n",
        "    return np.mod(dR + box_size * f32(0.5), box_size) - f32(0.5) * box_size\n",
        "\n",
        "  def shift_fn(R, dR, **unused_kwargs):\n",
        "    return np.mod(R + dR, box_size)\n",
        "\n",
        "  return displacement_fn, shift_fn\n",
        "  \n",
        "displacement, shift = setup_periodic_box(box_size)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q_vlEs5DEq-b"
      },
      "source": [
        "We now set up a function to calculate the total energy of the system. The JAX MD function ```smap.pair``` takes a given potential and promotes it to act on all particle pairs in a system. ```smap.pair``` does not actually return an energy, rather it returns a function that can be used to calculate the energy. \n",
        "\n",
        "For convenience and readability, we wrap ```smap.pair``` in a new function called ```harmonic_morse_pair```. For now, ignore the species keyword, we will return to this later."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y2KDd1x6AJPw"
      },
      "source": [
        "def harmonic_morse_pair(\n",
        "    displacement_or_metric, species=None, D0=5.0, alpha=10.0, r0=1.0, k=50.0): \n",
        "  D0 = np.array(D0, dtype=f32)\n",
        "  alpha = np.array(alpha, dtype=f32)\n",
        "  r0 = np.array(r0, dtype=f32)\n",
        "  k = np.array(k, dtype=f32)\n",
        "  return smap.pair(\n",
        "      harmonic_morse,\n",
        "      space.canonicalize_displacement_or_metric(displacement_or_metric),\n",
        "      species=species,\n",
        "      D0=D0,\n",
        "      alpha=alpha,\n",
        "      r0=r0,\n",
        "      k=k)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QYoPcRjXEfx-"
      },
      "source": [
        "Our helper function can be used to construct a function to compute the energy of the entire system as follows.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nT_ftIE8CQuA"
      },
      "source": [
        "# Create a function to calculate the total energy with specified parameters\n",
        "energy_fn = harmonic_morse_pair(displacement,D0=5.0,alpha=10.0,r0=1.0,k=500.0)\n",
        "\n",
        "# Use this to calculate the total energy\n",
        "print(energy_fn(R))\n",
        "\n",
        "# Use grad to calculate the net force\n",
        "force = -grad(energy_fn)(R)\n",
        "print(force[:5])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-gYTfzbXijOq"
      },
      "source": [
        "We are now in a position to use our energy function to manipulate the system. As an example, we perform energy minimization using JAX MD's implementation of the FIRE algorithm. \n",
        "\n",
        "We start by defining a function that takes an energy function, a set of initial positions, and a shift function and runs a specified number of steps of the minimization algorithm. The function returns the final set of positions and the maximum absolute value component of the force. We will use this function throughout this cookbook. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XX-Jdq_7maw_"
      },
      "source": [
        "def run_minimization(energy_fn, R_init, shift, num_steps=5000):\n",
        "  dt_start = 0.001\n",
        "  dt_max   = 0.004\n",
        "  init,apply=minimize.fire_descent(jit(energy_fn),shift,dt_start=dt_start,dt_max=dt_max)\n",
        "  apply = jit(apply)\n",
        "\n",
        "  @jit\n",
        "  def scan_fn(state, i):\n",
        "    return apply(state), 0.\n",
        "\n",
        "  state = init(R_init)\n",
        "  state, _ = lax.scan(scan_fn,state,np.arange(num_steps))\n",
        "\n",
        "  return state.position, np.amax(np.abs(-grad(energy_fn)(state.position)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "48YgfOdVnf5m"
      },
      "source": [
        "Now run the minimization with our custom energy function."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nY_hzbgIaZNI"
      },
      "source": [
        "Rfinal, max_force_component = run_minimization(energy_fn, R, shift)\n",
        "print('largest component of force after minimization = {}'.format(max_force_component))\n",
        "plot_system( Rfinal, box_size )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NKL3FcddI7GM"
      },
      "source": [
        "###Create a truncated potential"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tWsu6nvTJCsz"
      },
      "source": [
        "It is often desirable to have a potential that is strictly zero beyond a well-defined cutoff distance. In addition, MD simulations require the energy and force (i.e. first derivative) to be continuous. To easily modify an existing potential $V(r)$ to have this property, JAX MD follows the approach [taken by HOOMD Blue](https://hoomd-blue.readthedocs.io/en/stable/module-md-pair.html#hoomd.md.pair.pair). \n",
        "\n",
        "Consider the function \n",
        "\\begin{equation}\n",
        "S(r) =\n",
        "\\begin{cases}\n",
        "    1,& r<r_\\mathrm{on} \\\\\n",
        "    \\frac{(r_\\mathrm{cut}^2-r^2)^2 (r_\\mathrm{cut}^2 + 2r^2 - 3 r_\\mathrm{on}^2)}{(r_\\mathrm{cut}^2-r_\\mathrm{on}^2)^3},&  r_\\mathrm{on} \\leq r < r_\\mathrm{cut}\\\\\n",
        "    0,& r \\geq r_\\mathrm{cut}\n",
        "\\end{cases}\n",
        "\\end{equation}\n",
        "\n",
        "Here we plot both $S(r)$ and $\\frac{dS(r)}{dr}$, both of which are smooth and strictly zero above $r_\\mathrm{cut}$.\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bPs51LI_t684"
      },
      "source": [
        "dr = np.arange(0,3,0.01)\n",
        "S = energy.multiplicative_isotropic_cutoff(lambda dr: 1, r_onset=1.5, r_cutoff=2.0)(dr)\n",
        "ngradS = vmap(grad(energy.multiplicative_isotropic_cutoff(lambda dr: 1, r_onset=1.5, r_cutoff=2.0)))(dr)\n",
        "plt.plot(dr,S,label=r'$S(r)$')\n",
        "plt.plot(dr,ngradS,label=r'$\\frac{dS(r)}{dr}$')\n",
        "plt.legend()\n",
        "format_plot(r'$r$','')\n",
        "finalize_plot()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5Zmyipr-xaDj"
      },
      "source": [
        "We then use $S(r)$ to create a new function \n",
        "\\begin{equation}\\tilde V(r) = V(r) S(r),\n",
        "\\end{equation} \n",
        "which is exactly $V(r)$ below $r_\\mathrm{on}$, strictly zero above $r_\\mathrm{cut}$ and is continuous in its first derivative.\n",
        "\n",
        "This is implemented in JAX MD through ```energy.multiplicative_isotropic_cutoff```, which takes in a potential function $V(r)$ (e.g. our ```harmonic_morse``` function) and returns a new function $\\tilde V(r)$."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TCv8qIyeGA-k"
      },
      "source": [
        "harmonic_morse_cutoff = energy.multiplicative_isotropic_cutoff(\n",
        "    harmonic_morse, r_onset=1.5, r_cutoff=2.0)\n",
        "\n",
        "dr = np.arange(0,3,0.01)\n",
        "V = harmonic_morse(dr)\n",
        "V_cutoff = harmonic_morse_cutoff(dr)\n",
        "F = -vmap(grad(harmonic_morse))(dr)\n",
        "F_cutoff = -vmap(grad(harmonic_morse_cutoff))(dr)\n",
        "plt.plot(dr,V, label=r'$V(r)$')\n",
        "plt.plot(dr,V_cutoff, label=r'$\\tilde V(r)$')\n",
        "plt.plot(dr,F, label=r'$-\\frac{d}{dr} V(r)$')\n",
        "plt.plot(dr,F_cutoff, label=r'$-\\frac{d}{dr} \\tilde V(r)$')\n",
        "plt.legend()\n",
        "format_plot('$r$', '')\n",
        "plt.ylim(-13,5)\n",
        "finalize_plot()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AyHHLfQ1O6w0"
      },
      "source": [
        "As before, we can use ```smap.pair``` to promote this to act on an entire system."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "abP6O5atO6F-"
      },
      "source": [
        "def harmonic_morse_cutoff_pair(\n",
        "    displacement_or_metric, D0=5.0, alpha=5.0, r0=1.0, k=50.0,\n",
        "    r_onset=1.5, r_cutoff=2.0): \n",
        "  D0 = np.array(D0, dtype=f32)\n",
        "  alpha = np.array(alpha, dtype=f32)\n",
        "  r0 = np.array(r0, dtype=f32)\n",
        "  k = np.array(k, dtype=f32)\n",
        "  return smap.pair(\n",
        "      energy.multiplicative_isotropic_cutoff(\n",
        "          harmonic_morse, r_onset=r_onset, r_cutoff=r_cutoff),\n",
        "      space.canonicalize_displacement_or_metric(displacement_or_metric),\n",
        "      D0=D0,\n",
        "      alpha=alpha,\n",
        "      r0=r0,\n",
        "      k=k)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UgvzLddhBIhe"
      },
      "source": [
        "This is implemented as before"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CW_fOcqJKVlJ"
      },
      "source": [
        "# Create a function to calculate the total energy\n",
        "energy_fn = harmonic_morse_cutoff_pair(displacement, D0=5.0, alpha=10.0, r0=1.0, \n",
        "                                       k=500.0, r_onset=1.5, r_cutoff=2.0)\n",
        "\n",
        "# Use this to calculate the total energy\n",
        "print(energy_fn(R))\n",
        "\n",
        "# Use grad to calculate the net force\n",
        "force = -grad(energy_fn)(R)\n",
        "print(force[:5])\n",
        "\n",
        "# Minimize the energy using the FIRE algorithm\n",
        "Rfinal, max_force_component = run_minimization(energy_fn, R, shift)\n",
        "print('largest component of force after minimization = {}'.format(max_force_component))\n",
        "plot_system( Rfinal, box_size )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PAjtb9XZfLGp"
      },
      "source": [
        "##Specifying parameters"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HxXmuqQoXACk"
      },
      "source": [
        "###Dynamic parameters"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "korJBavMISgx"
      },
      "source": [
        "In the above examples, the strategy is to create a function ```energy_fn``` that takes a set of positions and calculates the energy of the system with all the parameters (e.g. ```D0```, ```alpha```, etc.) baked in. However, JAX MD allows you to override these baked-in values dynamically, i.e. when ```energy_fn``` is called. \n",
        "\n",
        "For example, we can print out the minimized energy and force of the above system with the truncated potential:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "osQ-l1cMQ7Xf"
      },
      "source": [
        "print(energy_fn(Rfinal))\n",
        "print(-grad(energy_fn)(Rfinal)[:5])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5JduIzmbRuXP"
      },
      "source": [
        "This uses the baked-in values of the 4 parameters: ```D0=5.0,alpha=10.0,r0=1.0,k=500.0```. If, for example, we want to dynamically turn off the attractive part of the potential, we simply pass ```D0=0``` to ```energy_fn```:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JlNgP8D0Q49c"
      },
      "source": [
        "print(energy_fn(Rfinal, D0=0))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SC0GqWcoCXoE"
      },
      "source": [
        "Since changing the potential moves the minimum, the force will not be zero:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7Fb3FDnuCW9M"
      },
      "source": [
        "print(-grad(energy_fn)(Rfinal, D0=0)[:5])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iDl-nqTV8VNR"
      },
      "source": [
        "This ability to dynamically pass parameters is very powerful. For example, if you want to shrink particles each step during a simulation, you can simply specify a different ```r0``` each step. \n",
        "\n",
        "This is demonstrated below, where we run a Brownian dynamics simulation at zero temperature with continuously decreasing ```r0```. The details of ```simulate.brownian``` are beyond the scope of this cookbook, but the idea is that we pass a new value of ```r0``` to the function ```apply``` each time it is called. The function ```apply``` takes a step of the simulation, and internally it passes any extra parameters like ```r0``` to ```energy_fn```."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2gxiNvCp7JHD"
      },
      "source": [
        "def run_brownian(energy_fn, R_init, shift, key, num_steps):\n",
        "  init, apply = simulate.brownian(energy_fn, shift, dt=0.00001, kT=0.0, gamma=0.1)\n",
        "  apply = jit(apply)\n",
        "\n",
        "  # Define how r0 changes for each step\n",
        "  r0_initial = 1.0\n",
        "  r0_final = .5\n",
        "  def get_r0(t):\n",
        "    return r0_final + (r0_initial-r0_final)*(num_steps-t)/num_steps\n",
        "\n",
        "  @jit\n",
        "  def scan_fn(state, t):\n",
        "    # Dynamically pass r0 to apply, which passes it on to energy_fn\n",
        "    return apply(state, r0=get_r0(t)), 0\n",
        "\n",
        "  key, split = random.split(key)\n",
        "  state = init(split, R_init)\n",
        "\n",
        "  state, _ = lax.scan(scan_fn,state,np.arange(num_steps))\n",
        "  return state.position, np.amax(np.abs(-grad(energy_fn)(state.position)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zc9_8W0YMkvt"
      },
      "source": [
        "If we use the previous result as the starting point for the Brownian Dynamics simulation, we find exactly what we would expect, the system contracts into a finite cluster, held together by the attractive part of the potential."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-nC5cDZ-7OJD"
      },
      "source": [
        "key, split = random.split(key)\n",
        "Rfinal2, max_force_component = run_brownian(energy_fn, Rfinal, shift, split, num_steps=6000)\n",
        "plot_system( Rfinal2, box_size )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fZ9aBlNTgoa7"
      },
      "source": [
        "###Particle-specific parameters"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AHDic52sg6iR"
      },
      "source": [
        "Our example potential has 4 parameters: ```D0```, ```alpha```, ```r0```, and ```k```. The usual way to pass these parameters is as a scalar (e.g. ```D0=5.0```), in which case that parameter is fixed for every particle pair. However, Python broadcasting allows for these parameters to be specified separately for every different particle pair by passing an $(N,N)$ array rather than a scalar. \n",
        "\n",
        "As an example, let's do this for the parameter ```r0```, which is an effective way of generating a system with continuous polydispersity in particle size. Note that the polydispersity disrupts the crystalline order after minimization."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iio81Poz1zfj"
      },
      "source": [
        "# Draw the radii from a uniform distribution\n",
        "key, split = random.split(key)\n",
        "radii = random.uniform(split, (N,), minval=1.0, maxval=2.0, dtype=f64)\n",
        "\n",
        "# Rescale to match the initial volume fraction\n",
        "radii = np.array([radii * np.sqrt(N/(4.*np.dot(radii,radii)))])\n",
        "\n",
        "# Turn this into a matrix of sums\n",
        "r0_matrix = radii+radii.transpose()\n",
        "\n",
        "# Create the energy function using r0_matrix\n",
        "energy_fn = harmonic_morse_pair(displacement, D0=5.0, alpha=10.0, r0=r0_matrix, \n",
        "                                       k=500.0)\n",
        "\n",
        "# Minimize the energy using the FIRE algorithm\n",
        "Rfinal, max_force_component = run_minimization(energy_fn, R, shift)\n",
        "print('largest component of force after minimization = {}'.format(max_force_component))\n",
        "plot_system( Rfinal, box_size )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7Ftp_f710xLl"
      },
      "source": [
        "In addition to standard Python broadcasting, JAX MD allows for the special case of additive parameters. If a parameter is passed as a (N,) array ```p_vector```, JAX MD will convert this into a (N,N) array ```p_matrix``` where ```p_matrix[i,j] = 0.5 (p_vector[i] + p_vector[j])```. This is a JAX MD specific ability and not a feature of Python broadcasting.\n",
        "\n",
        "As it turns out, our above polydisperse example falls into this category. Therefore, we could achieve the same result by passing ```r0=2.0*radii```."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4HnQK3fK1Wyp"
      },
      "source": [
        "# Create the energy function the radii array\n",
        "energy_fn = harmonic_morse_pair(displacement, D0=5.0, alpha=10.0, r0=2.*radii, \n",
        "                                       k=500.0)\n",
        "\n",
        "# Minimize the energy using the FIRE algorithm\n",
        "Rfinal, max_force_component = run_minimization(energy_fn, R, shift)\n",
        "print('largest component of force after minimization = {}'.format(max_force_component))\n",
        "plot_system( Rfinal, box_size )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cN-rwzbARInR"
      },
      "source": [
        "### Species"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fed3OOzMJmzu"
      },
      "source": [
        "It is often important to specify parameters differently for different particle pairs, but doing so with full ($N$,$N$) matrices is both inefficient and obnoxious. JAX MD allows users to create species, i.e. $N_s$ groups of particles that are identical to each other, so that parameters can be passed as much smaller ($N_s$,$N_s$) matrices.\n",
        "\n",
        "First, create an array that specifies which particles belong in which species. We will divide our system into two species."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zODRlLMhNQ2t"
      },
      "source": [
        "N_0 = N // 2  # Half the particles in species 0\n",
        "N_1 = N - N_0 # The rest in species 1\n",
        "species = np.array([0] * N_0 + [1] * N_1, dtype=np.int32)\n",
        "print(species)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aVFp3U1LBVRF"
      },
      "source": [
        "Next, create the $(2,2)$ matrix of ```r0```'s, which are set so that the overall volume fraction matches our monodisperse case. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FPylQCNNA1P5"
      },
      "source": [
        "rsmall=0.41099747 # Match the total volume fraction\n",
        "rlarge=1.4*rsmall\n",
        "r0_species_matrix = np.array([[2*rsmall, rsmall+rlarge],\n",
        "                              [rsmall+rlarge, 2*rlarge]])\n",
        "print(r0_species_matrix)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Mec9jBIqRZy3"
      },
      "source": [
        "energy_fn = harmonic_morse_pair(displacement, species=species, D0=5.0, \n",
        "                                alpha=10.0, r0=r0_species_matrix, k=500.0)\n",
        "\n",
        "Rfinal, max_force_component = run_minimization(energy_fn, R, shift)\n",
        "print('largest component of force after minimization = {}'.format(max_force_component))\n",
        "\n",
        "plot_system( Rfinal, box_size, species=species )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vvjaCB0cN8rA"
      },
      "source": [
        "###Dynamic Species"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jyYqW7dCDi0y"
      },
      "source": [
        "Just like standard parameters, the species list can be passed dynamically as well. However, unlike standard parameters, you have to tell `smap.pair` that the species will be specified dynamically. To do this, set  `species=2` be the total number of types of particles when creating your energy function.\n",
        "\n",
        "The following sets up an energy function where the attractive part of the interaction only exists between members of the first species, but where the species will be defined dynamically."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Q30CEsxaQGc7"
      },
      "source": [
        "D0_species_matrix = np.array([[ 5.0, 0.0],\n",
        "                              [0.0,  0.0]])\n",
        "\n",
        "energy_fn = harmonic_morse_pair(displacement, \n",
        "                                species=2, \n",
        "                                D0=D0_species_matrix, \n",
        "                                alpha=10.0,\n",
        "                                r0=0.5, \n",
        "                                k=500.0)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a2_Er0ttWcmu"
      },
      "source": [
        "Now we set up a finite temperature Brownian Dynamics simulation where, at every step, particles on the left half of the simulation box are assigned to species 0, while particles on the right half are assigned to species 1."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HlT0lZ0IO9rx"
      },
      "source": [
        "def run_brownian(energy_fn, R_init, shift, key, num_steps):\n",
        "  init, apply = simulate.brownian(energy_fn, shift, dt=0.00001, kT=1.0, gamma=0.1)\n",
        "  # apply = jit(apply)\n",
        "\n",
        "  # Define a function to recalculate the species each step\n",
        "  def get_species(R):\n",
        "    return np.where(R[:,0] < box_size / 2, 0, 1)\n",
        "\n",
        "  #@jit\n",
        "  def scan_fn(state, t):\n",
        "    # Recalculate the species list\n",
        "    species = get_species(state.position)\n",
        "    # Dynamically pass species to apply, which passes it on to energy_fn\n",
        "    return apply(state, species=species, species_count=2), 0\n",
        "\n",
        "  key, split = random.split(key)\n",
        "  state = init(split, R_init)\n",
        "\n",
        "  state, _ = lax.scan(scan_fn,state,np.arange(num_steps))\n",
        "  return state.position,np.amax(np.abs(-grad(energy_fn)(state.position,\n",
        "                                                        species=get_species(state.position), \n",
        "                                                        species_count=2)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L1Lso1xbXm10"
      },
      "source": [
        "When we run this, we see that particles on the left side form clusters while particles on the right side do not."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zshmmAiuROOP"
      },
      "source": [
        "key, split = random.split(key)\n",
        "Rfinal, max_force_component = run_brownian(energy_fn, R, shift, split, num_steps=10000)\n",
        "plot_system( Rfinal, box_size )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bvp_-L3cJYiM"
      },
      "source": [
        "##Efficeiently calculating neighbors"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5NWHAN8MD4Sh"
      },
      "source": [
        "The most computationally expensive part of most MD programs is calculating the force between all pairs of particles. Generically, this scales with $N^2$. However, for systems with isotropic pairwise interactions that are strictly zero beyond a cutoff, there are techniques to dramatically improve the efficiency. The two most common methods are cell list and neighbor lists.\n",
        "\n",
        "**Cell lists**\n",
        "\n",
        "The technique here is to divide space into small cells that are just larger than the largest interaction range in the system. Thus, if particle $i$ is in cell $c_i$ and particle $j$ is in cell $c_j$, $i$ and $j$ can only interact if $c_i$ and $c_j$ are neighboring cells. Rather than searching all $N^2$ combinations of particle pairs for non-zero interactions, you only have to search the particles in the neighboring cells. \n",
        "\n",
        "**Neighbor lists**\n",
        "\n",
        "Here, for each particle $i$, we make a list of *potential* neighbors: particles $j$ that are within some threshold distance $r_\\mathrm{threshold}$. If $r_\\mathrm{threshold} = r_\\mathrm{cutoff} + \\Delta r_\\mathrm{threshold}$ (where $r_\\mathrm{cutoff}$ is the largest interaction range in the system and $\\Delta r_\\mathrm{threshold}$ is an appropriately chosen buffer size), then all interacting particles will appear in this list as long as no particles moves by more than $\\Delta r_\\mathrm{threhsold}/2$. There is a tradeoff here: smaller $\\Delta r_\\mathrm{threhsold}$ means fewer particles to search over each MD step but the list must be recalculated more often, while larger $\\Delta r_\\mathrm{threhsold}$ means slower force calculates but less frequent neighbor list calculations. \n",
        "\n",
        "In practice, the most efficient technique is often to use cell lists to calculate neighbor lists. In JAX MD, this occurs under the hood, and so only calls to neighbor-list functionality are necessary."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c8sSH5ayU1N3"
      },
      "source": [
        "To implement neighbor lists, we need two functions: 1) a function to create and update the neighbor list, and 2) an energy function that uses a neighbor list rather than operating on all particle pairs. We create these functions with ```partition.neighbor_list``` and ```smap.pair_neighbor_list```, respectively. \n",
        "\n",
        "```partition.neighbor_list``` takes basic box information as well as the maximum interaction range ```r_cutoff``` and the buffer size ```dr_threshold```. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6QIkYw8wNZf4"
      },
      "source": [
        " def harmonic_morse_cutoff_neighbor_list(\n",
        "    displacement_or_metric,\n",
        "    box_size,\n",
        "    species=None,\n",
        "    D0=5.0, \n",
        "    alpha=5.0, \n",
        "    r0=1.0, \n",
        "    k=50.0,\n",
        "    r_onset=1.0,\n",
        "    r_cutoff=1.5, \n",
        "    dr_threshold=2.0,\n",
        "    **kwargs): \n",
        "\n",
        "  D0 = np.array(D0, dtype=np.float32)\n",
        "  alpha = np.array(alpha, dtype=np.float32)\n",
        "  r0 = np.array(r0, dtype=np.float32)\n",
        "  k = np.array(k, dtype=np.float32)\n",
        "  r_onset = np.array(r_onset, dtype=np.float32)\n",
        "  r_cutoff = np.array(r_cutoff, np.float32)\n",
        "  dr_threshold = np.float32(dr_threshold)\n",
        "\n",
        "  neighbor_fn = partition.neighbor_list(\n",
        "        displacement_or_metric, \n",
        "        box_size, \n",
        "        r_cutoff, \n",
        "        dr_threshold)\n",
        "\n",
        "  energy_fn = smap.pair_neighbor_list(\n",
        "    energy.multiplicative_isotropic_cutoff(harmonic_morse, r_onset, r_cutoff),\n",
        "    space.canonicalize_displacement_or_metric(displacement_or_metric),\n",
        "    species=species,\n",
        "    D0=D0,\n",
        "    alpha=alpha,\n",
        "    r0=r0,\n",
        "    k=k)\n",
        "\n",
        "  return neighbor_fn, energy_fn"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8AB0dK7hWNkt"
      },
      "source": [
        "To test this, we generate our new ```neighbor_fn``` and ```energy_fn```, as well as a comparison energy function using the default approach."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ND2eNJdFQwtr"
      },
      "source": [
        "r_onset  = 1.5\n",
        "r_cutoff = 2.0\n",
        "dr_threshold = 1.0\n",
        "\n",
        "neighbor_fn, energy_fn = harmonic_morse_cutoff_neighbor_list(\n",
        "    displacement, box_size, D0=5.0, alpha=10.0, r0=1.0, k=500.0,\n",
        "    r_onset=r_onset, r_cutoff=r_cutoff, dr_threshold=dr_threshold)\n",
        "\n",
        "energy_fn_comparison = harmonic_morse_cutoff_pair(\n",
        "    displacement, D0=5.0, alpha=10.0, r0=1.0, k=500.0,\n",
        "    r_onset=r_onset, r_cutoff=r_cutoff)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qufh0bh8WcwH"
      },
      "source": [
        "Next, we use ```neighbor_fn``` and the current set of positions to populate the neighbor list."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aFeZgIoVRGHj"
      },
      "source": [
        "nbrs = neighbor_fn(R)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b1UJwsAMWteb"
      },
      "source": [
        "To calculate the energy, we pass `nbrs` to `energy_fn`. The energy matches the comparison."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-d1gVB3ERFTg"
      },
      "source": [
        "print(energy_fn(R, neighbor=nbrs))\n",
        "print(energy_fn_comparison(R))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yJGBnaz2XGt-"
      },
      "source": [
        "Note that by default ```neighbor_fn``` uses a cell list internally to populate the neighbor list. This approach fails when the box size in any dimension is less than 3 times $r_\\mathrm{threhsold} = r_\\mathrm{cutoff} + \\Delta r_\\mathrm{threshold}$. In this case, ```neighbor_fn``` automatically turns off the use of cell lists, and instead searches over all particle pairs. This can also be done manually by passing ```disable_cell_list=True``` to ```partition.neighbor_list```. This can be useful for debugging or for small systems where the overhead of cell lists outweighs the benefit. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3PkSm7d4KTvb"
      },
      "source": [
        "###Updating neighbor lists"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5T_ViIa9Ksgn"
      },
      "source": [
        "The function ```neighbor_fn``` has two different usages, depending on how it is called. When used as above, i.e. ```nbrs = neighbor_fn(R)```, a new neighbor list is generated from scratch. Internally, JAX MD uses the given positions ```R``` to estimate a maximum capacity, i.e. the maximum number of neighbors any particle will have at any point during the use of the neighbor list. This estimate can be adjusted by passing a value of  ```capacity_multiplier``` to ```partition.neighbor_list```, which defaults to ```capacity_multiplier=1.25```.\n",
        "\n",
        "Since the maximum capacity is not known ahead of time, this construction of the neighbor list cannot be compiled. However, once a neighbor list is created in this way, repopulating the list with the same maximum capacity is a simpler operation that *can* be compiled. This is done by calling ```nbrs = neighbor_fn(R, nbrs)```. Internally, this checks if any particle has moved more than $\\Delta r_\\mathrm{threshold}/2$ and, if so, recomputes the neighbor list. If the new neighbor list exceeds the maximum capacity for any particle, the boolean variable ```nbrs.did_buffer_overflow``` is set to ```True```. \n",
        "\n",
        "These two uses together allow for safe and efficient neighbor list calculations. The example below demonstrates a typical simulation loop that uses neighbor lists. \n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9171zO3pQJnH"
      },
      "source": [
        "def run_brownian_neighbor_list(energy_fn, neighbor_fn, R_init, shift, key, num_steps):\n",
        "  nbrs = neighbor_fn(R_init)\n",
        "\n",
        "  init, apply = simulate.brownian(energy_fn, shift, dt=0.00001, kT=1.0, gamma=0.1)\n",
        "\n",
        "  def body_fn(state, t):\n",
        "    state, nbrs = state\n",
        "    nbrs = neighbor_fn(state.position, nbrs)\n",
        "    state = apply(state, neighbor=nbrs)\n",
        "    return (state, nbrs), 0\n",
        "\n",
        "  key, split = random.split(key)\n",
        "  state = init(split, R_init)\n",
        "\n",
        "  step = 0\n",
        "  step_inc=100\n",
        "  while step < num_steps/step_inc:\n",
        "    rtn_state, _ = lax.scan(body_fn, (state, nbrs), np.arange(step_inc))\n",
        "    new_state, nbrs = rtn_state\n",
        "    # If the neighbor list overflowed, rebuild it and repeat part of \n",
        "    # the simulation.\n",
        "    if nbrs.did_buffer_overflow:\n",
        "      print('Buffer overflow.')\n",
        "      nbrs = neighbor_fn(state.position)\n",
        "    else:\n",
        "      state = new_state\n",
        "      step += 1\n",
        "\n",
        "  return state.position"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C8rzGWg3JM67"
      },
      "source": [
        "To run this, we consider a much larger system than we have to this point. Warning: running this may take a few minutes."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F3BuxR1LKJ8H"
      },
      "source": [
        "Nlarge = 100*N\n",
        "box_size_large = 10*box_size\n",
        "displacement_large, shift_large = setup_periodic_box(box_size_large)\n",
        "\n",
        "key, split1, split2 = random.split(key,3)\n",
        "Rlarge = random.uniform(split1, (Nlarge,dimension), minval=0.0, maxval=box_size_large, dtype=f64) \n",
        "\n",
        "dr_threshold = 1.5\n",
        "neighbor_fn, energy_fn = harmonic_morse_cutoff_neighbor_list(\n",
        "    displacement_large, box_size_large, D0=5.0, alpha=10.0, r0=1.0, k=500.0,\n",
        "    r_onset=r_onset, r_cutoff=r_cutoff, dr_threshold=dr_threshold)\n",
        "energy_fn = jit(energy_fn)\n",
        "\n",
        "start_time = time.clock()\n",
        "Rfinal = run_brownian_neighbor_list(energy_fn, neighbor_fn, Rlarge, shift_large, split2, num_steps=4000)\n",
        "end_time = time.clock()\n",
        "print('run time = {}'.format(end_time-start_time))\n",
        "\n",
        "plot_system( Rfinal, box_size_large, ms=2 )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3LgpOW6bJeia"
      },
      "source": [
        "##Bonds"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u_xSakfJ8TFp"
      },
      "source": [
        "Bonds are a way of specifying potentials between specific pairs of particles that are \"on\" regardless of separation. For example, it is common to employ a two-sided spring potential between specific particle pairs, but JAX MD allows the user to specify arbitrary potentials with static or dynamic parameters. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u9hrkiNlQhyz"
      },
      "source": [
        "### Create and implement a bond potential"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lrugjvd5FKS1"
      },
      "source": [
        "We start by creating a custom potential that corresponds to a bistable spring, taking the form\n",
        "\\begin{equation}\n",
        "V(r) = a_4(r-r_0)^4 - a_2(r-r_0)^2.\n",
        "\\end{equation}\n",
        "$V(r)$ has two minima, at $r = r_0 \\pm \\sqrt{\\frac{a_2}{2a_4}}$."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fEUgo2lW_C-z"
      },
      "source": [
        "def bistable_spring(dr, r0=1.0, a2=2, a4=5, **kwargs):\n",
        "  return a4*(dr-r0)**4 - a2*(dr-r0)**2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d69FVa8uGimp"
      },
      "source": [
        "Plot $V(r)$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T9mAhlt9_W6g"
      },
      "source": [
        "drs = np.arange(0,2,0.01)\n",
        "U = bistable_spring(drs)\n",
        "plt.plot(drs,U)\n",
        "format_plot(r'$r$', r'$V(r)$')\n",
        "finalize_plot()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pbAhV1e6G99j"
      },
      "source": [
        "The next step is to promote this function to act on a set of bonds. This is done via ```smap.bond```, which takes our ```bistable_spring``` function, our displacement function, and a list of the bonds. It returns a function that calculates the energy for a given set of positions."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4aFl4IhnB2h8"
      },
      "source": [
        "def bistable_spring_bond(\n",
        "    displacement_or_metric, bond, bond_type=None, r0=1, a2=2, a4=5):\n",
        "  \"\"\"Convenience wrapper to compute energy of particles bonded by springs.\"\"\"\n",
        "  r0 = np.array(r0, f32)\n",
        "  a2 = np.array(a2, f32)\n",
        "  a4 = np.array(a4, f32)\n",
        "  return smap.bond(\n",
        "    bistable_spring,\n",
        "    space.canonicalize_displacement_or_metric(displacement_or_metric),\n",
        "    bond,\n",
        "    bond_type,\n",
        "    r0=r0,\n",
        "    a2=a2,\n",
        "    a4=a4)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_91khxTSIZ4d"
      },
      "source": [
        "However, in order to implement this, we need a list of bonds. We will do this by taking a system minimized under our original ```harmonic_morse``` potential:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "I6wmOB6YInKg"
      },
      "source": [
        "R_temp, max_force_component = run_minimization(harmonic_morse_pair(displacement,D0=5.0,alpha=10.0,r0=1.0,k=500.0), R, shift)\n",
        "print('largest component of force after minimization = {}'.format(max_force_component))\n",
        "plot_system( R_temp, box_size )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hVMOKqUaOT_-"
      },
      "source": [
        "We now place a bond between all particle pairs that are separated by less than 1.3. ```calculate_bond_data``` returns a list of such bonds, as well as a list of the corresponding current length of each bond.  "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ch34zyUT7vYP"
      },
      "source": [
        "bonds, lengths = calculate_bond_data(displacement, R_temp, 1.3)\n",
        "\n",
        "print(bonds[:5])   # list of particle index pairs that form bonds\n",
        "print(lengths[:5]) # list of the current length of each bond"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LfK-KuluBLjX"
      },
      "source": [
        "We use this length as the ```r0``` parameter, meaning that initially each bond is at the unstable local maximum $r=r_0$."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "p2GYzHgoAoIO"
      },
      "source": [
        "bond_energy_fn = bistable_spring_bond(displacement, bonds, r0=lengths)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H1QRVSpPPYQO"
      },
      "source": [
        "We now use our new ```bond_energy_fn``` to minimize the energy of the system. The expectation is that nearby particles should either move closer together or further apart, and the choice of which to do should be made collectively due to the constraint of constant volume. This is exactly what we see."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eLwRd59YJJCp"
      },
      "source": [
        "Rfinal, max_force_component = run_minimization(bond_energy_fn, R_temp, shift)\n",
        "print('largest component of force after minimization = {}'.format(max_force_component))\n",
        "plot_system( Rfinal, box_size )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ha811sZlRjbx"
      },
      "source": [
        "###Specifying bonds dynamically"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "66hlM5BDRnZ_"
      },
      "source": [
        "As with species or parameters, bonds can be specified dynamically, i.e. when the energy function is called. Importantly, note that this does NOT override bonds that were specified statically in ```smap.bond```."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Wg4Ye4RRSCyo"
      },
      "source": [
        "# Specifying the bonds dynamically ADDS additional bonds. \n",
        "#  Here, we dynamically pass the same bonds that were passed statically, which \n",
        "#  has the effect of doubling the energy\n",
        "print(bond_energy_fn(R))\n",
        "print(bond_energy_fn(R,bonds=bonds, r0=lengths))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5iYIcs-nStf6"
      },
      "source": [
        "We won't go thorugh a further example as the implementation is exactly the same as specifying species or parameters dynamically, but the ability to employ bonds both statically and dynamically is a very powerful and general framework."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lmnHa_pDSTCA"
      },
      "source": [
        "## Combining potentials \n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "762AxBtvTo7m"
      },
      "source": [
        "Most JAX MD functionality (e.g. simulations, energy minimizations) relies on a function that calculates energy for a set of positions. Importantly, while this cookbook focus on simple and robust ways of defining such functions, JAX MD is not limited to these methods; users can implement energy functions however they like. \n",
        "\n",
        "As an important example, here we consider the case where the energy includes both a pair potential and a bond potential. Specifically, we combine ```harmonic_morse_pair``` with ```bistable_spring_bond```. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KowW2fMuV_nY"
      },
      "source": [
        "# Note, the code in the \"Bonds\" section must be run prior to this.\n",
        "energy_fn = harmonic_morse_pair(displacement,D0=0.,alpha=10.0,r0=1.0,k=1.0)\n",
        "bond_energy_fn = bistable_spring_bond(displacement, bonds, r0=lengths)\n",
        "def combined_energy_fn(R):\n",
        "  return energy_fn(R) + bond_energy_fn(R)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LR8NW_NkZV0y"
      },
      "source": [
        "Here, we have set $D_0=0$, so the pair potential is just a one-sided repulsive harmonic potential. For particles connected with a bond, this raises the energy of the \"contracted\" minimum relative to the \"extended\" minimum."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vNd0A8CYXCuR"
      },
      "source": [
        "drs = np.arange(0,2,0.01)\n",
        "U = harmonic_morse(drs,D0=0.,alpha=10.0,r0=1.0,k=1.0)+bistable_spring(drs)\n",
        "plt.plot(drs,U)\n",
        "format_plot(r'$r$', r'$V(r)$')\n",
        "finalize_plot()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fMU9ethYZzry"
      },
      "source": [
        "This new energy function can be passed to the minimization routine (or any other JAX MD simulation routine) in the usual way."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "h67C8SRrWm8J"
      },
      "source": [
        "Rfinal, max_force_component = run_minimization(combined_energy_fn, R_temp, shift)\n",
        "print('largest component of force after minimization = {}'.format(max_force_component))\n",
        "plot_system( Rfinal, box_size )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7JkkgEY4H19L"
      },
      "source": [
        "##Specifying forces instead of energies"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RXaO8tDcb7aC"
      },
      "source": [
        "So far, we have defined functions that calculate the energy of the system, which we then pass to JAX MD. Internally, JAX MD uses automatic differentiation to convert these into functions that calculate forces, which are necessary to evolve a system under a given dynamics. However, JAX MD has the option to pass force functions directly, rather than energy functions. This creates additional flexibility because some forces cannot be represented as the gradient of a potential.\n",
        "\n",
        "As a simple example, we create a custom force function that zeros out the force of some particles. During energy minimization, where there is no stochastic noise, this has the effect of fixing the position of these particles.\n",
        "\n",
        "First, we break the system up into two species, as before."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RzFu0MnmHzp3"
      },
      "source": [
        "N_0 = N // 2  # Half the particles in species 0\n",
        "N_1 = N - N_0 # The rest in species 1\n",
        "species = np.array([0]*N_0 + [1]*N_1, dtype=np.int32)\n",
        "print(species)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Hry734J5Zh_F"
      },
      "source": [
        "Next, we we creat our custom force function. Starting with our ```harmonic_morse``` pair potential, we calculate the force manually (i.e. using built-in automatic differentiation), and then multiply the force by the species id, which has the desired effect. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yZRz8pxVG1O8"
      },
      "source": [
        "energy_fn = harmonic_morse_pair(displacement,D0=5.0,alpha=10.0,r0=1.0,k=500.0)\n",
        "force_fn = quantity.force(energy_fn)\n",
        "\n",
        "def custom_force_fn(R, **kwargs):\n",
        "  return vmap(lambda a,b: a*b)(force_fn(R),species)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PGpaSD0BZ-UT"
      },
      "source": [
        "Running simulations with custom forces is as easy as passing this force function to the simulation. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xhyWE6msitFX"
      },
      "source": [
        "def run_minimization_general(energy_or_force, R_init, shift, num_steps=5000):\n",
        "  dt_start = 0.001\n",
        "  dt_max   = 0.004\n",
        "  init,apply=minimize.fire_descent(jit(energy_or_force),shift,dt_start=dt_start,dt_max=dt_max)\n",
        "  apply = jit(apply)\n",
        "\n",
        "  @jit\n",
        "  def scan_fn(state, i):\n",
        "    return apply(state), 0.\n",
        "\n",
        "  state = init(R_init)\n",
        "  state, _ = lax.scan(scan_fn,state,np.arange(num_steps))\n",
        "\n",
        "  return state.position, np.amax(np.abs(quantity.canonicalize_force(energy_or_force)(state.position)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0vZJwbFuja0W"
      },
      "source": [
        "We run this as usual,"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "maSN2X6kFbn0"
      },
      "source": [
        "key, split = random.split(key)\n",
        "Rfinal, _ = run_minimization_general(custom_force_fn, R, shift)\n",
        "plot_system( Rfinal, box_size, species )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pDXzM-vIaZs_"
      },
      "source": [
        "After the above minimization, the blue particles have the same positions as they did initially:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ywwWkJW7JaMt"
      },
      "source": [
        "plot_system( R, box_size, species )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zT2peeMfKKp6"
      },
      "source": [
        "Note, this method for fixing particles only works when there is no stochastic noise (e.g. in Langevin or Brownian dynamics) because such noise affects partices whether or not they have a net force. A safer way to fix particles is to create a custom ```shift``` function."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i0OQMafrnArk"
      },
      "source": [
        "##Coupled ensembles"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0nbe8PfTQEpA"
      },
      "source": [
        "For a final example that demonstrates the flexibility within JAX MD, lets do something that is particularly difficult in most standard MD packages. We will create a \"coupled ensemble\" -- i.e. a set of two identical systems that are connected via a $Nd$ dimensional spring. An extension of this idea is used, for example, in the Doubly Nudged Elastic Band method for finding transition states. \n",
        "\n",
        "If the \"normal\" energy of each system is \n",
        "\\begin{equation}\n",
        "U(R) = \\sum_{i,j} V( r_{ij} ),\n",
        "\\end{equation}\n",
        "where $r_{ij}$ is the distance between the $i$th and $j$th particles in $R$ and the $V(r)$ is a standard pair potential, and if the two sets of positions, $R_0$ and $R_1$, are coupled via the potential\n",
        "\\begin{equation}\n",
        "U_\\mathrm{spr}(R_0,R_1) = \\frac 12 k_\\mathrm{spr} \\left| R_1 - R_0 \\right|^2,\n",
        "\\end{equation}\n",
        "so that the total energy of the system is \n",
        "\\begin{equation}\n",
        "U_\\mathrm{total} = U(R_0) + U(R_1) + U_\\mathrm{spr}(R_0,R_1).\n",
        "\\end{equation}\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uTtLhh3dKWAU"
      },
      "source": [
        "energy_fn = harmonic_morse_pair(displacement,D0=5.0,alpha=10.0,r0=0.5,k=500.0)\n",
        "def spring_energy_fn(Rall, k_spr=50.0, **kwargs):\n",
        "  metric = vmap(space.canonicalize_displacement_or_metric(displacement), (0, 0), 0)\n",
        "  dr = metric(Rall[0],Rall[1])\n",
        "  return 0.5*k_spr*np.sum((dr)**2)\n",
        "def total_energy_fn(Rall, **kwargs):\n",
        "  return np.sum(vmap(energy_fn)(Rall)) + spring_energy_fn(Rall)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tNy9BPXATOpU"
      },
      "source": [
        "We now have to define a new shift function that can handle arrays of shape $(2,N,d)$. In addition, we make two copies of our initial positions ```R```, one for each system. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rfuT7nC7TlLe"
      },
      "source": [
        "def shift_all(Rall, dRall, **kwargs):\n",
        "  return vmap(shift)(Rall, dRall)\n",
        "Rall = np.array([R,R])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3paT1E6zT4uS"
      },
      "source": [
        "Now, all we have to do is pass our custom energy and shift functions, as well as the $(2,N,d)$ dimensional initial position, to JAX MD, and proceed as normal. \n",
        "\n",
        "As a demonstration, we define a simple and general Brownian Dynamics simulation function, similar to the simulation routines above except without the special cases (e.g. chaning ```r0``` or species). "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iJqEQeF1kBXS"
      },
      "source": [
        "def run_brownian_simple(energy_or_force, R_init, shift, key, num_steps):\n",
        "  init, apply = simulate.brownian(energy_or_force, shift, dt=0.00001, kT=1.0, gamma=0.1)\n",
        "  apply = jit(apply)\n",
        "\n",
        "  @jit\n",
        "  def scan_fn(state, t):\n",
        "    return apply(state), 0\n",
        "\n",
        "  key, split = random.split(key)\n",
        "  state = init(split, R_init)\n",
        "\n",
        "  state, _ = lax.scan(scan_fn, state, np.arange(num_steps))\n",
        "  return state.position"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ug_49XmKkkVk"
      },
      "source": [
        "Note that nowhere in this function is there any indication that we are simulating an ensemble of systems. This comes entirely form the inputs: i.e. the energy function, the shift function, and the set of initial positions. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sbbX64tvMCJy"
      },
      "source": [
        "key, split = random.split(key)\n",
        "Rall_final = run_brownian_simple(total_energy_fn, Rall, shift_all, split, num_steps=10000)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RrdgrCPMUJzN"
      },
      "source": [
        "The output also has shape $(2,N,d)$. If we display the results, we see that the two systems are in similar, but not identical, positions, showing that we have succeeded in simulating a coupled ensemble. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m5vSaHU-PKO_"
      },
      "source": [
        "for Ri in Rall_final:\n",
        "  plot_system( Ri, box_size )\n",
        "finalize_plot((0.5,0.5))"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}