plot_fits.py

#!/usr/bin/env python
# Plots fits file generated by WSClean
import argparse
import glob
import os
import warnings

import astropy.units as u
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
from matplotlib.patches import Ellipse
import numpy as np
import sunpy.map

from multiprocessing import Pool

from astropy.coordinates import Angle, SkyCoord, EarthLocation
from mpl_toolkits.axes_grid1 import make_axes_locatable
from sunpy.coordinates import frames
from sunpy.time import TimeRange

from LOFAR_bf import LOFAR_BF
from manual_clean import convolve_model
from icrs_to_helio import icrs_to_helio
# from radec_to_hpc_map import radec_to_hpc

warnings.filterwarnings("ignore")


def plot_fits(fits_in, out_png=None):
    if out_png is None:
        out_png = fits_in[:-4] + 'png'
    smap = sunpy.map.Map(fits_in)
    smap.meta['wavelnth'] = smap.meta['crval3'] / 1e6
    smap.meta['waveunit'] = "MHz"

    fig = plt.figure()
    helio_smap = icrs_to_helio(fits_in)
    trange = TimeRange(helio_smap.date - 0.5 * u.s, helio_smap.date + 0.5 * u.s)
    # bf = LOFAR_BF(fits_in.split('_SB')[0] + '_B003_S0_P000_bf.h5', trange)
    # loc = np.where(abs(bf.freqs.to(u.MHz) - helio_smap.wavelength) == np.min(abs(bf.freqs.to(u.MHz) - helio_smap.wavelength)))[0][0]
    # dslice = np.mean(bf.data[:, loc - 8:loc + 8], axis=1)
    solar_PA = sunpy.coordinates.sun.P(smap.date).deg
    if fits_in[-8:] != "psf.fits":
        # gs = GridSpec(1, 1)
        # ax0 = fig.add_subplot(gs[0], projection=helio_smap)
        # ax1 = fig.add_subplot(gs[1])
        ax0 = fig.add_subplot(1, 1, 1, projection=helio_smap)
        helio_smap.plot_settings["title"] = str(
            np.round(helio_smap.meta['wavelnth'], 2)) + " MHz " + helio_smap.date.isot
        im = helio_smap.plot(axes=ax0,cmap='viridis')
        # ax1.plot(dslice)
        if fits_in[-10:] != "model.fits":
            helio_smap.draw_limb(color='r')

        b = Ellipse((200, 200), Angle(smap.meta['BMAJ'] * u.deg).arcsec / abs(smap.scale[0].to(u.arcsec / u.pix).value),
                    Angle(smap.meta['BMIN'] * u.deg).arcsec / abs(smap.scale[1].to(u.arcsec / u.pix).value),
                    angle=(90 + smap.meta['BPA']) - solar_PA, fill=False, color='w', ls='--')
        cbar = fig.colorbar(im, ax=ax0)
        cbar.set_label("Intensity (Jy/beam)")
    else:

        ax0 = fig.add_subplot(1, 1, 1, projection=smap)
        im = smap.plot(axes=ax0)

        b = Ellipse((smap.reference_pixel[0].value, smap.reference_pixel[1].value),
                    Angle(smap.meta['BMAJ'] * u.deg).arcsec / abs(smap.scale[0].to(u.arcsec / u.pix).value),
                    Angle(smap.meta['BMIN'] * u.deg).arcsec / abs(smap.scale[1].to(u.arcsec / u.pix).value),
                    angle=90 + smap.meta['BPA'], fill=False, color='w', ls='--')
        cbar = fig.colorbar(im, ax=ax0)
        cbar.set_label("PSF")
    ax0.add_patch(b)

    # plt.colorbar()
    print("Saving to {}".format(out_png))
    plt.savefig(out_png)
    # plt.close()


def plot_compare(fits_in, out_png=None):
    if out_png is None:
        out_png = fits_in[:-10] + 'compare.png'
    if fits_in[-10:] != "image.fits":
        print("Input fits must be clean image. e.g. wsclean-image.fits")
        return
    dirty_fits = fits_in[:-10] + "dirty.fits"
    helio_clean = icrs_to_helio(fits_in)
    helio_dirty = icrs_to_helio(dirty_fits)
    icrs_map = sunpy.map.Map(fits_in)
    solar_PA = sunpy.coordinates.sun.P(icrs_map.date).deg
    fig, ax = plt.subplots(1, 2, figsize=(10, 5), sharey=True)

    for smap, ax, label in zip([helio_dirty, helio_clean], ax, ['Dirty', 'Clean']):
        smap.plot_settings['title'] = str(np.round(smap.meta['wavelnth'], 2)) + " MHz " + smap.date.isot
        smap.plot_settings['cmap'] = 'viridis'
        # smap.plot_settings['norm'] = matplotlib.colors.Normalize(vmin=0, vmax=500)
        im = smap.plot(axes=ax)
        smap.draw_limb(color='r')
        divider = make_axes_locatable(ax)
        cax = divider.append_axes("right", size="5%", pad=0.05)
        fig.colorbar(im, cax=cax)
        ax.text(-3000, 3000, label, fontdict={'fontsize': 14, 'color': 'white'})

        b = Ellipse((-3000, -3000),
                    Angle(icrs_map.meta['BMAJ'] * u.deg).arcsec / abs(icrs_map.scale[0].to(u.arcsec / u.pix).value),
                    Angle(icrs_map.meta['BMIN'] * u.deg).arcsec / abs(icrs_map.scale[1].to(u.arcsec / u.pix).value),
                    angle=(90 + icrs_map.meta['BPA']) - solar_PA, fill=False, color='w', ls='--')
        ax.add_patch(b)
    plt.tight_layout()
    print("Saving to {}".format(out_png))
    plt.savefig(out_png)
    plt.close()


def plot_grid(fits_in, out_png=None):
    if out_png is None:
        out_png = fits_in[:-10] + 'grid_compare.png'
    if fits_in[-10:] != "image.fits":
        print("Input fits must be clean image. e.g. wsclean-image.fits")
        return
    dirty_fits = fits_in[:-10] + "dirty.fits"
    residual_fits = fits_in[:-10] + "residual.fits"
    model_fits = fits_in[:-10] + "model.fits"

    helio_clean = icrs_to_helio(fits_in)
    helio_dirty = icrs_to_helio(dirty_fits)
    conv_model = convolve_model(model_fits)
    #if not os.path.isfile(model_fits.replace('model.fits', 'convolved_model.fits')):
    conv_model.save(model_fits.replace('model.fits', 'convolved_model.fits'), overwrite=True)
    helio_residual = icrs_to_helio(residual_fits)
    helio_model = icrs_to_helio(model_fits.replace('model.fits', 'convolved_model.fits'))

    # icrs_map = sunpy.map.Map(fits_in)
    fig, ax = plt.subplots(2, 2, figsize=(10, 10), sharex=True, sharey=True)
    # fig = plt.figure(figsize=(10, 10))
    # ax1 = fig.add_subplot(2, 2, 1, projection=helio_clean)
    # ax2 = fig.add_subplot(2, 2, 2, sharey=ax1, projection=helio_clean)
    # ax3 = fig.add_subplot(2, 2, 3, sharex=ax1, projection=helio_clean)
    # ax4 = fig.add_subplot(2, 2, 4, sharex=ax1, sharey=ax1, projection=helio_clean)

    for smap, a, label in zip([helio_dirty, helio_clean, helio_residual, helio_model], ax.flatten(),
                              ['Dirty', 'Clean', 'Residual', 'Model']):
        smap.plot_settings['title'] = str(np.round(smap.meta['wavelnth'], 2)) + " MHz " + smap.date.isot
        smap.plot_settings['cmap'] = 'viridis'
        # smap.plot_settings['norm'] = matplotlib.colors.Normalize(vmin=0, vmax=500)
        im = smap.plot(axes=a, title='')
        smap.draw_limb()
        # core_ITRF = np.array((3826577.066, 461022.948, 5064892.786))
        # lofar_loc = EarthLocation.from_geocentric(*core_ITRF, u.m)
        # lofar_gcrs = SkyCoord(lofar_loc.get_gcrs(smap.date))
        # sun_centre_coord = sunpy.coordinates.sun.sky_position(t=smap.date, equinox_of_date=False)
        # sun_centre_coord = SkyCoord(sun_centre_coord[0], sun_centre_coord[1], sunpy.coordinates.sun.earth_distance(smap.date))
        # sun_centre_coord = sun_centre_coord.transform_to(frames.Helioprojective(observer=lofar_gcrs, obstime=smap.date))
        # sun_centre_coord = SkyCoord(sun_centre[0], sun_centre[1], sunpy.coordinates.sun.earth_distance(smap.date))
        # smap.draw_grid()
        # lon, lat = a.coords
        # lat.set_ticklabel_position('l')
        # lat.set_axislabel_position('l')
        # lon.set_ticklabel_position('b')
        # lon.set_axislabel_position('b')
        divider = make_axes_locatable(a)
        cax = divider.append_axes("right", size="5%", pad=0.05)
        # cax.grid(False)
        # cax_x = cax.coords[0]
        # cax_y = cax.coords[1]
        # cax_x.set_ticks_visible(False)
        # cax_x.set_ticklabel_visible(False)
        # cax_y.set_ticklabel_position('r')
        # cax_y.set_axislabel_position('r')
        # cax_y.set_ticks_position('r')
        # cax_y.set_axislabel('')
        fig.colorbar(im, cax=cax)
        a.text(-3000, 2900, label, fontdict={'fontsize': 12, 'color': 'white'})
        solar_PA = sunpy.coordinates.sun.P(smap.date).deg

        b = Ellipse((-2500, -2500),
                    Angle(smap.meta['BMAJ'] * u.deg).arcsec, #/ abs(smap.scale[0].to(u.arcsec / u.pix).value),
                    Angle(smap.meta['BMIN'] * u.deg).arcsec, # / abs(smap.scale[1].to(u.arcsec / u.pix).value),
                    angle=(90 + smap.meta['BPA']) - solar_PA, fill=False, color='w', ls='--')
        a.add_patch(b)
        # a.plot_coord(sun_centre_coord, '+', color='white')
        a.set_ylabel('Solar Y (arcsec)')
        a.set_xlabel('Solar X (arcsec)')

    # plt.setp(ax1.get_xticklabels(), visible=False)
    # plt.setp(ax2.get_xticklabels(), visible=False)
    # plt.setp(ax3.get_yticklabels(), visible=False)
    # plt.setp(ax4.get_yticklabels(), visible=False)
    fig.suptitle(helio_clean.date.isot)
    # plt.tight_layout()
    print("Saving to {}".format(out_png))
    plt.savefig(out_png)
    plt.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Plot a fits file generated by WSClean')
    parser.add_argument('fits_in', help='name of input fits file')
    parser.add_argument('-o', dest='out_png', help='name of output png.\
                        If not given defaults to same prefix as fits_in.',
                        default=None)
    parser.add_argument('-c', '--compare',
                        help='plot comparison between dirty and clean images.\
                        fits_in must be image.fits and have matching dirty.fits.',
                        action='store_true')
    parser.add_argument('-a', '--all',
                        help='plot all images of that match given wildcard.',
                        type=str, default=None)
    parser.add_argument('-g', '--grid',
                        help='plot comparison between dirty, clean, residuals and model.\
                        fits_in must be image.fits.',
                        action='store_true')
    args = parser.parse_args()
    fits_in = args.fits_in
    out_png = args.out_png
    compare = args.compare
    grid = args.grid
    plot_all = args.all
    if plot_all is None:
        if compare:
            plot_compare(fits_in, out_png)
        elif grid:
            plot_grid(fits_in, out_png)
        else:
            plot_fits(fits_in, out_png)
    else:
        mpl.use('Agg')
        all_files = glob.glob(plot_all)
        if len(all_files) == 0:
            print("No files found, exiting")

        all_files.sort()
        if compare:
            with Pool() as p:
                p.map(plot_compare, all_files)
        elif grid:
            with Pool() as p:
                p.map(plot_grid, all_files)
        else:
            with Pool() as p:
                p.map(plot_fits, all_files)