{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e465f900-0a94-46d4-b3c5-23cc0129557b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy\n",
    "\n",
    "from ipywidgets import IntSlider, interact\n",
    "from typing import Iterator, Tuple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b39df788-2971-4974-9ced-1518a47181e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ELM:\n",
    "    n_rows: int\n",
    "    n_cols: int\n",
    "    displacement: numpy.ndarray\n",
    "    velocity: numpy.ndarray\n",
    "    force: numpy.ndarray\n",
    "\n",
    "    # K\n",
    "    elastic_spring_constant: float\n",
    "    # c\n",
    "    bond_bending_constant: float\n",
    "    \n",
    "    def make(n_rows: int, n_cols: int):\n",
    "        return ELM(\n",
    "            n_rows,\n",
    "            n_cols,\n",
    "            numpy.zeros((n_rows, n_cols, 2)),\n",
    "            numpy.zeros((n_rows, n_cols, 2)),\n",
    "            numpy.zeros((n_rows, n_cols, 2))\n",
    "        )\n",
    "    \n",
    "    def __init__(\n",
    "        self,\n",
    "        n_rows: int,\n",
    "        n_cols: int,\n",
    "        displacement: numpy.ndarray,\n",
    "        velocity: numpy.ndarray,\n",
    "        force: numpy.ndarray\n",
    "    ) -> None:\n",
    "        self.n_rows = n_rows\n",
    "        self.n_cols = n_cols\n",
    "        self.displacement = displacement\n",
    "        self.velocity = velocity\n",
    "        self.force = force\n",
    "\n",
    "        self.elastic_spring_constant = 1.0\n",
    "        self.bond_bending_constant = 2.0\n",
    "    \n",
    "    # u_i(t)\n",
    "    def get_displacement(self, row: int, col: int) -> numpy.ndarray:\n",
    "        return self.displacement[row, col, :]\n",
    "\n",
    "    def set_displacement(self, row: int, col: int, displacement: numpy.ndarray) -> None:\n",
    "        self.displacement[row, col, :] = displacement\n",
    "\n",
    "    # x_i(t)\n",
    "    def get_lattice_position(self, row: int, col: int) -> numpy.ndarray:\n",
    "        return numpy.array([col, row], dtype=float)\n",
    "\n",
    "    # r_i(t)\n",
    "    def get_actual_position(self, row: int, col: int) -> numpy.ndarray:\n",
    "        return self.get_lattice_position(row, col) + self.get_displacement(row, col)\n",
    "        \n",
    "    # v_i(t)\n",
    "    def get_velocity(self, row: int, col: int) -> numpy.ndarray:\n",
    "        return self.velocity[row, col, :]\n",
    "\n",
    "    # F_i(t)\n",
    "    def get_force(self, row: int, col: int) -> numpy.ndarray:\n",
    "        return self.force[row, col, :]\n",
    "\n",
    "    def set_force(self, row: int, col: int, force: numpy.ndarray) -> None:\n",
    "        self.force[row, col, :] = force"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb5633c4-c2c0-4a9c-bde8-8a99c5db8a9d",
   "metadata": {},
   "source": [
    "(O'Brien, 2008)\n",
    "\n",
    "The force acting on node $i$ is given by\n",
    "\n",
    "$$\n",
    "\\mathbf{F}_i = \\sum^N_j \\mathbf{F}_{ij} \\frac{\\mathbf{r}_{ij}}{|\\mathbf{r}_{ij}|}\n",
    "$$\n",
    "\n",
    "Where $N = 8$ (for the 2D case) is the number of neighbors $j$.\n",
    "The force $\\mathbf{F}_{ij}$ on node $i$ from node $j$ is given by\n",
    "\n",
    "$$\n",
    "\\mathbf{F}_{ij} = -K_{ij} (\\mathbf{u}_{ij} \\cdot \\mathbf{x}_{ij}) + \\frac{c \\mathbf{u}_{ij}}{|\\mathbf{x}_{ij}|^2}\n",
    "$$\n",
    "\n",
    "where\n",
    "\n",
    "* $K_{ij} \\in \\mathbb{R}$ is the elastic spring constant\n",
    "* $\\mathbf{u}_{ij} = \\mathbf{u}_i - \\mathbf{u}_j \\in \\mathbb{R}^2$ is the displacement\n",
    "* $\\mathbf{x}_{ij} = \\mathbf{x}_i - \\mathbf{x}_j \\in \\mathbb{R}^2$ is the vector connecting nodes $i$ and $j$ on the undistorted lattice\n",
    "* $c \\in \\mathbb{R}$ is the bond-bending constant\n",
    "\n",
    "The $\\frac{c \\mathbf{u}_{ij}}{|\\mathbf{x}_{ij}|^2}$ part does not make sense to me since the first summand in $\\mathbf{F}_{ij}$ is scalar so the implementation below leaves it out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b4ed52bd-02aa-411a-9af4-be9bb28227bc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 13092 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def neighbors(elm: ELM, row: int, col: int) -> Iterator[Tuple[float, int, int]]:\n",
    "    for row_offset in [-1, 0, 1]:\n",
    "        for col_offset in [-1, 0, 1]:\n",
    "            nb_row = row + row_offset\n",
    "            nb_col = col + col_offset\n",
    "            \n",
    "            if nb_row < 0 or nb_row >= elm.n_rows or nb_col < 0 or nb_col >= elm.n_cols:\n",
    "                continue\n",
    "            if row_offset == 0 and col_offset == 0:\n",
    "                continue\n",
    "            \n",
    "            k = 1.0 if row_offset == 0 or col_offset == 0 else 2.0\n",
    "            yield (k, nb_row, nb_col)\n",
    "\n",
    "def calculate_force(elm: ELM, displacement: numpy.ndarray, row: int, col: int) -> numpy.ndarray:\n",
    "    K = elm.elastic_spring_constant\n",
    "    c = elm.bond_bending_constant\n",
    "\n",
    "    total = numpy.array([0.0, 0.0])\n",
    "    for k, nb_row, nb_col in neighbors(elm, row, col):\n",
    "        u_ij = displacement[row, col, :] - displacement[nb_row, nb_col, :]\n",
    "        x_ij = elm.get_lattice_position(row, col) - elm.get_lattice_position(nb_row, nb_col)\n",
    "        r_ij = x_ij + u_ij\n",
    "        \n",
    "        total -= K * k * numpy.dot(u_ij, x_ij) * (r_ij / numpy.linalg.norm(r_ij)) + (c * u_ij) / (numpy.linalg.norm(x_ij) ** 2)\n",
    "\n",
    "    return total\n",
    "\n",
    "def gaussian(mu, sigma_sq, x):\n",
    "    return numpy.exp(-(((x - mu) / numpy.sqrt(sigma_sq)) ** 2) / 2) / numpy.sqrt(2 * numpy.pi * sigma_sq)\n",
    "\n",
    "e = ELM.make(50, 50)\n",
    "n_plucked = 10\n",
    "base = e.n_rows // 2 - n_plucked // 2\n",
    "for row in range(n_plucked):\n",
    "    e.set_displacement(base + row, 0, numpy.array([\n",
    "        gaussian(n_plucked // 2, n_plucked / 3, row) * -5,\n",
    "        0.0\n",
    "    ]))\n",
    "\n",
    "x, y = numpy.meshgrid(\n",
    "    numpy.linspace(0, e.n_rows - 1, e.n_rows, dtype=int),\n",
    "    numpy.linspace(0, e.n_cols - 1, e.n_cols, dtype=int)\n",
    ")\n",
    "\n",
    "plt.scatter(\n",
    "    x + e.displacement[:, :, 0],\n",
    "    y + e.displacement[:, :, 1]\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dc183cc-5fa4-48f1-bc30-9cffff91be6a",
   "metadata": {},
   "source": [
    "A [velocity Verlet integration](https://en.wikipedia.org/wiki/Verlet_integration#Velocity_Verlet) scheme is used to solve the system.\n",
    "\n",
    "Our initial configuration defines for each of the $M$ nodes $i$:\n",
    "\n",
    "* $\\mathbf{u}_i(t_0) \\in \\mathbb{R}^M \\times 2$: Displacement\n",
    "* $\\mathbf{v}_i(t_0) \\in \\mathbb{R}^M \\times 2$: Velocity\n",
    "\n",
    "To calculate a time step, do for every node $i$:\n",
    "\n",
    "* Calculate the displacement $\\mathbf{u}_i(t + \\Delta t) = \\mathbf{u}_i(t) + \\mathbf{v}_i(t) \\Delta t + \\frac{1}{2} \\mathbf{F}_i(t) \\Delta t^2$\n",
    "* Calculate $\\mathbf{F}_i(t + \\Delta t)$ using $\\mathbf{u}_i(t + \\Delta t)$\n",
    "* Calculate $\\mathbf{v}_i(t + \\Delta t) = \\mathbf{v}_i(t) + \\frac{1}{2} (\\mathbf{F}_i(t) + \\mathbf{F}_i(t + \\Delta t)) \\Delta t$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ff7adc30-191e-4b7c-b146-2a533463d6ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "def velocity_verlet_step(elm: ELM, delta_t: float) -> ELM:\n",
    "    new_displacement = elm.displacement + elm.velocity * delta_t + 0.5 * elm.force * delta_t * delta_t\n",
    "    \n",
    "    new_force = numpy.zeros((elm.n_rows, elm.n_cols, 2))\n",
    "    for row in range(elm.n_rows):\n",
    "        for col in range(elm.n_cols):\n",
    "            new_force[row, col, :] = calculate_force(\n",
    "                elm,\n",
    "                new_displacement,\n",
    "                row,\n",
    "                col\n",
    "            )\n",
    "    \n",
    "    new_velocity = elm.velocity + 0.5 * (elm.force + new_force) * delta_t\n",
    "\n",
    "    return ELM(elm.n_rows, elm.n_cols, new_displacement, new_velocity, new_force)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bea964bf-6e94-4b72-a458-50c2964bbba1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "............................................................................................................................................................................................................................................................................................................"
     ]
    }
   ],
   "source": [
    "es = [e]\n",
    "for _ in range(300):\n",
    "    print(\".\", end=\"\")\n",
    "    es.append(velocity_verlet_step(es[-1], 0.1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "421c0fd2-2dba-4fd5-89ff-0968227b08b9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dc96f81e9f764b4fa58103f705f3e445",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=0, description='i', max=300), Output()), _dom_classes=('widget-interact'…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.do_plot(i: int) -> None>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def do_plot(i: int) -> None:\n",
    "    fig, ax = plt.subplots(figsize=(10, 5), ncols=2)\n",
    "    ax[0].scatter(\n",
    "        x + es[i].displacement[:, :, 0],\n",
    "        y + es[i].displacement[:, :, 1]\n",
    "    )\n",
    "    ax[0].quiver(\n",
    "        x + es[i].displacement[:, :, 0],\n",
    "        y + es[i].displacement[:, :, 1],\n",
    "        es[i].force[:, :, 0],\n",
    "        es[i].force[:, :, 1],\n",
    "        angles = \"xy\",\n",
    "        scale = 100\n",
    "    )\n",
    "    ax[1].imshow(numpy.linalg.norm(es[i].displacement, axis=2), cmap=\"hot\")\n",
    "    plt.show()\n",
    "\n",
    "interact(do_plot, i = IntSlider(min=0, max=len(es) - 1, value=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "96ab8cf1-8739-4faa-b6b5-70aa52a15a2a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 13124 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.animation as animation\n",
    "\n",
    "fps = 60\n",
    "\n",
    "fig, ax = plt.subplots(1, 1)\n",
    "im = plt.imshow(numpy.linalg.norm(es[0].displacement, axis=2), cmap=\"hot\")\n",
    "\n",
    "def frame(i):\n",
    "    im.set_data(numpy.linalg.norm(es[i].displacement, axis=2))\n",
    "    im.autoscale()\n",
    "    ax.set_title(f\"$i = {i} \\qquad \\sum |F_i| = {numpy.sum(numpy.linalg.norm(es[i].force)):.3f} \\qquad \\sum |v_i| = {numpy.sum(numpy.linalg.norm(es[i].velocity)):.3f} \\qquad \\sum |d_i| = {numpy.sum(numpy.linalg.norm(es[i].displacement)):.3f}$\")\n",
    "    return [im]\n",
    "\n",
    "anim = animation.FuncAnimation(\n",
    "    fig,\n",
    "    frame, \n",
    "    frames = len(es),\n",
    "    interval = 1000 / fps,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "adc4a988-310f-4fe7-90bd-e9cd643c469e",
   "metadata": {},
   "outputs": [],
   "source": [
    "anim.save('elastic.mp4', fps=fps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc2f7ad3-d825-498b-84ba-7fe3b3c5eaf5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
