Search code examples
pythonpandasanimationseaborncelluloid

Celluloid Animated Heatmap Issue with Axes and Missing Plots


I am trying to build an animated heatmap using celluloid. The x & y axis and color scale are the same but my code returns the weird output below.

enter image description here

My code uses seaborn, numpy, pandas, and celluloid and is simplified below:

from celluloid import Camera

## Set up celluloid
fig = plt.figure(figsize=(12, 9))
camera = Camera(fig)


## Loop to create figures
for item in range(len(df)):
   row = df.iloc[item]
   row = np.array(list(row))

   ## Create df from row
   shape = (8,12)
   df_row = pd.DataFrame(row.reshape(shape))

   ## Build seaborn heatmap
   ax = sns.heatmap(df_row, cmap="Greys", annot=False, vmin=0, vmax=1)
   ax.set_title(item)
   ax.xaxis.tick_top()
   for tick in ax.get_yticklabels():
      tick.set_rotation(0)
   
   ## Snap Celluloid
   camera.snap()

anim = camera.animate(interval=500)
anim.save("animation.mp4")

Solution

  • The problem is that seaborn constantly creates a new colorbar. To solve it, a fixed ax for the colorbar needs to be created at the start of the code.

    Here is the general setup, using celluloid's Camera. If you leave out cbar_ax=cbar_ax you'll see the strange behavior of an endless caravan of colorbars.

    import matplotlib.pyplot as plt
    import seaborn as sns
    import numpy as np
    from celluloid import Camera
    
    fig, (ax, cbar_ax) = plt.subplots(ncols=2, figsize=(12, 9), gridspec_kw={'width_ratios': [10, 1]})
    camera = Camera(fig)
    
    for _ in range(20):
        sns.heatmap(np.random.rand(8, 12), cmap="magma", annot=False, vmin=0, vmax=1,
                    ax=ax, cbar_ax=cbar_ax)
        ax.xaxis.tick_top()
        ax.tick_params(axis='y', labelrotation=0)
        camera.snap()
    
    anim = camera.animate(interval=500)
    anim.save("animation.mp4")
    

    The critical changes to your code would be:

    • replace fig = plt.figure(...) by fig, (ax, cbar_ax) = plt.subplots(...)
    • call sns.heatmap with ax=ax, cbar_ax=cbar_ax