{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2342d5ae-1839-4132-ace8-da8449223a1c",
   "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}\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}$ 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 Lamé constants are given by\n",
    "\n",
    "$\\lambda_\\text{2D} = K - \\frac{c}{\\Delta x^2}$\n",
    "\n",
    "and\n",
    "\n",
    "$\\mu_\\text{2D} = K + \\frac{c}{\\Delta x^2}$\n",
    "\n",
    "where $\\Delta x$ is the lattice grid spacing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "05d4e22c-3d7d-4ac2-bf22-59b78aa40b84",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "from ipywidgets import interact\n",
    "from typing import Tuple"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50ba2e84-bb40-44c8-bb86-8804fd26f674",
   "metadata": {},
   "source": [
    "$$\n",
    "\\mathbf{F}_{ij} = K_{ij} (\\mathbf{u}_{ij} \\cdot \\mathbf{x}_{ij}) + \\frac{c \\mathbf{u}_{ij}}{|\\mathbf{x}_{ij}|^2}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "82bb4a02-1ecf-4c84-8b9d-b753dee75337",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "j=(0, 0) K_ij=2.0000000000000004 u_ij=array([0.1, 0. ]) x_ij=array([1., 1.])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def calculate_force_elastic(\n",
    "    displacement: np.ndarray,\n",
    "    i: Tuple[int, int],\n",
    "    j: Tuple[int, int]\n",
    ") -> np.ndarray:\n",
    "    u_ij = displacement[i] - displacement[j]\n",
    "    x_ij = np.array(i, dtype=float) - np.array(j, dtype=float)\n",
    "    K_ij = np.linalg.norm(x_ij) ** 2\n",
    "    r_ij = x_ij + u_ij\n",
    "\n",
    "    print(f\"{j=} {K_ij=} {u_ij=} {x_ij=}\")\n",
    "\n",
    "    return 0\n",
    "    return -K_ij * np.dot(u_ij, x_ij) * r_ij / np.linalg.norm(r_ij)\n",
    "\n",
    "def calculate_force_bond(\n",
    "    displacement: np.ndarray,\n",
    "    i: Tuple[int, int],\n",
    "    j: Tuple[int, int]\n",
    ") -> np.ndarray:\n",
    "    u_ij = displacement[i] - displacement[j]\n",
    "    x_ij = np.array(i, dtype=float) - np.array(j, dtype=float)\n",
    "    c_ij = 1.0\n",
    "\n",
    "    return (-c_ij * u_ij) / (np.linalg.norm(x_ij) ** 2)\n",
    "\n",
    "def calculate_force(\n",
    "    displacement: np.ndarray,\n",
    "    i: Tuple[int, int],\n",
    "    j: Tuple[int, int]\n",
    ") -> np.ndarray:\n",
    "    return calculate_force_elastic(displacement, i, j) + calculate_force_bond(displacement, i, j)\n",
    "\n",
    "displacement = np.array([\n",
    "    [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]],\n",
    "    [[0.0, 0.0], [0.1, 0.0], [0.0, 0.0]],\n",
    "    [[0.0, 0.0], [0.0, 0.0], [0.0, 0.0]]\n",
    "])\n",
    "calculate_force_elastic(displacement, (1, 1), (0, 0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f7a019c2-8d20-4e84-9e0c-ef788781a07a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "10515f6b18554cd1bd86dd56891eed2b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(FloatSlider(value=0.0, description='dp_x', max=1.5, min=-1.5, step=0.05), FloatSlider(va…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.do_plot(dp_x: float, dp_y: float) -> None>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def do_plot(dp_x: float, dp_y: float) -> None:\n",
    "    i = (1, 1)\n",
    "    displacement[i] = np.array([dp_x, dp_y])\n",
    "    \n",
    "    width, height, _ = displacement.shape\n",
    "\n",
    "    xs = []\n",
    "    ys = []\n",
    "    elastic_xs = []\n",
    "    elastic_ys = []\n",
    "        \n",
    "    for x in range(width):\n",
    "        for y in range(height):\n",
    "            j = (x, y)\n",
    "            xs.append(x + displacement[j][0])\n",
    "            ys.append(y + displacement[j][1])\n",
    "\n",
    "            if j != i:\n",
    "                F_ij = calculate_force(displacement, i, j)\n",
    "                print(f\"{F_ij=}\")\n",
    "                elastic_xs.append(F_ij[0])\n",
    "                elastic_ys.append(F_ij[1])\n",
    "\n",
    "    pi = i + displacement[i]\n",
    "    plt.xlim(-1, 3)\n",
    "    plt.ylim(-1, 3)\n",
    "    plt.scatter(xs, ys)\n",
    "    plt.quiver(\n",
    "        [pi[0]] * len(elastic_xs),\n",
    "        [pi[1]] * len(elastic_ys),\n",
    "        elastic_xs,\n",
    "        elastic_ys,\n",
    "        angles = \"xy\",\n",
    "        pivot = \"tip\",\n",
    "        scale = 5,\n",
    "        width = 0.005\n",
    "    )\n",
    "    plt.quiver(\n",
    "        pi[0], pi[1],\n",
    "        sum(elastic_xs),\n",
    "        sum(elastic_ys),\n",
    "        pivot = \"tip\",\n",
    "        scale = 5,\n",
    "        angles = \"xy\",\n",
    "        color = \"red\"\n",
    "    )\n",
    "    \n",
    "\n",
    "interact(\n",
    "    do_plot,\n",
    "    dp_x = (-1.5, 1.5, 0.05),\n",
    "    dp_y = (-1.5, 1.5, 0.05)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4c55b68-903a-48de-8226-9cfbbbca4af6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f283012a-6b7c-48c3-8bdc-f45f5b458edc",
   "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
}
