def store_ejected(data_chunk, outfname=None):
import sys_to_si as s2s
import cPickle as cpk #use this as can store complete precision
f=open(outfname,'a')#open file
#dump data
cpk.dump((s2s.year(data_chunk[1]), num_ejected(data_chunk)), f, cpk.HIGHEST_PROTOCOL)
#close file
f.close()
print("Storing data point...")
def get_energy(data_chunk):
import sys
import sys_to_si as s2s
print("Finding energy of system...")
(pdata, time, N_p) = (data_chunk[0], data_chunk[1], data_chunk[2])
rockE = []
planetE=[]
for m, r, pos, vel, spn, col, idx in pdata:
if col!=3:
planetE.append(inst_energy(m, pos, vel))
#print(sum([posi*posi for posi in pos]))
else: rockE.append(inst_energy(m,pos,vel))
#print(time, sum(rockE), planetE[0])
return(s2s.year(time) ,sum(rockE)+sum(planetE), sum(rockE), sum(planetE))
def mVa_graph(data_chunk, dr=0.001, rmax=5, rmin=0, ymax=0.01, ymin=0,
outfname=None, FILL=True, BAR=True): #change as needed
import sys_to_si as s2s
import matplotlib.pyplot as plt
import numpy as np
import matplotlib as mpl
import math
mpl.figure.Figure.clear(plt.gcf())
Zs = bin_func(dr, rmin, rmax, data_chunk)
#Zs contains the Z value for each grid position, calculated by func3()
z_shape = Zs.shape
#print(z_shape)
print('Plotting image...')
if BAR:
plt.bar(np.arange(0,z_shape[0]), Zs, width=1, bottom=0, color='red', edgecolor='red')
else:
plt.plot(Zs, 'r-')
plt.axis([0, z_shape[0], ymin, ymax])
plt.xticks(range(0, z_shape[0], z_shape[0]/5),range(0,6,1))
#plt.ylim(0,0.01)
draw_resonances(data_chunk[0][0],dr, [(1,1),(1,2),(2,3),(1,3),(2,1),(3,2)])
xs = range(1,100)
#plt.plot(xs,[max(Zs)/math.pow(y*float(5)/100, 0.5) for y in xs], 'b--')
if FILL and not BAR:
plt.gca().fill_between(np.arange(0,z_shape[0]),0, Zs, facecolor='red')
plt.xlabel('Semi-major axis (AU)')
plt.ylabel('Mass (M$_\oplus$)')
#title2='Mass Evolution w.r.t Semi-Major Axis\n$T$ = %08.2f (yrs), $N$ = %07d, $d_r$=%02g (AU)' %(s2s.year(data_chunk[1]), data_chunk[2], dr)
title2 = '$T$ = %08.2f (yrs), $N$ = %07d, $d_r$=%02g (AU)' %(s2s.year(data_chunk[1]), data_chunk[2], dr)
plt.title(title2)
plt.legend()
print('Saving figure...')
if outfname==None:
plt.show()
return(0)
else:
plt.savefig(outfname)
return(0)
"bin_graph.py" Creates graphs of the "ss.dust" files outputted by PKDGRAV's
rubble planetesimals collision model.
bin_graph.py [dust_files] [output_folder]
[dust_files] = a list of "ss.dust" files to graph
[output_folder] = a folder to put the graphs in, will be named according to timestamp?
'''
if (len(sys.argv) == 1) or (sys.argv[1]=='-h') or (sys.argv[1]=='--help'):
print(usage)
sys.exit(0)
if len(sys.argv)>2:
outpath = sys.argv[-1]
max = -1
else:
outpath = None
max = 2
if outpath=='None' or outpath=='none' or outpath=='NONE':
outpath = None
for arg in sys.argv[1:max]:
data_chunk = unpackdust(arg)
(plist, time, n_particles) = unpackxdr_head(arg.rstrip('.dust'))
ident_string = 'ss.{0:06.5f}.dust_map'.format(s2s.year(time))
if outpath != None:
outname = os.path.join(outpath,ident_string+'.png')
else: outname = None
plot_dust((time, data_chunk), outfname=outname)
def plot_dust(data_chunk, outfname=None):
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.figure as fig
import sys_to_si as s2s
'''
data_chunk = [data, data2, data3...]
data = [bin, dust_mass, %from bin 1, %from bin 2, %from bin 3...]
'''
fig.Figure.clear(plt.gcf())
time = data_chunk[0]
data_chunk = data_chunk[1]
N_groups = len(data_chunk)
N_origins = len(data_chunk[0]) - 2
bin_norms = [data[1] for data in data_chunk]
bin_origins = [data[2:] for data in data_chunk]
#sys.exit() #for debugging
pic_array = make_pic(data_chunk)
pic_shape = pic_array.shape
the_cmap = mpl.cm.get_cmap('gist_heat') #Sets which colour map to use
#the_cmap.set_over('b')
im = plt.imshow(pic_array, cmap=the_cmap, origin='lower', zorder=0)
#plt.colorbar(extend='max') # add colour bar
#print(pic_array)
im.set_interpolation('bilinear')
#plt.xticks(range(pic_shape[1]/10, pic_shape[1], pic_shape[1]/5),range(-4,7,2))
#plt.yticks(range(pic_shape[0]/10, pic_shape[0], pic_shape[0]/5),range(-4,7,2))
#plt.xlabel('Displacement (AU)')
#plt.ylabel('Displacement (AU)')
#plt.show()
if outfname==None:
plt.show()
else:
pic_name = outfname.split(os.sep)
pic_name[-1] = pic_name[-1].split('.')
pic_name[-1][0] +='map'
pic_name[-1] = '_'.join(pic_name[-1][:-1])
pic_name[-1] += '.png'
pic_name = os.sep.join(pic_name)
print('\n'+pic_name)
plt.savefig(pic_name)
fig.Figure.clear(plt.gcf())
width = 0.5
bin_abs = []
for i in xrange(len(bin_norms)):
abs = []
for j in xrange(len(bin_origins[0])):
abs.append(bin_origins[i][j]*bin_norms[i])
bin_abs.append(abs)
bar_pos = range(1, len(data_chunk))
"""
col = ['r','g','b','k','c','m','y', 'w', (0.5,0.5,0.5), (0.7,0.5,0.5), (0.5,0.7,0.5),
(0.5,0.5,0.7), (0.7,0.5,0.7), (0.5,0.7,0.7), (0.7,0.7,0.5), (0.7,0.7,0.7)]
"""
#print(bin_abs)
plots = []
for i in xrange(len(bin_abs)):
plots.append(plt.bar(bar_pos, bin_abs[i], width = width, bottom = sumto(i, bin_abs[i]),
color = mpl.cm.spectral(float(i)*(1.0/float(len(bin_abs)))), log=False))
plt.xlim(0, len(data_chunk)+0.2*len(data_chunk))
plt.minorticks_on()
plt.title('Dust mass in each dust bin\nT = {0:08.2f} (years)'.format(s2s.year(time)))
plt.xlabel('Bin Number')
plt.ylabel('Mass in Bin')
plt.xticks([pos+width/2 for pos in bar_pos], range(1, len(bar_pos)+1) )
plt.legend([plot[0] for plot in plots], range(1,len(bar_pos)+1) )
if outfname==None:
plt.show()
return(0)
else:
plt.savefig(outfname)
return(0)
def surf_density(data_chunk, dx=0.1, dy=0.1, xmax=11, xmin=-11, ymax=11, ymin=-11,
colourmap='gnuplot', c_bar_lims=(0,0.2), autoda=False,
outfname=None):# change these as needed
import sys_to_si as s2s
import matplotlib.pyplot as plt
import numpy as np
import matplotlib as mpl
import matplotlib.figure as fig
"""
This function creates a surface density plot of a solar system simulation
Arguments:
data_chunk = [tuple] (pdata, time, n_particles) from either
'unpackxdr()' on readss or 'unpickle_chunk()'.
dx, dy = step in x, y dimension used for resolution, also
sets the lower limit for autoda. [number]
xmax, ymax = The upper x, y limits of the plot [number]
xmin, ymin = The lower x, y limits of the plot [number]
colourmap = The name of the colour map to use, see
http://matplotlib.sourceforge.net/examples/pylab_examples/show_colormaps.html
for a list of colourmaps.
c_bar_lims = The bottom and top limits of the colourbar, make sure
objects are visible in the range set.
autoda = This flag sets weather to auromatically determine
the dx, dy steps using the dx, dy values in the
arguments as lower limits. [Bool]
"""
fig.Figure.clear(plt.gcf())
inview=[dat[2] for dat in data_chunk[0] if r2(dat)<(xmin*ymin)]
x_range=max([x[0] for x in inview])-min([x[0] for x in inview])
y_range=max([x[1] for x in inview])-min([x[1] for x in inview])
n_particles = len(inview) #this many particles (total)
area=3.142*(y_range)*(x_range) #this much view space
#average distance between them, will be overestimate
#if some particles not in view area
smallest_dxy = np.sqrt(area/n_particles)
#print(smallest_dxy)
if autoda:
print('Automatically determining optimum binsize...')
if smallest_dxy>((dx+dy)/2):
if smallest_dxy>10*((dx+dy)/2):
print('Using 10* dx, dy as upper limits')
dx*=10
dy*=10
else:
dx = smallest_dxy
dy = smallest_dxy
else:
print('Using given dx, dy as lower limits...')
elif dx*dy<(smallest_dxy*smallest_dxy):
print('Warning: Resolution smaller than average inter-particle distance...')
x_bins = int(np.floor((xmax-xmin)/dx))
y_bins = int(np.floor((ymax-ymin)/dy))
Zs = bin_func(x_bins, y_bins, dx, dy, xmin, ymin,data_chunk)
#Zs contains the Z value for each grid position, calculated by func3()
z_shape = Zs.shape
print('Plotting image...')
#plt.hot()
the_cmap = mpl.cm.get_cmap(colourmap) #Sets which colour map to use
#the_cmap.set_over('b')
im = plt.imshow(Zs, cmap=the_cmap,
origin='lower', zorder=0)
# put the image on plot, using 'gnuplot' colour map
# origin at bottom left, draw this first
# list of colourmaps: http://matplotlib.sourceforge.net/examples/pylab_examples/show_colormaps.html
plt.colorbar(extend='max') # add colour bar
#set interpolation level
#im.set_interpolation('nearest')
#im.set_interpolation('bicubic')
im.set_interpolation('bilinear')
im.set_clim(*c_bar_lims) # Sets the lower and upper limits of the colour bar.
#plt.gca().invert_yaxis()
#plt.xticks(range(z_shape[1]/10, z_shape[1], z_shape[1]/5),range(-4,7,2))
#plt.yticks(range(z_shape[0]/10, z_shape[0], z_shape[0]/5),range(-4,7,2))
plt.xticks(range(z_shape[1]/22, z_shape[1], z_shape[1]/11),range(-10,11,2))
plt.yticks(range(z_shape[0]/22, z_shape[0], z_shape[0]/11),range(-10,11,2))
plt.xlabel('Displacement (AU)', fontsize=18)
plt.ylabel('Displacement (AU)', fontsize=18)
#title2='Surface density in M$_\oplus/$AU$^2$\n$T$ = %08.2f (yrs), $N$ = %07d, $d_{area}$=%04.2g (AU$^2$)' %(s2s.year(data_chunk[1]), data_chunk[2], dx*dy)
title2 = '$T$ = %08.2f (yrs)\n$N$ = %07d' %(s2s.year(data_chunk[1]), data_chunk[2])
plt.title(title2, fontsize=24)
plt.plot(int(np.floor((data_chunk[0][0][2][0]-xmin)/dx)-1),
int(np.floor((data_chunk[0][0][2][1]-ymin)/dy)-1), 'go', ms=20,
axes=plt.gca(), mfc='None', mec='g', mew=2)
# puts jupiter in correct place as a green circle
print('Saving figure...')
if outfname==None:
plt.show()
return(0)
else:
plt.savefig(outfname, bbox_inches="tight")
return(0)
def plot_RV(data_chunk, outfname=None):
import sys_to_si as s2s
import matplotlib.pyplot as plt
import matplotlib.figure as fig
(pdata, systime, particle_num) = data_chunk
fig.Figure.clear(plt.gcf())
print('Plotting data...')
#pdata format = mass, radius, pos(xyz), vel(xyz), spin(xyz), colour, original_identifier
#print([stuff[2] for stuff in pdata[:5]])
es = [inst_V(e[0], e[2],e[3]) for e in pdata] #get data for plot
#print('HERE', es[:2])
#print(sum([veli*veli for veli in pdata[0][3]]))
#print(sum([veli*veli for veli in pdata[1][3]]))
plt.plot([e[0] for e in es[1:]],[e[1] for e in es[1:]], 'r.',zorder=0, ms=0.2)
#plot the data, 'zorder=0' makes sure that it is drawn first and everything else is on top of it
#Draw vertical lines at resonances and add labels
plt.axvline(x=resonance(es[0][0],1,2), ymin=0, ymax=1, c='g', ls='--', label='Resonances')
plt.text(resonance(es[0][0],1,2), 0.8, '1:2', color='k', rotation=90, ha='center')
plt.axvline(x=resonance(es[0][0],1,3), ymin=0, ymax=1, c='g', ls='--', label='_nolegend_')
plt.text(resonance(es[0][0],1,3), 0.8, '1:3', color='k', rotation=90, ha='center')
plt.axvline(x=resonance(es[0][0],2,3), ymin=0, ymax=1, c='g', ls='--', label='_nolegend_')
plt.text(resonance(es[0][0],2,3), 0.8, '2:3', color='k', rotation=90, ha='center')
plt.axvline(x=resonance(es[0][0],2,1), ymin=0, ymax=1, c='g', ls='--', label='_nolegend_')
plt.text(resonance(es[0][0],2,1), 0.8, '2:1', color='k', rotation=90, ha='center')
plt.axvline(x=resonance(es[0][0],3,2), ymin=0, ymax=1, c='g', ls='--', label='_nolegend_')
plt.text(resonance(es[0][0],3,2), 0.8, '3:2', color='k', rotation=90, ha='center')
#plot planet data with error bars for hill radii
plt.errorbar(x=es[0][0], y=es[0][1], yerr=None, xerr=5*hill_rad(pdata[0][0],1,es[0][0]),
elinewidth=None, barsabove=True, marker='o',
ms=30*hill_rad(pdata[0][0],1,es[0][0])/hill_rad(10E-3,1,es[0][0]),
mec='b', mfc='None', mew=1)
plt.axis([0,5,-0.05,1])
#put in lables and titles etc
plt.minorticks_on()
plt.title('Velosity Magnitude vs Semi-major Axis:\nT = {0:08.2f} (years), N = {1:07}'.format(s2s.year(systime), particle_num))
plt.xlabel('Semi-major Axis (AU)')
plt.ylabel('Velocity (2PI AU per Year)')
plt.legend()
if outfname==None:
plt.show()
return(0)
else:
plt.savefig(outfname)
return(0)
def plot_eVa(data_chunk, outfname=None):
import sys_to_si as s2s
import matplotlib.pyplot as plt
import matplotlib.figure as fig
from own_cmaps import register_own_cmaps
register_own_cmaps()
(pdata, systime, particle_num) = data_chunk
fig.Figure.clear(plt.gcf())
print('Plotting data...')
#pdata format = mass, radius, pos(xyz), vel(xyz), spin(xyz), colour, original_identifier
#print([stuff[2] for stuff in pdata[:5]])
es = [inst_e(e[0], e[2],e[3]) for e in pdata] #get data for plot
es1 = sorted(es[1:], key=by_mass)
#print('HERE', es[:2])
#print(sum([veli*veli for veli in pdata[0][3]]))
#print(sum([veli*veli for veli in pdata[1][3]]))
#masses = [e[2] for e in es1[1:]]
masses = [e[2] for e in es1]
max_m = max(masses)
min_m = min(masses)
sizes = [2*(mass/min_m)**3 if mass >min_m else 5*0.125 for mass in masses ]
#sizes = [(mass/min_m) if mass >min_m else 0.125*mass/min_m for mass in masses ]
#print(min_m)
plt.scatter([e[0] for e in es1],[e[1] for e in es1], marker='o', zorder=0,
s=sizes,
cmap='own1', c = [e[2] for e in es1], edgecolors='none')
#plot the data, 'zorder=0' makes sure that it is drawn first and everything else is on top of it
#Draw vertical lines at resonances and add labels
plt.axvline(x=resonance(es[0][0],1,2), ymin=0, ymax=1, c='g', ls='--', label='Resonances')
plt.text(resonance(es[0][0],1,2), 0.8, '1:2', color='k', rotation=90, ha='center')
plt.axvline(x=resonance(es[0][0],1,3), ymin=0, ymax=1, c='g', ls='--', label='_nolegend_')
plt.text(resonance(es[0][0],1,3), 0.8, '1:3', color='k', rotation=90, ha='center')
plt.axvline(x=resonance(es[0][0],2,3), ymin=0, ymax=1, c='g', ls='--', label='_nolegend_')
plt.text(resonance(es[0][0],2,3), 0.8, '2:3', color='k', rotation=90, ha='center')
plt.axvline(x=resonance(es[0][0],2,1), ymin=0, ymax=1, c='g', ls='--', label='_nolegend_')
plt.text(resonance(es[0][0],2,1), 0.8, '2:1', color='k', rotation=90, ha='center')
plt.axvline(x=resonance(es[0][0],3,2), ymin=0, ymax=1, c='g', ls='--', label='_nolegend_')
plt.text(resonance(es[0][0],3,2), 0.8, '3:2', color='k', rotation=90, ha='center')
#plot planet data with error bars for hill radii
plt.errorbar(x=es[0][0], y=es[0][1], yerr=None, xerr=5*hill_rad(pdata[0][0],1,es[0][0]),
elinewidth=None, barsabove=True, marker='o',
ms=30*hill_rad(pdata[0][0],1,es[0][0])/hill_rad(10E-3,1,es[0][0]),
mec='b', mfc='None', mew=1)
plt.axis([0,11,-0.05,1]) #change as needed
#put in lables and titles etc
#plt.colorbar(extend='max') # add colour bar, if don't want colour bar just comment out
plt.clim(min_m, max_m)
#im.set_clim(max_m, min_m)
plt.minorticks_on()
#plt.title('Clearing and Resonances in e-a Space:\nT = {0:08.2f} (years), N = {1:07}'.format(s2s.year(systime), particle_num))
plt.title('T = {0:08.2f} (earth years), N={1:07}'.format(s2s.year(systime), particle_num))
plt.xlabel('Semi-major Axis (AU)')
plt.ylabel('Eccentricity')
plt.legend()
print("data plotted, showing...")
if outfname==None:
plt.show()
return(0)
else:
plt.savefig(outfname)
return(0)
print('Creating top level destination folder...')
os.makedirs(outfolder)
else:
print("Please choose an exsisting destination folder.")
sys.exit(0)
for label, function in graphlist:
if function in alltimedict and os.path.exists(os.path.join(outfolder,label,'temp')):
#delete old temporary storage space
os.system('rm '+os.path.join(outfolder,label,'temp'))
if '.dat' in sys.argv[1]:
for data_chunk in unpickle_chunk(sys.argv[1]):
#create identity string for each graph (timestamp)
#time from these files is in "system time" in units of years*2*pi
#therefore need to divide by 2*pi to make numbers correct
ident_string = 'ss.{0:08.3f}'.format(s2s.year(data_chunk[1]))
for label, function in graphlist:
#each set of graphs draws to different figures
plt.figure(graphlist.index((label,function)))
if not os.path.exists(os.path.join(outfolder,label)):
#make the second level folders if they don't exist
os.mkdir(os.path.join(outfolder,label))
if function in alltimedict:
#if graph is using data from entire time series,
#create a temporary storage file
outname=os.path.join(outfolder,label,'temp')
else:
#if graph is made new each time, stamp graph name with
#timestamp
outname = os.path.join(outfolder,label,ident_string+'.png')
#execute the function that creates the graph/datapoint
