In [ ]:
%matplotlib inline
import os
import sys
sys.path.append('/home/ychen/lib/util')
import util
import numpy as np
#import matplotlib
import yt
yt.mylog.setLevel("INFO")
import matplotlib
matplotlib.rcParams['figure.dpi'] = 150
from matplotlib import pyplot as plt
from plotSlices import plotSliceField
from tools import calcDen0 #,read_par, calcNozzleCoords
from particles.particle_filters import *
#def _jetmomentum_z(field, data):
# return data['gas', 'density'] * data['gas', 'velocity_z'] * data['flash', 'jet ']
#yt.add_field( 'jetmomentum_z', function=_jetmomentum_z, units='g/cm**2/s', take_log=False)
In [ ]:
#dir = '/home/ychen/data/0only_1204_M24_b01'
#dir = '/home/ychen/data/0only_1022_h1_10Myr'
#dir ='/home/ychen/d9/FLASH4/stampede/0314_L45_M10_b1_h1_nojiggle/'
#dir ='/home/ychen/data/0only_1106_M3_h1/'
#dir ='/home/ychen/d9/FLASH4/2015_production_runs/1022_L45_M10_b1_h1_10Myr/'
#dir ='/home/ychen/d9/FLASH4/Gravity_test/0627_tree_gravity/'
dir = '/home/ychen/data/0only_0529_h1/data/'
#dir = '/d/d5/ychen/2015_production_runs/1022_h1_10Myr/data'
#dir = '/home/ychen/Mount/engels2/data/1212_L45_M10_b1_h0_10Myr/'
fields_grid = ['density', 'velocity_z']
def rescan(printlist=False):
files = util.scan_files(dir, '*hdf5_chk_[0-9][0-9][0-9][0-9]', printlist=printlist, walk=True)
return files
files = rescan(True)
#for f in ds.derived_field_list:
#for f in ds.field_list:
# print f
ds = yt.load(files[0].fullpath)
#ds.index
In [ ]:
#ds.derived_field_list
box = ds.box(ds.domain_left_edge/16, ds.domain_right_edge/16)
plot = yt.SlicePlot(ds, 'x', fields=['current_density_z'], data_source=box).zoom(16)
plot.show()
In [ ]:
ds.derived_field_list
In [ ]:
plot.show()
In [ ]:
for f in files[-40:-39]:
ds = yt.load(f.fullpath)
plot = yt.SlicePlot(ds, 'x', fields=['mach_number'], width=(20, 'kpc'))
plot.set_zlim('mach_number', 2, 10)
plot.set_log('mach_number', False)
plot.show()
In [ ]:
plot.annotate_grids()
In [ ]:
ds.periodicity = (True, True, True)
ad = ds.all_data()
ptype = 'io'
ds.add_particle_filter(ptype)
xfield = (ptype, 'particle_position_y')
yfield = (ptype, 'particle_position_z')
cfield = (ptype, 'particle_type')
plot = yt.ParticlePhasePlot(ad, xfield, yfield, cfield)
plot.show()
In [ ]:
plot.zoom(16)
plot.show()
In [ ]:
ds1 = yt.load(files[3].fullpath)
ds2 = yt.load(files[-2].fullpath)
for key in ds1.parameters.iterkeys():
if ds1.parameters[key] != ds2.parameters[key]:
print '%20s %10.3e %10.3e' % (key, ds1.parameters[key], ds2.parameters[key])
#print key, ds1.parameters[key], ds2.parameters[key]
#key = 'nozzlevecy'
In [ ]:
ds = yt.load(files[-1].fullpath)
#box = ds.box(ds.domain_left_edge, ds.domain_right_edge)
#total_mass = sum(box['cell_mass'])
box = ds.box(ds.domain_left_edge/16, ds.domain_right_edge/16)
#total_mass_8 = sum(box['cell_mass'])
nbin = 40
vlim = (0,4.0E9)
null = plt.hist(box['velocity_magnitude'], bins=nbin, range=vlim, label='Bulk Velocity', alpha=0.3)
#null = plt.hist(np.abs(box['velocity_z']), bins=nbin, range=vlim, label='|Z Velocity|', alpha=0.3)
null = plt.hist(box['alfven_speed'], bins=nbin, range=vlim, label='Alfven Speed', alpha=0.5)
#null = plt.hist(box['sound_speed'], bins=nbin, range=vlim, label='Sound Speed', alpha=0.5)
plt.ylim(0, 5E5)
plt.legend()
In [ ]:
print max(box['sound_speed'])
print max(box['alfven_speed'])
print max(box['velocity_magnitude'])
In [ ]:
field = 'dens'
proj_axis = 'z'
center = (0,0,0)
ds = yt.load(files[6].fullpath)
print ds
ad = ds.all_data()
neg_partdens = ad['particle_dens'] < 0
xx = ad['particle_posx'][neg_partdens]
yy = ad['particle_posy'][neg_partdens]
zz = ad['particle_posz'][neg_partdens]
for x,y,z in zip(xx,yy,zz):
center = (x,y,z)
print center
plot=yt.SlicePlot(ds, proj_axis, field, center=center, origin='center-domain',\
width=((10,'kpc'), (10,'kpc')))
plot.annotate_marker(center, marker='o', plot_args={'color':'None', 'edgecolor':'white', 's':50, 'lw':1})
plot.show()
#plot.set_log(field, False)
#plot.set_zlim(field, -4E9, 4E9)
#plot.set_zlim(field, 1E-12, 5E-9)
#plot.set_cmap(field, 'seismic')
#plot.show()
In [ ]:
proj_axis = 'x'
ptype = 'jetp'
fields = ['magnetic_field_strength']
#center=([0.,2.,25.],'kpc')
center=(0.0,0.0,0.0)
LeftEdge = yt.YTArray([-80, -80, -160], 'kpc')
RightEdge = yt.YTArray([80, 80, 160], 'kpc')
#if False:
for file in files[-1:]:
ds = yt.load(file.fullpath)
ds.add_particle_filter(ptype)
ds.periodicity = (True, True, True)
#data = ds.box(LeftEdge, RightEdge)
#data = ds.all_data()
#nozzleCoords = calcNozzleCoords(pars, ds.current_time.value, proj_axis=proj_axis)
for field in fields:
if field in ['particle_gamc']:
xfield = (ptype, 'particle_position_y')
yfield = (ptype, 'particle_position_z')
cfield = (ptype, 'particle_gamc')
plot = yt.ParticlePlot(ds, xfield, yfield, cfield)#, width=((40,'kpc'),(80,'kpc')))
plot.set_zlim(cfield, 1E2, 1E5)
plot.set_cmap(cfield, 'algae')
else:
#plot=yt.SlicePlot(ds, proj_axis, field, center=center, width=((30,'kpc'),(50,'kpc')), origin='center-domain')
plot=yt.SlicePlot(ds, proj_axis, field, center=center, origin='center-domain')#,\
#width=((40,'kpc'), (80,'kpc')))
if field in ['magnetic_field_x', 'velocity_z']:
plot.set_log(field, False)
plot.set_cmap(field, 'seismic')
if field in ['magnetic_field_x', 'magnetic_field_z']:
plot.set_zlim(field, -5E-5, 5E-5)
if field == 'shok':
plot.set_log('shok', False)
plot.set_zlim('shok', 0, 1)
plot.set_cmap('shok', 'gist_heat_r')
if field == 'jet ':
plot.set_log('jet ', True)
plot.set_zlim('jet ', 1E-10, 1.0)
plot.set_cmap('jet ', 'gist_heat')
#else:
# plot.set_cmap(field, 'algae')
plot.zoom(4)
plot.set_figure_size(6)
#plot.annotate_grids()
#plot.set_cmap('density', 'gist_heat_r')
#plot.set_log(field, False)
#plot.set_zlim('velocity_z', -3E9, 3E9)
plot.annotate_timestamp(0.05, 0.95, time_format="{time:6.3f} {units}", time_unit='Myr', text_args={'color':'k'})
#plot.annotate_line_integral_convolution('magnetic_field_y', 'magnetic_field_z', \
# lim=(0.4,0.65), cmap='binary', alpha=0.8)#, const_alpha=True)
plot.set_zlim('magnetic_field_strength', 1E-6, 1E-5)
#plot.set_zlim('magnetic_pressure', 2.5e-14, 2.5e-11)
#plot.set_zlim('pressure', 1.0e-12, 2e-9)
#plot.set_zlim('density', 1.0e-28, 1e-25)
#plot.show()
#slab_width=3E21
#plot.annotate_particles(slab_width, col=ad['all', 'particle_gamc'], alpha=0.5)
#plot=plotSliceField(ds, proj_axis=proj_axis, field=field, center=center, zoom_fac=2,\
# plotgrid=True,plotvelocity=False,\
# savepath=None,\
# show=False)
#plot.annotate_quiver('velocity_y', 'velocity_z', 30, scale=1E10)
#plot.annotate_velocity(15, scale=1E10)
#plot.set_log(field, False)
#plot.zoom(8)
#plot.set_buff_size((2,4))
plot.show()
In [ ]:
plot.set_zlim(field, 1E-6, 1E-4)
plot.set_cmap(field, 'arbre')
In [ ]:
width = ds.domain_width[1:]/4
print(width.v)
fitsslice = yt.FITSSlice(ds, proj_axis, field, center=center, width=width.v, image_res=(1024,2048))
fitsslice.writeto('magnetic_field_strength_1380.fits', clobber=True)
In [ ]:
from yt_synchrotron_emissivity import *
nu = (150, 'MHz')
norm = yt.YTQuantity(*nu).in_units('GHz').value**0.5
print add_synchrotron_pol_emissivity(ds, ptype=ptype, nu=nu, proj_axis=proj_axis)
pol = 'i'
field = ('deposit', ('nn_emissivity_%s_%s_%%.1f%%s' % (pol, ptype)) % nu)
plot = yt.SlicePlot(ds, proj_axis, field, center=[0,0,0], width=((40,'kpc'),(80,'kpc')))
plot.annotate_grids(alpha=0.1)
plot.set_zlim(field,1E-30,1E-22)
plot.set_figure_size(6)
plot.show()
In [ ]:
proj = yt.ProjectionPlot(ds, proj_axis, field, center=[0,0,0], width=((40,'kpc'),(80,'kpc')))
cmap = plt.cm.hot
cmap.set_bad('k')
proj.set_cmap(field, cmap)
proj.set_figure_size(6)
proj.set_zlim(field, 1E-3/norm, 1E1/norm)
proj.annotate_timestamp(corner='upper_left', time_format="{time:6.3f} {units}", time_unit='Myr')
proj.show()
In [ ]:
proj.save('MHD_Jet_hdf5_plt_cnt_0640_Projection_x_nn_emissivity_i_lobe_150MHz.pdf')
In [ ]:
ad = ds.all_data()
filter = ad['index', 'z'] > 0
mass_weighted_z = sum(ad['flash', 'jet '][filter]*ad['cell_mass'][filter]*ad['index', 'z'][filter])\
/sum(ad['flash', 'jet '][filter]*ad['cell_mass'][filter])
print mass_weighted_z
In [ ]:
ds = yt.load(files[20].fullpath)
ad = ds.all_data()
tag = ad['all', 'particle_tag']
filter = np.logical_and((tag % 20 == 0), (ad['all', 'particle_shok']==0))
y = ad['all', 'particle_position_y'][filter]/3.08567758E21
z = ad['all', 'particle_position_z'][filter]/3.08567758E21
gamc = np.log10(np.abs(ad['all', 'particle_gamc'][filter]))
#plt.hist(tag, bins=140)
#plt.scatter(ad['all', 'particle_tadd'], ad['all', 'particle_tag'], s=1,linewidth=0)
fig=plt.figure(figsize=(5,8))
ax=fig.add_subplot(111)
ax.set_xlim(-30,30)
ax.set_ylim(-60,60)
ax.annotate('%6.3f Myr' % (float(ds.current_time)/3.15569E13),\
(1,0), xytext=(0.05, 0.96), textcoords='axes fraction',\
horizontalalignment='left', verticalalignment='center')
sc=ax.scatter(y,z,s=1,c=gamc,linewidth=0,cmap='algae',vmin=2,vmax=5,alpha=0.8)
cb=plt.colorbar(sc)
cb.set_label(u'log $\gamma_c$')
cb.set_ticks([2,3,4,5])
cb.set_ticklabels
plt.show()
In [ ]:
p = plot.plots['gas', 'density']
print p.image._A.shape
In [ ]:
#fields = ['dens', 'pres', 'temp', 'magp', 'magx', 'magy', 'magz', 'velx', 'vely', 'velz', 'jet ']
#fields = ['density', 'pressure', 'temperature', 'velocity_y', 'velocity_z', 'jet ', 'magnetic_field_x', \
# 'magnetic_field_y', 'magnetic_field_z', 'magnetic_pressure']
#fields = ['magnetic_field_x', 'magnetic_field_y', 'magnetic_field_z']
#fields = ['magx', 'magy', 'magz', 'magp']
#fields = ['magy']
#fields = ['pres', 'temp']
#fields = ['pressure']
fields = ['density']
#fields = ['velocity_y']
def plotfile(file):
try:
ds = yt.load(file.fullpath)
print "Plotting %s" % file.fullpath
plotSlices(ds, zoom_fac=1, center="c", fields=fields, proj_axes=['x'], drawnozzle=True,\
plotgrid=fields_grid, savepath=os.path.join(dir, 'figures'))
except KeyboardInterrupt:
print 'KeyboardInterrupt catched...'
raise Exception
nProcessor = multiprocessing.cpu_count()
pool = multiprocessing.Pool(nProcessor-2)
print nProcessor
files = rescan()
pool.map(plotfile, files[:])
#for f in files[-1:]:
# plotSlices(f)
pool.close()
pool.join()
In [ ]:
from mpl_toolkits.axes_grid1 import AxesGrid
fn = files[20].fullpath
ds = yt.load(fn) # load data
fig = plt.figure(figsize=(20,15))
# See http://matplotlib.org/mpl_toolkits/axes_grid/api/axes_grid_api.html
# These choices of keyword arguments produce two colorbars, both drawn on the
# right hand side. This means there are only two colorbar axes, one for Density
# and another for temperature. In addition, axes labels will be drawn for all
# plots.
grid = AxesGrid(fig, (0.075,0.075,0.85,0.85),
nrows_ncols = (2, 1),
axes_pad = 0.0,
label_mode = "L",
share_all = True,
cbar_location="right",
cbar_mode="edge",
cbar_size="5%",
cbar_pad="0%")
fields = ['density', 'temperature']
centers = [(0.0,0.0,20*3.086E21), (0.0,0.0,-20*3.086E21)]
for i, (field, center) in enumerate(zip(fields, centers)):
# Load the data and create a single plot
p = yt.SlicePlot(ds, 'x', field, center=center, width=(60*3.086E21, 40*3.086E21))#, origin='native')
p.set_cmap(field, 'algae')
#p.set_origin('native')
# This forces the ProjectionPlot to redraw itself on the AxesGrid axes.
plot = p.plots[field]
plot.figure = fig
plot.axes = grid[i].axes
# Since there are only two colorbar axes, we need to make sure we don't try
# to set the temperature colorbar to cbar_axes[4], which would if we used i
# to index cbar_axes, yielding a plot without a temperature colorbar.
# This unecessarily redraws the Density colorbar three times, but that has
# no effect on the final plot.
if field == 'density':
plot.cax = grid.cbar_axes[0]
elif field == 'temperature':
plot.cax = grid.cbar_axes[1]
# Finally, redraw the plot.
p._setup_plots()
plt.show()
In [ ]: