Pyplot Imshow Autozoom to cut out Irregular NaN padding

Question

I have the following code

import matplotlib.pyplot as plt
import numpy as np

array = np.pad(np.random.rand(300,300),10,'constant', constant_values =  nan)

fig, ax = plt.subplots()
l = ax.imshow(array, origin = 'lower')

plt.show()

As you can see, it plots an image with a border of NaNs around the edge. Is there a way to get Imshow to auto-crop or auto-zoom to the area of the plot without the NaNs?

There are a few caveats.

• Padding is Unequal: In this example, the padding is equal on all sides, but this is not true of my real data, so it cant just assume the same pad on each side.

• Cannot physically edit array: The real data is actually an astronomy FITS file with WCS coordinate data and so uses the WCSAxesSubplot system, with plt.fig.add_subplot(projection = wcs) command, and a FITS header file from astropy, but there is no way to nicely include that in my example code. (That said, I am confident if someone can show me how to do it in regular pyplot/matplotlib/imshow, I can transfer it over somehow.) The real problem with this is I therefore cannot edit the original data array to remove the NaNs as my header file (namely the reference pixels) would then be incorrect and my coordinates would be off.

• NaNs may not be aligned with image edge: Future data sets may be slightly rotated, or even have fields of NaNs between bits of the image. Therefore I cant have it slicing off corners or bits of my image just to nicely fit other bits. Some NaNs may still need to be visible after zooming. (This is less of an issue right now, as my current data sets all have nice, straight, axi-parallel edges).

Any help you could offer would be appreciated. I am happy to supply extra information. Also, once a zoom has been achieved automatically, if someone could point me to how to get the current zoom back out of the image, that would be great (as it solves an unrelated issue).

Answer 1

You can determine the indices of non-NaN elements as follows

bound = np.argwhere(~np.isnan(array))

and the xlim and ylim are simply the minima/maxima of these indices

plt.xlim(min(bound[:, 1]), max(bound[:, 1]))
plt.ylim(min(bound[:, 0]), max(bound[:, 0]))

To zoom back out, simply do the following

shape = np.shape(array)
plt.xlim(0, shape[1]-1)
plt.ylim(0, shape[0]-1)

Answer 2

This should work:

mask = ~np.isnan(array)
x = np.flatnonzero(np.any(mask, axis = 0))
x = np.arange(x.min(), x.max() + 1)[:, None]
y = np.flatnonzero(np.any(mask, axis = 1))
y = np.arange(y.min(), y.max() + 1)

plt.imshow(array[x,y], origin = 'lower')

Basically just building a fancy indexing of rows and columns with any non-nan values.