Allows to plot fits files produced as output

#!/usr/bin/python3

import numpy as np

import matplotlib.pyplot as plt

from astropy.io import fits

real_fits = "/home/emanuele/hpc_imaging/test_fits_real.fits"

imag_fits = "/home/emanuele/hpc_imaging/test_fits_img.fits"

nplanes = 1

with fits.open(real_fits) as hdu_real:

	img_hdu_real = np.array(hdu_real[0].data)

with fits.open(imag_fits) as hdu_imag:

	img_hdu_imag = np.array(hdu_imag[0].data)

xaxis = int(np.sqrt(img_hdu_real.size))

yaxes = xaxis

residual = np.vectorize(complex)(img_hdu_real, img_hdu_imag)

cumul2d = residual.reshape((xaxis,yaxes,nplanes), order='F')

for i in range(nplanes):

    gridded = np.squeeze(cumul2d[:,:,i])

    ax = plt.subplot()

    img = ax.imshow(np.abs(np.fft.fftshift(gridded)), aspect='auto', interpolation='none', origin='lower')

    ax.set_xlabel('cell')

    ax.set_ylabel('cell')

    cbar = plt.colorbar(img)

    cbar.set_label('norm(FFT)',size=18)

    figname='fits_image_' + str(i) + '.png'

    plt.savefig(figname)