Updating legend entry using imshow in Python3

Question

I'm attempting to add a legend to overlay an imshow() plot displaying an animated array of random numbers. I want the legend to update to display the step that we are viewing.

I attempted to follow the steps here , which shows how to create an animated legend for subplots() using FuncAnimation. I believe the only way to display animated arrays is using ArtistAnimation() and imshow(), but one or both of these is causing me an issue to follow the linked solution.

I've attached below the working code to generate the animated random array, with the legend solution (from link) double commented out.

Any help or advice to remedy would be enormously appreciated.

Thanks, C

import matplotlib.animation as animation
from matplotlib import colors
import matplotlib.pyplot as plt
import numpy as np

N=20
steps = 100
interval_pause = 100
repeat_pause = 1000

cmap = colors.ListedColormap(['white', 'black'])
bounds=[-1,0,1]
norm = colors.BoundaryNorm(bounds, cmap.N)
fig = plt.figure()
ax = plt.gca()
ax.axes.xaxis.set_ticklabels([])
ax.axes.yaxis.set_ticklabels([])
ax.axes.xaxis.set_ticks([])
ax.axes.yaxis.set_ticks([])
#plt.colorbar(img, cmap=cmap, norm=norm, boundaries=bounds, ticks=[-1,0,1])

array = 2*(np.random.rand(N,N,steps)-0.5)
state = np.zeros(steps)
ims = []    
##leg = ax.legend(loc='upper left',prop={'size':12})

for step in range(0,steps):
    state = array[:,:,step]
    im = plt.imshow(state,interpolation='nearest',cmap=cmap,norm=norm, animated=True)
    ##lab = 'step = '+str(step)
    ##leg.texts.set_text(lab)
    ims.append([im])##+leg])

ani = animation.ArtistAnimation(fig,ims,interval=interval_pause,repeat_delay=repeat_pause)
#ani.save('animate_evolution '+str(timer())+'.mp4')
plt.show()

Answer 1

As shown in the question you link to it is easier to use a FuncAnimation . This allows to simply update a single legend and imshow plot instead of creating several of those.

Because it's not really clear what the legend is supposed to show for an imshow plot, I just created a blue rectangle. You can of course replace it with whatever you like.

import matplotlib.animation as animation
from matplotlib import colors
import matplotlib.pyplot as plt
import numpy as np

N=20
steps = 100
interval_pause = 100
repeat_pause = 1000

cmap = colors.ListedColormap(['white', 'black'])
bounds=[-1,0,1]
norm = colors.BoundaryNorm(bounds, cmap.N)
fig = plt.figure()
ax = plt.gca()
ax.axes.xaxis.set_ticklabels([])
ax.axes.yaxis.set_ticklabels([])
ax.axes.xaxis.set_ticks([])
ax.axes.yaxis.set_ticks([])

array = 2*(np.random.rand(N,N,steps)-0.5)

leg = ax.legend([plt.Rectangle((0,0),1,1)],["step0"], loc='upper left',prop={'size':12})

img = ax.imshow(array[:,:,0],interpolation='nearest',cmap=cmap,norm=norm, animated=True)
fig.colorbar(img, cmap=cmap, norm=norm, boundaries=bounds, ticks=[-1,0,1])

def update(step):
    state = array[:,:,step]
    img.set_data(state)
    lab = 'step = '+str(step)
    leg.texts[0].set_text(lab)

ani = animation.FuncAnimation(fig,update,frames = steps, 
                              interval=interval_pause,repeat_delay=repeat_pause)
plt.show()