{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rössler attractor\n",
    "See https://en.wikipedia.org/wiki/R%C3%B6ssler_attractor\n",
    "\\begin{cases}  \\frac{dx}{dt} = -y - z \\\\ \\frac{dy}{dt} = x + ay \\\\ \\frac{dz}{dt} = b + z(x-c) \\end{cases}\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib ipympl\n",
    "import ipywidgets as widgets\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from scipy.integrate import solve_ivp\n",
    "from functools import lru_cache\n",
    "\n",
    "import mpl_interactions.ipyplot as iplt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define function to plot\n",
    "\n",
    "### Projecting on axes\n",
    "\n",
    "The Rossler attractor is a 3 dimensional system, but as 3D plots are not yet supported by `mpl_interactions` we will only visualize the `x` and `y` components.\n",
    "\n",
    "**Note:** Matplotlib supports 3D plots, but `mpl_interactions` does not yet support them. That makes this a great place to contribute to `mpl_interactions` if you're interested in doing so. If you want to have a crack at it feel free to comment on https://github.com/ianhi/mpl-interactions/issues/89 and `@ianhi` will be happy to help you through the process.\n",
    "\n",
    "\n",
    "### Caching\n",
    "One thing to note here is that `mpl_interactions` will cache function calls for a given set of parameters so that the same function isn't called multiple times if you are plotting it on multiple axes. However, that cache will not persist as the parameters are modified. So here we  build in our own cache to speed up execution\n",
    "\n",
    "### kwarg collisions\n",
    "\n",
    "\n",
    "We can't use the `c` argument to `f` as `c` is reserved by `plot` (and `scatter` and other functions) by matplotlib in order to control the colors of the plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_span = [0, 500]\n",
    "t_eval = np.linspace(0, 500, 1550)\n",
    "x0 = 0\n",
    "y0 = 0\n",
    "z0 = 0\n",
    "\n",
    "\n",
    "cache = {}\n",
    "def f(a, b, c_):\n",
    "    def deriv(t, cur_pos):\n",
    "        x, y, z = cur_pos\n",
    "        dxdt = -y - z\n",
    "        dydt = x + a * y\n",
    "        dzdt = b + z * (x - c_)\n",
    "        return [dxdt, dydt, dzdt]\n",
    "\n",
    "    id_ = (float(a), float(b), float(c_))\n",
    "    if id_ not in cache:\n",
    "        out = solve_ivp(deriv, t_span, y0=[x0, y0, z0], t_eval=t_eval).y[:2]\n",
    "        cache[id_] = out\n",
    "    else:\n",
    "        out = cache[id_]\n",
    "    return out.T  # requires shape (N, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "controls = iplt.plot(\n",
    "    f,\n",
    "    \".-\",\n",
    "    a=(0.05, 0.3, 1000),\n",
    "    b=0.2,\n",
    "    c_=(1, 20), # we can't use `c` because that is a kwarg for matplotlib that controls color\n",
    "    parametric=True,\n",
    "    alpha=0.5,\n",
    "    play_buttons=True,\n",
    "    play_button_pos=\"left\",\n",
    "    ylim=\"auto\",\n",
    "    xlim=\"auto\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Coloring by time point\n",
    "\n",
    "When we plot using `plot` we can't choose colors for individual points, so we can use the `scatter` function to color the points by the time point they have.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use a different argument for c because `c` is an argument to plt.scatter\n",
    "out = widgets.Output()\n",
    "display(out)\n",
    "\n",
    "\n",
    "def f(a, b, c_):\n",
    "    def deriv(t, cur_pos):\n",
    "        x, y, z = cur_pos\n",
    "        dxdt = -y - z\n",
    "        dydt = x + a * y\n",
    "        dzdt = b + z * (x - c_)\n",
    "        return [dxdt, dydt, dzdt]\n",
    "\n",
    "    id_ = (float(a), float(b), float(c_))\n",
    "    if id_ not in cache:\n",
    "        out = solve_ivp(deriv, t_span, y0=[0, 1, 0], t_eval=t_eval).y[:2]\n",
    "        cache[id_] = out\n",
    "    else:\n",
    "        out = cache[id_]\n",
    "    return out.T  # requires shape (N, 2)\n",
    "\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "controls = iplt.scatter(\n",
    "    f,\n",
    "    a=(0.05, 0.3, 1000),\n",
    "    b=0.2,\n",
    "    c_=(1, 20),\n",
    "    parametric=True,\n",
    "    alpha=0.5,\n",
    "    play_buttons=True,\n",
    "    play_button_pos=\"left\",\n",
    "    s=8,\n",
    "    c=t_eval,\n",
    ")\n",
    "controls = iplt.plot(\n",
    "    f,\n",
    "    \"-\",\n",
    "    controls=controls,\n",
    "    parametric=True,\n",
    "    alpha=0.5,\n",
    "    ylim=\"auto\",\n",
    "    xlim=\"auto\",\n",
    ")\n",
    "plt.colorbar().set_label(\"Time Point\")\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}