animate_learner2d.py

import itertools
import os

import adaptive
import holoviews.plotting.mpl
import matplotlib
import matplotlib.cm
import matplotlib.tri as mtri
import numpy as np
from matplotlib import animation
from matplotlib import pyplot as plt
from matplotlib.animation import FFMpegWriter
from tqdm import tqdm


class HistogramNormalize(matplotlib.colors.Normalize):
    def __init__(self, data, vmin=None, vmax=None, mixing_degree=1):
        self.mixing_degree = mixing_degree
        if vmin is not None:
            data = data[data >= vmin]
        if vmax is not None:
            data = data[data <= vmax]

        self.sorted_data = np.sort(data.flatten())
        matplotlib.colors.Normalize.__init__(self, vmin, vmax)

    def __call__(self, value, clip=None):
        hist_norm = np.ma.masked_array(
            np.searchsorted(self.sorted_data, value) / len(self.sorted_data)
        )
        linear_norm = super().__call__(value, clip)
        return self.mixing_degree * hist_norm + (1 - self.mixing_degree) * linear_norm


def learner_till(till, learner, data):
    new_learner = adaptive.Learner2D(None, bounds=learner.bounds)
    new_learner.data = {k: v for k, v in data[:till]}
    for x, y in learner._bounds_points:
        # always include the bounds
        new_learner.tell((x, y), learner.data[x, y])
    return new_learner


def n_grid_points(learner):
    from adaptive.learner.learner2D import areas
    from math import sqrt

    ip = learner.ip()
    # Calculate how many grid points are needed.
    # factor from A=√3/4 * a² (equilateral triangle)
    n = int(0.658 / sqrt(areas(ip).min()))
    return min(max(n, 10), 1000)


def gridded_learner_data(learner):
    n = n_grid_points(learner)
    return learner.plot(n).Image.I.data


def plot_tri(learner, ax):
    tri = learner.ip().tri
    triang = mtri.Triangulation(*tri.points.T, triangles=tri.vertices)
    return ax.triplot(triang, c="k", lw=0.2, alpha=0.8)


def get_new_artists(npoints, learner, data):
    new_learner = learner_till(npoints, learner, data)

    line1, line2 = plot_tri(new_learner, ax)

    data = gridded_learner_data(new_learner)
    norm = HistogramNormalize(data, mixing_degree=0.6)
    cmap = get_cmap(next(cycle))
    im = ax.imshow(data, norm=norm, extent=(-0.5, 0.5, -0.5, 0.5), cmap=cmap)
    return im, line1, line2


N = 128
cycle = itertools.cycle(range(N))


def get_cmap(roll, cmap=matplotlib.cm.inferno):
    import matplotlib.colors as mcolors

    to_roll = (
        np.linspace(0, 1, N // 2).tolist()[:-1]
        + np.linspace(0, 1, N // 2).tolist()[::-1]
    )
    colors = cmap(np.roll(to_roll, roll))
    rolled_cmap = mcolors.LinearSegmentedColormap.from_list("rolled_cmap", colors)
    return rolled_cmap


def bounds_from_saved_learner(fname):
    learner = adaptive.Learner2D(None, [(-1, 1), (-1, 1)])
    learner.load(fname)
    xs, ys = np.array(list(learner.data.keys())).T
    bounds = [(xs.min(), xs.max()), (ys.min(), ys.max())]
    return bounds


def learner_from_file(fname):
    if not os.path.exists(fname):
        raise Exception(
            "Download the data at https://www.dropbox.com/s/6y6mz6w0enahzhp/data.pickle?dl=0"
        )

    bounds = bounds_from_saved_learner(fname)
    learner = adaptive.Learner2D(None, bounds)
    learner.load(fname)
    return learner


if __name__ == "__main__":
    # check the following link to speedup the saving
    # https://stackoverflow.com/questions/30965355/speedup-matplotlib-animation-to-video-file
    learner = learner_from_file("data.pickle")

    data = list(learner.data.items())

    fig, ax = plt.subplots(figsize=(5, 5))
    fig.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=None, hspace=None)

    ax.set_xticks([])
    ax.set_yticks([])

    nseconds = 15
    npoints = (len(data) * np.linspace(0, 1, 24 * nseconds) ** 2).astype(int)

    artists = [get_new_artists(n, learner, data) for n in tqdm(npoints)]

    ani = animation.ArtistAnimation(fig, artists, blit=True)

    metadata = dict(
        artist="Bas Nijholt",
        title="Learning a Majorana phase diagram",
        genre="Quantum computing",
        subject="python-adaptive/adaptive",
        copyright="Bas Nijholt",
        comment="see http://adaptive.readthedocs.io",
    )

    writer = FFMpegWriter(fps=24, metadata=metadata, bitrate=2 * 1800)
    ani.save("movie.mp4", writer=writer, dpi=200)