Exemple #1
0
 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]
Exemple #2
0
 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
Exemple #3
0
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')
Exemple #4
0
 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]