def plot_compsite_map(vlafile, aiafile, outfile='', label='', pol=0, chans=[], chan_mask=None, x_range=[], y_range=[], levels=[0.9], plotstyle='centroid', figsize=(10, 8), figdpi=100, width=[5, 5], zorder=1, maponly=False, vlaonly=False, dspecdata={}, thrshd=None, cmap='jet', plt_clbar=True, aiaplt_args={ 'axes': None, 'cmap': None, 'vmax': None, 'vmin': None }, **kwargs): cmap = cm.get_cmap(cmap) if aiaplt_args['cmap'] is None: aiaplt_args.pop('cmap', None) if outfile: plt.ioff() if not vlaonly: if type(aiafile) != sunpy.map.sources.sdo.AIAMap and type( aiafile) != sunpy.map.compositemap.CompositeMap: aiamap = sunpy.map.Map(aiafile) else: aiamap = aiafile if x_range and y_range: try: aiamap = aiamap.submap(u.Quantity(x_range * u.arcsec), u.Quantity(y_range * u.arcsec)) except: from astropy.coordinates import SkyCoord bl = SkyCoord(x_range[0] * u.arcsec, y_range[0] * u.arcsec, frame=aiamap.coordinate_frame) tr = SkyCoord(x_range[1] * u.arcsec, y_range[1] * u.arcsec, frame=aiamap.coordinate_frame) aiamap = aiamap.submap(bl, tr) if type(vlafile) == dict: if dspecdata: if dspecdata['stack'] == 'Vstack': fig = plt.figure(figsize=(9, 12)) plt.subplot(2, 1, 1) elif dspecdata['stack'] == 'Hstack': fig = plt.figure(figsize=(20, 8)) plt.subplot(1, 2, 1) else: if aiaplt_args['axes'] is None: fig = plt.figure(figsize=figsize) plt.subplot() else: fig = plt.gcf() if not vlaonly: aiamap.plot(**aiaplt_args) ax1 = plt.gca() ax1.set_autoscale_on(False) else: ax1 = aiaplt_args['axes'] ax1.set_autoscale_on(False) if label: ax1.text(0.98, 0.98, label, horizontalalignment='right', verticalalignment='top', color='white', transform=ax1.transAxes, fontsize=14) clrange = vlafile['ColorMapper']['crange'] if not maponly: cargsort = np.argsort(np.array(vlafile['ColorMapper']['c'])) cargsort = cargsort[::zorder] vlafile['x'] = np.array(vlafile['x'])[cargsort] vlafile['y'] = np.array(vlafile['y'])[cargsort] vlafile['ColorMapper']['c'] = np.array( vlafile['ColorMapper']['c'])[cargsort] vlafile['s'] = np.array(vlafile['s'])[cargsort] im1 = ax1.scatter(vlafile['x'], vlafile['y'], c=vlafile['ColorMapper']['c'], s=vlafile['s'], vmin=clrange[0], vmax=clrange[-1], cmap=cmap, **kwargs) else: im1 = ax1.scatter([], [], c=[], s=[], vmin=clrange[0], vmax=clrange[-1], cmap=cmap, **kwargs) if not dspecdata: if plt_clbar: cb1 = plt.colorbar(im1, orientation='vertical', ax=ax1) if type(vlafile) == dict: if 'title' in vlafile['ColorMapper'].keys(): cb1.set_label(vlafile['ColorMapper']['title']) else: cb1.set_label('Frequency [GHz]') else: cb1.set_label('Frequency [GHz]') else: tim = dspecdata['time'] dt = np.mean(np.diff(tim)) cmapspec = cmap cmapspec.set_bad('white', 1.0) normspec = colors.Normalize(vmin=dspecdata['drange'][0], vmax=dspecdata['drange'][1]) if dspecdata['stack'] == 'Vstack': plt.subplot(2, 1, 2) elif dspecdata['stack'] == 'Hstack': plt.subplot(1, 2, 2) ax2 = plt.gca() im2 = plt.pcolormesh(tim, dspecdata['freq'], dspecdata['peak'], cmap=cmapspec, norm=normspec) ax2.add_patch( patches.Rectangle((vlafile['t'], dspecdata['freq'][0]), dt, dspecdata['freq'][-1] - dspecdata['freq'][0], facecolor='black', edgecolor='white', alpha=0.3)) ax2.set_xlim(tim[0], tim[-1]) ax2.set_ylim(dspecdata['freq'][0], dspecdata['freq'][-1]) ax2.set_title('Vector Dynamic spectrum') labels = ax2.get_xticks().tolist() newlabels = [ Time(lb / 24. / 3600., format='jd').iso.split(' ')[1] for lb in labels ] ax2.set_xticklabels(newlabels, rotation=45) if plt_clbar: cb1 = plt.colorbar(im1, orientation='vertical', ax=ax1) if type(vlafile) == dict: if 'title' in vlafile['ColorMapper'].keys(): cb1.set_label(vlafile['ColorMapper']['title']) else: cb1.set_label('Frequency [GHz]') else: cb1.set_label('Frequency [GHz]') cb2 = plt.colorbar(im2, orientation='vertical', ax=ax2) cb2.set_label('Max intensity [Jy/beam]') else: if aiaplt_args['axes'] is None: fig = plt.figure(figsize=figsize) plt.subplot() else: fig = plt.gcf() if not vlaonly: aiamap.plot(**aiaplt_args) ax1 = plt.gca() ax1.set_autoscale_on(False) else: ax1 = aiaplt_args['axes'] ax1.set_autoscale_on(False) # ax1.xaxis.set_ticks_position('top') # ax1.yaxis.set_ticks_position('right') if label: ax1.text(0.98, 0.98, label, horizontalalignment='right', verticalalignment='top', color='white', transform=ax1.transAxes, fontsize=14) # # Prevent the image from being re-scaled while overplotting. hdulist = fits.open(vlafile) hdu = hdulist[0] nfreq = hdu.data.shape[1] if not maponly: vladata = hdu.data[pol, 0, :, :] vlamap = sunpy.map.Map((vladata, hdu.header)) XX, YY = np.meshgrid(np.arange(vlamap.data.shape[1]), np.arange(vlamap.data.shape[0])) mapxy = vlamap.pixel_to_data(XX * u.pix, YY * u.pix) mapx, mapy = mapxy.Tx.value, mapxy.Ty.value clrange = (hdu.header['CRVAL3'] + np.arange(nfreq) * hdu.header['CDELT3']) / 1e9 if len(chans) == 0: chans = range(0, nfreq) if plotstyle == 'centroid': pltdata = {'x': [], 'y': [], 's': [], 'c': []} for idx, chan in enumerate(chans): vladata = hdu.data[pol, chan, :, :] if np.nanmax(vladata): vlamap.data = vladata maxxy = maxfit(vlamap, mapxy=[mapx, mapy], width=width) if maxxy: x, y = maxxy pltdata['x'].append(x) pltdata['y'].append(y) pltdata['c'].append( (hdu.header['CRVAL3'] + chan * hdu.header['CDELT3']) / 1e9) pltdata['s'].append(np.nanmax(vladata)) cargsort = np.argsort(np.array(pltdata['c'])) cargsort = cargsort[::zorder] pltdata['x'] = np.array(pltdata['x'])[cargsort] pltdata['y'] = np.array(pltdata['y'])[cargsort] pltdata['c'] = np.array(pltdata['c'])[cargsort] pltdata['s'] = np.array(pltdata['s'])[cargsort] im1 = ax1.scatter(pltdata['x'], pltdata['y'], c=pltdata['c'], s=pltdata['s'], vmin=clrange[0], vmax=clrange[-1], cmap=cmap, **kwargs) if plt_clbar: cb1 = plt.colorbar(im1, orientation='vertical', ax=ax1) cb1.set_label('Frequency [GHz]') elif plotstyle == 'contour': nchan = len(chans) for idx, chan in enumerate(chans): if chan_mask is not None: freq = (hdu.header['CRVAL3'] + chan * hdu.header['CDELT3']) / 1e9 if freq not in chan_mask: continue vladata = hdu.data[pol, chan, :, :] vlamap = sunpy.map.Map((vladata, hdu.header)) SRC_vlamap_contour = DButil.get_contour_data(mapx, mapy, vlamap.data, levels=levels) if SRC_vlamap_contour.data['xs']: for ii, xs in enumerate(SRC_vlamap_contour.data['xs']): x, y = xs, SRC_vlamap_contour.data['ys'][ii] if not thrshd or np.nanmax(vladata) >= thrshd: plt.plot(x, y, color=cmap( int(float(chan) / nfreq * 255)), zorder=nchan + zorder * idx, **kwargs) if plt_clbar: fig.subplots_adjust(right=0.8) cax1 = fig.add_axes([0.85, 0.1, 0.01, 0.8]) norm = colors.Normalize(vmin=clrange[0], vmax=clrange[-1]) cb1 = colorbar.ColorbarBase(cax1, cmap=cmap, norm=norm, orientation='vertical') cb1.set_label('Frequency [GHz]') try: cb1.ax.set_aspect(40) cb1.ax.tick_params(direction='in') except: pass try: cb2.ax.set_aspect(40) cb2.ax.tick_params(direction='in') except: pass if aiaplt_args['axes'] is None: fig.tight_layout(pad=3) if outfile: if outfile.endswith('.eps'): fig.savefig(outfile, format='eps', dpi=int(figdpi)) elif outfile.endswith('.pdf'): fig.savefig(outfile, format='pdf', dpi=int(figdpi)) elif outfile.endswith('.png'): fig.savefig(outfile, format='png', dpi=int(figdpi)) elif outfile.endswith('.jpeg'): fig.savefig(outfile, format='jpeg', dpi=int(figdpi)) elif outfile.endswith('.jpg'): fig.savefig(outfile, format='jpg', dpi=int(figdpi)) else: raise ValueError( 'Can not save {}. Provide a filename with extension of eps, pdf, png or jpeg.' .format(outfile))
def plotmap(vlafile, aiafile, outfile='', label='', pol=0, chans=[], x_range=[], y_range=[], levels=[0.9], plotstyle='centroid', figsize=(10, 8), figdpi=100, zorder=-1, maponly=False, dspecdata={}, **kwargs): import matplotlib.gridspec as gridspec if outfile: plt.ioff() if type(aiafile) != sunpy.map.sources.sdo.AIAMap: aiamap = sunpy.map.Map(aiafile) else: aiamap = aiafile if x_range and y_range: aiamap = aiamap.submap(u.Quantity(x_range * u.arcsec), u.Quantity(y_range * u.arcsec)) if type(vlafile) == dict: if dspecdata: if dspecdata['stack'] == 'Vstack': gs = gridspec.GridSpec(2, 2, width_ratios=[1, 0.02], height_ratios=[1, 0.5]) fig = plt.figure(figsize=(9, 12)) elif dspecdata['stack'] == 'Hstack': gs = gridspec.GridSpec(1, 4, width_ratios=[1, 0.02, 1.5, 0.02], height_ratios=[1]) fig = plt.figure(figsize=(18, 8)) ax1 = plt.subplot(gs[0], projection=aiamap) ax2 = plt.subplot(gs[1]) ax3 = plt.subplot(gs[2]) ax4 = plt.subplot(gs[3]) else: gs = gridspec.GridSpec(1, 2, width_ratios=[1, 0.05], height_ratios=[1]) fig = plt.figure(figsize=figsize) ax1 = plt.subplot(gs[0], projection=aiamap) ax2 = plt.subplot(gs[1]) aiamap.plot(axes=ax1) ax1.set_autoscale_on(False) if label: ax1.text(0.98, 0.98, label, horizontalalignment='right', verticalalignment='top', color='white', transform=ax1.transAxes, fontsize=14) freqs = vlafile['ColorMapper']['crange'] if not maponly: nchan = len(vlafile['ColorMapper']['c']) for idx, chan in enumerate(vlafile['ColorMapper']['c']): x, y = [vlafile['x'][idx], vlafile['y'][idx]] * u.arcsec s = vlafile['s'][idx] ax1.scatter(x.to(u.deg), y.to(u.deg), transform=ax1.get_transform('world'), facecolor=cm.jet( int((chan - freqs[0]) / (freqs[-1] - freqs[0]) * 255)), edgecolor='black', s=s, zorder=nchan + zorder * idx, **kwargs) if not dspecdata: cmap = cm.jet norm = colors.Normalize(vmin=freqs[0], vmax=freqs[-1]) cb1 = colorbar.ColorbarBase(ax2, cmap=cmap, norm=norm, orientation='vertical') if type(vlafile) == dict: cb1.set_label(vlafile['ColorMapper']['title']) else: cb1.set_label('Frequency [GHz]') else: tim = dspecdata['time'] dt = np.mean(np.diff(tim)) cmapspec = cm.jet cmapspec.set_bad('white', 1.0) normspec = colors.Normalize(vmin=dspecdata['drange'][0], vmax=dspecdata['drange'][1]) ax3.pcolormesh(tim, dspecdata['freq'], dspecdata['peak'], cmap=cmapspec, norm=normspec) ax3.add_patch( patches.Rectangle((vlafile['t'], dspecdata['freq'][0]), dt, dspecdata['freq'][-1] - dspecdata['freq'][0], facecolor='black', edgecolor='white', alpha=0.3)) ax3.set_xlim(tim[0], tim[-1] + dt) ax3.set_ylim(dspecdata['freq'][0], dspecdata['freq'][-1]) ax3.set_title('Vector Dynamic spectrum') labels = ax3.get_xticks().tolist() newlabels = [ Time(lb / 24. / 3600., format='jd').iso.split(' ')[1] for lb in labels ] ax3.set_xticklabels(newlabels, rotation=45) if dspecdata['stack'] == 'Vstack': fig.tight_layout(h_pad=3, pad=3) ax3.set_ylabel('Frequency [GHz]') pos1 = ax1.get_position() pos2 = ax2.get_position() pos4 = ax4.get_position() pos2new = [ pos2.x0 - (1.0 - pos1.x1) / 2, pos2.y0, pos2.width, pos2.height ] ax2.set_position(pos2new) pos4new = [ pos4.x0 - (1.0 - pos1.x1) / 2, pos4.y0, pos4.width, pos4.height ] ax4.set_position(pos4new) elif dspecdata['stack'] == 'Hstack': fig.tight_layout(w_pad=0, pad=5) ax3.set_yticklabels([]) pos2 = ax2.get_position() pos4 = ax4.get_position() pos2new = [pos2.x0 - 0.03, pos2.y0, pos2.width, pos2.height] ax2.set_position(pos2new) pos4new = [pos4.x0 - 0.01, pos4.y0, pos4.width, pos4.height] ax4.set_position(pos4new) cmap = cm.jet norm = colors.Normalize(vmin=freqs[0], vmax=freqs[-1]) cb1 = colorbar.ColorbarBase(ax2, cmap=cmap, norm=norm, orientation='vertical') cb2 = colorbar.ColorbarBase(ax4, cmap=cmapspec, norm=normspec, orientation='vertical') if type(vlafile) == dict: cb1.set_label(vlafile['ColorMapper']['title']) else: cb1.set_label('Frequency [GHz]') cb2.set_label('Max intensity [Jy/beam]') else: fig = plt.figure(figsize=figsize) ax1 = plt.subplot(projection=aiamap) aiamap.plot() if label: ax1.text(0.98, 0.98, label, horizontalalignment='right', verticalalignment='top', color='white', transform=ax1.transAxes, fontsize=14) # # Prevent the image from being re-scaled while overplotting. ax1.set_autoscale_on(False) hdulist = fits.open(vlafile) hdu = hdulist[0] if not maponly: vladata = hdu.data[pol, 0, :, :] vlamap = sunpy.map.Map((vladata, hdu.header)) XX, YY = np.meshgrid(np.arange(vlamap.data.shape[1]), np.arange(vlamap.data.shape[0])) mapx, mapy = vlamap.pixel_to_data(XX * u.pix, YY * u.pix) mapx, mapy = mapx.value, mapy.value nfreq = hdu.data[pol, :, :, :].shape[1] if len(chans) == 0: chans = range(0, nfreq) nchan = len(chans) if plotstyle == 'centroid': for idx, chan in enumerate(chans): vladata = hdu.data[pol, chan, :, :] if np.nanmax(vladata): vlamap.data = vladata maxxy = maxfit(vlamap, mapxy=[mapx, mapy]) if maxxy: x, y = maxxy * u.arcsec plt.scatter(x.to(u.deg), y.to(u.deg), transform=ax1.get_transform('world'), facecolor=cm.jet( int(float(chan) / nfreq * 255)), edgecolor='black', s=np.nanmax(vladata), zorder=nchan + zorder * idx, **kwargs) elif plotstyle == 'contour': for idx, chan in enumerate(chans): vladata = hdu.data[pol, chan, :, :] vlamap = sunpy.map.Map((vladata, hdu.header)) SRC_vlamap_contour = DButil.get_contour_data(mapx, mapy, vlamap.data, levels=levels) if SRC_vlamap_contour.data['xs']: xs = SRC_vlamap_contour.data['xs'][0] * u.arcsec ys = SRC_vlamap_contour.data['ys'][0] * u.arcsec plt.plot(xs.to(u.deg), ys.to(u.deg), transform=ax1.get_transform('world'), color=cm.jet(int(float(chan) / nfreq * 255)), zorder=nchan + zorder * idx, **kwargs) freqs = (hdu.header['CRVAL3'] + np.arange(nfreq) * hdu.header['CDELT3']) / 1e9 fig.subplots_adjust(right=0.8) cax1 = fig.add_axes([0.85, 0.1, 0.05, 0.8]) cmap = cm.jet norm = colors.Normalize(vmin=freqs[0], vmax=freqs[-1]) cb1 = colorbar.ColorbarBase(cax1, cmap=cmap, norm=norm, orientation='vertical') cb1.set_label('Frequency [GHz]') if outfile: if outfile.endswith('.eps'): fig.savefig(outfile, format='eps', dpi=int(figdpi)) elif outfile.endswith('.pdf'): fig.savefig(outfile, format='pdf', dpi=int(figdpi)) elif outfile.endswith('.png'): fig.savefig(outfile, format='png', dpi=int(figdpi)) elif outfile.endswith('.jpeg'): fig.savefig(outfile, format='jpeg', dpi=int(figdpi)) elif outfile.endswith('.jpg'): fig.savefig(outfile, format='jpg', dpi=int(figdpi)) else: raise ValueError( 'Can not save {}. Provide a filename with extension of eps, pdf, png or jpeg.' .format(outfile))