{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1af0d1e1-7b4d-4267-b336-e841303d19e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#You don't need to change anything in this block, although the modules need to be installed to run this notebook\n",
    "\n",
    "#We import numpy to handle vectors and some math\n",
    "import numpy as np\n",
    "\n",
    "#We import pandas to create a data frame of the experiment data\n",
    "#Such a table can later be used for plotting our results\n",
    "import pandas as pd\n",
    "\n",
    "# Import plotly, which is used for visualization\n",
    "import plotly.express as px\n",
    "import plotly.io as pio\n",
    "pio.renderers.default = 'iframe'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1616842d-eb1e-454f-9d30-868f1a0c8425",
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate N random particles in a 2d environment with a given pattern:\n",
    "def generate_pattern(N):\n",
    "    return np.random.rand(N, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eff15957-5a55-49de-91aa-875609fc0126",
   "metadata": {},
   "outputs": [],
   "source": [
    "# record data into dataframe (used for plotting)\n",
    "def make_df(data, t, type_name):\n",
    "    df = pd.DataFrame(data, columns=[\"x\", \"y\"])\n",
    "    df['t'] = t\n",
    "    df['type'] = type_name\n",
    "    df['pid'] = range(len(data))\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e4c8a1c-5538-40b0-9826-c20eb521a16c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# forces between a pair of particles i, j\n",
    "def force_ij(i, j, pattern, particles, k=0.3):\n",
    "    dist = np.linalg.norm(pattern[i, :] - pattern[j, :])\n",
    "    xi_minus_xj = particles[i, :] - particles[j, :]\n",
    "    f = -k * (np.linalg.norm(xi_minus_xj) - dist) * xi_minus_xj\n",
    "    return f\n",
    "\n",
    "# force for particle i\n",
    "def force_i(i, pattern, particles, k=0.3):\n",
    "    f = np.array([0.0, 0.0])\n",
    "    for j, _ in enumerate(particles):\n",
    "        if j == i:\n",
    "            continue\n",
    "        f = f + force_ij(i, j, pattern, particles, k=k)\n",
    "    return f\n",
    "\n",
    "# forces for all particles\n",
    "def force(pattern, particles, k=0.1):\n",
    "    return np.array([force_i(i, pattern, particles, k=k) for i, _ in enumerate(particles)])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66ff3a1b-ce35-4c4a-83a8-9dd394f17e6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# run an experiment with N random particles and a random pattern\n",
    "l = []\n",
    "data = []\n",
    "N = 5\n",
    "pattern = generate_pattern(N)\n",
    "x = generate_pattern(N)\n",
    "v = np.zeros_like(x)\n",
    "for t in range(30):\n",
    "    # record the current state at time t\n",
    "    data.append(make_df(x, t, \"particle\"))\n",
    "    data.append(make_df(pattern, t, \"pattern\"))\n",
    "  \n",
    "    # update v and x\n",
    "    v = force(pattern, x, k=0.3) + 0.3 * v\n",
    "    x = x + v\n",
    "df = pd.concat(data)\n",
    "\n",
    "# show the result\n",
    "fig = px.scatter(df, x=\"x\", y=\"y\", color=\"pid\", animation_frame=\"t\", animation_group=\"pid\", facet_col=\"type\")\n",
    "fig.update_layout(xaxis_range=(-0.5, 1.5), yaxis_range=(-0.5, 1.5), width=1600, height=800)\n",
    "fig.update_traces(marker={\"size\": 12})\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b0f2862-b831-4ccc-90ba-acd590a7f4a8",
   "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.9.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}