def x_inrange(self):
if self.dimensions==1:
xs=[]
for x in self.x:
ixrange=get_index_range(x,self.xlims,True)
xs=xs + [x[ixrange]]
return xs
if self.dimensions==2:
ixrange=get_index_range(self.x,self.xlims,True)
return self.x[ixrange]
def data_inrange(self):
#print self.dimensions
if self.dimensions==1:
#since a subplot can have 1D data of different length
data=[]
for x,d in zip(self.x,self.data):
ixrange=get_index_range(x,self.xlims,True)
data=data+[d[ixrange]]
if self.dimensions==2:
xlims=self.xlims
ylims=self.ylims
ixrange=numpy.array(get_index_range(self.x,xlims,True))
iyrange=numpy.array(get_index_range(self.y,ylims,True))
#print ixrange
#print iyrange
data=self.data[ixrange[:,numpy.newaxis],iyrange]
#print numpy.shape(data)
#print numpy.size(ixrange)
#print numpy.size(iyrange)
return data
def perfect_2d_plot(dirlist,attribs,xattr="psi",yattr="theta",normname="norms.namelist",speciesname="species",psiN_to_psiname="psiAHat.h5",cm=cm.rainbow,lg=True,xlims=[90,100],ylims=[0,2],species=True,sort_species=True,first=["D","He"],last=["e"],generic_labels=True,label_dict={"D":"i","He":"z","N":"z","e":"e"},vlines=None,hlines=None,share_scale=[],skip_species = [],simulList=None):
#dirlist: list of simulation directories
#attribs: list of fields to plot from simulation
#speciesname: species filename in the simuldir
#normname: norm filename in the simuldir
#lg: controls whether to interpret nearby simulations as being paired
if type(attribs) is not list:
attribs=[attribs]
if simulList == None:
normlist=[x + "/" + normname for x in dirlist]
specieslist=[x + "/" + speciesname for x in dirlist]
psiN_to_psiList=[x + "/" + psiN_to_psiname for x in dirlist]
simulList=perfect_simulations_from_dirs(dirlist,normlist,specieslist,psiN_to_psiList)
else:
"p_2d_plot: simulList specified externally, ignoring dirlist, etc."
#if we translate species labels to generic ion and impurity labels
#we still need to use original species for sorting preferences
if generic_labels:
for simul in simulList:
simul.species_list=generic_species_labels(simul.species_list,label_dict)
first=generic_species_labels(first,label_dict)
last=generic_species_labels(last,label_dict)
#get the total list of species found in each simulation
species_set=set([])
for simul in simulList:
# exclude species that are in the skip list
#nonexcluded = set(simul.species).difference(skip_species)
nonexcluded = [s for s in simul.species if s not in skip_species]
simul.species_list = nonexcluded
# add nonexcluded species to total species set
species_set = species_set | set(simul.species)
# sort species according to preferences
if sort_species:
species_set=sort_species_list(list(species_set),first,last)
print species_set
for attrib in attribs:
psp_list=[] #list of perfect subplot objects
attrib_sp_dep = is_attribute_species_dependent(simulList,attrib)
#we will assign data to the following attribute related groups
attrib_groupname=attrib
if attrib_sp_dep and species:
species_attrib_groupname="species_dependent"
this_species_set=species_set
gridspec=[len(simulList),len(species_set)]
else:
species_attrib_groupname="species_independent"
this_species_set=set([''])
gridspec=[len(simulList),1]
for i,simul in enumerate(simulList):
data0=getattr(simul,attrib) #need to split this into species data
#i_s will be coordinate of plot
for i_s,s in enumerate(this_species_set):
if attrib_sp_dep:
index=[i2 for i2,s2 in enumerate(simul.species) if s2 == s]
if len(index)==0:
#no data for this species in this simulation. Set data to zeros
data=numpy.zeros(data0[:,:,0].shape)
elif len(index)==1:
data=data0[:,:,index[0]]
else:
print "perfect_2d_plot: warning: more than one of the same species in the simulation. Will add contributions, but this is untested."
data=numpy.sum(data0[:,:,index],axis=2)
else:
data=data0
subplot_coordinates=(i,i_s)
#print subplot_coordinates
if i == 0:
show_zaxis_ticklabel=True
show_xaxis_ticklabel=False
title=s
else:
show_zaxis_ticklabel=False
title=''
if i == len(simulList)-1:
show_xaxis_ticklabel=True
else:
show_xaxis_ticklabel=False
if i_s == 0:
show_yaxis_ticklabel=True
else:
show_yaxis_ticklabel=False
#print simul.local
if simul.local:
gl_grp="local"
else:
gl_grp="global"
if (yattr == "theta") or (yattr == "theta_shifted"):
y_scale = 1/numpy.pi
else:
y_scale=1
if (xattr == "theta") or (xattr == "theta_shifted"):
x_scale = 1/numpy.pi
elif (xattr == "psi") or (xattr == "actual_psiN"):
x_scale=100
else:
x_scale=1
if xattr is not "psio":
x = getattr(simul,xattr)*x_scale
else:
print "WARNING: orbit width uses hardcoded index for deuterium and pedestal start and stop values!"
vlines2 = [0.94927395957025573,0.97463697978512787]
iD = 0
ped_start_index=get_index_range(simul.actual_psiN,[vlines2[0],vlines2[0]])[1]
ped_stop_index=get_index_range(simul.actual_psiN,[vlines2[1],vlines2[1]])[1]
ow_start=simul.orbit_width[ped_start_index,iD]
ow_stop=simul.orbit_width[ped_stop_index,iD]
x = (simul.actual_psiN-vlines2[1])/simul.orbit_width[:,iD]
vlines=[(vlines2[0]-vlines2[1])/ow_start,(vlines2[1]-vlines2[1])/ow_stop]
y = getattr(simul,yattr)*y_scale
psp_list.append(perfect_subplot(data,x=x,y=y,subplot_coordinates=subplot_coordinates,show_zaxis_ticklabel=show_zaxis_ticklabel,show_yaxis_ticklabel=show_yaxis_ticklabel,show_xaxis_ticklabel=show_xaxis_ticklabel,title=title,groups=[s,"sim"+str(i),gl_grp,"pair"+str(i/2),species_attrib_groupname],dimensions=2))
#end simulList loop
for psp in psp_list:
print psp.groups
species_groups=[perfect_subplot_group(psp_list,groups=[s,"species_dependent"],logic="and") for s in species_set if len(perfect_subplot_group(psp_list,groups=[s,"species_dependent"],logic="and").p_subplot_list)>0]
nospecies_group=perfect_subplot_group(psp_list,groups=["species_independent"])
sim_groups = [perfect_subplot_group(psp_list,groups=["sim"+str(i)]) for i in range(len(simulList))]
pair_groups = [perfect_subplot_group(psp_list,groups=["pair"+str(i)]) for i in range(len(simulList)/2)]
local_group = perfect_subplot_group(psp_list,groups=["local"])
global_group = perfect_subplot_group(psp_list,groups=["global"])
all_group=perfect_subplot_group(psp_list,groups='',get_all=True)
share_scale_group=perfect_subplot_group(psp_list,groups=share_scale,logic="or")
if lg==False:
color=iter(cm(numpy.linspace(0,1,len(simulList))))
for sim_group in sim_groups:
c=next(color)
sim_group.setattrs("border_color",c)
else:
color=iter(cm(numpy.linspace(0,1,len(pair_groups))))
for pair_group in pair_groups:
c=next(color)
pair_group.setattrs("border_color",c)
local_group.setattrs("border_linestyle","dashed")
all_group.setattrs("xlims",xlims)
#all_group.setattrs("ylims",[all_group.get_min("y",xlim=False,ylim=False),all_group.get_max("y",xlim=False,ylim=False)])
all_group.setattrs("ylims",ylims)
all_group.setattrs("vlines",vlines)
all_group.setattrs("hlines",hlines)
for species_group in species_groups:
species_group.setattrs("zlims",[species_group.get_min("data"),species_group.get_max("data")])
print [species_group.get_min("data"),species_group.get_max("data")]
if len(share_scale_group.p_subplot_list)>0:
share_scale_group.setattrs("zlims",[share_scale_group.get_min("data"),share_scale_group.get_max("data")])
if len(nospecies_group.p_subplot_list)>0:
nospecies_group.setattrs("zlims",[nospecies_group.get_min("data"),nospecies_group.get_max("data")])
for i,sim_group in enumerate(sim_groups):
if i != 0:
sim_group.setattrs("show_zaxis",False)
if xattr is not "psio":
global_xlabel=r"$100\psi_N$"
else:
global_xlabel=r"$\psi^\mathrm{o}$"
perfect_visualizer(psp_list,gridspec,global_xlabel=global_xlabel,dimensions=2,global_ylabel=r"$\theta/\pi$")
plt.savefig(attrib+'.pdf')
def y_inrange(self):
if self.dimensions==1:
print "warning-error: Cannot take yrange of 1D data."
if self.dimensions==2:
iyrange=get_index_range(self.y,self.ylims,True)
return self.y[iyrange]