Pyplot : imshow 3D array with a slider

Question

I have a 3d array

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider

A = np.random.rand(10,5,5)

and I would like to to visualize each 5x5 image in A

for Ai in A:
    plt.imshow(Ai)
    plt.show()

Except that rather than having 5 figures like in the example above, I would like to have a slider to switch between indexes for the first coordinate of A.

For the moment, I have tried the following :

idx0 = 3
l = plt.imshow(A[idx0])

axidx = plt.axes([0.25, 0.15, 0.65, 0.03])
slidx = Slider(axidx, 'index', 0, 9, valinit=idx0, valfmt='%d')

def update(val):
    idx = slidx.val
    l.set_data(A[idx])
    fig.canvas.draw_idle()
slidx.on_changed(update)

plt.show()

But when I use the slider, the image in the plot doesn't change, and I get the message

IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices

How can I get my slider to work with my 3d array ?

Answer 1

Turns out, the problem wasn't with the Slider at all, I just needed to convert the value returned by the slider into an int.

Replacing with

l.set_data(A[int(idx)])

does the trick