{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Tuple, Union\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.animation import FuncAnimation\n",
    "\n",
    "import adaptive\n",
    "from adaptive.learner.learner1D import default_loss, uniform_loss\n",
    "\n",
    "\n",
    "def set_style() -> None:\n",
    "    import mplcyberpunk\n",
    "\n",
    "    plt.style.use(\"cyberpunk\")\n",
    "    mplcyberpunk.make_lines_glow()\n",
    "\n",
    "\n",
    "def f(x: float, offset: float = 0.07357338543088588) -> float:\n",
    "    a = 0.01\n",
    "    return x + a**2 / (a**2 + (x - offset) ** 2)\n",
    "\n",
    "\n",
    "def init_axes(\n",
    "    ax: plt.Axes, title: str, bounds: Tuple[float, float], data: dict\n",
    ") -> Tuple[plt.Line2D]:\n",
    "    margin = 0.05 * (bounds[1] - bounds[0])\n",
    "    ax.set_xlim(bounds[0] - margin, bounds[1] + margin)\n",
    "    y_min = min(data.values())\n",
    "    y_max = max(data.values())\n",
    "    ax.set_ylim(y_min - 0.1 * (y_max - y_min), y_max + 0.1 * (y_max - y_min))\n",
    "    ax.set_title(title)\n",
    "    return (ax.plot([], [], marker=\"o\", markersize=4, linewidth=1)[0],)\n",
    "\n",
    "\n",
    "def update(\n",
    "    frame: int,\n",
    "    learners: List[adaptive.Learner1D],\n",
    "    lines: List[plt.Line2D],\n",
    "    texts: List[plt.Text],\n",
    ") -> List[Union[plt.Line2D, plt.Text]]:\n",
    "    for learner, line, text in zip(learners, lines, texts):\n",
    "        learner_n = learner_till(frame, learner)\n",
    "        xs, ys = learner_n.to_numpy().T\n",
    "        line.set_data(xs, ys)\n",
    "        text.set_text(f\"Points: {len(xs)}\")\n",
    "    return lines + texts\n",
    "\n",
    "\n",
    "def plot_learners(\n",
    "    learners: List[adaptive.Learner1D], titles: List[str], npoints: int = 100\n",
    ") -> None:\n",
    "    set_style()\n",
    "    n_learners = len(learners)\n",
    "    fig, axes = plt.subplots(1, n_learners, figsize=(5 * n_learners, 5))\n",
    "    if n_learners == 1:\n",
    "        axes = [axes]\n",
    "\n",
    "    lines = []\n",
    "    texts = []\n",
    "\n",
    "    for ax, learner, title in zip(axes, learners, titles):\n",
    "        (line,) = init_axes(ax, title, learner.bounds, learner.data)\n",
    "        lines.append(line)\n",
    "        text = ax.text(0.05, 0.95, \"\", transform=ax.transAxes, verticalalignment=\"top\")\n",
    "        texts.append(text)\n",
    "\n",
    "    ani = FuncAnimation(\n",
    "        fig,\n",
    "        update,\n",
    "        frames=range(1, npoints + 1),\n",
    "        init_func=lambda: tuple(lines),\n",
    "        fargs=(learners, lines, texts),\n",
    "        blit=True,\n",
    "    )\n",
    "    ani.save(\"_vs_\".join(titles) + \"_animation.mp4\", writer=\"ffmpeg\", fps=10)\n",
    "\n",
    "\n",
    "def plot_learner(learner: adaptive.Learner1D, title: str, npoints: int = 100) -> None:\n",
    "    plot_learners([learner], [title], npoints)\n",
    "\n",
    "\n",
    "def compare_learners(\n",
    "    learner1: adaptive.Learner1D,\n",
    "    learner2: adaptive.Learner1D,\n",
    "    title1: str,\n",
    "    title2: str,\n",
    "    npoints: int = 100,\n",
    ") -> None:\n",
    "    plot_learners([learner1, learner2], [title1, title2], npoints)\n",
    "\n",
    "\n",
    "def learner_till(npoints: int, learner: adaptive.Learner1D) -> adaptive.Learner1D:\n",
    "    new_learner = adaptive.Learner1D(\n",
    "        None, bounds=learner.bounds, loss_per_interval=learner.loss_per_interval\n",
    "    )\n",
    "    data = list(learner.data.items())\n",
    "    new_learner.data = dict(data[:npoints])\n",
    "    return new_learner\n",
    "\n",
    "\n",
    "learner = adaptive.Learner1D(f, bounds=(-1, 1), loss_per_interval=default_loss)\n",
    "learner_uni = adaptive.Learner1D(f, bounds=(-1, 1), loss_per_interval=uniform_loss)\n",
    "\n",
    "adaptive.runner.simple(learner, npoints_goal=100)\n",
    "adaptive.runner.simple(learner_uni, npoints_goal=100)\n",
    "\n",
    "compare_learners(learner, learner_uni, \"Adaptive\", \"Uniform sampling\")\n",
    "plot_learner(learner, \"Learner1D\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2D single"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import functools\n",
    "\n",
    "import matplotlib.tri as mtri\n",
    "import numpy as np\n",
    "from matplotlib import animation\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib.animation import FFMpegWriter\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "import adaptive\n",
    "\n",
    "\n",
    "def set_style() -> None:\n",
    "    import mplcyberpunk\n",
    "\n",
    "    plt.style.use(\"cyberpunk\")\n",
    "    mplcyberpunk.make_lines_glow()\n",
    "\n",
    "\n",
    "def learner_till(till, learner, data):\n",
    "    new_learner = adaptive.Learner2D(None, bounds=learner.bounds)\n",
    "    new_learner.data = dict(data[:till])\n",
    "    for x, y in learner._bounds_points:\n",
    "        # always include the bounds\n",
    "        new_learner.tell((x, y), learner.data[x, y])\n",
    "    return new_learner\n",
    "\n",
    "\n",
    "def plot_tri(learner, ax):\n",
    "    tri = learner.interpolator(scaled=True).tri\n",
    "    triang = mtri.Triangulation(*tri.points.T, triangles=tri.simplices)\n",
    "    return ax.triplot(triang, c=\"k\", lw=0.8, alpha=0.8)\n",
    "\n",
    "\n",
    "def get_new_artists(npoints, learner, data, ax):\n",
    "    new_learner = learner_till(npoints, learner, data)\n",
    "    line1, line2 = plot_tri(new_learner, ax)\n",
    "    data = np.rot90(new_learner.interpolated_on_grid()[-1])\n",
    "    im = ax.imshow(data, extent=(-0.5, 0.5, -0.5, 0.5), cmap=\"plasma\")\n",
    "    return im, line1, line2\n",
    "\n",
    "\n",
    "@functools.lru_cache\n",
    "def create_and_run_learner():\n",
    "    def ring(xy):\n",
    "        import numpy as np\n",
    "\n",
    "        x, y = xy\n",
    "        a = 0.2\n",
    "        return x + np.exp(-((x**2 + y**2 - 0.75**2) ** 2) / a**4)\n",
    "\n",
    "    learner = adaptive.Learner2D(ring, bounds=[(-1, 1), (-1, 1)])\n",
    "    adaptive.runner.simple(learner, loss_goal=0.005)\n",
    "    return learner\n",
    "\n",
    "\n",
    "def get_figure():\n",
    "    fig, ax = plt.subplots(figsize=(5, 5))\n",
    "    fig.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=None, hspace=None)\n",
    "    ax.set_xticks([])\n",
    "    ax.set_yticks([])\n",
    "    ax.spines[\"top\"].set_visible(False)\n",
    "    ax.spines[\"right\"].set_visible(False)\n",
    "    ax.spines[\"bottom\"].set_visible(False)\n",
    "    ax.spines[\"left\"].set_visible(False)\n",
    "    return fig, ax\n",
    "\n",
    "\n",
    "def setup(nseconds=15):\n",
    "    learner = create_and_run_learner()\n",
    "    data = list(learner.data.items())\n",
    "    fig, ax = get_figure()\n",
    "    npoints = (len(data) * np.linspace(0, 1, 24 * nseconds) ** 2).astype(int)\n",
    "    return npoints, learner, data, fig, ax\n",
    "\n",
    "\n",
    "def animate_mp4(fname=\"source/_static/logo_docs.mp4\", nseconds=15):\n",
    "    set_style()\n",
    "    npoints, learner, data, fig, ax = setup(nseconds)\n",
    "    artists = [get_new_artists(n, learner, data, ax) for n in tqdm(npoints)]\n",
    "    ani = animation.ArtistAnimation(fig, artists, blit=True)\n",
    "    ani.save(fname, writer=FFMpegWriter(fps=24))\n",
    "\n",
    "\n",
    "animate_mp4(\"Learner2D.mp4\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2D compare"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import functools\n",
    "import itertools\n",
    "from typing import List, Tuple, Union\n",
    "\n",
    "import matplotlib.tri as mtri\n",
    "import numpy as np\n",
    "from matplotlib import animation\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib.animation import FFMpegWriter\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "import adaptive\n",
    "\n",
    "\n",
    "def learner_till(\n",
    "    till: int, learner: adaptive.Learner2D, data: List\n",
    ") -> adaptive.Learner2D:\n",
    "    new_learner = learner.new()\n",
    "    new_learner.data = dict(data[:till])\n",
    "    for x, y in learner._bounds_points:\n",
    "        new_learner.tell((x, y), learner.data[x, y])\n",
    "    return new_learner\n",
    "\n",
    "\n",
    "def plot_tri(learner: adaptive.Learner2D, ax: plt.Axes) -> plt.Line2D:\n",
    "    tri = learner.interpolator(scaled=True).tri\n",
    "    triang = mtri.Triangulation(*tri.points.T, triangles=tri.simplices)\n",
    "    return ax.triplot(triang, c=\"k\", lw=0.8, alpha=0.8)\n",
    "\n",
    "\n",
    "def plot_grid(n_grid: int, ax: plt.Axes) -> List[Union[plt.Line2D, plt.Line2D]]:\n",
    "    xs = ys = np.linspace(-0.5, 0.5, n_grid)\n",
    "    h_lines = [ax.axhline(y, color=\"k\", lw=0.8, alpha=0.8) for y in ys]\n",
    "    v_lines = [ax.axvline(x, color=\"k\", lw=0.8, alpha=0.8) for x in xs]\n",
    "    return h_lines + v_lines\n",
    "\n",
    "\n",
    "def plot_image(learner, ax):\n",
    "    data = np.rot90(learner.interpolated_on_grid()[-1])\n",
    "    return ax.imshow(data, extent=(-0.5, 0.5, -0.5, 0.5), cmap=\"plasma\")\n",
    "\n",
    "\n",
    "def get_new_artists_adaptive(\n",
    "    npoints: int,\n",
    "    learner: adaptive.Learner2D,\n",
    "    data: List,\n",
    "    ax: plt.Axes,\n",
    "    is_top_panel: bool,\n",
    ") -> Tuple:\n",
    "    if is_top_panel:\n",
    "        learner = learner_till(npoints, learner, data)\n",
    "\n",
    "    im = plot_image(learner, ax)\n",
    "    if not is_top_panel:\n",
    "        tri_lines = plot_tri(learner, ax)\n",
    "        return learner, im, *tri_lines\n",
    "\n",
    "    return learner, im\n",
    "\n",
    "\n",
    "def get_new_artists_uniform(\n",
    "    npoints: int,\n",
    "    learner: adaptive.Learner2D,\n",
    "    ax: plt.Axes,\n",
    "    is_top_panel: bool,\n",
    ") -> Tuple:\n",
    "    n_grid = int(npoints**0.5)\n",
    "    if is_top_panel:\n",
    "        xs = ys = np.linspace(*learner.bounds[0], n_grid)\n",
    "        xys = list(itertools.product(xs, ys))\n",
    "        learner = learner.new()\n",
    "        learner.tell_many(xys, map(learner.function, xys))\n",
    "\n",
    "    im = plot_image(learner, ax)\n",
    "    if not is_top_panel:\n",
    "        grid = plot_grid(n_grid, ax)\n",
    "        return learner, im, *grid\n",
    "\n",
    "    return learner, im\n",
    "\n",
    "\n",
    "@functools.lru_cache\n",
    "def create_and_run_learner() -> adaptive.Learner2D:\n",
    "    def ring(xy: Tuple[float, float]) -> float:\n",
    "        x, y = xy\n",
    "        a = 0.2\n",
    "        return x + np.exp(-((x**2 + y**2 - 0.75**2) ** 2) / a**4)\n",
    "\n",
    "    learner = adaptive.Learner2D(ring, bounds=[(-1, 1), (-1, 1)])\n",
    "    adaptive.runner.simple(learner, loss_goal=0.005)\n",
    "    return learner\n",
    "\n",
    "\n",
    "def get_figure() -> Tuple[plt.Figure, plt.Axes]:\n",
    "    fig, axs = plt.subplots(2, 2, figsize=(10, 10), sharex=True, sharey=True)\n",
    "    fig.subplots_adjust(left=0, bottom=0, right=1, top=0.95, wspace=0.1, hspace=0.1)\n",
    "\n",
    "    for ax in axs.flatten():\n",
    "        ax.set_xticks([])\n",
    "        ax.set_yticks([])\n",
    "        ax.spines[\"top\"].set_visible(False)\n",
    "        ax.spines[\"right\"].set_visible(False)\n",
    "        ax.spines[\"bottom\"].set_visible(False)\n",
    "        ax.spines[\"left\"].set_visible(False)\n",
    "\n",
    "    axs[0, 0].set_title(\"Adaptive sampling\", fontsize=12)\n",
    "    axs[0, 1].set_title(\"Uniform sampling\", fontsize=12)\n",
    "    return fig, axs\n",
    "\n",
    "\n",
    "def setup(\n",
    "    nseconds: int = 15,\n",
    ") -> Tuple[np.ndarray, adaptive.Learner2D, List, plt.Figure, plt.Axes]:\n",
    "    adaptive_learner = create_and_run_learner()\n",
    "    data = list(adaptive_learner.data.items())\n",
    "    fig, axs = get_figure()\n",
    "    npoints_raw = np.linspace(0, 1, 24 * nseconds) ** 2\n",
    "    npoints = (len(data) * npoints_raw).astype(int)\n",
    "    npoints += 4\n",
    "    return npoints, adaptive_learner, data, fig, axs\n",
    "\n",
    "\n",
    "def text_npoints(ax, n):\n",
    "    return ax.text(\n",
    "        0.05, 0.95, f\"{n} points\", transform=ax.transAxes, verticalalignment=\"top\"\n",
    "    )\n",
    "\n",
    "\n",
    "def animate_mp4(fname: str = \"source/_static/logo_docs.mp4\", nseconds: int = 15):\n",
    "    npoints, adaptive_learner, data, fig, axs = setup(nseconds)\n",
    "\n",
    "    artists = []\n",
    "    for n in tqdm(npoints):\n",
    "        adaptive_artists = []\n",
    "        uniform_artists = []\n",
    "\n",
    "        new_learner, adaptive_im_top = get_new_artists_adaptive(\n",
    "            n, adaptive_learner, data, axs[0, 0], True\n",
    "        )\n",
    "        adaptive_artists.append(adaptive_im_top)\n",
    "\n",
    "        _, *adaptive_learner_bottom = get_new_artists_adaptive(\n",
    "            n, new_learner, data, axs[1, 0], False\n",
    "        )\n",
    "        adaptive_artists.extend(adaptive_learner_bottom)\n",
    "\n",
    "        new_learner, uniform_im_top = get_new_artists_uniform(\n",
    "            n, adaptive_learner, axs[0, 1], True\n",
    "        )\n",
    "        uniform_artists.append(uniform_im_top)\n",
    "\n",
    "        _, *uniform_learner_bottom = get_new_artists_uniform(\n",
    "            n, new_learner, axs[1, 1], False\n",
    "        )\n",
    "        uniform_artists.extend(uniform_learner_bottom)\n",
    "\n",
    "        adaptive_text = text_npoints(axs[0, 0], n)\n",
    "        uniform_text = text_npoints(axs[0, 1], n)\n",
    "\n",
    "        adaptive_artists.append(adaptive_text)\n",
    "        uniform_artists.append(uniform_text)\n",
    "\n",
    "        artists.append(adaptive_artists + uniform_artists)\n",
    "\n",
    "    ani = animation.ArtistAnimation(fig, artists, blit=True)\n",
    "    ani.save(fname, writer=FFMpegWriter(fps=24))\n",
    "\n",
    "\n",
    "plt.style.use(\n",
    "    \"https://github.com/dhaitz/matplotlib-stylesheets/raw/master/pitayasmoothie-dark.mplstyle\"\n",
    ")\n",
    "animate_mp4(\"Learner2D.mp4\")"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}