{
 "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 __init__(self, n_rows: int, n_cols: int):\n",
    "        self.n_rows = n_rows\n",
    "        self.n_cols = n_cols\n",
    "        self.displacement = numpy.zeros((n_rows, n_cols, 2))\n",
    "        self.velocity = numpy.zeros((n_rows, n_cols, 2))\n",
    "        self.force = numpy.zeros((n_rows, n_cols, 2))\n",
    "\n",
    "        self.elastic_spring_constant = 1.0\n",
    "        self.bond_bending_constant = 1.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": [],
   "source": [
    "def neighbors(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",
    "            if row_offset == 0 and col_offset == 0:\n",
    "                continue\n",
    "            else:\n",
    "                k = 1.0 if row_offset == 0 or col_offset == 0 else 2.0\n",
    "                yield (k, row + row_offset, col + col_offset)\n",
    "\n",
    "def calculate_force(elm: ELM, row: int, col: int) -> numpy.ndarray:\n",
    "    K = elm.elastic_spring_constant\n",
    "\n",
    "    total = numpy.array([0.0, 0.0])\n",
    "    for _, nb_row, nb_col in neighbors(row, col):\n",
    "        u_ij = elm.get_displacement(row, col) - elm.get_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 = elm.get_actual_position(row, col) - elm.get_actual_position(nb_row, nb_col)\n",
    "        \n",
    "        total -= K * numpy.dot(u_ij, x_ij) * (r_ij / numpy.linalg.norm(r_ij))\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(20, 20)\n",
    "for row in range(1, e.n_rows - 1):\n",
    "    e.set_displacement(row, 1, numpy.array([gaussian((e.n_rows - 1) / 2, 5.0, row) * -5, 0.0]))"
   ]
  },
  {
   "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_elm = ELM(elm.n_rows, elm.n_cols)\n",
    "\n",
    "    new_elm.displacement = elm.displacement + elm.velocity * delta_t + 0.5 * elm.force * delta_t * delta_t\n",
    "    for row in range(1, elm.n_rows - 1):\n",
    "        for col in range(1, elm.n_cols - 1):\n",
    "            new_elm.set_force(row, col, calculate_force(new_elm, row, col))\n",
    "    new_elm.velocity = elm.velocity + 0.5 * (elm.force + new_elm.force) * delta_t\n",
    "\n",
    "    return new_elm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "bea964bf-6e94-4b72-a458-50c2964bbba1",
   "metadata": {},
   "outputs": [],
   "source": [
    "es = [e]\n",
    "for _ in range(1000):\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": "32e0ea7e326e4b54badc0ba9b1e24e85",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(IntSlider(value=0, description='i', max=1000), 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": [
    "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",
    "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",
    "        scale = 10\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": null,
   "id": "d419c328-a466-42fa-a8d9-14807d7100e4",
   "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
}