In [ ]:
#%pdb
%matplotlib inline
import matplotlib.pyplot as plt
import os
import yt
yt.mylog.setLevel("INFO")
import numpy as np
from yt_synchrotron_emissivity import *
from mpl_toolkits.axes_grid1 import AxesGrid

In [ ]:
fname = '/d/d5/ychen/2015_production_runs/1022_h1_10Myr/data/MHD_Jet_10Myr_hdf5_plt_cnt_0910_synchrotron_peak_gc8'
#ds = yt.load(fname)
#setup_part_file(ds)
#ds.current_time.in_units('Myr')
#data = ds.all_data()

In [ ]:
ptype = 'lobe'
proj_axis = [1,0,2]
nus = [(150, 'MHz'), (1400, 'MHz')]
#nus = [(150, 'MHz')]
zoom_fac = 8
extend_cells=8
#pars = add_synchrotron_emissivity(ds, ptype=ptype, nu=nu)
#pars = add_synchrotron_dtau_emissivity(ds, ptype=ptype, nu=nu, proj_axis=proj_axis)
stokes = StokesFieldName(ptype, nus[0], proj_axis, field_type='flash')

In [ ]:
ds_sync = yt.load(fname)
ad = ds_sync.all_data()

In [ ]:
null = plt.hist(np.log10(ad[stokes.I]), range=(-30,-23))

In [ ]:
plot = yt.SlicePlot(ds_sync, 'x', stokes.I).zoom(8)
plot.set_zlim(stokes.I, 1E-28, 1E-23)
plot.show()

In [ ]:
#ds_sync = yt.load(synchrotron_file_name(ds, extend_cells=extend_cells))
ds_sync = yt.load(fname)
ds_sync.field_list
#ds_sync.coordinates.x_axis['x'] = 2
#ds_sync.coordinates.x_axis[0] = 2
#ds_sync.coordinates.y_axis['x'] = 1
#ds_sync.coordinates.y_axis[0] = 1
#ds_sync.field_units
beam_area = yt.YTQuantity(0.023, 'arcsec**2')
for nu in nus:
    stokes = StokesFieldName(ptype, nu, proj_axis, field_type='flash')
    for field in stokes.IQU:
        ds_sync.unit_registry.add('beam', float(beam_area.in_units('rad**2').v), dimensions=yt.units.dimensions.solid_angle)
        ds_sync.field_info[field].units = 'Jy/cm/arcsec**2'
        ds_sync.field_info[field].output_units = 'Jy/cm/beam'

In [ ]:
projs = {}
width = ds_sync.domain_width[[1,2]]/zoom_fac
#res = ds_sync.domain_dimensions[1:]*ds_sync.refine_by**ds_sync.index.max_level//zoom_fac//2
res = (1200, 2400)
for nu in nus:
    stokes = StokesFieldName(ptype, nu, proj_axis, field_type='flash')
    projs[nu] = yt.ProjectionPlot(ds_sync, proj_axis, stokes.IQU, center=[0,0,0], width=width).set_buff_size(res)

In [ ]:
stokes = StokesFieldName(ptype, nus[0], proj_axis, field_type='flash')
frb = projs[nus[0]].frb[stokes.I]
print(frb.mean())
beam_area = yt.YTQuantity(0.023, 'arcsec**2')

frb.units.registry.add('beam', float(beam_area.in_units('rad**2').v), dimensions=yt.units.dimensions.solid_angle)
print(frb.in_units('Jy/beam').mean())
prj = ds_sync.proj(stokes.I, 'x')

In [ ]:
prj[stokes.I].mean()

In [ ]:
prjfrb = prj.to_frb(width[0], (1200, 2400), height=width[1])
prjfrb[stokes.I].mean()

In [ ]:
frb_I = {}
for nu, proj in projs.items():
    stokes = StokesFieldName(ptype, nu, proj_axis, field_type='flash')
    norm = yt.YTQuantity(*nu).in_units('GHz').value**0.5
    #proj.set_figure_size(15)
    #proj.set_font_size(36)
    for field in [stokes.I]:
        proj.set_zlim(field, 1E-3/norm, 1E1/norm)
        cmap = plt.cm.hot
        cmap.set_bad('k')
        proj.set_cmap(field, cmap)

    for field in [stokes.Q, stokes.U]:
        proj.set_zlim(field, -5, 5)
        proj.set_cmap(field, 'seismic')
    proj._recreate_frb()
    frb_I[nu] = proj.frb.data[stokes.I].v
    frb_Q = proj.frb.data[stokes.Q].v
    frb_U = proj.frb.data[stokes.U].v
    #proj.annotate_timestamp(corner='upper_left', time_format="{time:6.3f} {units}",
    #                   time_unit='Myr', draw_inset_box=True)
    proj.annotate_clear()
    #proj.annotate_polline(frb_I[nu], frb_Q, frb_U, factor=25, scale=15)
    #proj.annotate_contour(stokes.I, clim=(1E-3/norm, 4E0/norm))
    proj.set_axes_unit('kpc')
    #proj.set_figure_size(15)
    proj.show()

In [ ]:
from scipy.ndimage import gaussian_filter
import matplotlib
matplotlib.rcParams['font.family'] = 'stixgeneral'
sigma = 1
nu1, nu2 = nus
I1 = gaussian_filter(frb_I[nu1], sigma)
I2 = gaussian_filter(frb_I[nu2], sigma)
alpha = np.log10(I2/I1)/np.log10(1400/150)
alpha = np.ma.masked_where(I2<1E-3, np.array(alpha))
ext = ds.arr([-0.5*width[0], 0.5*width[0], -0.5*width[1], 0.5*width[1]])
plt.figure(figsize=(6,10), dpi=150)
cmap = plt.cm.jet
cmap.set_bad('navy')
plt.imshow(alpha, cmap=cmap, vmin=-1.5, vmax=-0.5, extent=ext.in_units('kpc'), origin='lower', aspect='equal')
plt.xlabel('z (kpc)')
plt.ylabel('y (kpc)')
plt.axes().tick_params(direction='in')
cb = plt.colorbar(fraction=0.10, pad=0, aspect=50)
cb.set_label('Spectral Index (%s) (1.4GHz/150MHz)' % ptype)
cb.ax.tick_params(direction='in')

In [ ]:
width/2 - width

In [ ]:
from yt.utilities.fits_image import FITSProjection
fieldi = ('nn_emissivity_i_%s_%%.1f%%s' % ptype) % nu
#fits_proj = FITSProjection(ds, proj_axis, [fieldi, fieldq, fieldu], center=[0,0,0], width=((40,'kpc'),(80,'kpc')))
width = ds.domain_width[1:]/8
res = ds.domain_dimensions[1:]*ds.refine_by**ds.index.max_level//8
fits_proj = FITSProjection(ds, proj_axis, fieldi, center=[0,0,0], width=width, image_res=res)

In [ ]:
fits_proj.writeto('emissivity_i_600_1.4GHz.fits')

In [ ]:
frb = proj.frb.data[stokes.U]
fig = plt.figure(figsize=(8,16))
plt.imshow(gaussian_filter(frb, 8), cmap='seismic', vmin=-5, vmax=5)
print(frb.shape)

In [ ]:
yt.visualization.fixed_resolution_filters.filter_registry

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 [ ]:
fig = plt.figure(figsize=(16,4))

ax1 = fig.add_subplot(131)
null = ax1.hist(frac[I_bin.nonzero()].flatten()*100, range=(0,80), bins=40,
                                  weights=I_bin[I_bin.nonzero()].flatten() )
ax1.set_xlim(0,80)
ax1.set_xlabel('Polarization fraction (%)')
#ax1.set_xlim(0, 80)

ax2  = fig.add_subplot(132)
null = ax2.hist(psi[I_bin.nonzero()].flatten(), bins=50, range=(-0.5*np.pi, 0.5*np.pi),
                                  weights=I_bin[I_bin.nonzero()].flatten() )
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.add_subplot(133)
null = ax3.hist(np.abs(psi[I_bin.nonzero()].flatten()), bins=25, range=(0.0, 0.5*np.pi))
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:])

In [ ]:
fig = plt.figure(figsize=(15,20))
plt.imshow(psi)
plt.colorbar()

In [ ]:
field = fieldi
xx0, xx1 = plot.xlim
print xx0, xx1
yy0, yy1 = plot.ylim
print plot.plots

factor = 10.0
nx = plot.plots[field].image._A.shape[0] / factor
ny = plot.plots[field].image._A.shape[1] / factor
X,Y = np.meshgrid(np.linspace(xx0,xx1,nx,endpoint=True),\
                  np.linspace(yy0,yy1,ny,endpoint=True))

#frb_U = X
#frb_Q = Y
#frb_I = np.sqrt(X*X+Y*Y)

psi = 0.5*np.arctan(frb_U[::factor,::factor]/frb_Q[::factor,::factor])
frac = np.sqrt(frb_Q[::factor,::factor]**2+frb_U[::factor,::factor]**2)/frb_I[::factor,::factor]
#null = plt.hist(frac.flatten(), range=(0,5), bins=100)

In [ ]:
convolved_image = frb_I
convolved_image_q = frb_Q
convolved_image_u = frb_U
# First, we plot the background image
fig = plt.figure(figsize=(8,16))
i_plot = fig.add_subplot(111)
i_plot.imshow(np.log10(convolved_image+1e-3), vmin=-1, vmax=1, origin='lower')

# ranges of the axis
xx0, xx1 = i_plot.get_xlim()
yy0, yy1 = i_plot.get_ylim()

# binning factor
factor = [32, 32]

# re-binned number of points in each axis
nx_new = convolved_image.shape[1] // factor[0]
ny_new = convolved_image.shape[0] // factor[1]

# These are the positions of the quivers
X,Y = np.meshgrid(np.linspace(xx0,xx1,nx_new,endpoint=True),
                  np.linspace(yy0,yy1,ny_new,endpoint=True))

# bin the data
I_bin = convolved_image.reshape(nx_new, factor[0], ny_new, factor[1]).sum(3).sum(1)
Q_bin = convolved_image_q.reshape(nx_new, factor[0], ny_new, factor[1]).sum(3).sum(1)
U_bin = convolved_image_u.reshape(nx_new, factor[0], ny_new, factor[1]).sum(3).sum(1)

# polarization angle
psi = 0.5*np.arctan2(U_bin, Q_bin)

# polarization fraction
frac = np.sqrt(Q_bin**2+U_bin**2)/I_bin

# mask for low signal area
mask = I_bin < 0.1

frac[mask] = 0
psi[mask] = 0

pixX = frac*np.cos(psi) # X-vector 
pixY = frac*np.sin(psi) # Y-vector

# keyword arguments for quiverplots
quiveropts = dict(headlength=0, headwidth=1, pivot='middle')
i_plot.quiver(X, Y, pixX, pixY, scale=8, **quiveropts)
plt.show()

In [ ]:
plot.save()

In [ ]:
p = plot.plots[fieldi]
quiveropts = dict(headlength=code0, headwidth=1, pivot='middle')#, scale=3, 
    #linewidth=.5, units='xy', width=.05, headwidth=1)v
p.axes.cla()
p.axes.quiver(X,Y,frac*np.sin(psi),frac*np.cos(psi), **quiveropts)
#p.show()
plot.show()

In [ ]:
code0 = ds.quan(0, 'code_length')

In [ ]:
plot.zoom(64)
plot.annotate_grids()
plot.annotate_particles((0.5, 'kpc'), col='white', ptype='jetp')
plot.set_zlim(field, 1E-1, 1E3)
plot.set_cmap(field, cmap)
plot.show()

In [ ]:
ptype = 'all'
#B = np.array([data[(ptype, 'particle_magx')], data[(ptype, 'particle_magy')], data[(ptype, 'particle_magz')]])\
B_cut = 1E-20
Bvec = np.array([data[('flash', 'magx')], data[('flash', 'magy')], data[('flash', 'magz')]])\
    *np.sqrt(4.0*np.pi)
print Bvec.shape
#B = B[:,np.nonzero(B2)][:,0,:]
#print B.shape

In [ ]:
B[:,1]
B2 = np.sum(B*B, axis=0)
print B2.shape
flag = B2 > B_cut**2
print flag.shape
print B.shape
print B[:,flag].shape
#print B[:,np.nonzero(B2)].shape
#print B[:,np.nonzero(B2)][:,0,:].shape

In [ ]:
import matplotlib.pyplot as plt
plt.hist(np.log10(B2[B2>0.0]), bins=40)#, range=(2E-8, 1E-7))

#print cross
#print np.sum(cross*cross, axis=1)
#print np.sum(B*B, axis=0)

In [ ]:
for proj_axis in ['x', 'y', 'z']:
    if proj_axis=='x': los = [1.,0.,0.]
    elif proj_axis=='y': los = [0.,1.,0.]
    elif proj_axis=='z': los = [0.,0.,1.]

    cross = np.cross(los, B, axisb=0)
    sina = np.sqrt(np.sum(cross*cross, axis=1)/np.sum(B*B, axis=0))
    #print sina.shape
    #print "sin(alpha)=", sina
    a = np.arcsin(sina)/np.pi*180
    #print 'alpha=', a
    dummy = plt.hist(a, bins=40, alpha=0.5, label=proj_axis)
plt.xlabel('B position angle (deg)')
plt.legend(loc=2)

In [ ]:
from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

xs, ys, zs = zip([0,0,0], LOS)
ax.plot(xs, ys, zs)
xs, ys, zs = zip([0,0,0], B0/np.sqrt(np.sum(B0*B0)))
ax.plot(xs, ys, zs)



ax.set_xlim(-1,1)
ax.set_ylim(-1,1)
ax.set_zlim(-1,1)

In [ ]:
Bvec = np.array([data[('flash', 'magx')], data[('flash', 'magy')], data[('flash', 'magz')]])\
    *np.sqrt(4.0*np.pi)
print Bvec.shape
los = np.array([1.,0.,0.])
print los.shape

In [ ]:
B2 = np.sum(Bvec*Bvec, axis=0)
mask = B2>0.0
Bvec = Bvec[:,mask]
print Bvec.shape
print np.inner(Bvec.transpose(),los)

In [ ]:
inner = np.inner(Bvec.transpose(), los)
print inner.shape
inner.reshape((inner.shape[0], 1))
print inner
print (inner.reshape((-1, 1))*los).shape

#Bproj = Bvec - *los

In [ ]:
Bvec.transpose() - np.inner(Bvec.transpose(), los).reshape((-1,1))*los

In [ ]:
from mpl_toolkits.axes_grid1 import AxesGrid
dirs = ['/home/ychen/data/0only_0605_hinf/',\
        '/home/ychen/data/0only_0529_h1/',\
        '/home/ychen/data/0only_0605_h0/']

labels = ['toroidal', 'helical', 'poloidal', 'hydro']
ptype = 'lobe'
pol = 'i'
nu = (150, 'MHz')
proj_axis = 'x'
norm = yt.YTQuantity(*nu).in_units('GHz').value**0.5
field = ('deposit', ('nn_emissivity_%s_%s_%%.1f%%s' % (pol, ptype)) % nu)
center=(0.0,0.0,0.0)

fig = plt.figure()

grid = AxesGrid(fig, (0.075,0.075,0.85,0.85),
                nrows_ncols = (1, 4),
                axes_pad = 0.05,
                label_mode = "L",
                share_all = True,
                cbar_location="right",
                cbar_mode="single",
                cbar_size="5%",
                cbar_pad="0%")

for i, dir in enumerate(dirs):
    # Load the data and create a single plot
    ds = yt.load(os.path.join(dir, 'MHD_Jet_hdf5_plt_cnt_0640'))# load data
    pars = add_synchrotron_pol_emissivity(ds, ptype=ptype, nu=nu, proj_axis=proj_axis)
    
    p=yt.ProjectionPlot(ds, proj_axis, field, center=center, origin='center-domain',\
                              width=((40,'kpc'), (80,'kpc')),)
    # Ensure the colorbar limits match for all plots
    p.set_zlim(field, 1E-3/norm, 1E1/norm)
    cmap = plt.cm.hot
    cmap.set_bad('k')
    p.set_cmap(field, cmap)
    p.annotate_text((0.75, 0.95), labels[i], coord_system='axis', text_args={'color':'grey'})
    p.set_buff_size((200,400))
    p.set_font_size(9)
    
    if i == 0:
        p.annotate_timestamp(0.05, 0.95, time_format="{time:6.3f} {units}", time_unit='Myr', text_args={'color':'grey'})
    

    # This forces the ProjectionPlot to redraw itself on the AxesGrid axes.
    plot = p.plots[field]
    plot.figure = fig
    plot.axes = grid[i].axes
    plot.cax = grid.cbar_axes[i]

    # Finally, this actually redraws the plot.
    p._setup_plots()
    p.plots[field].axes.set_xticks([-10,0,10,20])
    p.plots[field].axes.set_xticklabels([-10,0,10,20])

fig.set_figwidth(10)
fig.set_figheight(5)
fig.savefig('compare_4_synchrotron_lobe_I_64.pdf')

In [ ]: