In [ ]:

    

#%pdb
%matplotlib inline
import matplotlib
matplotlib.rcParams['figure.dpi'] = 200
import matplotlib.pyplot as plt
import os
import yt
yt.mylog.setLevel("WARNING")
import numpy as np
from yt_synchrotron_emissivity import *
import pyfits
from scipy.ndimage import gaussian_filter
from mpl_toolkits.axes_grid1 import AxesGrid


    

In [ ]:

    

#dir = '/d/d5/ychen/2015_production_runs/1022_h1_10Myr/'
#ds = yt.load(dir + 'data/MHD_Jet_10Myr_hdf5_plt_cnt_0910')
dir = '/d/d5/ychen/2015_production_runs/0529_h1_00/'
ds = yt.load(dir + 'data/MHD_Jet_hdf5_plt_cnt_0600')

proj_axis = 'x'
nus = [(100, 'MHz'), (1400, 'MHz')]

ptype = 'lobe'
extend_cells = 8
res = (512, 256)
sigma = 1

print(ds.basename, ds.current_time.in_units('Myr'))

fitsname = synchrotron_fits_filename(ds, dir, ptype, proj_axis)
hdulist = pyfits.open(fitsname)
frb_I, frb_Q, frb_U = {},{},{}
for nu in nus:
    stokes = StokesFieldName(ptype, nu, proj_axis, field_type='flash')
    frb_I[nu] = hdulist[stokes.I[1]].data
    frb_Q[nu] = hdulist[stokes.Q[1]].data
    frb_U[nu] = hdulist[stokes.U[1]].data
header = hdulist[stokes.I[1]].header
xr = -header['CRPIX1']*header['CDELT1'] + header['CRVAL1']
xl = (header['NAXIS1'] - header['CRPIX1'])*header['CDELT1'] + header['CRVAL1']
yr = -header['CRPIX2']*header['CDELT2'] + header['CRVAL2']
yl = (header['NAXIS2'] - header['CRPIX2'])*header['CDELT2'] + header['CRVAL2']
ext = ds.arr([yr, yl, xr, xl], input_units='cm').in_units('kpc')


I_bin, Q_bin, U_bin = {}, {}, {}
psi, frac = {}, {}

for nu in nus:
    I_bin[nu] = gaussian_filter(frb_I[nu], sigma)
    Q_bin[nu] = gaussian_filter(frb_Q[nu], sigma)
    U_bin[nu] = gaussian_filter(frb_U[nu], sigma)
    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 [ ]:

    

nu1, nu2 = nus

fig = plt.figure(figsize=(8,3))
ax = fig.add_subplot(111)
cmap = plt.cm.hot
norm = yt.YTQuantity(*nu1).in_units('GHz').value**0.5
im = ax.imshow(np.log10(I_bin[nu1]).transpose(), cmap=cmap, 
               vmin=np.log10(1E-3/norm), vmax=np.log10(1E1/norm), 
               extent=ext, 
               origin='lower', aspect='equal')
ax.set_facecolor('k')
cbar = fig.colorbar(im)
cbar.ax.tick_params(direction='in')

timestamp = '%.1f Myr' % ds.current_time.in_units('Myr')
ax.annotate(timestamp, (0.04, 0.75) , xycoords='axes fraction', color='k')


    

In [ ]:

    

pol = np.sqrt(Q_bin[nu1]**2 + I_bin[nu1]**2)
weight = pol/pol.max()


    

In [ ]:

    

nu1, nu2 = nus
# I1 = gaussian_filter(frb_I[nu1], sigma)
# I2 = gaussian_filter(frb_I[nu2], sigma)
#alpha = np.log10(I2/I1)/np.log10(nu2[0]/nu1[0])
#alpha = np.ma.masked_where(I2<1E-7, np.array(alpha))

fig = plt.figure(figsize=(8,3))
ax = fig.add_subplot(111)
cmap = plt.cm.RdBu
im = ax.imshow((psi[nu1]).transpose(), cmap=cmap, vmin=-1.5, vmax=1.5, extent=ext, origin='lower', aspect='equal')
ax.set_facecolor('w')
cbar = fig.colorbar(im)
cbar.ax.tick_params(direction='in', )
#cbar.ax.set_yticks([-0.5, -0])
#ax.annotate(labels[i % len(dir)], (0.65, 0.75) , xycoords='axes fraction', color='white')
timestamp = '%.1f Myr' % ds.current_time.in_units('Myr')
ax.annotate(timestamp, (0.04, 0.85) , xycoords='axes fraction', color='k')


    

In [ ]:

    

mask = frb_I[nu1] > 0.0
null = plt.hist(psi[nu1][mask].flatten(), bins=100)


    

In [ ]:

    

fig = plt.figure(figsize=(8,3))
ax = fig.add_subplot(111)
cmap = plt.cm.gist_heat_r

#weight = I_bin[nu2]*frac[nu2]/I_bin[nu2].flatten().max()/frac[nu2].flatten().max()
weight = 1
image = np.abs(psi[nu2]-psi[nu1])*weight
image = np.where(image < np.pi/2, image, np.pi-image)
im = ax.imshow(image.transpose(), cmap=cmap, vmin=0, vmax=np.pi/2, extent=ext, origin='lower', aspect='equal')
ax.set_facecolor('w')
cbar = fig.colorbar(im)
cbar.ax.tick_params(direction='in')
#cbar.ax.set_yticks([-0.5, -0])
timestamp = '%.1f Myr' % ds.current_time.in_units('Myr')
ax.annotate(timestamp, (0.04, 0.85) , xycoords='axes fraction', color='k')


    

In [ ]:

    

nu1, nu2 = nus
# I1 = gaussian_filter(frb_I[nu1], sigma)
# I2 = gaussian_filter(frb_I[nu2], sigma)
#alpha = np.log10(I2/I1)/np.log10(nu2[0]/nu1[0])
#alpha = np.ma.masked_where(I2<1E-7, np.array(alpha))
#ext = ds.arr([-0.5*width[0], 0.5*width[0], -0.5*width[1], 0.5*width[1]]).in_units('kpc')

fig = plt.figure(figsize=(8,3))
ax = fig.add_subplot(111)
cmap = plt.cm.viridis
im = ax.imshow(frac[nu1].transpose(), cmap=cmap, vmin=0, vmax=0.7, extent=ext, origin='lower', aspect='equal')
ax.set_facecolor('w')
cbar = fig.colorbar(im)
cbar.ax.tick_params(direction='in')
#cbar.ax.set_yticks([-0.5, -0])


#ax.annotate(labels[i % len(dir)], (0.65, 0.75) , xycoords='axes fraction', color='white')

timestamp = '%.1f Myr' % ds.current_time.in_units('Myr')
ax.annotate(timestamp, (0.04, 0.85) , xycoords='axes fraction', color='k')


    

In [ ]:

    

nu1, nu2 = nus
# I1 = gaussian_filter(frb_I[nu1], sigma)
# I2 = gaussian_filter(frb_I[nu2], sigma)
#alpha = np.log10(I2/I1)/np.log10(nu2[0]/nu1[0])
#alpha = np.ma.masked_where(I2<1E-7, np.array(alpha))
#ext = ds.arr([-0.5*width[0], 0.5*width[0], -0.5*width[1], 0.5*width[1]]).in_units('kpc')

fig = plt.figure(figsize=(8,3))
ax = fig.add_subplot(111)
cmap = plt.cm.RdBu
im = ax.imshow((frac[nu2]-frac[nu1]).transpose(), cmap=cmap, vmin=-0.5, vmax=0.5, extent=ext, origin='lower', aspect='equal')
ax.set_facecolor('w')
cbar = fig.colorbar(im)
cbar.ax.tick_params(direction='in')
#cbar.ax.set_yticks([-0.5, -0])


#ax.annotate(labels[i % len(dir)], (0.65, 0.75) , xycoords='axes fraction', color='white')

timestamp = '%.1f Myr' % ds.current_time.in_units('Myr')
ax.annotate(timestamp, (0.04, 0.85) , xycoords='axes fraction', color='k')


    

In [ ]:

    

dir = '/d/d5/ychen/2015_production_runs/0529_h1_00/'
ds = yt.load(dir + 'data/MHD_Jet_hdf5_plt_cnt_0600')
proj_axis = [1,0,2]
zoom_fac = 6
nus = [(150, 'MHz'), (1400, 'MHz')]

ptype = 'lobe'
extend_cells = 8
res = (512, 256)
sigma = 1

print(ds.basename, ds.current_time.in_units('Myr'))
width = ds.domain_width[[2,1]]/zoom_fac

fitsname = synchrotron_fits_filename(ds, dir, ptype, proj_axis)
hdulist = pyfits.open(fitsname)
frb_I, frb_Q, frb_U = {},{},{}
for nu in nus:
    stokes = StokesFieldName(ptype, nu, proj_axis, field_type='flash')
    frb_I[nu] = hdulist[stokes.I[1]].data
    frb_Q[nu] = hdulist[stokes.Q[1]].data
    frb_U[nu] = hdulist[stokes.U[1]].data
header = hdulist[stokes.I[1]].header
xr = -header['CRPIX1']*header['CDELT1'] + header['CRVAL1']
xl = (header['NAXIS1'] - header['CRPIX1'])*header['CDELT1'] + header['CRVAL1']
yr = -header['CRPIX2']*header['CDELT2'] + header['CRVAL2']
yl = (header['NAXIS2'] - header['CRPIX2'])*header['CDELT2'] + header['CRVAL2']
ext = ds.arr([yr, yl, xr, xl], input_units='cm').in_units('kpc')


I_bin, Q_bin, U_bin = {}, {}, {}
psi, frac = {}, {}

for nu in nus:
    I_bin[nu] = gaussian_filter(frb_I[nu], sigma)
    Q_bin[nu] = gaussian_filter(frb_Q[nu], sigma)
    U_bin[nu] = gaussian_filter(frb_U[nu], sigma)
    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 [ ]:

    

nu1, nu2 = nus
# I1 = gaussian_filter(frb_I[nu1], sigma)
# I2 = gaussian_filter(frb_I[nu2], sigma)
#alpha = np.log10(I2/I1)/np.log10(nu2[0]/nu1[0])
#alpha = np.ma.masked_where(I2<1E-7, np.array(alpha))
#ext = ds.arr([-0.5*width[0], 0.5*width[0], -0.5*width[1], 0.5*width[1]]).in_units('kpc')

fig = plt.figure(figsize=(8,3))
ax = fig.add_subplot(111)
cmap = plt.cm.RdBu
im = ax.imshow(psi[nu1].transpose(), cmap=cmap, vmin=-1.5, vmax=1.5, extent=ext, origin='lower', aspect='equal')
ax.set_facecolor('w')
cbar = fig.colorbar(im)
cbar.ax.tick_params(direction='in')
#cbar.ax.set_yticks([-0.5, -0])


#ax.annotate(labels[i % len(dir)], (0.65, 0.75) , xycoords='axes fraction', color='white')

timestamp = '%.1f Myr' % ds.current_time.in_units('Myr')
ax.annotate(timestamp, (0.04, 0.85) , xycoords='axes fraction', color='k')


    

In [ ]:

    

fig = plt.figure(figsize=(8,3))
ax = fig.add_subplot(111)
cmap = plt.cm.gist_heat_r
image = np.abs(psi[nu2]-psi[nu1])*weight
image = np.where(image < np.pi/2, image, np.pi-image)
im = ax.imshow(image.transpose(), cmap=cmap, vmin=0, vmax=np.pi/2, extent=ext, origin='lower', aspect='equal')
ax.set_facecolor('w')
cbar = fig.colorbar(im)
cbar.ax.tick_params(direction='in')
#cbar.ax.set_yticks([-0.5, -0])
timestamp = '%.1f Myr' % ds.current_time.in_units('Myr')
ax.annotate(timestamp, (0.04, 0.85) , xycoords='axes fraction', color='k')


    

In [ ]:

    

nu1, nu2 = nus

fig = plt.figure(figsize=(8,3))
ax = fig.add_subplot(111)
cmap = plt.cm.viridis
im = ax.imshow(frac[nu2].transpose(), cmap=cmap, vmin=0, vmax=0.7, extent=ext, origin='lower', aspect='equal')
ax.set_facecolor('w')
cbar = fig.colorbar(im)
cbar.ax.tick_params(direction='in')
timestamp = '%.1f Myr' % ds.current_time.in_units('Myr')
ax.annotate(timestamp, (0.04, 0.85) , xycoords='axes fraction', color='k')


    

In [ ]:

    

frb_I = proj.frb.data[stokes.I].v
frb_Q = proj.frb.data[stokes.Q].v
frb_U = proj.frb.data[stokes.U].v
#null = plt.hist(np.log10(arri.flatten()), range=(-15,3), bins=100)

sigma = 8

frb_I = gaussian_filter(frb_I, 8)
frb_Q = gaussian_filter(frb_Q, 8)
frb_U = gaussian_filter(frb_U, 8)

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

I_bin = frb_I.reshape(ny, factor, nx, factor).sum(3).sum(1)
Q_bin = frb_Q.reshape(ny, factor, nx, factor).sum(3).sum(1)
U_bin = frb_U.reshape(ny, factor, nx, 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
#print len(psi.flatten()[psi.flatten().nonzero()])
#null = plt.hist(np.abs(psi.flatten())[psi.flatten().nonzero()], bins=50)
#plt.xlim(0, 0.5*np.pi)


    

In [ ]: