Python yt.load函数代码示例

def load_prof(directory,filament_num,var,keep_list):
    #~~~~~~~~~~~~~~~~#
    # NOT  TESTED    #
    #                #
    #~~~~~~~~~~~~~~~~#

#Function which the profile data for one variable from disk
    #Only returns the array of variable bins, and no other profile metadata - for that use load_all_profs
    #Will also filter out any segments of filaments deemed to be within a halo
    
    #Init list to return

    var_profs = [[] for i in range(filament_num)]

    

    #Gather list of filenames, ergo gathering list of fils and segs
    import os
 
    filelist = sorted(os.listdir(directory))
    #At this stage we can split the filelist into two depending on variable required
    #This is due to density and temp/velocity profiles being stored seperately.
    vardict = {'density':0,'dark_matter_density':0,"temperature":1,"cylindrical_radial_velocity":1}
    profnum = len(filelist) / 2
    filelist = filelist[profnum * vardict[var]:profnum * (vardict[var]+ 1)]
    
    #Parallize the import of the data to speed this process up
    #Use yt's easy to use parallel_objects implementation
    
    storage = {}

    #Gather x data from one profile, incase this is also needed
    x = yt.load(''.join([directory,'/',filelist[0]])).data['x']

    filelist=filelist[:10]
    for sto, file_in_dir in yt.parallel_objects(filelist, storage=storage):
        #Determine filament and segment number
        
        filnum = int(file_in_dir[7:10])
        segnum = int(file_in_dir[13:16])
        
        #Check to see if segment is within a halo, if so, disregard
        #if keep_list[filnum][segnum] == True:
             
        prof = yt.load(''.join([directory,'/',file_in_dir])).data[var]
        sto.result = prof
        sto.result_id = "%d_%d" %(filnum,segnum)
        
        
            

    for (fil,seg),prof in sorted(storage.items()):
        var_prof[fil].append(prof)
        
               
    
    return var_profs,x

def test_store():
    ds = yt.load(G30)
    store = ds.parameter_filename + '.yt'
    field = "density"
    if os.path.isfile(store):
        os.remove(store)

    proj1 = ds.proj(field, "z")
    sp = ds.sphere(ds.domain_center, (4, 'kpc'))
    proj2 = ds.proj(field, "z", data_source=sp)

    proj1_c = ds.proj(field, "z")
    yield assert_equal, proj1[field], proj1_c[field]

    proj2_c = ds.proj(field, "z", data_source=sp)
    yield assert_equal, proj2[field], proj2_c[field]

    def fail_for_different_method():
        proj2_c = ds.proj(field, "z", data_source=sp, method="mip")
        return (proj2[field] == proj2_c[field]).all()
    yield assert_raises, YTUnitOperationError, fail_for_different_method

    def fail_for_different_source():
        sp = ds.sphere(ds.domain_center, (2, 'kpc'))
        proj2_c = ds.proj(field, "z", data_source=sp, method="integrate")
        return assert_equal(proj2_c[field], proj2[field])
    yield assert_raises, AssertionError, fail_for_different_source

def generateAbsorbers(pfname, pfdir, filepath, i = 0):
    """
    pfname -> name of data dump.... like "RD0020"
    pfdir  -> parent directory to dump ==>  pfdir + "/" + pfname + "/" pfname
                                            should be the data dump itself
    filepath -> path to absorption results... should contain observation result
                text files.
    """
    
    #i = 0
    halos = np.genfromtxt(pfdir + "/" + pfname + "_halos.txt", names=True)
    center = np.array([halos['x'][i],halos['y'][i],halos['z'][i]])
    R_vir  = halos['virial_radius'][i]
    
    ds = yt.load(pfdir + "/" + pfname + "/" + pfname)
    
    cluster = GalCluster(ds, center, R_vir)
    
    line_file = filepath + "/QSO_data_absorbers.out"
    data = np.genfromtxt(line_file, names=True)
    all_points = plot_gas_3D(cluster, filepath, data)
    
    if not os.path.isdir(filepath + "/absorbers"):
        os.mkdir(filepath + "/absorbers")
    
    saveAbsorberList(all_points, filepath + "/absorbers/absorberList.pickle")

    return all_points

def generate_container_data(test_dir):

    file_to_write = os.path.join(test_dir, "containers.h5")

    ds = yt.load("enzo_tiny_cosmology/DD0046/DD0046")

    for key, value in cont_dict.items():
        cont = getattr(ds, camelcase_to_underscore(value[0]), None)
        if cont is not None:
            c = cont(*value[1], **value[2])
            c["density"].write_hdf5(file_to_write, dataset_name=key)
            
    # Special handling

    sp1 = ds.sphere(*cont_dict["sp1"][1])
    prj = ds.proj("density", 1, data_source=sp1)
    prj["density"].write_hdf5(file_to_write, dataset_name="prj3")

    sp2 = ds.sphere(*cont_dict["sp2"][1])
    conditions = ["obj['kT'] > 0.5"]
    cr = sp2.cut_region(conditions)
    cr["density"].write_hdf5(file_to_write, dataset_name="cr")

    for i, grid in enumerate(ds.index.grids):
        grid["density"].write_hdf5(file_to_write, dataset_name="grid_%04d" % (i+1))

def extract_DensVelVort(fil,args,vars):
    lev=0
    if args.ytVersion2: 
       ds=load(fil)
       lbox= float(ds.domain_width[0])
       cube= ds.h.covering_grid(level=lev,left_edge=ds.domain_left_edge,
                          dims=ds.domain_dimensions) 
    else:
        ds=yt.load(fil)
        lbox= float(ds.domain_width[0])
        cube= ds.covering_grid(level=lev,left_edge=ds.domain_left_edge,
                          dims=ds.domain_dimensions)
    print "gathering 3d data arrays"
    #3D arrays
    rho_x = np.array(cube["density"])
    vx = np.array(cube["X-momentum"])/rho_x
    vy = np.array(cube["Y-momentum"])/rho_x
    vz = np.array(cube["Z-momentum"])/rho_x
    print "3d arrays extracted, now calculating vorticity and single numbers like 3D Mach"
    rho_inf= rho_x.flatten().mean()
    (vmag_x,Mach_3D)= getVrms_weighted3DVrms(vx,vy,vz,rho_x,args)
    wmag_x= Sim_curl.getCurl(vx,vy,vz,  ds) #3D array
    del vx,vy,vz,cube
    #store variables
    vars.rho.append(rho_x)
    vars.vmag.append(vmag_x)
    vars.wmag.append(wmag_x)
    vars.basename.append(ds.basename)
    vars.lbox.append(lbox)
    vars.mach3d.append(Mach_3D)
    del rho_x,vmag_x,wmag_x

    def __init__(self, pfname, pfdir = '', i = 0, res_type='cluster'):
        self.pfname = pfname
        self.pf = yt.load(pfdir + pfname + "/" + pfname)
        self.res_type = res_type

        if res_type == 'cluster':
            halos = np.genfromtxt(pfdir + pfname +"_halos.txt", names=True)
            self.cluster_center = np.array([halos['x'][i],halos['y'][i],halos['z'][i]])
            self.mass   = halos['virial_mass'][i]
            self.R_vir  = halos['virial_radius'][i] 
            self.ray_length = 5.0*self.R_vir # *****************************************

            

                
        elif res_type == 'igm':
            self.ray_length = self.pf.domain_width.convert_to_units('Mpc')[0].value * 0.99
            
            
        self.z_center = self.ray_length/2.0 * self.pf.hubble_constant*100.0/CONST_c
        

        
        


        self.defineFieldLabels()

def all_direction_slices(timestep):
    from mpl_toolkits.axes_grid1 import AxesGrid
    ds = yt.load("mhd_sphere_hdf5_chk_{}".format(str(timestep).zfill(4)))
    fig = plt.figure()
    grid = AxesGrid(fig, ( (0, 0, 0.8, 0.8)),
                    nrows_ncols = (1, 3),
                    axes_pad = 1.0,
                    label_mode = "1",
                    share_all = True,
                    cbar_location="right",
                    cbar_mode="each",
                    cbar_size="3%",
                    cbar_pad="0%")
    direction = ['x','y','z']
    physical_quantity='magnetic_field_strength'
    for i, direc in enumerate(direction):
        slc = yt.SlicePlot(ds,direc, physical_quantity)
        slc.set_axes_unit('pc')
        slc.set_font_size(12)
        slc.annotate_magnetic_field()
        plot = slc.plots[physical_quantity]
        plot.figure = fig
        slc.set_cmap(physical_quantity,"rainbow")
        plot.axes = grid[i].axes
        plot.cax = grid.cbar_axes[i]
        slc._setup_plots()
    plt.savefig(SAVE_PATH+"{0}_alldir_Bstrength.png".format(ds))

def make_gas_profiles(dslist = [], outname = 'gas_profile_evolution.h5', overwrite = False):

    if os.path.exists(outname):
        print "Output file exists - do not overwrite"
        return dd.io.load(outname)

    all_data_dict = {}

    all_data_dict['t'] = np.zeros(len(dslist))
    all_data_dict['surface_density'] = [None]*len(dslist)
#    all_data_dict['column_density']  = [None]*len(dslist)

    i = 0
    for d in dslist:
        ds = yt.load(d)
        com = profiles.center_of_mass(ds)
        r, sigma = profiles.generate_gas_profile(ds, ds.all_data(), com = com)


        all_data_dict['r'] = r
        all_data_dict['t'][i] = ds.current_time.convert_to_units('Myr').value
        all_data_dict['surface_density'][i] = sigma

        i = i + 1
#        all_data_dict['column_density'][i]  = N


    dd.io.save(outname, all_data_dict)

    return all_data_dict

def pretty_pic(data_hdf5,args):
	ds = yt.load(data_hdf5) # load data
	ds.add_field("KaysVx", function=_KaysVx, units="cm/s")
	ds.add_field("KaysVy", function=_KaysVy, units="cm/s")
	ds.add_field("KaysVz", function=_KaysVz, units="cm/s")
	ds.add_field("KaysBx", function=_KaysBx, units="gauss")
	ds.add_field("KaysBy", function=_KaysBy, units="gauss")
	ds.add_field("KaysBz", function=_KaysBz, units="gauss")

	msink=13./32
	cs=1.
	mach=5.
	rho0=1.e-2
	rBH= msink/cs**2/(mach**2+1)
	sink= InitSink(args.data_sink)
	sink.coord= dict(x=0.,y=0.,z=0.) #override it 
	center = [0,0,0]
	width = (float(ds.domain_width[0]), 'cm') 
	Npx= 1000
	res = [Npx, Npx] # create an image with 1000x1000 pixels
	fig,ax= plt.subplots() #sharey=True,sharex=True)
	#fig.set_size_inches(15, 10)
	cut_dir='z'
	if args.twod == 'slice': 
		proj = ds.slice(axis=cut_dir,coord=sink.coord[cut_dir])
	else: 
		proj = ds.proj(field="density",axis=cut_dir)
	frb = proj.to_frb(width=width, resolution=res, center=center)
	img= np.array(frb['density'])[::-1]/rho0
	xlab,ylab='x/rBH','y/rBH'
	im= ax.imshow(np.log10(img),vmin=np.log10(args.clim[0]),vmax=np.log10(args.clim[1]))
	ax.autoscale(False)
	#sink marker
	ax.scatter((sink.x+1)*Npx/2,(sink.y+1)*Npx/2,s=50,c='black',marker='o') #all sinks
	#user defined ticks
	xticks=(Npx*np.array([0.,0.25,0.5,0.75,1.])).astype('int')
	# xticks_labs= (xticks-500)*rBH/Npx
	ax.set_xticks(xticks)
	ax.set_xticklabels(((xticks-Npx/2)*float(args.width_rBH)/Npx).astype('int'))  #ax.xaxis.get_ticklocs()/100.)
	ax.set_yticks(xticks)
	ax.set_yticklabels(((xticks[::-1]-Npx/2)*float(args.width_rBH)/Npx).astype('int')) #ax.yaxis.get_ticklocs()/100.)
	ax.set_xlabel(xlab)
	ax.set_ylabel(ylab)
	ax.set_title(cut_dir+' '+args.twod)

	cax = fig.add_axes([0.83, 0.15, 0.02, 0.7])
	cbar= fig.colorbar(im, cax=cax)    

	# cbar = fig.colorbar(cax) # ticks=[-1,0,1])
	cbar_labels = [item.get_text() for item in cbar.ax.get_yticklabels()]
	cbar_labels = ['']*len(cbar_labels)
	cbar_labels[0]= args.clim[0]
	cbar_labels[-1]= args.clim[1]
	cbar.ax.set_yticklabels(cbar_labels)
	if args.twod == 'slice': lab= r'$\mathbf{ \rho/\bar{\rho} }$'
	else: lab= r'$\mathbf{ \Sigma/\bar{\Sigma} }$'
	cbar.set_label(r'$\mathbf{ \rho/\bar{\rho} }$',fontsize='xx-large')
	add_name= 'prettypic'
	plt.savefig(os.path.join(args.outdir,args.twod+'_'+add_name+'_'+os.path.basename(data_hdf5)+'_.png'))

def create_profile(infile):
    ds = yt.load(infile)
    ad = ds.all_data()
    profile = yt.create_profile(ad, args.axis, fields=[args.field],
                                weight_field=None, n_bins=args.nbins,
                                logs={args.axis: args.axis_log,
                                      args.field: args.field_log})
    return profile

def project( filein, fileout = 'junk.png', sinkfile=None) :
	pf = yt.load(file)
	print moviefile
	p = yt.ProjectionPlot(pf, "z", "density")
	p.set_zlim("density", 1e-3,1e0)
	if( not sinkfile == None) : 
		annotate_particles(p, sinkfile) 
		
        p.save(fileout)

def load_halos(ds,location):
    #Merely loads an already saved halocatalog
    #Calls a fn to gather positions and radii of halos over a mass threshold
    halo_ds = yt.load(location)
    hc = HaloCatalog(data_ds=ds,halos_ds=halo_ds)
    hc.add_filter('quantity_value','particle_mass','>',1E12,'Msun')
    halo_ds = hc.halos_ds.all_data()

    radii, positions = get_halo_pos(ds,halo_ds)
    return hc, radii, positions

def get_halo_bounds(nth_most_massive, pdf, hdf):

    particles_midx = pdf + ".midx10"
    pds = yt.load(pdf, midx_filename=particles_midx)

    halos_midx = hdf + ".midx7"
    hds = yt.load(hdf, midx_filename=halos_midx)

    halo = get_halo_params(nth_most_massive)

    sds = SuperData(pds, hds)

    hx, hy, hz, hr = halo['x'], halo['y'], halo['z'], halo['r200b']

    px, py, pz = sds.particle_position(sds.halo_dataset.arr([hx, hy, hz], 'code_length')).d
    pr = sds.conv_arr(sds.halo_dataset.arr([hr], 'code_length'), sds.full_particle_dataset).d
    center = np.array([px, py, pz])
    left = center - 2*pr
    right = center + 2*pr
    return halo, center, left, right, pr[0]

def plot_var(i,physical_quantity,cut="z",velocity=False,grid=False,zmin ="",zmax="",particle=False):
    ds = yt.load("mhd_sphere_hdf5_chk_{}".format(str(i).zfill(4)))
    slc = yt.SlicePlot(ds, cut,physical_quantity)#,center=(0.5,0.5,0.5))
    slc.set_figure_size(5)
    if grid: slc.annotate_grids()
    if velocity: slc.annotate_velocity()
    magnetic = True
    if magnetic: slc.annotate_magnetic_field()
    slc.set_cmap("all","rainbow")
    if zmin!="" and zmax!="": slc.set_zlim(physical_quantity,zmin,zmax)
    #slc.show()
    slc.save(SAVE_PATH+"{0}_{1}_{2}.png".format(ds,physical_quantity,cut))

def test_particle_filter():
    """Test the particle_filter decorator"""

    @particle_filter(filtered_type='all', requires=['creation_time'])
    def stars(pfilter, data):
        filter_field = (pfilter.filtered_type, "creation_time")
        return (data.ds.current_time - data[filter_field]) > 0

    ds = yt.load(iso_galaxy)
    ds.add_particle_filter('stars')
    ad = ds.all_data()
    ad['deposit', 'stars_cic']
    assert True