###This is a script to show the polarization information of the selected regions to be tested
###The data used in this script is obtained from the MPIFR's survey sampler: https://www3.mpifr-bonn.mpg.de/survey.html
### This data have been corrected with WMAP data.
%matplotlib inline
import aplpy
import matplotlib as mpl
import matplotlib.pyplot as plt
from astropy.io import fits
# initialize plot settings
mpl.rcParams['font.size'] = 16
mpl.rcParams['lines.linewidth'] = 2
mpl.rcParams['patch.linewidth'] = 2
mpl.rcParams['lines.markeredgewidth'] = 2
mpl.rcParams['lines.markersize'] = 10
mpl.rcParams['axes.linewidth'] = 2
mpl.rcParams['xtick.major.width'] = 2
mpl.rcParams['xtick.minor.width'] = 2
mpl.rcParams['xtick.major.size'] = 10
mpl.rcParams['xtick.major.pad'] = 5
mpl.rcParams['xtick.minor.size'] = 5
mpl.rcParams['xtick.minor.pad'] = 5
mpl.rcParams['ytick.major.width'] = 2
mpl.rcParams['ytick.minor.width'] = 2
mpl.rcParams['ytick.major.size'] = 10
mpl.rcParams['ytick.minor.size'] = 5
mpl.rcParams['figure.figsize'] = (10, 10/1.6)
data = fits.getdata("w3.fits")
hd = fits.getheader("w3.fits")
dat = data[0,:,:]
hd.remove('CRVAL3')
hd.remove('CRPIX3')
hd.remove('CDELT3')
hd.remove('CTYPE3')
hdu = fits.PrimaryHDU(dat, hd)
f1 = aplpy.FITSFigure(hdu)
f1.show_colorscale(cmap='jet',vmin=0, vmax=3.2e6,interpolation='bicubic')
f1.show_contour(levels=[2e5, 4e5, 8e5, 1.6e6, 3.2e6])
f1.set_tick_labels_xformat('dd')
f1.set_tick_labels_yformat('dd')
#f1.recenter(34.5, -0.5 , width=2, height=2)
f1.add_label(136, 3, "W3",size=40, color='white')
f1.show_colorbar()
f1.add_beam(9.5/60,9.5/60,0) ##9.5 arcmin
f1.beam.set_color('white')
########polarized intensity ########################
rawpol = fits.getdata('w3-polI.fits')
phd = fits.getheader('w3-polI.fits')
pol = rawpol[0,:,:]
phd.remove('CRVAL3')
phd.remove('CRPIX3')
phd.remove('CDELT3')
phd.remove('CTYPE3')
polhdu = fits.PrimaryHDU(pol, phd)
fpol = aplpy.FITSFigure(polhdu)
fpol.show_colorscale(cmap='jet')
fpol.set_tick_labels_xformat('dd')
fpol.set_tick_labels_yformat('dd')
#f1.recenter(34.5, -0.5 , width=2, height=2)
fpol.add_label(136, 3, "W3",size=40, color='white')
fpol.show_colorbar()
fpol.add_beam(9.5/60,9.5/60,0) ##9.5 arcmin
fpol.beam.set_color('white')
####### linear polarization #######################
rawang = fits.getdata('w3-pa.fits')
anghd = fits.getheader('w3-pa.fits')
ang = rawang[0,:,:]
anghd.remove('CRVAL3')
anghd.remove('CRPIX3')
anghd.remove('CDELT3')
anghd.remove('CTYPE3')
anghdu = fits.PrimaryHDU(ang, anghd)
poldeg = pol/dat
mask = pol > 20000 ## test values
poldegv2 = poldeg* mask
polhduv2 = fits.PrimaryHDU(poldegv2, phd)
#fits.writeto('try.fits',poldegv2, phd)
#fpol.show_vectors(polhdu, anghdu,cutoff=20000,step=2)
fpol.show_vectors(polhdu, anghdu,cutoff=15000,step=2,scale=5e-5,color='black')
#fpol.show_vectors(polhduv2, anghdu,cutoff=0,step=2,scale=0.05,color='black')