{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Custom Animations\n",
    "\n",
    "These examples show some basic functions of PySwarming, like target, aggregation, and repulsion.\n",
    "\n",
    "For more examples, you can see the [notebooks](https://github.com/mrsonandrade/pyswarming/tree/main/notebooks) and [examples](https://github.com/mrsonandrade/pyswarming/tree/main/examples) directories of PySwarming."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Target behavior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "051ef3c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "# importing the swarming behaviors\n",
    "import pyswarming.behaviors as ps\n",
    "\n",
    "# importing numpy to work with arrays\n",
    "import numpy as np\n",
    "\n",
    "# importing matplotlib to plot the animation\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation\n",
    "\n",
    "# importing functools.partial to use in the animation\n",
    "from functools import partial\n",
    "\n",
    "# define each robot (x, y, z) position\n",
    "robot_poses = np.asarray([0., 0., 0.])\n",
    "\n",
    "# set the robot linear velocity\n",
    "robot_speed = 0.025\n",
    "\n",
    "# define a target (x, y, z) position\n",
    "swarm_target = np.asarray([8., 8., 0.])\n",
    "\n",
    "# First set up the figure, the axis, and the plot element we want to animate\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.set_xlim([-2,10])\n",
    "ax.set_ylim([-2,10])\n",
    "ax.set_xlabel('X(m)')\n",
    "ax.set_ylabel('Y(m)')\n",
    "ax.grid()\n",
    "ax.set_aspect('equal')\n",
    "ax.set_title('Target behavior')\n",
    "\n",
    "robot, = ax.plot([], [], marker='o', lw=0)\n",
    "target = ax.plot(swarm_target[0], swarm_target[1], marker='x', color='red', lw=0)\n",
    "\n",
    "# initialization function: plot the background of each frame\n",
    "def init():\n",
    "    robot.set_data([], [])\n",
    "    return (robot,)\n",
    "\n",
    "# animation function. This is called sequentially\n",
    "def animate(i, robot_poses):\n",
    "    r_i = robot_poses\n",
    "    r_i += robot_speed*ps.target(r_i, swarm_target)\n",
    "    robot.set_data(r_i[0], r_i[1])\n",
    "    return (robot,)\n",
    "\n",
    "# call the animator. blit=True means only re-draw the parts that \n",
    "# have changed.\n",
    "anim = animation.FuncAnimation(fig, partial(animate, robot_poses=robot_poses), init_func=init,\n",
    "                               frames=480, interval=1, blit=True)\n",
    "\n",
    "from IPython.display import HTML\n",
    "HTML(anim.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35bcb062",
   "metadata": {},
   "source": [
    "## Aggregation behavior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32e10bb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "# importing pyswarming behaviors\n",
    "import pyswarming.behaviors as ps\n",
    "\n",
    "# importing numpy to work with arrays\n",
    "import numpy as np\n",
    "\n",
    "# importing matplotlib to plot the animation\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation\n",
    "\n",
    "# importing functools.partial to use in the animation\n",
    "from functools import partial\n",
    "\n",
    "# define each robot (x, y, z) position\n",
    "robot_poses = np.asarray([[8., 8., 0.],\n",
    "                          [-8., 8., 0.],\n",
    "                          [8., -8., 0.],\n",
    "                          [-8., -8., 0.]])\n",
    "\n",
    "# set the robot linear velocity\n",
    "robot_speed = 0.025\n",
    "\n",
    "# First set up the figure, the axis, and the plot element we want to animate\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.set_xlim([-10,10])\n",
    "ax.set_ylim([-10,10])\n",
    "ax.set_xlabel('X(m)')\n",
    "ax.set_ylabel('Y(m)')\n",
    "ax.grid()\n",
    "ax.set_aspect('equal')\n",
    "ax.set_title('Aggregation behavior')\n",
    "\n",
    "robot1, = ax.plot([], [], marker='o', lw=0)\n",
    "robot2, = ax.plot([], [], marker='o', lw=0)\n",
    "robot3, = ax.plot([], [], marker='o', lw=0)\n",
    "robot4, = ax.plot([], [], marker='o', lw=0)\n",
    "\n",
    "# initialization function: plot the background of each frame\n",
    "def init():\n",
    "    robot1.set_data([], [])\n",
    "    robot2.set_data([], [])\n",
    "    robot3.set_data([], [])\n",
    "    robot4.set_data([], [])\n",
    "    return (robot1,robot2,robot3,robot4,)\n",
    "\n",
    "# animation function. This is called sequentially\n",
    "def animate(i, robot_poses):\n",
    "    for r_ind in range(len(robot_poses)):\n",
    "        r_i = robot_poses[r_ind]\n",
    "        r_j = np.delete(robot_poses, np.array([r_ind]), axis=0)\n",
    "        robot_poses[r_ind] += robot_speed*ps.aggregation(r_i, r_j)\n",
    "    robot1.set_data(robot_poses[0][0], robot_poses[0][1])\n",
    "    robot2.set_data(robot_poses[1][0], robot_poses[1][1])\n",
    "    robot3.set_data(robot_poses[2][0], robot_poses[2][1])\n",
    "    robot4.set_data(robot_poses[3][0], robot_poses[3][1])\n",
    "    return (robot1,robot2,robot3,robot4,)\n",
    "\n",
    "# call the animator. blit=True means only re-draw the parts that \n",
    "# have changed.\n",
    "anim = animation.FuncAnimation(fig, partial(animate, robot_poses=robot_poses), init_func=init,\n",
    "                               frames=480, interval=1, blit=True)\n",
    "\n",
    "from IPython.display import HTML\n",
    "HTML(anim.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d6c573c",
   "metadata": {},
   "source": [
    "## Repulsion behavior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "014abcee",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "# importing pyswarming behaviors\n",
    "import pyswarming.behaviors as ps\n",
    "\n",
    "# importing numpy to work with arrays\n",
    "import numpy as np\n",
    "\n",
    "# importing matplotlib to plot the animation\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation\n",
    "\n",
    "# importing functools.partial to use in the animation\n",
    "from functools import partial\n",
    "\n",
    "# define each robot (x, y, z) position\n",
    "robot_poses = np.asarray([[1., 1., 0.],\n",
    "                          [-1., 1., 0.],\n",
    "                          [1., -1., 0.],\n",
    "                          [-1., -1., 0.]])\n",
    "\n",
    "# set the robot linear velocity\n",
    "robot_speed = 0.025\n",
    "\n",
    "# First set up the figure, the axis, and the plot element we want to animate\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.set_xlim([-10,10])\n",
    "ax.set_ylim([-10,10])\n",
    "ax.set_xlabel('X(m)')\n",
    "ax.set_ylabel('Y(m)')\n",
    "ax.grid()\n",
    "ax.set_aspect('equal')\n",
    "ax.set_title('Repulsion behavior')\n",
    "\n",
    "robot1, = ax.plot([], [], marker='o', lw=0)\n",
    "robot2, = ax.plot([], [], marker='o', lw=0)\n",
    "robot3, = ax.plot([], [], marker='o', lw=0)\n",
    "robot4, = ax.plot([], [], marker='o', lw=0)\n",
    "\n",
    "# initialization function: plot the background of each frame\n",
    "def init():\n",
    "    robot1.set_data([], [])\n",
    "    robot2.set_data([], [])\n",
    "    robot3.set_data([], [])\n",
    "    robot4.set_data([], [])\n",
    "    return (robot1,robot2,robot3,robot4,)\n",
    "\n",
    "# animation function. This is called sequentially\n",
    "def animate(i, robot_poses):\n",
    "    for r_ind in range(len(robot_poses)):\n",
    "        r_i = robot_poses[r_ind]\n",
    "        r_j = np.delete(robot_poses, np.array([r_ind]), axis=0)\n",
    "        robot_poses[r_ind] += robot_speed*ps.repulsion(r_i, r_j, 5.0)\n",
    "    robot1.set_data(robot_poses[0][0], robot_poses[0][1])\n",
    "    robot2.set_data(robot_poses[1][0], robot_poses[1][1])\n",
    "    robot3.set_data(robot_poses[2][0], robot_poses[2][1])\n",
    "    robot4.set_data(robot_poses[3][0], robot_poses[3][1])\n",
    "    return (robot1,robot2,robot3,robot4,)\n",
    "\n",
    "# call the animator. blit=True means only re-draw the parts that \n",
    "# have changed.\n",
    "anim = animation.FuncAnimation(fig, partial(animate, robot_poses=robot_poses), init_func=init,\n",
    "                               frames=480, interval=1, blit=True)\n",
    "\n",
    "from IPython.display import HTML\n",
    "HTML(anim.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50eca133",
   "metadata": {},
   "source": [
    "## Aggregation + Repulsion behaviors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "983c7cd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "# importing pyswarming behaviors\n",
    "import pyswarming.behaviors as ps\n",
    "\n",
    "# importing numpy to work with arrays\n",
    "import numpy as np\n",
    "\n",
    "# importing matplotlib to plot the animation\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation\n",
    "\n",
    "# importing functools.partial to use in the animation\n",
    "from functools import partial\n",
    "\n",
    "# define each robot (x, y, z) position\n",
    "robot_poses = np.asarray([[8., 8., 0.],\n",
    "                          [-8., 8., 0.],\n",
    "                          [8., -8., 0.],\n",
    "                          [-8., -8., 0.]])\n",
    "\n",
    "# set the robot linear velocity\n",
    "robot_speed = 0.025\n",
    "\n",
    "# First set up the figure, the axis, and the plot element we want to animate\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.set_xlim([-10,10])\n",
    "ax.set_ylim([-10,10])\n",
    "ax.set_xlabel('X(m)')\n",
    "ax.set_ylabel('Y(m)')\n",
    "ax.grid()\n",
    "ax.set_aspect('equal')\n",
    "ax.set_title('Aggregation + Repulsion behaviors')\n",
    "\n",
    "robot1, = ax.plot([], [], marker='o', lw=0)\n",
    "robot2, = ax.plot([], [], marker='o', lw=0)\n",
    "robot3, = ax.plot([], [], marker='o', lw=0)\n",
    "robot4, = ax.plot([], [], marker='o', lw=0)\n",
    "\n",
    "# initialization function: plot the background of each frame\n",
    "def init():\n",
    "    robot1.set_data([], [])\n",
    "    robot2.set_data([], [])\n",
    "    robot3.set_data([], [])\n",
    "    robot4.set_data([], [])\n",
    "    return (robot1,robot2,robot3,robot4,)\n",
    "\n",
    "# animation function. This is called sequentially\n",
    "def animate(i, robot_poses):\n",
    "    for r_ind in range(len(robot_poses)):\n",
    "        r_i = robot_poses[r_ind]\n",
    "        r_j = np.delete(robot_poses, np.array([r_ind]), axis=0)\n",
    "        robot_poses[r_ind] += robot_speed*(ps.aggregation(r_i, r_j) + ps.repulsion(r_i, r_j, 5.0))\n",
    "    robot1.set_data(robot_poses[0][0], robot_poses[0][1])\n",
    "    robot2.set_data(robot_poses[1][0], robot_poses[1][1])\n",
    "    robot3.set_data(robot_poses[2][0], robot_poses[2][1])\n",
    "    robot4.set_data(robot_poses[3][0], robot_poses[3][1])\n",
    "    return (robot1,robot2,robot3,robot4,)\n",
    "\n",
    "# call the animator. blit=True means only re-draw the parts that \n",
    "# have changed.\n",
    "anim = animation.FuncAnimation(fig, partial(animate, robot_poses=robot_poses), init_func=init,\n",
    "                               frames=720, interval=1, blit=True)\n",
    "\n",
    "from IPython.display import HTML\n",
    "HTML(anim.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b3ce0a4",
   "metadata": {},
   "source": [
    "## Area Coverage behavior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fca4ef1",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "# importing pyswarming behaviors\n",
    "import pyswarming.behaviors as ps\n",
    "\n",
    "# importing numpy to work with arrays\n",
    "import numpy as np\n",
    "\n",
    "# importing matplotlib to plot the animation\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation\n",
    "\n",
    "# importing functools.partial to use in the animation\n",
    "from functools import partial\n",
    "\n",
    "# define each robot (x, y, z) position\n",
    "robot_poses = np.asarray([[36., 35., 0.],\n",
    "                          [35., 36., 0.],\n",
    "                          [36., 36., 0.],\n",
    "                          [35., 35., 0.]])\n",
    "\n",
    "# set the robot linear velocity\n",
    "robot_speed = 0.20\n",
    "\n",
    "# First set up the figure, the axis, and the plot element we want to animate\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.set_xlim([-10,50])\n",
    "ax.set_ylim([-10,50])\n",
    "ax.set_xlabel('X(m)')\n",
    "ax.set_ylabel('Y(m)')\n",
    "ax.grid()\n",
    "ax.set_aspect('equal')\n",
    "ax.set_title('Area Coverage')\n",
    "\n",
    "robot1, = ax.plot([], [], marker='o', lw=0)\n",
    "robot2, = ax.plot([], [], marker='o', lw=0)\n",
    "robot3, = ax.plot([], [], marker='o', lw=0)\n",
    "robot4, = ax.plot([], [], marker='o', lw=0)\n",
    "\n",
    "# sphere function: used as region to be filled\n",
    "sphere = lambda x: x[0]**2 + x[1]**2 + x[2]**2 - 4.0\n",
    "\n",
    "# initialization function: plot the background of each frame\n",
    "def init():\n",
    "    robot1.set_data([], [])\n",
    "    robot2.set_data([], [])\n",
    "    robot3.set_data([], [])\n",
    "    robot4.set_data([], [])\n",
    "    return (robot1,robot2,robot3,robot4,)\n",
    "\n",
    "# animation function. This is called sequentially\n",
    "def animate(i, robot_poses):\n",
    "    for r_ind in range(len(robot_poses)):\n",
    "        r_i = robot_poses[r_ind]\n",
    "        r_j = np.delete(robot_poses, np.array([r_ind]), axis=0)\n",
    "        robot_poses[r_ind] += robot_speed*ps.area_coverage(r_i, r_j, sphere, 3.0, 3)\n",
    "    robot1.set_data(robot_poses[0][0], robot_poses[0][1])\n",
    "    robot2.set_data(robot_poses[1][0], robot_poses[1][1])\n",
    "    robot3.set_data(robot_poses[2][0], robot_poses[2][1])\n",
    "    robot4.set_data(robot_poses[3][0], robot_poses[3][1])\n",
    "    return (robot1,robot2,robot3,robot4,)\n",
    "\n",
    "# call the animator. blit=True means only re-draw the parts that \n",
    "# have changed.\n",
    "anim = animation.FuncAnimation(fig, partial(animate, robot_poses=robot_poses), init_func=init,\n",
    "                               frames=480, interval=1, blit=True)\n",
    "\n",
    "from IPython.display import HTML\n",
    "HTML(anim.to_jshtml())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7932e07f",
   "metadata": {},
   "source": [
    "## Collective Navigation behavior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77393de9",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "# importing pyswarming behaviors\n",
    "import pyswarming.behaviors as ps\n",
    "\n",
    "# importing numpy to work with arrays\n",
    "import numpy as np\n",
    "\n",
    "# importing matplotlib to plot the animation\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.animation as animation\n",
    "\n",
    "# importing functools.partial to use in the animation\n",
    "from functools import partial\n",
    "\n",
    "# define each robot (x, y, z) position\n",
    "robot_poses = np.asarray([[1., 1., 0.],\n",
    "                [-1., 1., 0.],\n",
    "                [1., -1., 0.],\n",
    "                [-1., -1., 0.]])\n",
    "\n",
    "# set the robot linear velocity\n",
    "robot_speed = 0.10\n",
    "\n",
    "# define a target (x, y, z) position\n",
    "swarm_target = np.asarray([35., 35., 0.])\n",
    "\n",
    "# First set up the figure, the axis, and the plot element we want to animate\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.set_xlim([-10,50])\n",
    "ax.set_ylim([-10,50])\n",
    "ax.set_xlabel('X(m)')\n",
    "ax.set_ylabel('Y(m)')\n",
    "ax.grid()\n",
    "ax.set_aspect('equal')\n",
    "ax.set_title('Collective Navigation')\n",
    "\n",
    "robot1, = ax.plot([], [], marker='o', lw=0)\n",
    "robot2, = ax.plot([], [], marker='o', lw=0)\n",
    "robot3, = ax.plot([], [], marker='o', lw=0)\n",
    "robot4, = ax.plot([], [], marker='o', lw=0)\n",
    "\n",
    "target = ax.plot(swarm_target[0], swarm_target[1], marker='x', color='red', lw=0)\n",
    "\n",
    "# initialization function: plot the background of each frame\n",
    "def init():\n",
    "    robot1.set_data([], [])\n",
    "    robot2.set_data([], [])\n",
    "    robot3.set_data([], [])\n",
    "    robot4.set_data([], [])\n",
    "    return (robot1,robot2,robot3,robot4,)\n",
    "\n",
    "# animation function. This is called sequentially\n",
    "def animate(i, robot_poses):\n",
    "    for r_ind in range(len(robot_poses)):\n",
    "        r_i = robot_poses[r_ind]\n",
    "        r_j = np.delete(robot_poses, np.array([r_ind]), axis=0)\n",
    "        robot_poses[r_ind] += robot_speed*ps.collective_navigation(r_i, r_j, swarm_target, 5.0)\n",
    "    robot1.set_data(robot_poses[0][0], robot_poses[0][1])\n",
    "    robot2.set_data(robot_poses[1][0], robot_poses[1][1])\n",
    "    robot3.set_data(robot_poses[2][0], robot_poses[2][1])\n",
    "    robot4.set_data(robot_poses[3][0], robot_poses[3][1])\n",
    "    return (robot1,robot2,robot3,robot4,)\n",
    "\n",
    "# call the animator. blit=True means only re-draw the parts that \n",
    "# have changed.\n",
    "anim = animation.FuncAnimation(fig, partial(animate, robot_poses=robot_poses), init_func=init,\n",
    "                               frames=960, interval=1, blit=True)\n",
    "\n",
    "from IPython.display import HTML\n",
    "HTML(anim.to_jshtml())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.0 64-bit",
   "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.0"
  },
  "vscode": {
   "interpreter": {
    "hash": "7e1998ff7f8aa20ada591c520b972326324e5ea05489af9e422744c7c09f6dad"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}