{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "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": 2,
   "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, dim=2):\n",
    "    return np.random.rand(N, dim)*2.0 - 1.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "eff15957-5a55-49de-91aa-875609fc0126",
   "metadata": {},
   "outputs": [],
   "source": [
    "# record data into dataframe (used for plotting)\n",
    "def make_df(data, t, type_name, **kwargs):\n",
    "    df = pd.DataFrame(data, columns=[f\"x{i}\" for i,_ in enumerate(data[0])])\n",
    "    df['t'] = t\n",
    "    df['type'] = type_name\n",
    "    df['pid'] = range(len(data))\n",
    "    for k, v in kwargs.items():\n",
    "        df[k] = v\n",
    "    return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "27323c25-9f75-4895-8de7-7fef52dd33ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f1(x):\n",
    "    return np.linalg.norm(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3e4c8a1c-5538-40b0-9826-c20eb521a16c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def random_phi(dim):\n",
    "    return np.random.rand()\n",
    "\n",
    "def force_i(i, particles, x_p, c1=2.0, c2=2.0):\n",
    "    return c2 * random_phi(dim) * (x_p[i] - particles[i])\n",
    "    \n",
    "# forces for all particles\n",
    "def force(particles, x_p, c1=2.0, c2=2.0):\n",
    "    return np.array([force_i(i, particles, x_p, c1=c1, c2=c2) for i, _ in enumerate(particles)])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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 = 10\n",
    "dim = 2\n",
    "w = 0.4\n",
    "c1 = 2.0\n",
    "c2 = 2.0\n",
    "\n",
    "def fitness(x):\n",
    "    return np.array([f1(x_i) for x_i in x])\n",
    "\n",
    "# use 31 runs, when testing more which type runs better\n",
    "for run in range(1):\n",
    "    \n",
    "    x = generate_pattern(N, dim=dim)\n",
    "    v = np.zeros_like(x)\n",
    "\n",
    "    # fitness of population\n",
    "    f = fitness(x)\n",
    "    # fitness of previous best\n",
    "    f_p = fitness(x)\n",
    "    # fitness of local best (global best for fully connected)\n",
    "    f_l = min(fitness(x))\n",
    "    x_p = x\n",
    "    \n",
    "    for t in range(100):\n",
    "        # record the current state at time t\n",
    "        data.append(make_df(x.copy(), t, \"particle\", fitness=f, best_fitness=f_l, previous_bestf=f_p, run=run))\n",
    "        # update v and x\n",
    "        v = force(x, x_p, c1=c1, c2=c2) + w * v\n",
    "        x = x + v\n",
    "        # compute fitness\n",
    "        f = fitness(x)\n",
    "        # update x_p\n",
    "        x_p[f < f_p] = x[f < f_p]\n",
    "        # update f_p\n",
    "        f_p[f < f_p] = f[f < f_p]\n",
    "        # update f_l\n",
    "        if f_l > min(f):\n",
    "            f_l = min(f)\n",
    "df = pd.concat(data)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3b0f2862-b831-4ccc-90ba-acd590a7f4a8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<iframe\n",
       "    scrolling=\"no\"\n",
       "    width=\"100%\"\n",
       "    height=\"545px\"\n",
       "    src=\"iframe_figures/figure_7.html\"\n",
       "    frameborder=\"0\"\n",
       "    allowfullscreen\n",
       "></iframe>\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = px.line(df, x='t', y='best_fitness', color='run')\n",
    "fig.update_layout(yaxis_range=(0, None))\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2103480f-e6ef-4102-842c-d0eb2357ab08",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<iframe\n",
       "    scrolling=\"no\"\n",
       "    width=\"100%\"\n",
       "    height=\"545px\"\n",
       "    src=\"iframe_figures/figure_8.html\"\n",
       "    frameborder=\"0\"\n",
       "    allowfullscreen\n",
       "></iframe>\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = px.box(df, x='t', y='best_fitness')\n",
    "fig.update_layout(yaxis_range=(0, None))\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c8a2beb9-daa7-4eec-a3a0-75ea16b410f1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<iframe\n",
       "    scrolling=\"no\"\n",
       "    width=\"820px\"\n",
       "    height=\"820\"\n",
       "    src=\"iframe_figures/figure_9.html\"\n",
       "    frameborder=\"0\"\n",
       "    allowfullscreen\n",
       "></iframe>\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# show the result\n",
    "fig = px.scatter(df.loc[df.run.eq(0)], x=\"x0\", y=\"x1\", color=\"pid\", animation_frame=\"t\", animation_group=\"pid\")\n",
    "fig.update_layout(xaxis_range=(-1, 1), yaxis_range=(-1, 1), width=800, height=800)\n",
    "fig.update_traces(marker={\"size\": 12})\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41e65fc1-ae21-4d8d-ae45-b826d8729bbd",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f12fdecf-eb65-4364-b85e-6499379202f4",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}