def iterative_RH(planet_pos,sph_pos,dust_pos,M_star,MP,M_sph,M_dust):
'''Iteratively calculate the Hill Radius'''
sph_pos = sph_pos-planet_pos
dust_pos = dust_pos-planet_pos
planet_r = np.sqrt(planet_pos[0]**2+planet_pos[1]**2)
sph_rel_P = np.sort(np.sqrt(sph_pos[:,0]**2+sph_pos[:,1]**2+sph_pos[:,2]**2))
dust_rel_P = np.sort(np.sqrt(dust_pos[:,0]**2+dust_pos[:,1]**2+dust_pos[:,2]**2))
print 'MP', MP
mass_temp = MP.copy()
i = 0
ri = sph_rel_P[i]
RH_temp = b.hill_radius(planet_r,mass_temp,M_star)
bRH = RH_temp.copy()
while ri < RH_temp:
RH_temp = b.hill_radius(planet_r,mass_temp,M_star)
mass_temp += M_sph
i += 1
ri = sph_rel_P[i]
RH = RH_temp
MH = mass_temp
ind = np.argmin((sph_rel_P-RH/2)**2)
indd = np.argmin((dust_rel_P-RH/2)**2)
MH_2 = ind*M_sph
MH_2_dust = indd*M_dust
return RH,MH,MH_2,MH_2_dust
def res_hists(xs, ys, sph_h, sph_u, M_star, MP, planet_p, it, time):
'''Produce histograms of smoothing length ratios'''
xs, ys = xs * code_L, ys * code_L #cm
Rs = np.sqrt(xs**2 + ys**2) #cm
hs = sph_h * code_L #cm
us = sph_u * code_L**2 / code_time**2 #cm^2/s^-2
Om_Ks = b.v_kepler(M_star * code_M, Rs) / (Rs) #/s
Hs = np.sqrt((c.gamma_mono - 1) * us) / Om_Ks #cm
h_Hs = hs / Hs
planet_R = np.sqrt(planet_p[0]**2 + planet_p[1]**2) * code_L #cm
RH = b.hill_radius(planet_R, MP, M_star) #cm
dR = np.sqrt((xs - planet_p[0])**2 + (ys - planet_p[1])**2) #cm
h_RH = hs[dR < RH] / RH
print 'h_RH', h_RH
#=== Plotting ===#
txt_ys = [0.9, 0.85, 0.8, 0.75, 0.7]
bins = np.linspace(0, 2, 100)
fig14 = plt.figure(14, facecolor='white')
ax1 = fig14.add_axes([0.15, 0.1, 0.37, 0.85])
ax2 = fig14.add_axes([0.6, 0.1, 0.37, 0.85])
ax1.hist(h_Hs,
bins=bins,
edgecolor=cols[it],
lw=3,
facecolor='none',
histtype='step')
ax1.set_xlabel(r'$h / H$')
ax1.set_ylabel('Particle Number')
ax2.hist(h_RH,
bins=bins,
edgecolor=cols[it],
lw=3,
facecolor='none',
histtype='step')
ax2.set_xlabel(r'$h / R_H$')
ax2.text(0.6,
txt_ys[it],
str(int(time)) + ' Years',
color=cols[it],
transform=ax2.transAxes)
return
def render_SPH(runfolder,
args=gas_args,
snapprefix='snapshot_',
snap=-1,
inc=0,
azi=0,
inc2=0,
render_mode='Sigma',
vel_field=False,
upto=0,
zoom=False,
polyzoom=False,
overlay=False,
write=False,
dpi=1000,
h_hist=False,
it=0,
rerun=False,
amin=0,
amax=0,
tracking=True,
track_R=3):
'''Make rendered animation or single image of an SPH plot. Option to save.
field_mode: Sigma - Surface Density [gcm^{-2}]
tau - Dimensionless stopping time []
Q - Toomre Q parameter []
set snap to find snapshot. snap=-1 -> movie
If inc != 0 or 180, plot rho instead of Sigma
overlay - pixellated overlay of dust particles. Scatter too expensive.
vel_field - overlay of render data velocity structure
zoom - zooms on first planet
polyzoom - zooms on max rho SPH particles
upto - limits movie frames
tracking - track gas particles that are initially close to the planet
'''
#==== Load Arguments ====#
mode_name, p_type, map_col = args[0], args[1], args[2]
vmin, vmax, field_mode = args[3][0], args[3][1], args[5]
zoom_str, overlay_str, angle_str = '', '', ''
movie = False
box_lim = XYdim * AU_scale
if polyzoom == True:
zoom = True
if zoom == True:
zoom_str = '_zoom'
box_lim = XYdim_zoom * AU_scale
vmin, vmax = args[4][0], args[4][1]
if overlay == True:
overlay_str = '_dustgrid'
if inc != 0:
angle_str += '_i' + str(inc)
if azi != 0:
angle_str += '_a' + str(azi)
if field_mode == 'rho':
render_mode = 'rho_grid'
#==== Build grids ====#
bins = np.linspace(-box_lim, box_lim, ngrid + 1)
bin_mids = (bins[1:] + bins[:-1]) / 2
Bin_Xs, Bin_Ys = np.meshgrid(bin_mids, bin_mids)
#================ Find Time data and num snaps ==================#
time1 = time.time()
try:
time_zero = readheader(filepath + runfolder + snapprefix + '001',
'time')
try:
time_one = readheader(filepath + runfolder + snapprefix + '002',
'time')
time_dt = time_one - time_zero
except:
'Only one snapshot!'
time_dt = 0
except:
print 'No snapshots found!'
time_dt = 0
snap_dt = time_dt * code_time / c.sec_per_year
if snap == -1:
movie = True
print savedir + runfolder
if upto != 0:
num_snaps = upto
else:
num_snaps = find_snap_num(runfolder)
snapids = np.arange(num_snaps)
else:
num_snaps = 1
snapids = [snap]
if tracking == True:
track_IDs = ID_selection(runfolder,
snapprefix=snapprefix,
track_R=track_R,
zoom='planet')
else:
track_IDs = []
#----------- Load saved files or rerun rendering routine ----------#
try:
if (rerun == True) or (movie == False):
1 / 0
else:
print savedir + runfolder.rstrip(
'//'
) + '_' + p_type + zoom_str + overlay_str + angle_str + '_movie.npy'
save_array = np.load(savedir + runfolder.rstrip('//') + '_' +
p_type + zoom_str + overlay_str + angle_str +
'_movie.npy')
except:
print 'No file currently exists, must compute simulation render!'
save_array = np.zeros((num_snaps, 5, ngrid + 1, ngrid))
for i in range(len(snapids)):
idn = snapids[i]
print 'i', idn
store, dust_store = render(runfolder,
snapprefix=snapprefix,
snapid=idn,
p_type=p_type,
inc=inc,
azi=azi,
inc2=inc2,
render_mode=render_mode,
zoom=zoom,
track_IDs=track_IDs,
polyzoom=polyzoom,
overlay=overlay,
h_hist=h_hist,
it=it)
save_array[i, 0, 0, :] = store[0, 0, :] #Header info
save_array[
i, 0,
1:, :] = store[0, 1:, :] * code_M / code_L**2 #Sigma [g/cm^2]
save_array[i, 1, 1:, :] = store[
1, 1:, :] * code_L**2 / code_time**2 #u [cm^2/s^2]
save_array[i, 2,
1:, :] = store[2,
1:, :] * code_L / code_time #vx [cm/s]
save_array[i, 3,
1:, :] = store[3,
1:, :] * code_L / code_time #vy [cm/s]
#==== Add image floor to remove low density errors ====#
save_array[i, 0, 1:, :][save_array[i, 0, 1:, :] > roof] = floor
save_array[i, 0, 1:, :][save_array[i, 0, 1:, :] < floor] = floor
print 'Min/Max: ', np.min(save_array[i, 0, 1:, :]), np.max(
save_array[i, 0, 1:, :])
print '\n'
#==== Add dust_overlay info ====#
if overlay == True:
save_array[i, 0,
1:, :] = array_dust_overlay(save_array[i, 0, 1:, :],
dust_store * AU_scale,
box_lim)
save_array[i, 1,
1:, :] = array_dust_overlay(save_array[i, 1, 1:, :],
dust_store * AU_scale,
box_lim)
#==== Save output ====#
if snap == -1:
try:
print 'try removing file'
print savedir + runfolder.rstrip(
'//'
) + '_' + p_type + zoom_str + overlay_str + angle_str + '_movie.npy'
os.remove(savedir + runfolder.rstrip('//') + '_' + p_type +
zoom_str + overlay_str + angle_str + '_movie.npy')
print 'file removed'
except:
print 'pass'
print 'Starting save!'
np.save(
savedir + runfolder.rstrip('//') + '_' + p_type + zoom_str +
overlay_str + angle_str + '_movie', save_array)
print 'Loading Complete'
#------------------======== Load Header info =========----------------#
M_star = save_array[:, 0, 0, 0] #code_M
MP = save_array[:, 0, 0, 2] #code_M
NP = save_array[:, 0, 0, 1] #Number of planets
p_pos = save_array[:, 0, 0, 3:6] * AU_scale #AU
p_sep = np.sqrt(p_pos[:, 0]**2 + p_pos[:, 1]**2 + p_pos[:, 2]**2)
RH = b.hill_radius(p_sep, MP, M_star) #AU
Sigma = save_array[:, 0, 1:, :] + floor
bonus_field = save_array[:, 1, 1:, :]
vx_field = save_array[:, 2, 1:, :]
vy_field = save_array[:, 3, 1:, :]
NP_max = np.max(NP)
print 'Max Number of planets: ', NP_max
p_poss = save_array[:, 0, 0, 2:6 + 4 * (NP_max - 1)] * AU_scale
p_poss = np.reshape(p_poss, (num_snaps, NP_max, 4))
p_poss[p_poss == 0] = 1e5
#====================== Determine plot field ========================#
if field_mode == 'Sigma':
plot_field = Sigma
cax_label = r'$\Sigma \quad \rm{[g cm^{-2}}]$'
if field_mode == 'rho':
plot_field = Sigma / code_L
cax_label = r'$\rho \quad \rm{[g cm^{-3}}]$'
else:
Om_K, Rs = calc_Om_K_field(M_star, zoom=zoom, p_pos=p_pos)
if field_mode == 'tau':
plot_field = calc_tau(u=bonus_field,
Om_K=Om_K,
Sigma=Sigma,
render_mode=render_mode)
cax_label = r'$\tau$'
elif field_mode == 'Q':
if zoom == True:
Om_K, Rs = calc_Om_K_field(MP,
zoom=zoom,
p_pos=[0, 0],
render_mode=render_mode)
plot_field = calc_Q(u=bonus_field, Om_K=Om_K, Sigma=Sigma, Rs=Rs)
cax_label = 'Toomre Q'
elif field_mode == 'temp':
plot_field = calc_temp(u=bonus_field,
Sigma=Sigma,
Om_K=Om_K,
render_mode=render_mode)
cax_label = 'T [K]'
#========================= Plotting Code ========================#
#-=========================-------------========================-#
fig1 = plt.figure(it, facecolor='white', figsize=(8, 8))
ax1 = fig1.add_axes([0.15, 0.2, 0.76, 0.76])
cax = fig1.add_axes([0.15, 0.08, 0.76, 0.02])
smoothed = ax1.imshow(plot_field[0],
interpolation='none',
norm=LogNorm(vmin=vmin, vmax=vmax),
cmap=map_col,
extent=[-box_lim, box_lim, -box_lim, box_lim],
origin='lower')
plt.colorbar(smoothed, cax=cax, orientation='horizontal')
ax1.set_xlabel('x [AU]')
ax1.set_ylabel('y [AU]')
ax1.set_xlim(-box_lim, box_lim)
ax1.set_ylim(-box_lim, box_lim)
if inc == 90:
ax1.set_ylabel('z [AU]')
#==== Plot sinks ====#
print np.shape(p_poss)
planets = ax1.scatter(p_poss[0, :, 1], p_poss[0, :, 2])
#Get parameter plotting strings
beta_string = beta_stringify(runfolder)
grain_string = grain_stringify(runfolder) + ' cm'
tbox = dict(facecolor='white')
ax1.set_rasterized(True)
#ax1.text(0.04,1.015,mode_name + r', $\beta$ = ' + beta_string, transform=ax1.transAxes)
ax1.text(0.06,
0.9,
r'$\beta$ = ' + beta_string + '\n' + mode_name,
transform=ax1.transAxes,
bbox=tbox)
timetext = ax1.text(0.75,
1.015,
'Time: {:.0f} '.format(snapids[0] * snap_dt) + ' Yrs',
transform=ax1.transAxes)
cax.set_xlabel(cax_label)
if vel_field == True:
n_v = 32 #int(ngrid/50)
dv = int(ngrid / n_v)
vel_grid = np.linspace(-box_lim, box_lim, n_v)
Quivers = ax1.quiver(vel_grid,
vel_grid,
vx_field[0, ::dv, ::dv],
vy_field[0, ::dv, ::dv],
angles='xy')
if zoom == True:
Hill_radius = plt.Circle((0, 0), RH[0], fc='none', ec='#ffffff')
Half_Hill_radius = plt.Circle((0, 0),
RH[0] / 2,
fc='none',
ec='#ffffff',
ls='--')
ax1.add_patch(Hill_radius)
ax1.add_patch(Half_Hill_radius)
#==== Movie making code ====#
if movie == True:
def animate(i):
#print '2222', contours
returns = []
smoothed.set_array(plot_field[i])
timetext.set_text('Time: {:.0f} '.format(int(i * snap_dt)) +
' Yrs')
returns.append([smoothed, timetext])
planets.set_offsets(p_poss[i, :, 1:3])
returns.append(planets)
if vel_field == True:
Quivers.set_UVC(vx_field[i, ::dv, ::dv],
vy_field[i, ::dv, ::dv])
if zoom == True:
Hill_radius.set_radius(RH[i])
Half_Hill_radius.set_radius(RH[i] / 2)
returns.append([Hill_radius, Half_Hill_radius])
return returns
print 'reached animation'
ani = animation.FuncAnimation(fig1,
animate,
interval=200,
frames=num_snaps,
blit=False,
repeat=True)
if write == True:
print 'Writing savefile'
writer = animation.writers['ffmpeg'](fps=5)
ani.save(runfolder.strip('//') + '_' + p_type + zoom_str +
overlay_str + angle_str + field_mode + '.mp4',
writer=writer,
dpi=dpi)
plt.show()
#==== If no movie, save image ====#
else:
if write == True:
fig1.savefig(runfolder.strip('//') + '_' + p_type + zoom_str +
overlay_str + angle_str + field_mode + '_' +
str(snap).zfill(3) + '.pdf',
dpi=dpi)
print 'Runtime: ', time.time() - time1
return
def render_SPH(runfolder,
args=gas_args,
snapprefix='snapshot_',
snap=-1,
inc=0,
azi=0,
inc2=0,
vel_field=False,
Hill_rad=False,
Slice=False,
upto=0,
zoom=False,
polyzoom=False,
overlay=False,
write=False,
hammer=False,
RZ_plot=False,
roffset=0.01,
dpi=100,
h_hist=False,
it=0,
rerun=False,
amin=0,
amax=0,
tracking=False,
track_R=3,
soffset=0):
'''Make rendered animation or single image of an SPH plot. Option to save.
field_mode:
---- Vertically averaged modes ----
Sigma - Surface Density [gcm^{-2}]
---------- Slice modes ----------
temp - Temperature [K]
tau - Dimensionless stopping time []
Q - Toomre Q parameter []
vz - zcoord (or rcoord) velocity [cm/s]
set snap to find snapshot. snap=-1 -> movie
If inc != 0 or 180, plot rho instead of Sigma
overlay - pixellated overlay of dust particles. Scatter too expensive.
vel_field - overlay of render data velocity structure
zoom - zooms on first planet
polyzoom - zooms on max rho SPH particles
upto - limits movie frames
tracking - track gas particles that are initially close to the planet
'''
#==== Load Arguments ====#
mode_name, p_type, map_col = args[0], args[1], args[2]
vmin, vmax, field_mode = args[3][0], args[3][1], args[5]
zoom_str, overlay_str, angle_str, hammer_str = '', '', '', ''
movie = False
box_lim = XYdim * AU_scale
if polyzoom == True:
zoom = True
if zoom == True:
zoom_str = '_zoom'
box_lim = XYdim_zoom * AU_scale
vmin, vmax = args[4][0], args[4][1]
if inc != 0:
angle_str += '_i' + str(inc)
if mode_name == 'vz':
mode_name = r'$v_{\phi}$'
if azi != 0:
angle_str += '_a' + str(azi)
if field_mode == 'Sigma':
render_mode = 'Sigma'
if Slice == False:
if (field_mode == 'rho') or (field_mode == 'temp') or (
field_mode == 'tau') or (field_mode == 'vz') or (field_mode
== 'h'):
render_mode = 'rho_grid'
else:
render_mode = 'Sigma'
if field_mode == 'vz':
A = 'vz'
angle_str += '_vz'
elif field_mode == 'h':
A = 'h'
else:
A = 'u'
box_limx, box_limy = box_lim, box_lim
if hammer == True:
angle_str += 'hammer'
box_limx = np.pi
box_limy = 1
#==== Build grids ====#
binsx = np.linspace(-box_limx, box_limx, ngrid + 1)
binx_mids = (binsx[1:] + binsx[:-1]) / 2
binsy = np.linspace(-box_limx, box_limx, ngrid + 1)
biny_mids = (binsy[1:] + binsy[:-1]) / 2
Bin_Xs, Bin_Ys = np.meshgrid(binx_mids, biny_mids)
#================ Find Time data and num snaps ==================#
globbed = glob.glob(filepath + runfolder + 'snapshot*')
num_snaps = len(globbed)
snaps = []
slen = len(globbed[0].split('.hdf')[0].split('snapshot_')[1])
print 'slen', slen
time1 = time.time()
try:
time_zero = readheader(
filepath + runfolder + snapprefix + '001'.zfill(slen), 'time')
try:
time_one = readheader(
filepath + runfolder + snapprefix + '002'.zfill(slen), 'time')
time_dt = time_one - time_zero
except:
'Only one snapshot!'
time_dt = 0
except:
print 'No snapshots found!'
time_dt = 0
snap_dt = time_dt * code_time / c.sec_per_year
if snap == -1:
movie = True
print savedir + runfolder
if upto != 0:
num_snaps = upto
min_snap = 0
else:
for i in range(len(globbed)):
snaps.append(int(globbed[i][-5 - slen:][:slen]))
min_snap = min(snaps)
print 'Running movie for ', num_snaps, ' frames'
snapids = np.arange(num_snaps) + min_snap
else:
num_snaps = 1
snapids = [snap]
#======== Set up overlay data ========#
if tracking == True:
coarse = 4
track_IDs = b.ID_selection(runfolder,
snapprefix=snapprefix,
track_R=track_R,
zoom='planet')
ngrid_coarse = int(ngrid / coarse)
#track_store = np.zeros((num_snaps,ngrid_coarse,ngrid_coarse))
else:
track_IDs = []
#======== Load saved files or rerun rendering routine ========#
try:
if (rerun == True) or (movie == False):
1 / 0
else:
print savedir + runfolder.rstrip(
'//'
) + '_' + p_type + zoom_str + overlay_str + angle_str + render_mode + '_movie.npy'
save_array = np.load(savedir + runfolder.rstrip('//') + '_' +
p_type + zoom_str + overlay_str + angle_str +
render_mode + '_movie.npy')
except:
print 'No file currently exists, must compute simulation render!'
save_array = np.zeros((num_snaps, 6, ngrid + 1, ngrid))
for i in range(len(snapids)):
idn = snapids[i]
print 'i', idn + 1, ' out of', num_snaps
print 'render_mode', render_mode
store, dust_pos, track_pos = render(runfolder,
snapprefix=snapprefix,
snapid=idn,
slen=slen,
p_type=p_type,
inc=inc,
azi=azi,
inc2=inc2,
hammer=hammer,
RZ_plot=RZ_plot,
A=A,
roffset=roffset,
render_mode=render_mode,
zoom=zoom,
track_IDs=track_IDs,
polyzoom=polyzoom,
overlay=overlay,
h_hist=h_hist,
it=it)
print 'A', store[1, 1:, :]
save_array[i, 0, 0, :] = store[0, 0, :] #Header info
save_array[
i, 0,
1:, :] = store[0, 1:, :] * code_M / code_L**2 #Sigma [g/cm^2]
save_array[i, 2,
1:, :] = store[2,
1:, :] * code_L / code_time #vx [cm/s]
save_array[i, 3,
1:, :] = store[3,
1:, :] * code_L / code_time #vy [cm/s]
dust_pos *= AU_scale
track_pos *= AU_scale
if A == 'u':
save_array[i, 1, 1:, :] = store[
1, 1:, :] * code_L**2 / code_time**2 #u [cm^2/s^2]
elif A == 'h':
save_array[i, 1, 1:, :] = store[1, 1:, :] * AU_scale #h in AU
elif A == 'vz':
save_array[i, 1,
1:, :] = store[1, 1:, :] * code_L / code_time #cm/s
#==== Add image floor to remove low density errors ====#
save_array[i, 0, 1:, :][save_array[i, 0, 1:, :] > roof] = floor
save_array[i, 0, 1:, :][save_array[i, 0, 1:, :] < floor] = floor
print 'Min/Max: ', np.min(save_array[i, 0, 1:, :]), np.max(
save_array[i, 0, 1:, :])
print '\n'
#==== Add dust_overlay info ====#
save_array[i, 4, 1:, :] = track_overlay(dust_pos, box_lim, ngrid)
if tracking == True:
save_array[i, 5,
1:ngrid_coarse + 1, :ngrid_coarse] = track_overlay(
track_pos, box_lim, ngrid_coarse)
#==== Save output ====#
if snap == -1:
try:
print 'try removing file'
print savedir + runfolder.rstrip(
'//'
) + '_' + p_type + zoom_str + overlay_str + angle_str + render_mode + '_movie.npy'
os.remove(savedir + runfolder.rstrip('//') + '_' + p_type +
zoom_str + overlay_str + angle_str + render_mode +
'_movie.npy')
print 'file removed'
except:
print 'pass'
print 'Starting save!'
np.save(
savedir + runfolder.rstrip('//') + '_' + p_type + zoom_str +
overlay_str + angle_str + render_mode + '_movie', save_array)
print 'Loading Complete'
#------------------======== Load Header info =========----------------#
Time = save_array[:, 0, 0,
0] #snap_dt*np.arange(num_snaps)#save_array[:,0,0,0] #Years
M_star = save_array[:, 0, 0, 1] #code_M
NP = save_array[:, 0, 0, 2] #Number of planets
MP = save_array[:, 0, 0, 3] #code_M
p_pos = save_array[:, 0, 0, 4:7] * AU_scale #AU
p_sep = np.sqrt(p_pos[:, 0]**2 + p_pos[:, 1]**2 + p_pos[:, 2]**2)
RH = b.hill_radius(p_sep, MP, M_star) #AU
Sigma = save_array[:, 0, 1:, :] + floor
bonus_field = save_array[:, 1, 1:, :]
vx_field = save_array[:, 2, 1:, :]
vy_field = save_array[:, 3, 1:, :]
dust_store = save_array[:, 4, 1:, :]
if tracking == True:
track_store = save_array[:, 5, 1:ngrid_coarse + 1, :ngrid_coarse]
NP_max = int(np.max(NP))
print 'Max Number of planets: ', NP_max
p_poss = save_array[:, 0, 0, 3:7 + 4 * (NP_max - 1)] * AU_scale #AU
p_poss = np.reshape(p_poss, (num_snaps, NP_max, 4))
if hammer == True:
p_poss[:, :, 1] *= np.pi / box_lim
p_poss[:, :, 2] *= 1 / box_lim
#====================== Determine plot field ========================#
if field_mode == 'Sigma':
plot_field = Sigma
cax_label = r'$\Sigma \quad \rm{[g cm^{-2}}]$'
if field_mode == 'rho':
plot_field = Sigma / code_L
cax_label = r'$\rho \quad \rm{[g cm^{-3}}]$'
else:
Om_K, Rs = calc_Om_K_field(M_star, zoom=zoom, p_pos=p_pos)
if field_mode == 'tau':
plot_field = calc_tau(u=bonus_field,
Om_K=Om_K,
Sigma=Sigma,
render_mode=render_mode)
cax_label = r'$\tau$'
elif field_mode == 'Q':
if zoom == True:
Om_K, Rs = calc_Om_K_field(MP,
zoom=zoom,
p_pos=[0, 0],
render_mode=render_mode)
plot_field = calc_Q(u=bonus_field, Om_K=Om_K, Sigma=Sigma, Rs=Rs)
cax_label = 'Toomre Q'
elif field_mode == 'temp':
plot_field = calc_temp(u=bonus_field,
Sigma=Sigma,
Om_K=Om_K,
render_mode=render_mode)
cax_label = 'T [K]'
elif field_mode == 'vz':
plot_field = bonus_field / code_L * code_time
cax_label = r'$v_r [cms^{-1}]$'
if inc != 0:
cax_label = r'$v_{\phi}$ $[cms^{-1}]$'
elif field_mode == 'h':
plot_field = bonus_field #h in AU
cax_label = r'Smoothing length h'
#========================= Plotting Code ========================#
#-=========================-------------========================-#
fig1 = plt.figure(it, facecolor='white', figsize=(8, 8))
ax1 = fig1.add_axes([0.15, 0.2, 0.76, 0.76])
cax = fig1.add_axes([0.15, 0.08, 0.76, 0.02])
norm = LogNorm(vmin=vmin, vmax=vmax)
if field_mode == 'vz':
norm = SymLogNorm(linthresh=vmin, vmin=-vmax, vmax=vmax)
map_col = 'RdYlBu_r'
elif (field_mode == 'temp') and (zoom == True):
norm = colors.Normalize(vmin=vmin, vmax=vmax)
elif field_mode == 'h':
norm = colors.Normalize(vmin=vmin, vmax=vmax)
smoothed = ax1.imshow(plot_field[0],
interpolation='none',
norm=norm,
aspect='auto',
cmap=map_col,
extent=[-box_limx, box_limx, -box_limy, box_limy],
origin='lower')
plt.colorbar(smoothed, cax=cax, orientation='horizontal')
'''
#==== 3D test ====#
dim = np.shape(plot_field[0])[0]
Xs = np.linspace(-box_lim,box_lim,dim)
X,Y = np.meshgrid(Xs,Xs)
ax1 = fig1.gca(projection='3d')
ax1.plot_surface(X, Y, plot_field[0], rstride=8, cstride=8, alpha=0.3)
'''
print 'Max dust store', np.max(dust_store)
if overlay == True:
dust = ax1.imshow(dust_store[0],
interpolation='none',
norm=LogNorm(1, 5e4),
aspect='auto',
cmap='gray_r',
extent=[-box_limx, box_limx, -box_limy, box_limy],
origin='lower',
alpha=0.7)
if tracking == True:
tracker = ax1.imshow(track_store[0],
interpolation='none',
norm=LogNorm(1, 100),
alpha=0.5,
aspect='auto',
cmap='RdPu',
extent=[-box_limx, box_limx, -box_limy, box_limy],
origin='lower')
ax1.set_xlabel('x [AU]')
ax1.set_ylabel('y [AU]')
if hammer == True:
ax1.set_xlabel(r'$\phi$')
ax1.set_ylabel(r'cos($\theta$)')
ax1.set_xlim(-box_limx, box_limx)
if inc == 90:
ax1.set_ylabel('z [AU]')
#==== Plot sinks ====#
print np.shape(p_poss)
planet_col = '#581845' #'#22ddff'
if args[0] == 'Dust':
planet_col = '#dd00cc'
planets = ax1.scatter(p_poss[0, :, 1],
p_poss[0, :, 2],
s=30,
color=planet_col) #'#ff1493')#,edgecolors='purple')
#Get parameter plotting strings
beta_string = beta_stringify(runfolder)
grain_string = grain_stringify(runfolder) + ' cm'
tbox = dict(facecolor='white')
ax1.set_rasterized(True)
#ax1.text(0.04,1.015,mode_name + r', $\beta$ = ' + beta_string, transform=ax1.transAxes)
ax1.text(0.06,
0.9,
r'$\beta$ = ' + beta_string + '\n' + mode_name,
transform=ax1.transAxes,
bbox=tbox)
cax.set_xlabel(cax_label)
timetext = ax1.text(0.75,
1.015,
'Time: {:.0f} '.format(Time[0]) + ' Yrs',
transform=ax1.transAxes)
#timetext = ax1.text(0.75,1.015,'Time: {:.0f} '.format(int(snap_dt*soffset)) + ' Yrs', transform=ax1.transAxes)
if vel_field == True:
n_v = 16 #int(ngrid/50)
dv = int(ngrid / n_v)
vel_gridx = np.linspace(-box_limx, box_limx, n_v)
vel_gridy = np.linspace(-box_limy, box_limy, n_v)
Quivers = ax1.quiver(vel_gridx,
vel_gridy,
vx_field[0, ::dv, ::dv],
vy_field[0, ::dv, ::dv],
angles='xy')
if (zoom == True) and (Hill_rad == True):
Hill_radius = plt.Circle((0, 0), RH[0], fc='none', ec='#ffffff')
Half_Hill_radius = plt.Circle((0, 0),
RH[0] / 2,
fc='none',
ec='#ffffff',
ls='--')
ax1.add_patch(Hill_radius)
ax1.add_patch(Half_Hill_radius)
#==== Movie making code ====#
if movie == True:
def animate(i):
returns = []
smoothed.set_array(plot_field[i])
timetext.set_text('Time: {:.0f} '.format(int(Time[i])) + ' Yrs')
#timetext.set_text('Time: {:.0f} '.format(int(snap_dt*(i+soffset))) + ' Yrs')
returns.append([smoothed, timetext])
planets.set_offsets(p_poss[i, :, 1:3])
returns.append(planets)
if vel_field == True:
Quivers.set_UVC(vx_field[i, ::dv, ::dv],
vy_field[i, ::dv, ::dv])
if (zoom == True) and (polyzoom == False) and (Hill_rad == True):
Hill_radius.set_radius(RH[i])
Half_Hill_radius.set_radius(RH[i] / 2)
returns.append([Hill_radius, Half_Hill_radius])
if overlay == True:
dust.set_array(dust_store[i])
returns.append([dust])
if tracking == True:
tracker.set_array(track_store[i])
returns.append([tracker])
return returns
print 'reached animation'
interval = 200 #10#50#100
print 'N snaps', num_snaps
anim = player.Player(fig1,
animate,
interval=interval,
maxi=num_snaps,
blit=False,
repeat=True,
save_count=num_snaps - 1)
#anim = FuncAnimation(fig1, animate, interval=interval, blit=False, repeat=True,save_count=num_snaps)
if write == True:
print 'Writing savefile'
anim = FuncAnimation(fig1,
animate,
interval=interval,
blit=False,
repeat=True,
save_count=num_snaps)
writer = animation.writers['ffmpeg'](fps=20)
anim.save(runfolder.strip('//') + '_' + p_type + zoom_str +
overlay_str + angle_str + field_mode + '.mp4',
writer=writer,
dpi=dpi)
plt.show()
#==== If no movie, save image ====#
else:
if write == True:
fig1.savefig(runfolder.strip('//') + '_' + p_type + zoom_str +
overlay_str + angle_str + field_mode + '_' +
str(snap).zfill(3) + '.pdf',
dpi=dpi)
print 'Runtime: ', time.time() - time1
return