figures.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Adaptive paper figures"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## imports and function definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import functools\n",
    "import itertools\n",
    "import os\n",
    "import pickle\n",
    "\n",
    "import holoviews.plotting.mpl\n",
    "import matplotlib\n",
    "\n",
    "matplotlib.use(\"agg\")\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.tri as mtri\n",
    "import numpy as np\n",
    "from scipy import interpolate\n",
    "\n",
    "import adaptive\n",
    "\n",
    "\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'svg'\n",
    "\n",
    "golden_mean = (np.sqrt(5) - 1) / 2  # Aesthetic ratio\n",
    "fig_width_pt = 246.0  # Columnwidth\n",
    "inches_per_pt = 1 / 72.27  # Convert pt to inches\n",
    "fig_width = fig_width_pt * inches_per_pt\n",
    "fig_height = fig_width * golden_mean  # height in inches\n",
    "fig_size = [fig_width, fig_height]\n",
    "\n",
    "params = {\n",
    "    \"backend\": \"ps\",\n",
    "    \"axes.labelsize\": 13,\n",
    "    \"font.size\": 13,\n",
    "    \"legend.fontsize\": 10,\n",
    "    \"xtick.labelsize\": 10,\n",
    "    \"ytick.labelsize\": 10,\n",
    "    \"text.usetex\": True,\n",
    "    \"figure.figsize\": fig_size,\n",
    "    \"font.family\": \"serif\",\n",
    "    \"font.serif\": \"Computer Modern Roman\",\n",
    "    \"legend.frameon\": True,\n",
    "    \"savefig.dpi\": 300,\n",
    "}\n",
    "\n",
    "plt.rcParams.update(params)\n",
    "plt.rc(\"text.latex\", preamble=[r\"\\usepackage{xfrac}\", r\"\\usepackage{siunitx}\"])\n",
    "\n",
    "\n",
    "class HistogramNormalize(matplotlib.colors.Normalize):\n",
    "    def __init__(self, data, vmin=None, vmax=None, mixing_degree=1):\n",
    "        self.mixing_degree = mixing_degree\n",
    "        if vmin is not None:\n",
    "            data = data[data >= vmin]\n",
    "        if vmax is not None:\n",
    "            data = data[data <= vmax]\n",
    "\n",
    "        self.sorted_data = np.sort(data.flatten())\n",
    "        matplotlib.colors.Normalize.__init__(self, vmin, vmax)\n",
    "\n",
    "    def __call__(self, value, clip=None):\n",
    "        hist_norm = np.ma.masked_array(\n",
    "            np.searchsorted(self.sorted_data, value) / len(self.sorted_data)\n",
    "        )\n",
    "        linear_norm = super().__call__(value, clip)\n",
    "        return self.mixing_degree * hist_norm + (1 - self.mixing_degree) * linear_norm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1D funcs\n",
    "def f(x, offset=0.123):\n",
    "    a = 0.02\n",
    "    return x + a ** 2 / (a ** 2 + (x - offset) ** 2)\n",
    "\n",
    "\n",
    "def g(x):\n",
    "    return np.tanh(x * 40)\n",
    "\n",
    "\n",
    "def h(x):\n",
    "    return np.sin(100 * x) * np.exp(-x ** 2 / 0.1 ** 2)\n",
    "\n",
    "\n",
    "funcs_1D = [\n",
    "    dict(function=f, bounds=(-1, 1), title=\"peak\"),\n",
    "    dict(function=g, bounds=(-1, 1), title=\"tanh\"),\n",
    "    dict(function=h, bounds=(-0.3, 0.3), title=\"wave packet\"),\n",
    "]\n",
    "\n",
    "# 2D funcs\n",
    "def ring(xy, offset=0.123):\n",
    "    a = 0.2\n",
    "    x, y = xy\n",
    "    return x + np.exp(-(x ** 2 + y ** 2 - 0.75 ** 2) ** 2 / a ** 4)\n",
    "\n",
    "\n",
    "@functools.lru_cache()\n",
    "def phase_diagram_setup(fname):\n",
    "    data = adaptive.utils.load(fname)\n",
    "    points = np.array(list(data.keys()))\n",
    "    values = np.array(list(data.values()), dtype=float)\n",
    "    bounds = [\n",
    "        (points[:, 0].min(), points[:, 0].max()),\n",
    "        (points[:, 1].min(), points[:, 1].max()),\n",
    "    ]\n",
    "    ll, ur = np.reshape(bounds, (2, 2)).T\n",
    "    inds = np.all(np.logical_and(ll <= points, points <= ur), axis=1)\n",
    "    points, values = points[inds], values[inds].reshape(-1, 1)\n",
    "    return interpolate.LinearNDInterpolator(points, values), bounds\n",
    "\n",
    "\n",
    "def phase_diagram(xy, fname):\n",
    "    ip, _ = phase_diagram_setup(fname)\n",
    "    zs = ip(xy)\n",
    "    return np.round(zs)\n",
    "\n",
    "\n",
    "def density(x, eps=0):\n",
    "    e = [0.8, 0.2]\n",
    "    delta = [0.5, 0.5, 0.5]\n",
    "    c = 3\n",
    "    omega = [0.02, 0.05]\n",
    "\n",
    "    H = np.array(\n",
    "        [\n",
    "            [e[0] + 1j * omega[0], delta[0], delta[1]],\n",
    "            [delta[0], e[1] + c * x + 1j * omega[1], delta[1]],\n",
    "            [delta[1], delta[2], e[1] - c * x + 1j * omega[1]],\n",
    "        ]\n",
    "    )\n",
    "    H += np.eye(3) * eps\n",
    "    return np.trace(np.linalg.inv(H)).imag\n",
    "\n",
    "\n",
    "def level_crossing(xy):\n",
    "    x, y = xy\n",
    "    return density(x, y) + y\n",
    "\n",
    "\n",
    "funcs_2D = [\n",
    "    dict(function=ring, bounds=[(-1, 1), (-1, 1)], npoints=33, title=\"ring\"),\n",
    "    dict(\n",
    "        function=functools.partial(phase_diagram, fname=\"phase_diagram.pickle\"),\n",
    "        bounds=phase_diagram_setup(\"phase_diagram.pickle\")[1],\n",
    "        npoints=100,\n",
    "        fname=\"phase_diagram.pickle\",\n",
    "        title=\"phase diagram\",\n",
    "    ),\n",
    "    dict(function=level_crossing, bounds=[(-1, 1), (-3, 3)], npoints=50, title=\"level crossing\"),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interval and loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(1)\n",
    "xs = np.array([0.1, 0.3, 0.35, 0.45])\n",
    "f = lambda x: x ** 3\n",
    "ys = f(xs)\n",
    "means = lambda x: np.convolve(x, np.ones(2) / 2, mode=\"valid\")\n",
    "xs_means = means(xs)\n",
    "ys_means = means(ys)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=fig_size)\n",
    "ax.scatter(xs, ys, c=\"k\")\n",
    "ax.plot(xs, ys, c=\"k\")\n",
    "\n",
    "ax.annotate(\n",
    "    s=r\"$L_{1,2} = \\sqrt{\\Delta x^2 + \\Delta y^2}$\",\n",
    "    xy=(np.mean([xs[0], xs[1]]), np.mean([ys[0], ys[1]])),\n",
    "    xytext=(xs[0] + 0.05, ys[0] - 0.05),\n",
    "    arrowprops=dict(arrowstyle=\"->\"),\n",
    "    ha=\"center\",\n",
    "    zorder=10,\n",
    ")\n",
    "\n",
    "for i, (x, y) in enumerate(zip(xs, ys)):\n",
    "    sign = [1, -1][i % 2]\n",
    "    ax.annotate(\n",
    "        s=fr\"$x_{i+1}, y_{i+1}$\",\n",
    "        xy=(x, y),\n",
    "        xytext=(x + 0.01, y + sign * 0.04),\n",
    "        arrowprops=dict(arrowstyle=\"->\"),\n",
    "        ha=\"center\",\n",
    "    )\n",
    "\n",
    "ax.scatter(xs, ys, c=\"green\", s=5, zorder=5, label=\"existing data\")\n",
    "losses = np.hypot(xs[1:] - xs[:-1], ys[1:] - ys[:-1])\n",
    "ax.scatter(\n",
    "    xs_means, ys_means, c=\"red\", s=300 * losses, zorder=8, label=\"candidate points\"\n",
    ")\n",
    "xs_dense = np.linspace(xs[0], xs[-1], 400)\n",
    "ax.plot(xs_dense, f(xs_dense), alpha=0.3, zorder=7, label=\"function\")\n",
    "\n",
    "ax.legend()\n",
    "ax.axis(\"off\")\n",
    "plt.savefig(\"figures/loss_1D.pdf\", bbox_inches=\"tight\", transparent=True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learner1D vs grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(2, len(funcs_1D), figsize=(fig_width, 1.5 * fig_height))\n",
    "n_points = 50\n",
    "for i, ax in enumerate(axs.T.flatten()):\n",
    "    ax.xaxis.set_ticks([])\n",
    "    ax.yaxis.set_ticks([])\n",
    "    kind = \"homogeneous\" if i % 2 == 0 else \"adaptive\"\n",
    "    index = i // 2 if kind == \"homogeneous\" else (i - 1) // 2\n",
    "    d = funcs_1D[index]\n",
    "    bounds = d[\"bounds\"]\n",
    "    f = d[\"function\"]\n",
    "\n",
    "    if kind == \"homogeneous\":\n",
    "        xs = np.linspace(*bounds, n_points)\n",
    "        ys = f(xs)\n",
    "        ax.set_title(rf\"\\textrm{{{d['title']}}}\")\n",
    "    elif kind == \"adaptive\":\n",
    "        loss = adaptive.learner.learner1D.curvature_loss_function()\n",
    "        learner = adaptive.Learner1D(f, bounds=bounds, loss_per_interval=loss)\n",
    "        adaptive.runner.simple(learner, goal=lambda l: l.npoints >= n_points)\n",
    "        xs, ys = zip(*sorted(learner.data.items()))\n",
    "    xs_dense = np.linspace(*bounds, 1000)\n",
    "    ax.plot(xs_dense, f(xs_dense), c=\"red\", alpha=0.3, lw=0.5)\n",
    "    ax.scatter(xs, ys, s=0.5, c=\"k\")\n",
    "\n",
    "axs[0][0].set_ylabel(r\"$\\textrm{homogeneous}$\")\n",
    "axs[1][0].set_ylabel(r\"$\\textrm{adaptive}$\")\n",
    "plt.savefig(\"figures/Learner1D.pdf\", bbox_inches=\"tight\", transparent=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learner2D vs grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(2, len(funcs_2D), figsize=(fig_width, 1.5 * fig_height))\n",
    "\n",
    "plt.subplots_adjust(hspace=-0.1, wspace=0.1)\n",
    "\n",
    "with_tri = False\n",
    "\n",
    "for i, ax in enumerate(axs.T.flatten()):\n",
    "    label = \"abcdef\"[i]\n",
    "    ax.text(\n",
    "        0.5,\n",
    "        1.05,\n",
    "        f\"$\\mathrm{{({label})}}$\",\n",
    "        transform=ax.transAxes,\n",
    "        horizontalalignment=\"center\",\n",
    "        verticalalignment=\"bottom\",\n",
    "    )\n",
    "    ax.xaxis.set_ticks([])\n",
    "    ax.yaxis.set_ticks([])\n",
    "    kind = \"homogeneous\" if i % 2 == 0 else \"adaptive\"\n",
    "    d = funcs_2D[i // 2] if kind == \"homogeneous\" else funcs_2D[(i - 1) // 2]\n",
    "    bounds = d[\"bounds\"]\n",
    "    npoints = d[\"npoints\"]\n",
    "    f = d[\"function\"]\n",
    "    fname = d.get(\"fname\")\n",
    "    if fname is not None:\n",
    "        f = functools.partial(f, fname=fname)\n",
    "\n",
    "    if kind == \"homogeneous\":\n",
    "        xs, ys = [np.linspace(*bound, npoints) for bound in bounds]\n",
    "        data = {xy: f(xy) for xy in itertools.product(xs, ys)}\n",
    "        learner = adaptive.Learner2D(f, bounds=bounds)\n",
    "        learner.data = data\n",
    "        d[\"learner_hom\"] = learner\n",
    "    elif kind == \"adaptive\":\n",
    "        learner = adaptive.Learner2D(f, bounds=bounds)\n",
    "        if fname is not None:\n",
    "            learner.load(fname)\n",
    "        learner.data = {\n",
    "            k: v for i, (k, v) in enumerate(learner.data.items()) if i <= npoints ** 2\n",
    "        }\n",
    "        adaptive.runner.simple(learner, goal=lambda l: l.npoints >= npoints ** 2)\n",
    "        d[\"learner\"] = learner\n",
    "\n",
    "    if with_tri:\n",
    "        tri = learner.ip().tri\n",
    "        triang = mtri.Triangulation(*tri.points.T, triangles=tri.vertices)\n",
    "        ax.triplot(triang, c=\"w\", lw=0.2, alpha=0.8)\n",
    "\n",
    "    values = np.array(list(learner.data.values()))\n",
    "    plot_data = learner.plot(npoints if kind == \"homogeneous\" else None).Image.I.data\n",
    "    im = ax.imshow(\n",
    "        plot_data,\n",
    "        extent=(-0.5, 0.5, -0.5, 0.5),\n",
    "        interpolation=\"none\",\n",
    "    )\n",
    "    ax.set_xticks([])\n",
    "    ax.set_yticks([])\n",
    "\n",
    "    if i in [2, 3]:\n",
    "        norm = HistogramNormalize(plot_data, mixing_degree=0.6)\n",
    "        im.set_norm(norm)\n",
    "\n",
    "axs[0][0].set_ylabel(r\"$\\textrm{homogeneous}$\")\n",
    "axs[1][0].set_ylabel(r\"$\\textrm{adaptive}$\")\n",
    "\n",
    "plt.savefig(\"figures/Learner2D.pdf\", bbox_inches=\"tight\", transparent=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "learner = adaptive.Learner1D(funcs_1D[0][\"function\"], bounds=funcs_1D[0][\"bounds\"])\n",
    "times = []\n",
    "for i in range(10000):\n",
    "    t_start = time.time()\n",
    "    points, _ = learner.ask(1)\n",
    "    t_end = time.time()\n",
    "    times.append(t_end - t_start)\n",
    "    learner.tell(points[0], learner.function(points[0]))\n",
    "plt.plot(np.cumsum(times))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Algorithm explaination: Learner1D step-by-step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 1, figsize=(fig_width, 3 * fig_height))\n",
    "\n",
    "\n",
    "def f(x, offset=0.123):\n",
    "    a = 0.2\n",
    "    return a ** 2 / (a ** 2 + (x - offset) ** 2)\n",
    "\n",
    "\n",
    "learner = adaptive.Learner1D(\n",
    "    f, (0, 2), loss_per_interval=adaptive.learner.learner1D.curvature_loss_function()\n",
    ")\n",
    "learner._recompute_losses_factor = 0.1\n",
    "xs_dense = np.linspace(*learner.bounds, 400)\n",
    "ys_dense = f(xs_dense)\n",
    "step = 0.4\n",
    "\n",
    "for i in range(11):\n",
    "    offset = -i * step\n",
    "\n",
    "    x = learner.ask(1)[0][0]\n",
    "    y = f(x)\n",
    "    learner.tell(x, y)\n",
    "    xs, ys = map(np.array, zip(*sorted(learner.data.items())))\n",
    "    ys = ys + offset\n",
    "    if i >= 1:\n",
    "        axs.plot(xs_dense, ys_dense + offset, c=\"k\", alpha=0.3, zorder=0)\n",
    "        axs.plot(xs, ys, zorder=1, c=\"k\")\n",
    "        axs.scatter(xs, ys, alpha=1, zorder=2, c=\"k\")\n",
    "        (x_left, x_right), loss = list(learner.losses.items())[\n",
    "            0\n",
    "        ]  # it's an ItemSortedDict\n",
    "        (y_left, y_right) = [\n",
    "            learner.data[x_left] + offset,\n",
    "            learner.data[x_right] + offset,\n",
    "        ]\n",
    "        axs.scatter([x_left, x_right], [y_left, y_right], c=\"r\", s=10, zorder=3)\n",
    "        x_mid = np.mean((x_left, x_right))\n",
    "        y_mid = np.interp(x_mid, (x_left, x_right), (y_left, y_right))\n",
    "        axs.scatter(x_mid, y_mid, zorder=4, marker=\"x\", c=\"green\")\n",
    "\n",
    "axs.text(\n",
    "    -0.1,\n",
    "    0.5,\n",
    "    (r\"$\\mathrm{iterations}$\" + \"\\n\" + \"$\\longleftarrow$\"),\n",
    "    transform=axs.transAxes,\n",
    "    horizontalalignment=\"center\",\n",
    "    verticalalignment=\"center\",\n",
    "    rotation=90,\n",
    "    fontsize=18,\n",
    ")\n",
    "\n",
    "\n",
    "# legend\n",
    "\n",
    "import matplotlib.patches as mpatches\n",
    "import matplotlib.lines as mlines\n",
    "\n",
    "\n",
    "class LargestInterval:\n",
    "    pass\n",
    "\n",
    "\n",
    "class Interval:\n",
    "    pass\n",
    "\n",
    "\n",
    "class Function:\n",
    "    pass\n",
    "\n",
    "\n",
    "class IntervalHandler:\n",
    "    def __init__(self, with_inner=True, length=20, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.with_inner = with_inner\n",
    "        self.length = length\n",
    "\n",
    "    def legend_artist(self, legend, orig_handle, fontsize, handlebox):\n",
    "        x0, y0 = handlebox.xdescent, handlebox.ydescent\n",
    "        offsets = [0, self.length]\n",
    "        line = mlines.Line2D((0, offsets[-1]), (0, 0), zorder=0, c=\"k\")\n",
    "        handlebox.add_artist(line)\n",
    "\n",
    "        for offset in offsets:\n",
    "            circle1 = mpatches.Circle(\n",
    "                [x0 + offset, y0], 4, facecolor=\"k\", lw=3, zorder=1\n",
    "            )\n",
    "            handlebox.add_artist(circle1)\n",
    "            if self.with_inner:\n",
    "                circle2 = mpatches.Circle(\n",
    "                    [x0 + offset, y0], 3, facecolor=\"red\", lw=3, zorder=1\n",
    "                )\n",
    "                handlebox.add_artist(circle2)\n",
    "\n",
    "\n",
    "class FunctionHandler:\n",
    "    def __init__(self, xs, ys, *args, **kwargs):\n",
    "        super().__init__(*args, **kwargs)\n",
    "        self.xs = xs / xs.ptp() * 20\n",
    "        self.ys = ys - ys.mean()\n",
    "\n",
    "    def legend_artist(self, legend, orig_handle, fontsize, handlebox):\n",
    "        x0, y0 = handlebox.xdescent, handlebox.ydescent\n",
    "\n",
    "        line = mlines.Line2D(self.xs, self.ys * 10, zorder=0, c=\"k\", alpha=0.3)\n",
    "\n",
    "        handlebox.add_artist(line)\n",
    "\n",
    "\n",
    "plt.legend(\n",
    "    [\n",
    "        Function(),\n",
    "        mlines.Line2D([], [], marker=\"o\", lw=0, c=\"k\"),\n",
    "        LargestInterval(),\n",
    "        Interval(),\n",
    "        mlines.Line2D([], [], marker=\"x\", lw=0, c=\"green\"),\n",
    "    ],\n",
    "    [\n",
    "        \"original function\",\n",
    "        \"known point\",\n",
    "        \"interval\",\n",
    "        \"largest loss interval\",\n",
    "        \"next candidate point\",\n",
    "    ],\n",
    "    handler_map={\n",
    "        LargestInterval: IntervalHandler(False),\n",
    "        Interval: IntervalHandler(True),\n",
    "        Function: FunctionHandler(xs, ys),\n",
    "    },\n",
    "    bbox_to_anchor=(0.25, 0.9, 1.0, 0.0),\n",
    "    ncol=1,\n",
    ")\n",
    "\n",
    "# On grid, uncomment for homogeneous plot\n",
    "# axs.plot(learner.bounds, [-(i + 0.5) * step, -(i + 0.5) * step], c='k', ls='--')\n",
    "# xs_hom = np.linspace(*learner.bounds, i)\n",
    "# ys_hom = f(xs_hom) - (i + 3) * step\n",
    "# axs.plot(xs_hom, ys_hom, zorder=1, c=\"k\")\n",
    "# axs.scatter(xs_hom, ys_hom, alpha=1, zorder=2, c=\"k\")\n",
    "\n",
    "axs.axis(\"off\")\n",
    "plt.savefig(\"figures/algo.pdf\", bbox_inches=\"tight\", transparent=True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Line loss visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.patches import Polygon\n",
    "\n",
    "fig, axs = plt.subplots(5, 1, figsize=(fig_width, 1.5 * fig_height))\n",
    "f = lambda x: np.sin(np.array(x)) ** 2\n",
    "xs = np.array([0, 1.3, 3, 5, 7, 8])\n",
    "ys = f(xs)\n",
    "\n",
    "\n",
    "def plot(xs, ax):\n",
    "    ys = f(xs)\n",
    "    xs_dense = np.linspace(xs[0], xs[-1], 300)\n",
    "    ys_dense = f(xs_dense)\n",
    "    ax.plot(xs_dense, ys_dense, alpha=0.3, c=\"k\")\n",
    "    ax.plot(xs, ys, c=\"k\")\n",
    "    ax.scatter(xs, ys, zorder=10, s=14, c=\"k\")\n",
    "\n",
    "\n",
    "plot(xs, axs[0])\n",
    "plot(xs, axs[1])\n",
    "\n",
    "for i, ax in enumerate(axs):\n",
    "    ax.axis(\"off\")\n",
    "    ax.set_ylim(-1.5, 1.5)\n",
    "    label = \"abcdefgh\"[i]\n",
    "    ax.text(\n",
    "        0.5,\n",
    "        0.9,\n",
    "        f\"$\\mathrm{{({label})}}$\",\n",
    "        transform=ax.transAxes,\n",
    "        horizontalalignment=\"center\",\n",
    "        verticalalignment=\"bottom\",\n",
    "    )\n",
    "\n",
    "\n",
    "def plot_tri(xs, ax, colors):\n",
    "    ys = f(xs)\n",
    "    for i in range(len(xs)):\n",
    "        if i == 0 or i == len(xs) - 1:\n",
    "            continue\n",
    "        verts = [(xs[i - 1], ys[i - 1]), (xs[i], ys[i]), (xs[i + 1], ys[i + 1])]\n",
    "        poly = Polygon(verts, facecolor=colors[xs[i]], alpha=0.4)\n",
    "        ax.add_patch(poly)\n",
    "        ax.scatter([xs[i]], [ys[i]], c=colors[xs[i]], s=6, zorder=11)\n",
    "        ax.plot(\n",
    "            [xs[i - 1], xs[i + 1]], [ys[i - 1], ys[i + 1]], c=\"k\", ls=\"--\", alpha=0.3\n",
    "        )\n",
    "\n",
    "\n",
    "learner = adaptive.Learner1D(\n",
    "    f,\n",
    "    (xs[0], xs[-1]),\n",
    "    loss_per_interval=adaptive.learner.learner1D.curvature_loss_function(),\n",
    ")\n",
    "learner.tell_many(xs, ys)\n",
    "\n",
    "for i, ax in enumerate(axs[1:]):\n",
    "    if i != 0:\n",
    "        x_new = learner.ask(1)[0][0]\n",
    "        learner.tell(x_new, f(x_new))\n",
    "        xs, ys = zip(*sorted(learner.data.items()))\n",
    "    plot(xs, ax)\n",
    "    colors = {x: f\"C{i}\" for i, x in enumerate(learner.data.keys())}\n",
    "    plot_tri(xs, ax, colors=colors)\n",
    "\n",
    "plt.savefig(\"figures/line_loss.pdf\", bbox_inches=\"tight\", transparent=True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Line loss L1-error(N)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import collections\n",
    "from adaptive import Learner1D, LearnerND\n",
    "from scipy import interpolate\n",
    "\n",
    "\n",
    "def err_1D(xs, ys, xs_rand, f):\n",
    "    ip = interpolate.interp1d(xs, ys)\n",
    "    abserr = np.abs(ip(xs_rand) - f(xs_rand))\n",
    "    return np.average(abserr ** 2) ** 0.5\n",
    "\n",
    "\n",
    "def get_err_1D(learner, N):\n",
    "    xs_rand = np.random.uniform(*learner.bounds, size=100_000)\n",
    "\n",
    "    xs = np.linspace(*learner.bounds, N)\n",
    "    ys = learner.function(xs)\n",
    "    err_lin = err_1D(xs, ys, xs_rand, learner.function)\n",
    "\n",
    "    xs, ys = zip(*learner.data.items())\n",
    "    xs, ys = xs[:N], ys[:N]\n",
    "    ip = interpolate.interp1d(xs, ys)\n",
    "    err_learner = err_1D(xs, ys, xs_rand, learner.function)\n",
    "    return err_lin, err_learner\n",
    "\n",
    "\n",
    "def err_2D(zs, zs_rand):\n",
    "    abserr = np.abs(zs - zs_rand)\n",
    "    return np.average(abserr ** 2) ** 0.5\n",
    "\n",
    "\n",
    "def get_err_2D(learner, N):\n",
    "    xys_rand = np.vstack(\n",
    "        [\n",
    "            np.random.uniform(*learner.bounds[0], size=int(100_000)),\n",
    "            np.random.uniform(*learner.bounds[1], size=int(100_000)),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    xys_hom = np.array(\n",
    "        [\n",
    "            (x, y)\n",
    "            for x in np.linspace(*learner.bounds[0], int(N ** 0.5))\n",
    "            for y in np.linspace(*learner.bounds[1], int(N ** 0.5))\n",
    "        ]\n",
    "    )\n",
    "    N = len(xys_hom)\n",
    "\n",
    "    try:\n",
    "        # Vectorized\n",
    "        zs_hom = learner.function(xys_hom.T)\n",
    "        zs_rand = learner.function(xys_rand)\n",
    "    except:\n",
    "        # Non-vectorized\n",
    "        zs_hom = np.array([learner.function(xy) for xy in xys_hom])\n",
    "        zs_rand = np.array([learner.function(xy) for xy in xys_rand.T])\n",
    "\n",
    "    ip = interpolate.LinearNDInterpolator(xys_hom, zs_hom)\n",
    "    zs = ip(xys_rand.T)\n",
    "    err_lin = err_2D(zs, zs_rand)\n",
    "\n",
    "    ip = interpolate.LinearNDInterpolator(learner.points[:N], learner.values[:N])\n",
    "    zs = ip(xys_rand.T)\n",
    "    err_learner = err_2D(zs, zs_rand)\n",
    "\n",
    "    return err_lin, err_learner\n",
    "\n",
    "\n",
    "recalculate = False\n",
    "N_max = 10_000\n",
    "Ns = np.geomspace(4, N_max, 50).astype(int)\n",
    "fname = \"error_line_loss.pickle\"\n",
    "\n",
    "if not os.path.exists(fname) and not recalculate:\n",
    "    loss_1D = adaptive.learner.learner1D.curvature_loss_function()\n",
    "    loss_2D = adaptive.learner.learnerND.curvature_loss_function()\n",
    "\n",
    "    err = collections.defaultdict(dict)\n",
    "    for i, (funcs, loss, Learner, get_err) in enumerate(\n",
    "        zip(\n",
    "            [funcs_1D, funcs_2D],\n",
    "            [loss_1D, loss_2D],\n",
    "            [Learner1D, LearnerND],\n",
    "            [get_err_1D, get_err_2D],\n",
    "        )\n",
    "    ):\n",
    "        for d in funcs:\n",
    "            learner = Learner(d[\"function\"], d[\"bounds\"], loss)\n",
    "            adaptive.runner.simple(learner, goal=lambda l: l.npoints >= N_max)\n",
    "            errs = [get_err(learner, N) for N in Ns]\n",
    "            err_hom, err_adaptive = zip(*errs)\n",
    "            err[i][d[\"title\"]] = (err_hom, err_adaptive)\n",
    "\n",
    "    with open(fname, \"wb\") as f:\n",
    "        pickle.dump(err, f)\n",
    "else:\n",
    "    with open(fname, \"rb\") as f:\n",
    "        err = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(2, 1, figsize=(fig_width, 1.6 * fig_height))\n",
    "plt.subplots_adjust(hspace=0.3)\n",
    "\n",
    "axs[1].set_xlabel(\"$N$\")\n",
    "\n",
    "for i, ax in enumerate(axs):\n",
    "    ax.set_ylabel(r\"$\\text{Err}_{1}(\\tilde{f})$\")\n",
    "\n",
    "    for j, (title, (err_hom, err_adaptive)) in enumerate(err[i].items()):\n",
    "        color = f\"C{j}\"\n",
    "        label = \"abc\"[j]\n",
    "        label = f\"$\\mathrm{{({label})}}$ {title}\"\n",
    "        ax.loglog(Ns, err_hom, ls=\"--\", c=color)\n",
    "        ax.loglog(Ns, err_adaptive, label=label, c=color)\n",
    "#         error = np.array(err_hom) / np.array(err_adaptive)\n",
    "#         if i == 0:\n",
    "#             ax.loglog(Ns[:36], error[:36], c=color, label=label)\n",
    "#         else:\n",
    "#             ax.loglog(Ns, error, c=color, label=label)\n",
    "        ax.legend()\n",
    "\n",
    "plt.savefig(\"figures/line_loss_error.pdf\", bbox_inches=\"tight\", transparent=True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Error reduction\n",
    "print(\"1D\")\n",
    "for title, (err_hom, err_adaptive) in err[0].items():\n",
    "    print(title, err_hom[-1] / err_adaptive[-1])\n",
    "\n",
    "print(\"\\n2D\")\n",
    "for title, (err_hom, err_adaptive) in err[1].items():\n",
    "    print(title, err_hom[-1] / err_adaptive[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## iso-line plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import lru_cache\n",
    "import numpy as np\n",
    "import scipy.linalg\n",
    "import scipy.spatial\n",
    "import kwant\n",
    "\n",
    "\n",
    "@lru_cache()\n",
    "def create_syst(unit_cell):\n",
    "    lat = kwant.lattice.Polyatomic(unit_cell, [(0, 0, 0)])\n",
    "    syst = kwant.Builder(kwant.TranslationalSymmetry(*lat.prim_vecs))\n",
    "    syst[lat.shape(lambda _: True, (0, 0, 0))] = 6\n",
    "    syst[lat.neighbors()] = -1\n",
    "    return kwant.wraparound.wraparound(syst).finalized()\n",
    "\n",
    "\n",
    "def get_brillouin_zone(unit_cell):\n",
    "    syst = create_syst(unit_cell)\n",
    "    A = get_A(syst)\n",
    "    neighbours = kwant.linalg.lll.voronoi(A)\n",
    "    lattice_points = np.concatenate(([[0, 0, 0]], neighbours))\n",
    "    lattice_points = 2 * np.pi * (lattice_points @ A.T)\n",
    "    vor = scipy.spatial.Voronoi(lattice_points)\n",
    "    brillouin_zone = vor.vertices[vor.regions[vor.point_region[0]]]\n",
    "    return scipy.spatial.ConvexHull(brillouin_zone)\n",
    "\n",
    "\n",
    "def momentum_to_lattice(k, syst):\n",
    "    A = get_A(syst)\n",
    "    k, residuals = scipy.linalg.lstsq(A, k)[:2]\n",
    "    if np.any(abs(residuals) > 1e-7):\n",
    "        raise RuntimeError(\n",
    "            \"Requested momentum doesn't correspond to any lattice momentum.\"\n",
    "        )\n",
    "    return k\n",
    "\n",
    "\n",
    "def get_A(syst):\n",
    "    B = np.asarray(syst._wrapped_symmetry.periods).T\n",
    "    return np.linalg.pinv(B).T\n",
    "\n",
    "\n",
    "def energies(k, unit_cell):\n",
    "    syst = create_syst(unit_cell)\n",
    "    k_x, k_y, k_z = momentum_to_lattice(k, syst)\n",
    "    params = {\"k_x\": k_x, \"k_y\": k_y, \"k_z\": k_z}\n",
    "    H = syst.hamiltonian_submatrix(params=params)\n",
    "    energies = np.linalg.eigvalsh(H)\n",
    "    return min(energies)\n",
    "\n",
    "\n",
    "from functools import partial\n",
    "\n",
    "from ipywidgets import interact_manual\n",
    "\n",
    "\n",
    "# Define the lattice vectors of some common unit cells\n",
    "lattices = dict(\n",
    "    hexegonal=((0, 1, 0), (np.cos(-np.pi / 6), np.sin(-np.pi / 6), 0), (0, 0, 1)),\n",
    "    simple_cubic=((1, 0, 0), (0, 1, 0), (0, 0, 1)),\n",
    "    fcc=((0, 0.5, 0.5), (0.5, 0.5, 0), (0.5, 0, 0.5)),\n",
    "    bcc=((-0.5, 0.5, 0.5), (0.5, -0.5, 0.5), (0.5, 0.5, -0.5)),\n",
    ")\n",
    "\n",
    "\n",
    "from adaptive.learner.learnerND import volume\n",
    "\n",
    "\n",
    "def isoline_loss_function(level, priority):\n",
    "    def loss(simplex, values, value_scale):\n",
    "        values = np.array(values)\n",
    "        which_side = np.sign(level * value_scale - values)\n",
    "        crosses_isoline = np.any(np.diff(which_side))\n",
    "        return volume(simplex)* (1 + priority * crosses_isoline)\n",
    "    return loss\n",
    "\n",
    "level = 0.1\n",
    "loss_per_simplex = isoline_loss_function(level, 0.5)\n",
    "\n",
    "learners = []\n",
    "fnames = []\n",
    "for name, unit_cell in lattices.items():\n",
    "    hull = get_brillouin_zone(unit_cell)\n",
    "    learner = adaptive.LearnerND(\n",
    "        partial(energies, unit_cell=unit_cell), hull, loss_per_simplex=loss_per_simplex\n",
    "    )\n",
    "    fnames.append(name)\n",
    "    learners.append(learner)\n",
    "\n",
    "mapping = dict(zip(fnames, learners))\n",
    "\n",
    "learner = adaptive.BalancingLearner(learners, strategy=\"npoints\")\n",
    "\n",
    "\n",
    "def select(name, learners, fnames):\n",
    "    return learners[fnames.index(name)]\n",
    "\n",
    "\n",
    "def iso(unit_cell, level=8.5):\n",
    "    l = select(unit_cell, learners, fnames)\n",
    "    adaptive.runner.simple(l, goal=lambda l: l.npoints > 1000)\n",
    "    return l.plot_isosurface(level=level)\n",
    "\n",
    "\n",
    "interact_manual(iso, level=(-6, 9, 0.1), unit_cell=lattices.keys())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(xy):\n",
    "    x, y = xy\n",
    "    return x ** 2 + y ** 3\n",
    "\n",
    "\n",
    "bounds = [(-1, 1), (-1, 1)]\n",
    "npoints = 12 ** 2\n",
    "\n",
    "\n",
    "level = 0.1\n",
    "loss = isoline_loss_function(level, priority=50)\n",
    "\n",
    "with_tri = True\n",
    "fig, axs = plt.subplots(1, 2, figsize=(1.5*fig_width, 1.5*fig_height))\n",
    "plt.subplots_adjust(wspace=0.3)\n",
    "\n",
    "for i, ax in enumerate(axs.flatten()):\n",
    "    learner = adaptive.LearnerND(f, bounds, loss_per_simplex=loss)\n",
    "    label = \"ab\"[i]\n",
    "    ax.text(\n",
    "        0.5,\n",
    "        1.05,\n",
    "        f\"$\\mathrm{{({label})}}$\",\n",
    "        transform=ax.transAxes,\n",
    "        horizontalalignment=\"center\",\n",
    "        verticalalignment=\"bottom\",\n",
    "    )\n",
    "    ax.xaxis.set_ticks([])\n",
    "    ax.yaxis.set_ticks([])\n",
    "    kind = \"homogeneous\" if i == 0 else \"adaptive\"\n",
    "\n",
    "    if kind == \"homogeneous\":\n",
    "        xs, ys = [np.linspace(*bound, int(npoints ** 0.5)) for bound in bounds]\n",
    "        data = {xy: f(xy) for xy in itertools.product(xs, ys)}\n",
    "        learner.data = data\n",
    "    elif kind == \"adaptive\":\n",
    "        learner.data = {\n",
    "            k: v for i, (k, v) in enumerate(learner.data.items()) if i <= npoints\n",
    "        }\n",
    "        adaptive.runner.simple(learner, goal=lambda l: l.npoints >= npoints)\n",
    "\n",
    "    if with_tri:\n",
    "        xy = np.array([learner.tri.get_vertices(s) for s in learner.tri.simplices])\n",
    "        print(f\"{len(xy)} triangles\")\n",
    "        triang = mtri.Triangulation(*xy.reshape(-1, 2).T)\n",
    "        ax.triplot(triang, c=\"w\", lw=0.2, alpha=0.8)\n",
    "\n",
    "    # Isolines\n",
    "    vertices, lines = learner._get_iso(level, which=\"line\")\n",
    "    paths = np.array([[vertices[i], vertices[j]] for i, j in lines]).T\n",
    "    print(\"{} line segments\".format(len(paths.T)))\n",
    "    ax.plot(*paths, c=\"k\", lw=0.7)\n",
    "\n",
    "    values = np.array(list(learner.data.values()))\n",
    "    ax.imshow(\n",
    "        learner.plot(npoints if kind == \"homogeneous\" else None).Image.I.data,\n",
    "        extent=(-1, 1, -1, 1),\n",
    "        interpolation=\"none\",\n",
    "    )\n",
    "    ax.set_xticks([])\n",
    "    ax.set_yticks([])\n",
    "\n",
    "axs[0].set_ylabel(r\"$\\textrm{homogeneous}$\")\n",
    "axs[1].set_ylabel(r\"$\\textrm{adaptive}$\")\n",
    "plt.savefig(\"figures/isoline.pdf\", bbox_inches=\"tight\", transparent=True)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}