In [ ]:

    

%matplotlib inline

# Load local version of yabox
import sys
sys.path.insert(0, '../')

from yabox import DE, PDE
import numpy as np

# Imports required for 3d animations
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import animation, rcParams
from IPython.display import HTML


    

In [ ]:

    

# There is a bug in matplotlib that prevents the use of ffmpeg and avconv!
# https://github.com/matplotlib/matplotlib/pull/8743

avail_writers = matplotlib.animation.writers.list()
    
# Functions to generate 3d animations

def display_animation(anim):
    plt.close(anim._fig)
    return HTML(anim.to_html5_video())


def generate_video(problem, algorithm, figsize=(12, 8), frames=100, interval=100):
    # Try to use tqdm to show progress
    use_tqdm = True
    try:
        from tqdm.auto import tqdm
    except:
        use_tqdm = False
    
    figure = plt.figure(figsize=figsize)
    ax = Axes3D(figure)
    problem.plot3d(ax3d=ax)
    minz = min(ax.get_zlim())
    it = algorithm.geniterator()
    if use_tqdm:
        it = iter(tqdm(it, total=frames))

    def animate(i):
        ax.clear()
        ax.autoscale(enable=False)
        problem.plot3d(ax3d=ax)
        status = next(it)
        population = status.population
        P = algorithm.denormalize(population)
        fitness = status.fitness
        idx = status.best_idx
        PT = P.T
        # Individuals
        ax.scatter(PT[0], PT[1], fitness, s=30, c='#930039', marker='o', depthshade=False, zorder=999)
        # Shadow projections
        ax.scatter(PT[0], PT[1], np.full_like(PT[0], minz), alpha=0.5, s=50, c='black', marker='o', 
                   edgecolors='none', depthshade=False, zorder=999)
        
    anim = animation.FuncAnimation(figure, animate, frames=frames, interval=interval, blit=False)
    return anim


    

In [ ]:

    

from yabox.problems import Ackley, Levy
problem = Ackley()


    

In [ ]:

    

algorithm = DE(problem, problem.bounds)
anim = generate_video(problem, algorithm, figsize=(12, 8), frames=10, interval=100)
display_animation(anim)