In [ ]:
%matplotlib inline
import matplotlib.pyplot as plt
import yt
import logging
logging.getLogger('yt').setLevel(logging.ERROR)
import numpy as np
from yt import derived_field

nus = [(1500, 'MHz'), (150, 'MHz')]
proj_axis = [1,0,2]

prj = {}
frb_I, frb_Q, frb_U = {}, {}, {}
I_bin, Q_bin, U_bin = {}, {}, {}
psi, frac = {}, {}

nu0 = nus[0]
nu1 = nus[1]
nus_str = ['%.1f %s' % nu0, '%.1f %s' % nu1]

postfix = ('_synchrotron_gc0')

dir = '/home/ychen/d9/FLASH4/2015_production_runs/0529_L45_M10_b1_h1'
fname = dir + '/MHD_Jet_hdf5_plt_cnt_0630' + postfix
ds = yt.load(fname)
ptype = 'lobe'


for nu in nus:
    stokes = StokesFieldName(ptype, nu, proj_axis, field_type='flash')
    north_vector = [0,0,1]
    if proj_axis == 'x':
        prj[nu] = yt.ProjectionPlot(ds, proj_axis, fields, width=(80, 'kpc'))
    else:
        prj[nu] = yt.OffAxisProjectionPlot(ds, proj_axis, fields,width=(60, 'kpc'),
                                   north_vector=north_vector)
    frb_I[nu] = prj[nu].frb.data[fields[0]].v
    frb_Q[nu] = prj[nu].frb.data[fields[1]].v
    frb_U[nu] = prj[nu].frb.data[fields[2]].v

In [ ]:
def plot_polarization_histogram(frac, psi, I_bin, fig=None, label=None):
    
    if not fig:
        fig = plt.figure(figsize=(16,4))
    
    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 (%)')

    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 [ ]:
# Binning
for nu in nus:
    factor = 1
    nx = 800//factor
    ny = 800//factor

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

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

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

for nu in reversed(nus):
    nu_str = '%.0f%s' % nu
    fig = plot_polarization_histogram(frac[nu], psi[nu], I_bin[nu], fig=fig, label=nu_str)

In [ ]:
ns, bins = {}, {}
for nu in nus:
    frac_nu = frac[nu]
    ns[nu], bins[nu], patches = plt.hist(frac_nu[frac_nu.nonzero()]*100, range=(0,80), bins=40, \
                              weights=I_bin[nu][frac_nu.nonzero()], normed=True)

plt.cla()
plt.scatter((bins[nu0][1:]+bins[nu0][:-1])/2, (ns[nu0]-ns[nu1])/(ns[nu0]+ns[nu1]))
plt.hlines(0, 0, 80, linestyle=':')
plt.ylim(-0.3,1)

In [ ]:
fig = plt.figure(figsize=(20,8))
ax1 = fig.add_subplot(121)
ax1.imshow(frac[nu0], vmin=0, vmax=0.7)
plt.title(nus_str[0])
ax2 = fig.add_subplot(122)
img2 = ax2.imshow(frac[nu1], vmin=0, vmax=0.7)
plt.title(nus_str[1])
cb = plt.colorbar(img2, pad=0)
cb.set_label('polarization fraction')

In [ ]:
fig = plt.figure(figsize=(10,8))
ax = fig.add_subplot(111)
img = ax.imshow((frac[nu0]-frac[nu1]), cmap='seismic', vmin=-0.5, vmax=0.5)
cb = plt.colorbar(img, pad=0)

cb.set_label('polarization fraction difference (%s - %s)' % tuple(nus_str))

In [ ]:
fig = plt.figure(figsize=(22,8))
ax1 = fig.add_subplot(121)
ax1.imshow(psi[nu0], vmin=0, vmax=np.pi/2)
plt.title(nus_str[0])
ax2 = fig.add_subplot(122)
img2 = ax2.imshow(psi[nu1], vmin=0, vmax=np.pi/2)
plt.title(nus_str[1])
cb = plt.colorbar(img2)
cb.set_label('polarization angle')

In [ ]:
fig = plt.figure(figsize=(10,8))
ax = fig.add_subplot(111)
img = plt.imshow((psi[nu0]-psi[nu1]), cmap='seismic', vmin=-1.5, vmax=1.5)
cb = plt.colorbar(img, pad=0)

cb.set_label('polarization angle difference (%s - %s)' % tuple(nus_str))

In [ ]:
#frb_I = proj.frb.data[fields[0]].v
#frb_Q = proj.frb.data[fields[1]].v
#frb_U = proj.frb.data[fields[2]].v

#proj.annotate_polline(frb_I, frb_Q, frb_U, factor=16)
#proj.show()