def apply_negative_radial_fit(name,Ncores): KK=np.load(name+'_KK.npy') KT=KK[0] KD=KK[1] Nbins=len(KT) my_storage = {} maps=[] print('computing optimization') indices=range(Nbins) for sto, kt in yt.parallel_objects(set(KT), Ncores, storage = my_storage): output = temporal_turbulence(name,kt) sto.result_id = output sto.result = kt print('applying optimization') my_storage2 = {} for sto, ij in yt.parallel_objects(indices, Ncores, storage = my_storage2): mapa=apply_temp(name,KT,KD,ij,Nbins) sto.result_id = ij sto.result = mapa if yt.is_root(): for fn, vals in sorted(my_storage2.items()): if ij==0: mapa=vals else: mapa+=vals np.save(name+'_negative_fit_%02d' %Nbins,mapa)
def Circulation_negative_turbulence(name,Nbins,start,kmin,kmax,Nk,Ncores): ktarray=np.exp(np.linspace(np.log(kmin),np.log(kmax),Nk)) my_storage = {} kdmatrix=[] errmatrix=[] for sto, kt in yt.parallel_objects(ktarray, Ncores, storage = my_storage): output = Circulation_negative_fits(name,kt,Nbins,start) sto.result_id = kt sto.result = output if yt.is_root(): for fn, vals in sorted(my_storage.items()): kdmatrix.append(vals[0]) errmatrix.append(vals[1]) rcen = vals[2] kdmatrix=np.array(kdmatrix) errmatrix=np.array(errmatrix) KT=np.zeros(Nbins) KD=np.zeros(Nbins) for j in range(Nbins): optimo=np.where(errmatrix[:,j]==errmatrix[:,j].min()) KD[j]=kdmatrix[optimo,j] KT[j]=ktarray[optimo] KK=np.array([KT,KD]) np.save(name+'_KK',KK) return KT,KD
def sep_fils(filament,radii,positions,ds): #Routine which checks to see if segments of a filament reside within a halo #Init vars storage = {} fils = np.array(filaments) keep_arr = [] #Loop across all filaments, utilizing parellization, split into pool of 4 jobs, this is arbritrary for sto,index in yt.parallel_objects(range(len(fils)),storage = storage,njobs=4): #init some new vars, keeplist will be a 'mask' array for filaments keeplist = np.empty(fils[index].shape[0],dtype=bool) segstor = {} #loop over all segments of a filament, again utlize parellization. for stor,(i,section) in yt.parallel_objects(enumerate(fils[index]),storage=segstor): keep = True #Check if segment is within halo or not for pos,rad in zip(positions,radii): if np.linalg.norm(section[0:3] - pos) < rad: keep = False break stor.result = keep stor.result_id = i for i in range(fils[index].shape[0]): keeplist[i] = segstor[i] sto.result = keeplist sto.result_id = index #Aggragate arrays from all parallel jobs keep_arr = [ lists for (key,lists) in sorted(storage.items())] return keep_arr
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 save_all_profs(profiles,dsname,ds): #Save ALL profiles to disk, where profiles are supplied in a list of 2 variable types ( those weighted by mass and those by volume), containing lists of filaments containing the list of segments. #Init useful vars, important varshort is 4 chars long for loading from disk later varshort = ['dens','kine'] #Split job into 2, one for the different weight fields for jobnumber in yt.parallel_objects(range(2),njobs=2): #Parallelize the workload efficiently for i,fil in yt.parallel_objects(enumerate(profiles)): for j,seg in enumerate(fil): #Determine which variable we are dealing with if jobnumber % 2 == 0: segment = seg[0] else: segment = seg[1] #Save segment of filament profile to disk save_prof(segment,dsname,varshort[jobnumber],(i,j))
def load_all_profs(directory,filament_num): #Loads profiles from disk. #Profiles are stored as a 3D list of segments listed within filaments listen in var type #This means there are three keys needed to identify one profile, its var type, its filament number and its segment number #I.E. They are accessed via profiles[ variable number ] [ filament number] [ segment number ] #Variables are in the order dens, temp, dark , velo - later changing this to a dict would be much more conveniant and elegant #Init lists of filaments, and dict to help address this vardict = {'dens':0,'kine':1} profiles = [ [ [] for x in range(filament_num) ] for i in range(2)] #Determine total list of files filelist = sorted(os.listdir(directory)) #For each file in this list, load data, and append to correct part of profiles structure filelen = len(filelist) denslist = filelist[0:filelen/2] kinslist = filelist[filelen/2:] stordict = {} for stor,(i,file_in_dir) in yt.parallel_objects(enumerate(filelist),storage=stordict): #prof = yt.load(''.join([directory,'/',file_in_dir])) key_id = (vardict[file_in_dir[0:4]],int(file_in_dir[7:10]),int(file_in_dir[13:16])) stor.result = yt.load(''.join([directory,'/',file_in_dir])) stor.result_id = key_id print key_id for (key0,key1,key2),value in sorted(stordict): profiles[key0][key1].append(values) return profiles
def Circulation_negative_optimum_mpi(name,kmin,kmax,Nk,Ncores): ktarray=np.exp(np.linspace(np.log(kmin),np.log(kmax),Nk)) my_storage = {} kdarray=[] erarray=[] for sto, kt in yt.parallel_objects(ktarray, Ncores, storage = my_storage): output = Circulation_negative_optimum_thread(name,kt) sto.result_id = kt sto.result = output if yt.is_root(): for fn, vals in sorted(my_storage.items()): kdarray.append(vals[0]) erarray.append(vals[1]) kdarray=np.array(kdarray) erarray=np.array(erarray) optimo=np.where(erarray==erarray.min()) return ktarray[optimo],kdarray[optimo]
def make_save_profiles(fils,dsname,ds,check=False): #Routine which generates temperature, density and velocity profiles for all segments of a filament #inits useful vars var = [('density'),('dark_matter_density'),('temperature'),('cylindrical_radial_velocity')] weights = ("cell_volume","cell_mass") profiles = [[] for i in range(len(fils))] filsize = len(fils) #check to see if some profiles already exist on disk, in whcih case we dont need to make the profiles if check == True: checkfils = [] directory = ''.join(['~/data',dsname.upper(),'/profiles/']) filelist = os.listdir(directory) for files in filelist: checkfils.append((int(files[7:10]),int(files[13:16]))) else: checkfils = [] #Loops over all filaments for i in yt.parallel_objects(range(filsize),dynamic=check): profs = [] filaments = fils[i] for section in (range(len(filaments) - 1)): #Generate useful profiles for all segments of filament if (i,section) in checkfils: var_prof = profiling.det_fil_profile(filaments[section],filaments[section + 1],var[0:2],weights[0],2,ds) var_prof2 = profiling.det_fil_profile(filaments[section],filaments[section + 1],var[2:],weights[1],2,ds) #Save these to disk save_prof(var_prof,dsname,'dens',(i,section)) save_prof(var_prof2,dsname,'kine',(i,section))
if args.make_frames_only == 'False': verbatim = False if rank == 0: verbatim = True usable_files = mym.find_files(m_times, files, sink_form_time,sink_id, verbatim=False) del sink_form_time del files sys.stdout.flush() CW.Barrier() if args.make_frames_only == 'False': #Trying yt parallelism file_int = -1 for fn in yt.parallel_objects(usable_files, njobs=int(size/6)): if size > 1: file_int = usable_files.index(fn) else: file_int = file_int + 1 if usable_files[file_int] == usable_files[file_int-1]: os.system('cp '+ save_dir + "movie_frame_" + ("%06d" % frames[file_int-1]) + ".pkl " + save_dir + "movie_frame_" + ("%06d" % frames[file_int]) + ".pkl ") make_pickle = False if args.plot_time is None: pickle_file = save_dir + "movie_frame_" + ("%06d" % frames[file_int]) + ".pkl" else: pickle_file = save_dir + "time_" + str(args.plot_time) +".pkl" if os.path.isfile(pickle_file) == False: make_pickle = True elif os.path.isfile(pickle_file) == True: if os.stat(pickle_file).st_size == 0:
if __name__ == "__main__": es = yt.load("rs_normal_bg1.h5") fns = es.data["filename"].astype(str) value = 1e-3 halo_dir = "halo_2170858" # halo_dir = "halo_2171203" data_dir = os.path.join(halo_dir, "clumps_%.0e_inner" % value) ofn = os.path.join(halo_dir, "bh_clump_distance_edge_%.0e_inner.h5" % value) contained_info = {} my_storage = {} for sto, (i, fn) in yt.parallel_objects(enumerate(fns), storage=my_storage): my_dmin = [] contained = 0 cfns = glob.glob(os.path.join(data_dir, os.path.dirname(fn), "*.h5")) pbar = yt.get_pbar( "%s (z = %f) - Calculating distances" % (os.path.dirname(fn), es.data["redshift"][i]), len(cfns)) for cfn in cfns: clump_dmin = [] mylog.setLevel(40) ds = yt.load(cfn) mylog.setLevel(llevel) if not hasattr(ds, "tree") or \ ds.tree.children is None or \ len(ds.tree.children) == 0:
# Enable MPI yt.enable_parallelism() # Load data ds = yt.load(str(SNAP)) ds.index hc = pd.read_csv(CAND) if yt.is_root(): print(len(hc), 'halos to measure') # Measure local environment density storage = {} for sto, (index, row) in yt.parallel_objects( list(hc.to_dict(orient='index').items()), storage=storage, dynamic=True ): cache = BASE_DIR / f'data/halo/menv-{index}' if cache.exists(): m_env = struct.unpack('d', cache.read_bytes())[0] print(index, 'from cache') else: print(index, 'begin') center = ds.arr([row['Xc'], row['Yc'], row['Zc']], 'kpccm/h') sp = ds.sphere(center, (R_ENV, 'Mpc')) m_env = sp['all', 'particle_mass'].sum().to_value('Msun') cache.write_bytes(struct.pack('d', m_env)) print(index, 'end') sto.result_id = index sto.result = m_env
#experimenting import numpy as np import matplotlib.pyplot as plt import yt import glob fpath = '/zang/msoares/JLopez/' ffiles= '50_hdf5_plt_cnt_[0-9][0-9][0][0]' num_procs=20 # get the files to loop through my_filenames = glob.glob(fpath+ffiles) my_filenames.sort() for f in yt.parallel_objects(my_filenames,num_procs): print f pf=yt.load(f) tf=yt.ColorTransferFunction((-29,-24)) tf.add_layers(8,w=0.05, colormap="Hue Sat Value 2") cam = pf.camera([0.5, 0.5, 0.5], [1.0, 1.0, 1.0], (.5, 'pc'), 512, tf, fields=["d\ ensity"]) frame = 0 # Do a rotation over 5 frames for i, snapshot in enumerate(cam.rotation(np.pi, 5, clip_ratio=8.0)): snapshot.write_png('%s_camera_movement_%04i.png' % (pf,frame)) frame += 1
usable_files = [args.specific_file] m_times = mym.generate_frame_times(files, args.time_step, presink_frames=0, end_time=args.end_time, form_time=sink_form_time) if rank == 0: print("Usable_files=", usable_files) no_frames = len(m_times) frames = list(range(args.start_frame, no_frames)) del sink_form_time del files sys.stdout.flush() CW.Barrier() #Trying yt parallelism file_int = -1 for fn_it in yt.parallel_objects(range(len(usable_files)), njobs=int(size/(8*4))): #8 projection and 4 fields. fn = usable_files[fn_it] print("File", fn, "is going to rank", rank) if size > 1: file_int = usable_files.index(fn) else: file_int = file_int + 1 #This was implemented for the specific cases where dt < simulation dump frequency, and the same file might be used for multiple times. if args.plot_time != None: if os.path.exists(save_dir + "time_" + (str(int(args.plot_time)))) == False: try: os.makedirs(save_dir + "time_" + (str(int(args.plot_time)))) except: print(save_dir + "time_" + (str(int(args.plot_time))), "Already exists") else: if os.path.exists(save_dir + "movie_frame_" + ("%06d" % frames[file_int])) == False:
import yt import numpy as np import matplotlib.pyplot as plt from scipy.integrate import simps from scipy.interpolate import interp1d from astropy.table import Table , Column ,vstack,hstack from Vorticity_Difussion import * yt.enable_parallelism() import sys import os cmdargs = sys.argv Ncores=2 directory='Examples' N=1200 narr=np.arange(0.0,1.95,0.1) Sarr=np.arange(1,1004,50) SDiff=[0.1,0.2,0.5,1,2,3,4,5,6,8,10,12,14,16,20,25,50,75,100] N=1200 my_storage = {} for sto, nd in yt.parallel_objects(narr, Ncores, storage = my_storage): output = grid_table_n(N,nd,snarr,sdarr,directory) sto.result_id = 'n%.1f_'%nd sto.result = output
usable_files = files del files gc.collect() sys.stdout.flush() CW.Barrier() ''' rit = -1 for fn_it in range(len(usable_files)): rit = rit + 1 if rit == size: rit = 0 if rank == rit: ''' if args.make_projection == "True": for fn_it in yt.parallel_objects(range(len(usable_files)), njobs=int(size / (100))): pickle_file = save_dir + "movie_frame_" + ("%06d" % fn_it) + "_proj.pkl" if os.path.exists(pickle_file) == False: fn = usable_files[fn_it] ds = yt.load(fn, units_override=units_override) time_real = ds.current_time.in_units('yr') time_val = np.round(time_real.in_units('yr')) proj = yt.ProjectionPlot(ds, args.perp_axis, ("ramses", "Density"), width=units['length_unit'], method='integrate') proj_array = np.array(proj.frb.data[("ramses", "Density")] / units['length_unit'].in_units('cm'))
@yt.derived_field(name="IonizedHydrogen", units="", display_name=r"\frac{\rho_{HII}}{\rho_H}") def IonizedHydrogen(field, data): return data["HII_Density"] / (data["HI_Density"] + data["HII_Density"]) ts = yt.DatasetSeries("SED800/DD*/*.index", parallel=8) ionized_z = np.zeros(ts[0].domain_dimensions, dtype="float32") t1 = time.time() for ds in ts.piter(): z = ds.current_redshift for g in yt.parallel_objects(ds.index.grids, njobs=16): i1, j1, k1 = g.get_global_startindex() # Index into our domain i2, j2, k2 = g.get_global_startindex() + g.ActiveDimensions # Look for the newly ionized gas newly_ion = (g["IonizedHydrogen"] > 0.999) & (ionized_z[i1:i2, j1:j2, k1:k2] < z) ionized_z[i1:i2, j1:j2, k1:k2][newly_ion] = z g.clear_data() print("Iteration completed %0.3e" % (time.time() - t1)) comm = communication_system.communicators[-1] for i in range(ionized_z.shape[0]): ionized_z[i, :, :] = comm.mpi_allreduce(ionized_z[i, :, :], op="max") print("Slab % 3i has minimum z of %0.3e" % (i, ionized_z[i, :, :].max())) t2 = time.time() print("Completed. %0.3e" % (t2 - t1))
sys.stdout.flush() CW.Barrier() #Trying yt parallelism file_int = -1 field_list = [ proj_field, ('gas', 'Radial_Velocity'), ('gas', 'Proj_x_velocity'), ('gas', 'Proj_y_velocity') ] N_proj_vectors = 8 if len(usable_files) == 1: njobs = 1 else: njobs = int(size / 8 * 4) for fn_it in yt.parallel_objects(range(len(usable_files)), njobs=njobs): #8 projection and 4 fields. fn = usable_files[fn_it] print("File", fn, "is going to rank", rank) if size > 1: file_int = usable_files.index(fn) else: file_int = file_int + 1 #This was implemented for the specific cases where dt < simulation dump frequency, and the same file might be used for multiple times. if args.plot_time != None: if os.path.exists(save_dir + "time_" + (str(int(args.plot_time)))) == False: try: os.makedirs(save_dir + "time_" + (str(int(args.plot_time)))) except: print(save_dir + "time_" + (str(int(args.plot_time))), "Already exists")
XMax = 25 YMin = -25 YMax = 25 R_grid, THETA_grid = np.meshgrid(R[NG:-NG], THETA[NG:-NG]) X = R_grid*np.sin(THETA_grid) Y = R_grid*np.cos(THETA_grid) alpha = np.sqrt(-gCon[:,:,0,0])**(-1) PolGamma = 4./3 minorLocatorX = pl.FixedLocator(np.linspace(0,60,13)) minorLocatorY = pl.FixedLocator(np.linspace(-20,20,9)) majorLocatorX = pl.FixedLocator(np.linspace(0,60,7)) majorLocatorY = pl.FixedLocator(np.linspace(-20,20,5)) for file_number, dump_file in yt.parallel_objects(enumerate(data_files)): print "File number = ", file_number frame_index = file_number data_file = h5py.File(dump_file, "r") primVars = data_file['primVars'] elem = returnFluidElement(primVars) A_plot = returnMagneticVectorPotential(elem) N_start = 1 fig, axes = pl.subplots(nrows=1, ncols=3) fig.set_size_inches((24, 12))
from scipy.interpolate import interp1d from astropy.table import Table , Column ,vstack,hstack from Vorticity_Difussion import * yt.enable_parallelism() import sys import os cmdargs = sys.argv Ncores=2 directory='Examples' N=1200 #narr=[0.1,0.25,0.5,0.75,1.0,1.25,1.5] narr=np.linspace(0.0,1.99,20) Sarr=np.arange(50,600,50) SDiff=[1,2,3,4,5,6,8,10,12,14,16,20,25] #SDiff='None' par=[] for n in narr: for S in Sarr: par.append([n,S]) my_storage = {} for sto, li in yt.parallel_objects(par, Ncores, storage = my_storage): output = save_example(N,li[0],li[1],SDiff,directory) sto.result_id = 'n%.1f_'%li[0]+'S%.1f'%li[1] sto.result = output
def run_sightlines(outputfilename,save_after_num,parallel,\ simulation_dest = None,run = 'default',throwerrors = 'warn'): if run not in ['default', 'test']: print('unknown option for "run" %s.'%run+\ ' Please restart with "run = default"'+\ ' or "run = test".') #do not print out anything from yt (it prints plenty) yt.funcs.mylog.setLevel(50) if parallel: yt.enable_parallelism() readvalsoutput = simulation_quasar_sphere.read_values(outputfilename) #by creating a QuasarSphere, it knows all its metadata and other #information from simparams and scanparams (first lines of file at #'filename') q = simulation_quasar_sphere.SimQuasarSphere( start_up_info_packet=readvalsoutput) if q.simparams[6] is None: if simulation_dest: q.simparams[6] = simulation_dest else: raise NoSimulationError('Simulation file location unknown, '+\ 'run with "simulation_dest" to process') else: simulation_dest = q.simparams[6] ds,fields_to_keep = code_specific_setup.load_and_setup(simulation_dest,\ q.fullname,ions = q.ions,\ redshift = q.redshift) set_up_general(ds, q.code, q.center, q.bulk_velocity, q.Rvir) code_specific_setup.check_redshift(ds, outputfilename=outputfilename) num_bin_vars = q.gasbins.get_length() #Can start at a position further than 0 if reached starting_point = q.length_reached bins = np.append(np.arange(starting_point, q.length, save_after_num), q.length) #first for loop is non-parallel. If 32 processors available, it will break up #into bins of size 32 at a time for example. At end, saves data from all 32. #this is (~12 bins) in usual circumstances for i in range(0, len(bins) - 1): current_info = q.info[bins[i]:bins[i + 1]] if yt.is_root(): tprint("%s-%s /%s" % (bins[i], bins[i + 1], len(q.info))) my_storage = {} #2nd for loop is parallel. Each vector goes to a different processor, and creates #a separate trident sightline (~32 sightlines [in a bin]). #the longest processing step is ray = trident.make_simple_ray, and it's #the only step which actually takes any time (below two for loops go by fast) for sto, in_vec in yt.parallel_objects(current_info, storage=my_storage): vector = np.copy(in_vec) index = vector[0] toprint = "line %s, (r = %.0f) densities " % (str( int(index)), vector[3]) tprint("<line %d, starting process> " % index) ident = str(index) start = ds.arr(tuple(vector[5:8]), 'unitary') end = ds.arr(tuple(vector[8:11]), 'unitary') try: ray = trident.make_simple_ray(ds, start_position=start, end_position=end, data_filename="ray" + ident + ".h5", fields=fields_to_keep, ftype='gas') except KeyboardInterrupt: print('skipping sightline %s ...' % index) print('Interrupt again within 5 seconds to *actually* end') time.sleep(5) continue except ValueError as e: throw_errors_if_allowed( e, throwerrors, 'ray has shape %s - %s, but size 0' % (start, end)) except Exception as e: throw_errors_if_allowed(e, throwerrors, 'problem with making ray') continue trident.add_ion_fields(ray, q.ions) field_data = ray.all_data() dl = field_data['gas', 'dl'] #3rd for loop is for processing each piece of info about each ion #including how much that ion is in each bin according to gasbinning #here just process topline data (column densities and ion fractions) #(~10 ions) for j in range(len(q.ions)): ion = q.ions[j] ionfield = field_data["gas", ion_to_field_name(ion)] cdens = np.sum((ionfield * dl).in_units('cm**-2')).value vector[11 + j * (num_bin_vars + 2)] = cdens total_nucleus = np.sum(ionfield[ionfield>0]/\ field_data["gas",ion_to_field_name(ion,'ion_fraction')][ionfield>0]\ * dl[ionfield>0]) vector[11 + j * (num_bin_vars + 2) + 1] = cdens / total_nucleus #4th for loop is processing each gasbin for the current ion #(~20 bins) for k in range(num_bin_vars): try: variable_name, edges, units = q.gasbins.get_field_binedges_for_num( k, ion) if variable_name is None: vector[11 + j * (num_bin_vars + 2) + k + 2] = np.nan elif variable_name in ray.derived_field_list: if units: data = field_data[variable_name].in_units( units) else: data = field_data[variable_name] abovelowerbound = data > edges[0] belowupperbound = data < edges[1] withinbounds = np.logical_and( abovelowerbound, belowupperbound) coldens_in_line = (ionfield[withinbounds]) * ( dl[withinbounds]) coldens_in_bin = np.sum(coldens_in_line) vector[11 + j * (num_bin_vars + 2) + k + 2] = coldens_in_bin / cdens else: print( str(variable_name) + " not in ray.derived_field_list") except Exception as e: throw_errors_if_allowed( e, throwerrors, 'Could not bin into %s with edges %s' % (variable_name, edges)) toprint += "%s:%e " % (ion, cdens) #gets some more information from the general sightline. #metallicity, average density (over the whole sightline) #mass-weighted temperature try: if ('gas', "H_nuclei_density") in ray.derived_field_list: Z = np.sum(field_data['gas',"metal_density"]*dl)/ \ np.sum(field_data['gas',"H_nuclei_density"]*mh*dl) else: Z = np.sum(field_data['gas',"metal_density"]*dl)/ \ np.sum(field_data['gas',"number_density"]*mh*dl) vector[-1] = Z except Exception as e: throw_errors_if_allowed(e, throwerrors, 'problem with average metallicity') try: n = np.sum( field_data['gas', 'number_density'] * dl) / np.sum(dl) vector[-2] = n except Exception as e: throw_errors_if_allowed(e, throwerrors, 'problem with average density') try: T = np.average(field_data['gas','temperature'],\ weights=field_data['gas','density']*dl) vector[-3] = T except Exception as e: throw_errors_if_allowed(e, throwerrors, 'problem with average temperature') try: os.remove("ray" + ident + ".h5") except: pass tprint(toprint) #'vector' now contains real data, not just '-1's sto.result_id = index sto.result = vector #save all parallel sightlines after they finish (every 32 lines are saved at once) if yt.is_root(): keys = my_storage.keys() for key in keys: q.info[int(key)] = my_storage[key] q.scanparams[6] += (bins[i + 1] - bins[i]) q.length_reached = q.scanparams[6] if run != 'test': outputfilename = q.save_values(oldfilename=outputfilename) tprint("file saved to " + outputfilename + ".")
def calc_ang_mom_and_fluxes(halo, foggie_dir, output_dir, run, **kwargs): outs = kwargs.get("outs", "all") trackname = kwargs.get("trackname", "halo_track") ### set up the table of all the stuff we want data = Table(names=('redshift', 'radius', 'nref_mode', \ 'net_mass_flux', 'net_metal_flux', \ 'mass_flux_in', 'mass_flux_out', \ 'metal_flux_in', 'metal_flux_out', \ 'net_kinetic_energy_flux','net_thermal_energy_flux','net_entropy_flux',\ 'kinetic_energy_flux_in','kinetic_energy_flux_out',\ 'thermal_energy_flux_in','thermal_energy_flux_out',\ 'entropy_flux_in','entropy_flux_out',\ 'net_cold_mass_flux', 'cold_mass_flux_in', 'cold_mass_flux_out', \ 'net_cool_mass_flux', 'cool_mass_flux_in', 'cool_mass_flux_out', \ 'net_warm_mass_flux', 'warm_mass_flux_in', 'warm_mass_flux_out', \ 'net_hot_mass_flux', 'hot_mass_flux_in', 'hot_mass_flux_out', \ 'annular_ang_mom_gas_x', 'annular_ang_mom_gas_y','annular_ang_mom_gas_z', \ 'annular_spec_ang_mom_gas_x', 'annular_spec_ang_mom_gas_y','annular_spec_ang_mom_gas_z',\ 'annular_ang_mom_dm_x', 'annular_ang_mom_dm_y','annular_ang_mom_dm_z', \ 'annular_spec_ang_mom_dm_x', 'annular_spec_ang_mom_dm_y', 'annular_spec_ang_mom_dm_z', \ 'outside_ang_mom_gas_x', 'outside_ang_mom_gas_y', 'outside_ang_mom_gas_z', \ 'outside_spec_ang_mom_gas_x', 'outside_spec_ang_mom_gas_y', 'outside_spec_ang_mom_gas_z', \ 'outside_ang_mom_dm_x', 'outside_ang_mom_dm_y','outside_ang_mom_dm_z',\ 'outside_spec_ang_mom_dm_x', 'outside_spec_ang_mom_dm_y', 'outside_spec_ang_mom_dm_z', \ 'inside_ang_mom_stars_x', 'inside_ang_mom_stars_y', 'inside_ang_mom_stars_z', \ 'inside_spec_ang_mom_stars_x', 'inside_spec_ang_mom_stars_y', 'inside_spec_ang_mom_stars_z'), dtype=('f8', 'f8', 'i8','f8', 'f8', 'f8', 'f8', 'f8', 'f8','f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8' )) data2 = Table(names=('redshift', 'radius', 'net_mass_flux', 'net_metal_flux', \ 'mass_flux_in', 'mass_flux_out', \ 'metal_flux_in', 'metal_flux_out', \ 'net_cold_mass_flux', 'cold_mass_flux_in', 'cold_mass_flux_out', \ 'net_cool_mass_flux', 'cool_mass_flux_in', 'cool_mass_flux_out', \ 'net_warm_mass_flux', 'warm_mass_flux_in', 'warm_mass_flux_out', \ 'net_hot_mass_flux', 'hot_mass_flux_in', 'hot_mass_flux_out', \ 'annular_ang_mom_gas_x', 'annular_ang_mom_gas_y','annular_ang_mom_gas_z', \ 'annular_spec_ang_mom_gas_x', 'annular_spec_ang_mom_gas_y','annular_spec_ang_mom_gas_z',\ 'outside_ang_mom_gas_x', 'outside_ang_mom_gas_y', 'outside_ang_mom_gas_z', \ 'outside_spec_ang_mom_gas_x', 'outside_spec_ang_mom_gas_y', 'outside_spec_ang_mom_gas_z'), \ dtype=('f8', 'f8', 'f8','f8', 'f8', 'f8', 'f8', 'f8', 'f8','f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8')) print(foggie_dir) track_name = foggie_dir + 'halo_00' + str(halo) + '/' + run + trackname if args.system == "pleiades": track_name = foggie_dir + "halo_008508/nref11f_refine200kpc_z4to2/halo_track" print("opening track: " + track_name) track = Table.read(track_name, format='ascii') track.sort('col1') ## default is do allll the snaps in the directory ## want to add flag for if just one run_dir = foggie_dir + 'halo_00' + str(halo) + '/' + run if halo == "8508": prefix = output_dir + 'plots_halo_008508/' + run + '/' else: prefix = output_dir + 'other_halo_plots/' + str(halo) + '/' + run + '/' if not (os.path.exists(prefix)): os.system("mkdir " + prefix) if outs == "all": print("looking for outputs in ", run_dir) outs = glob.glob(os.path.join(run_dir, '?D????/?D????')) else: print("outs = ", outs) new_outs = [glob.glob(os.path.join(run_dir, snap)) for snap in outs] print("new_outs = ", new_outs) new_new_outs = [snap[0] for snap in new_outs] outs = new_new_outs #print("outs:") #print (outs) #dataset_series = yt.load(outs) storage = {} #for sto,ds in dataset_series.piter(storage=storage): for sto, snap in yt.parallel_objects(outs, 2, storage=storage): # load the snapshot #print('opening snapshot '+ snap) ds = yt.load(snap) # add the particle filters ds.add_particle_filter('stars') ds.add_particle_filter('dm') # create all the regions zsnap = ds.get_parameter('CosmologyCurrentRedshift') #proper_box_size = get_proper_box_size(ds) # another option than the function: proper_box_size = ds.quan(1., 'code_length').to('kpc') refine_box, refine_box_center, refine_width_code = get_refine_box( ds, zsnap, track) refine_width = refine_width_code * proper_box_size # center is trying to be the center of the halo halo_center, halo_velocity = get_halo_center(ds, refine_box_center) ### OK, now want to set up some spheres of some sizes and get the stuff radii = refine_width * 0.5 * np.arange(0.9, 0.1, -0.01) # 0.5 because radius small_sphere = ds.sphere(halo_center, 0.05 * refine_width_code) # R=10ckpc/h big_sphere = ds.sphere(halo_center, 0.45 * refine_width_code) # we want to subtract the bulk velocity from the radial velocities bulk_velocity = big_sphere.quantities["BulkVelocity"]() # find number of cells for the FRB # by default, it uses nref10 as the cell size for the frb # then create the 3D FRB for calculating the fluxes cell_size = np.unique(big_sphere['dx'].in_units('kpc'))[1] box_width = ds.quan(0.9 * refine_width, 'kpc') nbins = int(np.ceil(box_width / cell_size).value) halo_center = ds.arr(halo_center, 'code_length') xL, xR = halo_center[0] - box_width / 2., halo_center[ 0] + box_width / 2. yL, yR = halo_center[1] - box_width / 2., halo_center[ 1] + box_width / 2. zL, zR = halo_center[2] - box_width / 2., halo_center[ 2] + box_width / 2. jnbins = complex(0, nbins) box = ds.r[xL:xR:jnbins, yL:yR:jnbins, zL:zR:jnbins] box.set_field_parameter("center", halo_center) box.set_field_parameter("bulk_velocity", bulk_velocity) ### OK, now want to call the fields that we'll need for the fluxes ### otherwise, the code crashes when trying to select subsets of the data ## GAS FIELDS temperature = box['Temperature'].flatten() cell_mass = box['cell_mass'].to("Msun").flatten() metal_mass = box[('gas', 'metal_mass')].to("Msun").flatten() radius = box['radius'].to("kpc").flatten() radial_velocity = box['radial_velocity'].to('kpc/yr').flatten() cell_volume = box['cell_volume'].flatten() #grid_levels = box['index', 'grid_level'] gas_ang_mom_x = box[('gas', 'angular_momentum_x')].flatten() gas_ang_mom_y = box[('gas', 'angular_momentum_y')].flatten() gas_ang_mom_z = box[('gas', 'angular_momentum_z')].flatten() gas_spec_ang_mom_x = box[('gas', 'specific_angular_momentum_x')].flatten() gas_spec_ang_mom_y = box[('gas', 'specific_angular_momentum_y')].flatten() gas_spec_ang_mom_z = box[('gas', 'specific_angular_momentum_z')].flatten() kinetic_energy = box['gas', 'kinetic_energy'].flatten() kinetic_energy = (kinetic_energy * cell_volume / cell_mass).to('erg/g') thermal_energy = box['gas', 'thermal_energy'].flatten() entropy = box['entropy'].flatten() hden = box['H_nuclei_density'].flatten() ## STAR PARTICLE FIELDS star_ang_mom_x = big_sphere['stars', 'particle_angular_momentum_x'].flatten() star_ang_mom_y = big_sphere['stars', 'particle_angular_momentum_y'].flatten() star_ang_mom_z = big_sphere['stars', 'particle_angular_momentum_z'].flatten() star_spec_ang_mom_x = big_sphere[ 'stars', 'particle_specific_angular_momentum_x'].flatten() star_spec_ang_mom_y = big_sphere[ 'stars', 'particle_specific_angular_momentum_y'].flatten() star_spec_ang_mom_z = big_sphere[ 'stars', 'particle_specific_angular_momentum_z'].flatten() star_distance = np.sqrt( (big_sphere['stars', 'particle_position_x'] - halo_center[0])**2. + (big_sphere['stars', 'particle_position_y'] - halo_center[1])**2. + (big_sphere['stars', 'particle_position_z'] - halo_center[2])**2.).to("kpc") ## DM PARTICLE FIELDS dm_ang_mom_x = big_sphere['dm', 'particle_angular_momentum_x'].flatten() dm_ang_mom_y = big_sphere['dm', 'particle_angular_momentum_y'].flatten() dm_ang_mom_z = big_sphere['dm', 'particle_angular_momentum_z'].flatten() dm_spec_ang_mom_x = big_sphere[ 'dm', 'particle_specific_angular_momentum_x'].flatten() dm_spec_ang_mom_y = big_sphere[ 'dm', 'particle_specific_angular_momentum_y'].flatten() dm_spec_ang_mom_z = big_sphere[ 'dm', 'particle_specific_angular_momentum_z'].flatten() dm_distance = np.sqrt( (big_sphere['dm', 'particle_position_x'] - halo_center[0])**2. + (big_sphere['dm', 'particle_position_y'] - halo_center[1])**2. + (big_sphere['dm', 'particle_position_z'] - halo_center[2])**2.).to("kpc") table1 = np.zeros((len(radii), len(data.keys()))) table2 = np.zeros((len(radii), len(data2.keys()))) for rad in parallel_objects(radii): print(rad) if rad != np.max(radii): idI = np.where(radii == rad)[0] if rad == radii[-1]: minrad, maxrad = ds.quan(0., 'kpc'), rad else: maxrad, minrad = rad, radii[idI[0] + 1] # some radius / geometry things dr = maxrad - minrad rad_here = (minrad + maxrad) / 2. # find the indices that I'm going to need idR = np.where((radius >= minrad) & (radius < maxrad))[0] idCd = np.where((radius >= minrad) & (radius < maxrad) & (temperature <= 1e4))[0] idCl = np.where((radius >= minrad) & (radius < maxrad) & (temperature > 1e4) & (temperature <= 1e5))[0] idW = np.where((radius >= minrad) & (radius < maxrad) & (temperature > 1e5) & (temperature <= 1e6))[0] idH = np.where((radius >= minrad) & (radius < maxrad) & (temperature >= 1e6)) idRdm = np.where((dm_distance >= minrad) & (dm_distance < maxrad))[0] big_annulusGAS = np.where(radius >= rad_here)[0] big_annulusDM = np.where(dm_distance >= rad_here)[0] inside = np.where(star_distance < rad_here)[0] # most common refinement level #nref_mode = stats.mode(grid_levels[idR]) nref_mode = 10. ## FIX FOR NOW! # mass fluxes gas_flux = (np.sum(cell_mass[idR] * radial_velocity[idR]) / dr).to("Msun/yr") metal_flux = (np.sum(metal_mass[idR] * radial_velocity[idR]) / dr).to("Msun/yr") kinetic_energy_flux = ( np.sum(kinetic_energy[idR] * radial_velocity[idR]) / dr).to("erg/(g*yr)") thermal_energy_flux = ( np.sum(thermal_energy[idR] * radial_velocity[idR]) / dr).to("erg/(g*yr)") entropy_flux = (np.sum(entropy[idR] * radial_velocity[idR]) / dr) ## also filter based off radial velocity idVin = np.where(radial_velocity[idR] <= 0.)[0] idVout = np.where(radial_velocity[idR] > 0.)[0] gas_flux_in = (np.sum( cell_mass[idR][idVin] * radial_velocity[idR][idVin]) / dr).to("Msun/yr") gas_flux_out = (np.sum( cell_mass[idR][idVout] * radial_velocity[idR][idVout]) / dr).to("Msun/yr") metal_flux_in = (np.sum( metal_mass[idR][idVin] * radial_velocity[idR][idVin]) / dr).to("Msun/yr") metal_flux_out = (np.sum( metal_mass[idR][idVout] * radial_velocity[idR][idVout]) / dr).to("Msun/yr") kinetic_energy_flux_in = (np.sum( kinetic_energy[idR][idVin] * radial_velocity[idR][idVin]) / dr).to("erg/(g*yr)") kinetic_energy_flux_out = ( np.sum(kinetic_energy[idR][idVout] * radial_velocity[idR][idVout]) / dr).to("erg/(g*yr)") thermal_energy_flux_in = (np.sum( thermal_energy[idR][idVin] * radial_velocity[idR][idVin]) / dr).to("erg/(g*yr)") thermal_energy_flux_out = ( np.sum(thermal_energy[idR][idVout] * radial_velocity[idR][idVout]) / dr).to("erg/(g*yr)") entropy_flux_in = ( np.sum(entropy[idR][idVin] * radial_velocity[idR][idVin]) / dr) entropy_flux_out = (np.sum( entropy[idR][idVout] * radial_velocity[idR][idVout]) / dr) ## and filter on temperature! and velocity! woo! idVin = np.where(radial_velocity[idH] <= 0.)[0] idVout = np.where(radial_velocity[idH] > 0.)[0] hot_gas_flux = (np.sum(cell_mass[idH] * radial_velocity[idH]) / dr).to("Msun/yr") hot_gas_flux_in = (np.sum( cell_mass[idH][idVin] * radial_velocity[idH][idVin]) / dr).to("Msun/yr") hot_gas_flux_out = (np.sum( cell_mass[idH][idVout] * radial_velocity[idH][idVout]) / dr).to("Msun/yr") idVin = np.where(radial_velocity[idW] <= 0.)[0] idVout = np.where(radial_velocity[idW] > 0.)[0] warm_gas_flux = ( np.sum(cell_mass[idW] * radial_velocity[idW]) / dr).to("Msun/yr") warm_gas_flux_in = (np.sum( cell_mass[idW][idVin] * radial_velocity[idW][idVin]) / dr).to("Msun/yr") warm_gas_flux_out = (np.sum( cell_mass[idW][idVout] * radial_velocity[idW][idVout]) / dr).to("Msun/yr") idVin = np.where(radial_velocity[idCl] <= 0.)[0] idVout = np.where(radial_velocity[idCl] > 0.)[0] cool_gas_flux = ( np.sum(cell_mass[idCl] * radial_velocity[idCl]) / dr).to("Msun/yr") cool_gas_flux_in = (np.sum( cell_mass[idCl][idVin] * radial_velocity[idCl][idVin]) / dr).to("Msun/yr") cool_gas_flux_out = (np.sum( cell_mass[idCl][idVout] * radial_velocity[idCl][idVout]) / dr).to("Msun/yr") idVin = np.where(radial_velocity[idCd] <= 0.)[0] idVout = np.where(radial_velocity[idCd] > 0.)[0] cold_gas_flux = ( np.sum(cell_mass[idCd] * radial_velocity[idCd]) / dr).to("Msun/yr") cold_gas_flux_in = (np.sum( cell_mass[idCd][idVin] * radial_velocity[idCd][idVin]) / dr).to("Msun/yr") cold_gas_flux_out = (np.sum( cell_mass[idCd][idVout] * radial_velocity[idCd][idVout]) / dr).to("Msun/yr") ## GAS angular momentum! annular_ang_mom_gas_x = np.sum(gas_ang_mom_x[idR]) annular_ang_mom_gas_y = np.sum(gas_ang_mom_y[idR]) annular_ang_mom_gas_z = np.sum(gas_ang_mom_z[idR]) annular_spec_ang_mom_gas_x = np.mean(gas_spec_ang_mom_x[idR]) annular_spec_ang_mom_gas_y = np.mean(gas_spec_ang_mom_y[idR]) annular_spec_ang_mom_gas_z = np.mean(gas_spec_ang_mom_z[idR]) outside_ang_mom_gas_x = np.sum(gas_ang_mom_x[big_annulusGAS]) outside_ang_mom_gas_y = np.sum(gas_ang_mom_y[big_annulusGAS]) outside_ang_mom_gas_z = np.sum(gas_ang_mom_z[big_annulusGAS]) outside_spec_ang_mom_gas_x = np.mean( gas_spec_ang_mom_x[big_annulusGAS]) outside_spec_ang_mom_gas_y = np.mean( gas_spec_ang_mom_y[big_annulusGAS]) outside_spec_ang_mom_gas_z = np.mean( gas_spec_ang_mom_z[big_annulusGAS]) ## PARTICLE angular momentum calculations inside_ang_mom_stars_x = np.sum(star_ang_mom_x[inside]) inside_ang_mom_stars_y = np.sum(star_ang_mom_y[inside]) inside_ang_mom_stars_z = np.sum(star_ang_mom_z[inside]) inside_spec_ang_mom_stars_x = np.sum( star_spec_ang_mom_x[inside]) inside_spec_ang_mom_stars_y = np.sum( star_spec_ang_mom_y[inside]) inside_spec_ang_mom_stars_z = np.sum( star_spec_ang_mom_z[inside]) annular_ang_mom_dm_x = np.sum(dm_ang_mom_x[idRdm]) annular_ang_mom_dm_y = np.sum(dm_ang_mom_y[idRdm]) annular_ang_mom_dm_z = np.sum(dm_ang_mom_z[idRdm]) annular_spec_ang_mom_dm_x = np.mean(dm_spec_ang_mom_x[idRdm]) annular_spec_ang_mom_dm_y = np.mean(dm_spec_ang_mom_y[idRdm]) annular_spec_ang_mom_dm_z = np.mean(dm_spec_ang_mom_z[idRdm]) outside_ang_mom_dm_x = np.sum(dm_ang_mom_x[big_annulusDM]) outside_ang_mom_dm_y = np.sum(dm_ang_mom_y[big_annulusDM]) outside_ang_mom_dm_z = np.sum(dm_ang_mom_z[big_annulusDM]) outside_spec_ang_mom_dm_x = np.mean( dm_spec_ang_mom_x[big_annulusDM]) outside_spec_ang_mom_dm_y = np.mean( dm_spec_ang_mom_y[big_annulusDM]) outside_spec_ang_mom_dm_z = np.mean( dm_spec_ang_mom_z[big_annulusDM]) table1[idI,:] = [ds.current_redshift, rad, nref_mode, gas_flux, metal_flux, \ gas_flux_in, gas_flux_out, metal_flux_in, metal_flux_out, \ kinetic_energy_flux, thermal_energy_flux, entropy_flux,\ kinetic_energy_flux_in,kinetic_energy_flux_out,\ thermal_energy_flux_in,thermal_energy_flux_out,\ entropy_flux_in,entropy_flux_out,\ cold_gas_flux, cold_gas_flux_in, cold_gas_flux_out, \ cool_gas_flux, cool_gas_flux_in, cool_gas_flux_out, \ warm_gas_flux, warm_gas_flux_in, warm_gas_flux_out, \ hot_gas_flux, hot_gas_flux_in, hot_gas_flux_out, annular_ang_mom_gas_x, annular_ang_mom_gas_y,annular_ang_mom_gas_z, \ annular_spec_ang_mom_gas_x, annular_spec_ang_mom_gas_y,annular_spec_ang_mom_gas_z,\ annular_ang_mom_dm_x, annular_ang_mom_dm_y,annular_ang_mom_dm_z, \ annular_spec_ang_mom_dm_x, annular_spec_ang_mom_dm_y, annular_spec_ang_mom_dm_z, \ outside_ang_mom_gas_x, outside_ang_mom_gas_y, outside_ang_mom_gas_z, \ outside_spec_ang_mom_gas_x, outside_spec_ang_mom_gas_y, outside_spec_ang_mom_gas_z, \ outside_ang_mom_dm_x, outside_ang_mom_dm_y,outside_ang_mom_dm_z,\ outside_spec_ang_mom_dm_x, outside_spec_ang_mom_dm_y, outside_spec_ang_mom_dm_z, \ inside_ang_mom_stars_x, inside_ang_mom_stars_y, inside_ang_mom_stars_z, \ inside_spec_ang_mom_stars_x, inside_spec_ang_mom_stars_y, inside_spec_ang_mom_stars_z] ####### REPEATING THE CALCUATION WITHOUT THE DENSEST ISM-LIKE GAS ######### idR = np.where((radius >= minrad) & (radius < maxrad) & (hden < 0.1))[0] idCd = np.where((radius >= minrad) & (radius < maxrad) & (temperature <= 1e4) & (hden < 0.1))[0] idCl = np.where((radius >= minrad) & (radius < maxrad) & (temperature > 1e4) & (temperature <= 1e5) & (hden < 0.1))[0] idW = np.where((radius >= minrad) & (radius < maxrad) & (temperature > 1e5) & (temperature <= 1e6) & (hden < 0.1))[0] idH = np.where((radius >= minrad) & (radius < maxrad) & (temperature >= 1e6) & (hden < 0.1)) big_annulusGAS = np.where((radius >= rad_here) & (hden < 0.1))[0] # mass fluxes gas_flux = (np.sum(cell_mass[idR] * radial_velocity[idR]) / dr).to("Msun/yr") metal_flux = (np.sum(metal_mass[idR] * radial_velocity[idR]) / dr).to("Msun/yr") ## also filter based off radial velocity idVin = np.where(radial_velocity[idR] <= 0.)[0] idVout = np.where(radial_velocity[idR] > 0.)[0] gas_flux_in = (np.sum( cell_mass[idR][idVin] * radial_velocity[idR][idVin]) / dr).to("Msun/yr") gas_flux_out = (np.sum( cell_mass[idR][idVout] * radial_velocity[idR][idVout]) / dr).to("Msun/yr") metal_flux_in = (np.sum( metal_mass[idR][idVin] * radial_velocity[idR][idVin]) / dr).to("Msun/yr") metal_flux_out = (np.sum( metal_mass[idR][idVout] * radial_velocity[idR][idVout]) / dr).to("Msun/yr") ## and filter on temperature! and velocity! woo! idVin = np.where(radial_velocity[idH] <= 0.)[0] idVout = np.where(radial_velocity[idH] > 0.)[0] hot_gas_flux = (np.sum(cell_mass[idH] * radial_velocity[idH]) / dr).to("Msun/yr") hot_gas_flux_in = (np.sum( cell_mass[idH][idVin] * radial_velocity[idH][idVin]) / dr).to("Msun/yr") hot_gas_flux_out = (np.sum( cell_mass[idH][idVout] * radial_velocity[idH][idVout]) / dr).to("Msun/yr") idVin = np.where(radial_velocity[idW] <= 0.)[0] idVout = np.where(radial_velocity[idW] > 0.)[0] warm_gas_flux = ( np.sum(cell_mass[idW] * radial_velocity[idW]) / dr).to("Msun/yr") warm_gas_flux_in = (np.sum( cell_mass[idW][idVin] * radial_velocity[idW][idVin]) / dr).to("Msun/yr") warm_gas_flux_out = (np.sum( cell_mass[idW][idVout] * radial_velocity[idW][idVout]) / dr).to("Msun/yr") idVin = np.where(radial_velocity[idCl] <= 0.)[0] idVout = np.where(radial_velocity[idCl] > 0.)[0] cool_gas_flux = ( np.sum(cell_mass[idCl] * radial_velocity[idCl]) / dr).to("Msun/yr") cool_gas_flux_in = (np.sum( cell_mass[idCl][idVin] * radial_velocity[idCl][idVin]) / dr).to("Msun/yr") cool_gas_flux_out = (np.sum( cell_mass[idCl][idVout] * radial_velocity[idCl][idVout]) / dr).to("Msun/yr") idVin = np.where(radial_velocity[idCd] <= 0.)[0] idVout = np.where(radial_velocity[idCd] > 0.)[0] cold_gas_flux = ( np.sum(cell_mass[idCd] * radial_velocity[idCd]) / dr).to("Msun/yr") cold_gas_flux_in = (np.sum( cell_mass[idCd][idVin] * radial_velocity[idCd][idVin]) / dr).to("Msun/yr") cold_gas_flux_out = (np.sum( cell_mass[idCd][idVout] * radial_velocity[idCd][idVout]) / dr).to("Msun/yr") ## GAS angular momentum! annular_ang_mom_gas_x = np.sum(gas_ang_mom_x[idR]) annular_ang_mom_gas_y = np.sum(gas_ang_mom_y[idR]) annular_ang_mom_gas_z = np.sum(gas_ang_mom_z[idR]) annular_spec_ang_mom_gas_x = np.mean(gas_spec_ang_mom_x[idR]) annular_spec_ang_mom_gas_y = np.mean(gas_spec_ang_mom_y[idR]) annular_spec_ang_mom_gas_z = np.mean(gas_spec_ang_mom_z[idR]) outside_ang_mom_gas_x = np.sum(gas_ang_mom_x[big_annulusGAS]) outside_ang_mom_gas_y = np.sum(gas_ang_mom_y[big_annulusGAS]) outside_ang_mom_gas_z = np.sum(gas_ang_mom_z[big_annulusGAS]) outside_spec_ang_mom_gas_x = np.mean( gas_spec_ang_mom_x[big_annulusGAS]) outside_spec_ang_mom_gas_y = np.mean( gas_spec_ang_mom_y[big_annulusGAS]) outside_spec_ang_mom_gas_z = np.mean( gas_spec_ang_mom_z[big_annulusGAS]) table2[idI,:] = [zsnap, rad,gas_flux, metal_flux, \ gas_flux_in, gas_flux_out, metal_flux_in, metal_flux_out, \ cold_gas_flux, cold_gas_flux_in, cold_gas_flux_out, \ cool_gas_flux, cool_gas_flux_in, cool_gas_flux_out, \ warm_gas_flux, warm_gas_flux_in, warm_gas_flux_out, \ hot_gas_flux, hot_gas_flux_in, hot_gas_flux_out, annular_ang_mom_gas_x, annular_ang_mom_gas_y,annular_ang_mom_gas_z, \ annular_spec_ang_mom_gas_x, annular_spec_ang_mom_gas_y,annular_spec_ang_mom_gas_z,\ outside_ang_mom_gas_x, outside_ang_mom_gas_y, outside_ang_mom_gas_z, \ outside_spec_ang_mom_gas_x, outside_spec_ang_mom_gas_y, outside_spec_ang_mom_gas_z] print("#################") print("combining threads") print("#################") comm = communication_system.communicators[-1] for i in range(table1.shape[0]): table1[i, :] = comm.mpi_allreduce(table1[i, :], op="sum") for i in range(table2.shape[0]): table2[i, :] = comm.mpi_allreduce(table2[i, :], op="sum") sto.result = (table1, table2) sto.result_id = str(ds) if yt.is_root(): print("####################") print("In final table build") print("####################") for key in storage.keys(): table1, table2 = storage[key] for i in range(table1.shape[0] - 1): data.add_row(table1[i + 1, :]) data2.add_row(table2[i + 1, :]) data = set_table_units(data) data2 = set_table_units(data2) tablename = run_dir + '/' + args.run + '_angular_momenta_and_fluxes.hdf5' data.write(tablename, path='all_data', serialize_meta=True, overwrite=True) data2.write(tablename, path='noISM_data', serialize_meta=True, overwrite=False, append=True) return "whooooo angular momentum wheeeeeeee"
""" import os import sys import yt yt.enable_parallelism() from yt.extensions.p2p import \ add_p2p_particle_filters if __name__ == "__main__": es = yt.load(sys.argv[1]) fns = es.data['filename'].astype(str)[::-1] data_dir = es.directory for fn in yt.parallel_objects(fns, njobs=-1, dynamic=True): ds = yt.load(os.path.join(data_dir, fn)) output_file = os.path.join("pop3", f"{ds.basename}.h5") if os.path.exists(output_file): continue add_p2p_particle_filters(ds) region = ds.box(ds.parameters["RefineRegionLeftEdge"], ds.parameters["RefineRegionRightEdge"]) fields = [ "particle_mass", "particle_index", "particle_type", "particle_position_x", "particle_position_y", "particle_position_z", "particle_velocity_x", "particle_velocity_y", "particle_velocity_z", "creation_time",
ds.current_time + ds.domain_width[0]/ds.quan(2.998*10**8, 'm/s')) lambda_shift = 1215.67 lambda_min = lambda_shift * (z - dz + 1) lambda_max = lambda_shift * (z + 1) locals = range(rank, nSamples, size) LR = trident.LightRay(ds) LocalID = range(rank, nSamples, size) count = 0 keepFrac = float(sys.argv[5]) / 1024**2 keepPerRank = float(sys.argv[5]) // size # # Make light-rays first and save to disk # for i in yt.parallel_objects(range(nSamples)): if i % 5000 == 0: print('%d/%d' % (i, nSamples)) if np.random.uniform(0, 1) > keepFrac: continue # want to cut down and NOT make 1e6 spectra... sFile ='%s/rayData/%s/RD%04d/spectra/rank%d_%d.h5'\ %(dataRepo, simname, d, rank, count) rFile = '%s/rayData/%s/RD%04d/rays/rank%d_%d.h5'\ %(dataRepo,simname, d, rank, count) if os.path.exists(sFile): count += 1 continue x = (i // dim + 0.5) * dx y = (i % dim + 0.5) * dx start = [x, y, 0] end = [x, y, 1] ray = LR.make_light_ray(start_position = start,\
import numpy as np import glob import pickle f = open('../Bhspos_r8.pickle') bhspos = pickle.load(f) f.close() DD = glob.glob('../DD*/*.hierarchy') DD.sort() DD = DD[0:102] allhaloidx = np.loadtxt('../allhalos1.txt') my_storage = {} I = range(len(DD)) for sto, i in yt.parallel_objects(I, 16, storage = my_storage): pf = load(DD[i]) sto.result_id = DD[i] sto.result = [] sp = pf.sphere(allhaloidx[i,17:20]/28.4, (8*allhaloidx[i,11], 'kpccm/h')) proj = ProjectionPlot(pf,'y','Density', center = allhaloidx[i,17:20]/28.4, width = 20*allhaloidx[i,11]/28400, data_source = sp) for j in range(bhspos[i].shape[0]): proj.annotate_marker([bhspos[i][j][0],bhspos[i][j][1],bhspos[i][j][2]],marker='o') proj.save()
avgdmdens = np.zeros((numsims, nbins)) #the global containers dm dens avgstarsdens = np.zeros((numsims, nbins)) #the global containers dm dens avgodens = np.zeros((numsims, nbins)) #the global containers for overdesnity avgdmodens = np.zeros((numsims, nbins)) #the global containers for overdesnity #Standard deviation stddens = np.zeros((numsims, nbins)) stdtemp = np.zeros((numsims, nbins)) stddmdens = np.zeros((numsims, nbins)) stdstarsdens = np.zeros((numsims, nbins)) stdodens = np.zeros((numsims, nbins)) stddmodens = np.zeros((numsims, nbins)) ############################################################### #Run over the simulations in parallel for sto, path in yt.parallel_objects(datapath, num_procs, container): print datetime.now().strftime('%H:%M:%S'), 'Loading Sim #', sto.result_id #Load data ds = yt.load(path) #pick the halos to be used in this simulation according to the mass limits simid = simtype[sto.result_id][3] #sto.result_id works like enumerate halos = ds.arr( np.loadtxt('Halos/%s.txt' % simid, skiprows=1, unpack=False, usecols=(1, 2, 3)), 'Mpc/h') virial = ds.arr( np.loadtxt('Halos/%s.txt' % simid, skiprows=1,
filepath = os.getcwd() + "/dumps" moment_files = np.sort(glob.glob(filepath + '/moment*.bin')) lagrange_multiplier_files = \ np.sort(glob.glob(filepath+'/lagrange_multipliers*.bin')) print("moment files : ", moment_files.size) print("lagrange multiplier files : ", lagrange_multiplier_files.size) dt = params.dt dump_interval = params.dump_steps time_array = np.loadtxt("dump_time_array.txt") io = PetscBinaryIO.PetscBinaryIO() for file_number, dump_file in yt.parallel_objects(enumerate(moment_files)): print("file number = ", file_number, "of ", moment_files.size) moments = io.readBinaryFile(dump_file) moments = moments[0].reshape(N_q2, N_q1, 3) density = moments[:, :, 0] j_x = moments[:, :, 1] j_y = moments[:, :, 2] lagrange_multipliers = \ io.readBinaryFile(lagrange_multiplier_files[file_number]) lagrange_multipliers = lagrange_multipliers[0].reshape(N_q2, N_q1, 7) #h5f = h5py.File(lagrange_multiplier_files[file_number], 'r') #lagrange_multipliers = h5f['lagrange_multipliers'][:]
np.sort(glob.glob(filepath+'/lagrange_multipliers*.h5')) distribution_function_files = np.sort(glob.glob(filepath+'/f_*.h5')) dt = params.dt dump_interval = params.dump_steps time_array = np.loadtxt("dump_time_array.txt") h5f = h5py.File(distribution_function_files[0], 'r') dist_func_background = h5f['distribution_function'][:] h5f.close() q1_position = 10 q2_position = 100 for file_number, dump_file in yt.parallel_objects(enumerate(distribution_function_files)): print("file number = ", file_number, "of ", distribution_function_files.size) h5f = h5py.File(distribution_function_files[file_number], 'r') dist_func = h5f['distribution_function'][:] h5f.close() f_at_desired_q = np.reshape((dist_func - dist_func_background)[q2_position, q1_position, :], [N_p1, N_p2] ) pl.contourf(p1_meshgrid, p2_meshgrid, f_at_desired_q, 100, cmap='bwr') pl.title(r'Time = ' + "%.2f"%(time_array[file_number]) + " ps") pl.xlabel('$p_x$') pl.ylabel('$p_y$')
# Loading datasets fn_list = open(sys.argv[1], 'r') fns = fn_list.readlines() # FIRE with AMIGA centroids if AMIGA: # Read in halo information from amiga output if present amiga_data = get_amiga_data(sys.argv[2]) # Smooth the data to remove jumps in centroid amiga_data = smooth_amiga(amiga_data) # Tempest with Rockstar centroids elif Rockstar: rockstar_data = get_rockstar_data(sys.argv[2], int(sys.argv[3])) for fn in yt.parallel_objects(fns): fn = fn.strip() # Get rid of trailing \n fn_head = fn.split('/')[-1] cdens_fn = "%s_cdens.h5" % fn_head # Define ions we care about ions = [] ion_fields = [] full_ion_fields = [] ions.append('H I') ion_fields.append('H_number_density') full_ion_fields.append(('gas', 'H_number_density')) #ions.append('Mg II') #ion_fields.append('Mg_p1_number_density') #full_ion_fields.append(('gas', 'Mg_p1_number_density')) #ions.append('Si II')
h5f.close() density.append(moments[:, :, 0]) edge_density.append(density[file_number][0, sensor_1_left_indices]) density = np.array(density) edge_density = np.array(edge_density) mean_density = np.mean(density) max_density = np.max(density) min_density = np.min(density) np.savetxt("edge_density.txt", edge_density) print("Dumping data...") for file_number in yt.parallel_objects(range(density.shape[0])): print("File number = ", file_number, ' of ', moment_files.size) pl.semilogy( q2[sensor_1_left_indices], density[file_number][0, sensor_1_left_indices], ) #pl.title(r'Time = ' + "%.2f"%(time_array[file_number]) + " ps") pl.title(r'Time = ' + "%.2f" % (file_number * dt * dump_interval) + " ps") pl.xlim([sensor_1_left_start, sensor_1_left_end]) #pl.ylim([min_density-mean_density, max_density-mean_density]) #pl.ylim([0., np.log(max_density)]) #pl.gca().set_aspect('equal')
pfields = [("black_hole", field) for field in [ "particle_mass", "particle_type", "particle_index", "creation_time", "metallicity_fraction", "particle_velocity_x", "particle_velocity_y", "particle_velocity_z" ]] cfield = "BH_to_Eddington" a = ytree.load(sys.argv[2]) cfield_min = float(sys.argv[3]) data_dir = os.path.join( a.filename.split("/")[-3], "clumps_%.0e" % cfield_min) last_snap_idx = int(a[0]["Snap_idx"]) idx_offset = fns.size - last_snap_idx - 1 for i, fn in yt.parallel_objects(enumerate(fns)): if not os.path.exists(fn): continue ds = yt.load(fn) c_min = ds.quan(cfield_min, "1/Msun") output_dir = os.path.join(data_dir, ds.directory) ensure_dir(output_dir) setup_ds = True search_idx = i - idx_offset sel = 'tree["tree", "Snap_idx"] == %d' % search_idx my_halos = a.select_halos(sel, fields=["Snap_idx"]) yt.mylog.info("Finding clumps for %d halos." % my_halos.size)
def gen_profs(ds,fils,dsname,keep=True,totalratio=None,mask=False): x = np.zeros(10) storage = {} #Check to see if totalratio has been set, else determine it. if totalratio == None: #Determine simulation-wide ratio of baryon to Total matter, done in parallel to speed up for stor,i in yt.parallel_objects(range(1),storage = storage): denstotal = ds.all_data().quantities.weighted_average_quantity("density","cell_volume") dmtotal = ds.all_data().quantities.weighted_average_quantity("dark_matter_density","cell_volume") totalratio = denstotal / (dmtotal + denstotal) del dmtotal stor.result = totalratio,denstotal.in_units('g/cm**3') stor.result_id = 1 totalratio,denstotal = storage[1] print totalratio # print to stdout, allows us to check ratio makes sense/debug del storage #Print to a file, allows checking of progress when in job queue if yt.is_root(): with open("Testinfo.dat",'w') as f: f.write('Ratio determined') print("All Ratio determined") x = [] storage = {} #Gather list of all profiles on disk filelist = sorted(os.listdir(''.join(['/shome/mackie/data',dsname,'/profiles']))) #Include only density profiles filelist = filelist[:len(filelist)/2] #Load a numpy mask array if required if keep == True: keeplist = np.load(''.join(['/shome/mackie/data',dsname,'/filkeep.npy'])) for stor,file_in_dir in yt.parallel_objects(filelist,storage=storage): filnum = int(file_in_dir[7:10]) segnum = int(file_in_dir[13:16]) prof = yt.load(''.join(['/shome/mackie/data',dsname,'/profiles/',file_in_dir])) dm = prof.data['dark_matter_density'] dens = prof.data['density'].in_units('g/cm**3') result = (dens /( dm + dens)) - totalratio dens = dens/denstotal result = result.v if keep == True: if keeplist[filnum][segnum] == True: stor.result = result,dens stor.result_id = file_in_dir else: stor.result = result,dens stor.result_id = file_in_dir del prof,dm,dens,filnum,segnum results = [] denresult = [] for keys,values in storage.items(): results.append(values[0]) denresult.append(values[1]) results = np.array(results) denresult = np.array(denresult) if mask: results = np.ma.masked_array(results,mask=~mask,fill_value=np.nan) #Get x bins x = yt.load(''.join(['/shome/mackie/data',dsname,'/profiles/densfil000seg000.h5'])).data['x'] if yt.is_root(): with open("Testinfo.dat",'a') as f: f.write("returning results to plot") print results return results,denresult,x