In [ ]:
%matplotlib inline
import os
import yt
yt.mylog.setLevel("WARNING")
import numpy as np
import matplotlib
matplotlib.rcParams['font.family'] = 'stixgeneral'
matplotlib.rcParams['figure.dpi'] = 200
import matplotlib.pyplot as plt
import util
from synchrotron.yt_synchrotron_emissivity import StokesFieldName,\
                                                synchrotron_fits_filename
from mpl_toolkits.axes_grid1 import AxesGrid, ImageGrid
import pyfits

#dirs = ['/home/ychen/data/0only_0529_h1',\
#        '/home/ychen/data/0only_0605_hinf',\
#        '/home/ychen/data/0only_0605_h0']
#subdir = 'dtau_synchrotron_QU_nn_lobe/fits'

#regex = 'synchrotron_i_lobe_0?00_150MHz.fits'

#def rescan(dir, printlist=True):
#    fitsdir = os.path.join(dir,subdir)
#    files = util.scan_files(fitsdir, regex=regex, printlist=printlist, reverse=False)
#    return files

#fitsfiles = rescan(dirs[1])

In [ ]:
def plot_pol_frac_histogram(ax, frac, I_bin, label):
    ns, bins, patches = ax.hist(frac[I_bin.nonzero()].flatten()*100, range=(0,80),
                    bins=40, alpha=0.8, histtype='step',
                    weights=I_bin[I_bin.nonzero()].flatten(),
                    normed=True, label=label)
    #ax.set_xlabel('Polarization fraction (%)')
    ax.set_xlim(0, 80)
    return ns, bins

def plot_pol_angle_histogram(ax, psi, I_bin, label):
    ns, bins, patches = ax.hist(psi[I_bin.nonzero()].flatten(), range=(-0.5*np.pi, 0.5*np.pi),
                   bins=40, alpha=0.8, histtype='step',
                   weights=I_bin[I_bin.nonzero()].flatten(),
                   normed=True, label=label)
    
    x_tick = np.linspace(-0.5, 0.5, 5, endpoint=True)
    x_label = [r"$-\frac{\pi}{2}$", r"$-\frac{\pi}{4}$", r"$0$", r"$+\frac{\pi}{4}$", r"$+\frac{\pi}{2}$"]
    ax.set_xlim(-0.5*np.pi, 0.5*np.pi)
    ax.set_xticks(x_tick*np.pi)
    ax.set_xticklabels(x_label)
    return ns, bins

In [ ]:
#dirs = ['/home/ychen/data/00only_0605_hinf/',\
#        '/home/ychen/data/00only_0529_h1/',\
#        '/home/ychen/data/00only_0605_h0/',]

def plot_synchrotron_histogram_grids(proj_axis, nus):
    dirs = ['/d/d5/ychen/2015_production_runs/0204_hinf_10Myr/',\
            '/d/d5/ychen/2015_production_runs/1022_h1_10Myr/',\
            #'/d/d5/ychen/2015_production_runs/0204_h0_10Myr/'
            '/d/d5/ychen/2016_production_runs/1212_h0_10Myr/'
            #'/d/d5/ychen/2015_production_runs/0529_h1_00/'
           ]

    labels = ['toroidal', 'helical', 'poloidal']

    #filenumbers = [200, 400, 600, 800, 1200, 1300]
    filenumbers = [200, 600, 910, 1050]

    iterator = []
    for filenumber in filenumbers:
        for dir in dirs:
            iterator.append((filenumber, dir))

    labels = ['toroidal', 'helical', 'poloidal']

    ptype = 'lobe'
    nu0 = nus[0]
    nu1 = nus[1]

    fig, axes = plt.subplots(len(filenumbers), len(dirs), sharex=True, sharey=True, figsize=(8,6))

    for j, dir in enumerate(dirs):
        for i, filenumber in enumerate(filenumbers):
            # Load the data and create a single plot
            ds = yt.load(os.path.join(dir, 'data/MHD_Jet_10Myr_hdf5_plt_cnt_%04d' % filenumber))

            fitsname = synchrotron_fits_filename(ds, dir, ptype, proj_axis)
            if not os.path.isfile(fitsname): continue
            hdulist = pyfits.open(fitsname)
            print(fitsname)
            #for hdu in hdulist:
            #    print(hdu.name)

            ns, bins = {}, {}
            for nu in nus:
                #norm = yt.YTQuantity(*nu).in_units('GHz').value**0.5
                stokes = StokesFieldName(ptype, nu, proj_axis, field_type='fits')
                frb_I = hdulist[stokes.I[1]].data
                frb_Q = hdulist[stokes.Q[1]].data
                frb_U = hdulist[stokes.U[1]].data

                factor = 1
                nx = frb_I.shape[1]//factor
                ny = frb_I.shape[0]//factor

                I_bin = frb_I.reshape(nx, factor, ny, factor).sum(3).sum(1)
                Q_bin = frb_Q.reshape(nx, factor, ny, factor).sum(3).sum(1)
                U_bin = frb_U.reshape(nx, factor, ny, factor).sum(3).sum(1)

                psi = 0.5*np.arctan2(U_bin, Q_bin)
                frac = np.sqrt(Q_bin**2+U_bin**2)/I_bin

                ax = axes[i,j]

                #plot_pol_frac_histogram(ax, frac, I_bin)
                nu_str = '%i %s' % nu
                ns[nu], bins[nu] = plot_pol_frac_histogram(ax, frac, I_bin, nu_str)
                #plot_pol_angle_histogram(ax, psi, I_bin, nu_str)
            ax.set_ylim(-0.02, 0.04)

            #ax.scatter(frac.flatten()[::10], np.log10((I_bin*frac).flatten()[::10]), s=1, lw=0)
            #ax.set_aspect(3, adjustable='box')
            plt.subplots_adjust(hspace = 0, wspace = 0.1)
            if j == 0:
                timestamp = '%.1f Myr' % ds.current_time.in_units('Myr')
                ax.annotate(timestamp, (0.04, 0.80) , xycoords='axes fraction')

            if i == 0:
                ax.annotate(labels[j], (0.70, 0.80) , xycoords='axes fraction')

            # Ploting the difference in fraction
            #ax2 = ax.twinx()
            ax2 = ax
            #ax2.plot((bins[nu0][1:]+bins[nu0][:-1])/2, (ns[nu1]-ns[nu0])/(ns[nu1]+ns[nu0]), '--', lw=1, alpha=0.7, c='green')
            ax2.plot((bins[nu0][1:]+bins[nu0][:-1])/2, (ns[nu1]-ns[nu0]), '--', lw=1, alpha=0.7, c='green', label='difference')
            ax2.hlines(0, 0, 80, colors='k', lw=1, linestyles=':')
            #ax2.set_ylim(-0.01, 0.02)
            ax2.set_xlim(0, 80)
            if j in [0,1]:
                ax2.get_yaxis().set_visible(False)
            if i == 1 and j == 2:
                ax.legend()
    return fig
        #if j == 0:
        #    break
            #if i == 4 and j == 1:
            #    ax.set_xlabel('Polarization Angle')

In [ ]:
nus = [(150, 'MHz'), (8000, 'MHz')]
fig = plot_synchrotron_histogram_grids('x', nus)
fig.savefig('synchrotron_polarization_x_frac_histogram.pdf')

In [ ]:
nus = [(150, 'MHz'), (8000, 'MHz')]
fig = plot_synchrotron_histogram_grids([1,0,2], nus)
fig.savefig('synchrotron_polarization_102_frac_histogram.pdf')

In [ ]:
fitsname = '/d/d5/ychen/2015_production_runs/1022_h1_10Myr/cos_synchrotron_QU_nn_lobe/fits/synchrotron_lobe_0910.fits'
# Load the FITS file into the program
hdulist = pyfits.open(fitsname)

# Load table data as tbdata
#tbdata = hdulist[1].data
for hdu in hdulist:
    print(hdu.name)

In [ ]:
ptype = 'lobe'
proj_axis = 'x'
nu = (150, 'MHz')
zoom_fac = 6
extend_cells = 32

stokes = StokesFieldName(ptype, nu, proj_axis, field_type='fits')
frb_I = hdulist[stokes.I[1]].data
frb_Q = hdulist[stokes.Q[1]].data
frb_U = hdulist[stokes.U[1]].data

In [ ]:
#fields = []
fields += stokes.IQU
print(fields)

In [ ]:
def find_fwhm(arr_x, arr_y):

    difference = max(arr_y) - min(arr_y)
    HM = difference / 2

    pos_extremum = arr_y.argmax()  # or in your case: arr_y.argmin()

    nearest_above = (np.abs(arr_y[pos_extremum:-1] - HM)).argmin()
    nearest_below = (np.abs(arr_y[0:pos_extremum] - HM)).argmin()

    FWHM = (np.mean(arr_x[nearest_above + pos_extremum]) - 
        np.mean(arr_x[nearest_below]))
    return FWHM

In [ ]:
def plot_polarization_histogram(frac, psi, I_bin, fig=None, label=None):

    if not fig:
        fig = plt.figure(figsize=(12,12))

    ax1 = fig.axes[0]
    null = ax1.hist(frac[I_bin.nonzero()].flatten()*100, range=(0,80),
                    bins=40, alpha=0.5,
                    weights=I_bin[I_bin.nonzero()].flatten(),
                    normed=True)
    ax1.set_xlabel('Polarization fraction (%)')
    ax1.set_xlim(0, 80)

    ax2 = fig.axes[1]
    null = ax2.hist(psi[I_bin.nonzero()].flatten(), bins=50,
                    range=(-0.5*np.pi, 0.5*np.pi), alpha=0.5,
                    weights=I_bin[I_bin.nonzero()].flatten(),
                    normed=True)
    x_tick = np.linspace(-0.5, 0.5, 5, endpoint=True)

    x_label = [r"$-\pi/2$", r"$-\pi/4$", r"$0$", r"$+\pi/4$", r"$+\pi/2$"]
    ax2.set_xlim(-0.5*np.pi, 0.5*np.pi)
    ax2.set_xticks(x_tick*np.pi)
    ax2.set_xticklabels(x_label)
    #ax2.set_title(ds.basename + '  %.1f %s' % nu)

    ax3 = fig.axes[2]
    null = ax3.hist(np.abs(psi[I_bin.nonzero()].flatten()), bins=25,
                    range=(0.0, 0.5*np.pi), alpha=0.5,
                    label=label)
    ax3.legend()
    ax3.set_xlim(0.0, 0.5*np.pi)
    ax3.set_xticks([x_tick[2:]*np.pi])
    ax3.set_xticks(x_tick[2:]*np.pi)
    ax3.set_xticklabels(x_label[2:])

    return fig

In [ ]:
factor = 1
nx = frb_I.shape[1]//factor
ny = frb_I.shape[0]//factor

I_bin = frb_I.reshape(nx, factor, ny, factor).sum(3).sum(1)
Q_bin = frb_Q.reshape(nx, factor, ny, factor).sum(3).sum(1)
U_bin = frb_U.reshape(nx, factor, ny, factor).sum(3).sum(1)

psi = 0.5*np.arctan2(U_bin, Q_bin)
frac = np.sqrt(Q_bin**2+U_bin**2)/I_bin

fig = plt.figure(figsize=(16,4))
ax1 = fig.add_subplot(131)
ax2 = fig.add_subplot(132)
ax3 = fig.add_subplot(133)

fig = plot_polarization_histogram(frac, psi, I_bin, fig=fig)