Setting correct limits with imshow if image data shape changes

Question

I have a 3D array, of which the first two dimensions are spatial, so say (x,y). The third dimension contains point-specific information.

print H.shape  # --> (200, 480, 640)  spatial extents (200,480)

Now, by selecting a certain plane in the third dimension, I can display an image with

imdat = H[:,:,100]    # shape (200, 480)
img = ax.imshow(imdat, cmap='jet',vmin=imdat.min(),vmax=imdat.max(), animated=True, aspect='equal')

I want to now rotate the cube, so that I switch from (x,y) to (y,x).

    H = np.rot90(H)          # could also use H.swapaxes(0,1) or H.transpose((1,0,2)) 
    print H.shape    # --> (480, 200, 640)

Now, when I call:

imdat = H[:,:,100]   # shape (480,200)
img.set_data(imdat)
ax.relim()
ax.autoscale_view(tight=True)

I get weird behavior. The image along the rows displays the data till 200th row, and then it is black until the end of the y-axis (480). The x-axis extends from 0 to 200 and shows the rotated data. Now on, another rotation by 90-degrees, the image displays correctly (just rotated 180 degrees of course)

It seems to me like after rotating the data, the axis limits, (or image extents?) or something is not refreshing correctly. Can somebody help?

PS: to indulge in bad hacking, I also tried to regenerate a new image (by calling ax.imshow) after each rotation, but I still get the same behavior.

Answer 1

Below I include a solution to your problem. The method resetExtent uses the data and the image to explicitly set the extent to the desired values. Hopefully I correctly emulated the intended outcome.

import matplotlib.pyplot as plt
import numpy as np

def resetExtent(data,im):
    """
    Using the data and axes from an AxesImage, im, force the extent and 
    axis values to match shape of data.
    """
    ax = im.get_axes()
    dataShape = data.shape

    if im.origin == 'upper':
        im.set_extent((-0.5,dataShape[0]-.5,dataShape[1]-.5,-.5))
        ax.set_xlim((-0.5,dataShape[0]-.5))
        ax.set_ylim((dataShape[1]-.5,-.5))
    else:
        im.set_extent((-0.5,dataShape[0]-.5,-.5,dataShape[1]-.5))
        ax.set_xlim((-0.5,dataShape[0]-.5))
        ax.set_ylim((-.5,dataShape[1]-.5))

def main():
    fig = plt.gcf()
    ax = fig.gca()

    H = np.zeros((200,480,10))
    # make distinguishing corner of data
    H[100:,...] = 1
    H[100:,240:,:] = 2

    imdat = H[:,:,5]
    datShape = imdat.shape

    im = ax.imshow(imdat,cmap='jet',vmin=imdat.min(),
                    vmax=imdat.max(),animated=True,
                    aspect='equal',
                    #                origin='lower'
                    )

    resetExtent(imdat,im)

    fig.savefig("img1.png")

    H = np.rot90(H)

    imdat = H[:,:,0]
    im.set_data(imdat)
    resetExtent(imdat,im)

    fig.savefig("img2.png")

if __name__ == '__main__':
  main()

This script produces two images: First un-rotated: Then rotated:

I thought just explicitly calling set_extent would do everything resetExtent does, because it should adjust the axes limits if 'autoscle' is True. But for some unknown reason, calling set_extent alone does not do the job.