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 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 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 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 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 __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 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 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
def doit(filename):
ds = yt.load(filename)
ad = ds.all_data()
profile = yt.create_profile(ad, args.binfields, args.fields,
n_bins=tuple(args.nbins),
weight_field=None)
print(profile.keys())
print(profile["enucdot"].shape)
print(profile["y"].shape)
for y, e in zip(profile["y"], profile["enucdot"]):
print('({}, {})'.format(y, e))
exit()
return ds.time, emax, elocmax, efwhm
def main():
args = parse_args()
filename = args.filename
print filename
print "Loading yt dataset " + filename.split("/")[-1] + "..."
ds = yt.load(filename)
#establish a moab instance for use
mb = core.Core()
yt2moab_uniform_gridgen(mb,ds)
#write file
mb.write_file("test.h5m")
def test_radmc3d_exporter_continuum():
"""
This test is simply following the description in the docs for how to
generate the necessary output files to run a continuum emission map from
dust for one of our sample datasets.
"""
ds = yt.load(ISO_GAL)
# Make up a dust density field where dust density is 1% of gas density
dust_to_gas = 0.01
def _DustDensity(field, data):
return dust_to_gas * data["density"]
ds.add_field(("gas", "dust_density"), function=_DustDensity, units="g/cm**3")
yield RadMC3DValuesTest(ds, ("gas", "dust_density"))
def plot_stuff(i):
print "Plotting timestep ",i
ds = yt.load("orszag_mhd_2d_hdf5_chk_{}".format(str(i).zfill(4)))
# os.chdir(SAVE_PATH)
xdir=0
ydir=1
slicedirection=2
sl = ds.slice(slicedirection,0) ##Get the Slice
w = [ds.domain_left_edge[xdir],ds.domain_right_edge[xdir],ds.domain_left_edge[ydir],ds.domain_right_edge[ydir]]
frb1 = FixedResolutionBuffer(sl,w,(128,128)) #Create FixedResolution Buffer
plt.figure()
plt.imshow(np.array(frb1["density"]))
plt.colorbar()
plt.savefig(SAVE_PATH+"{0}_density.png".format(ds))
plt.figure()
plt.imshow(np.array(frb1["magnetic_field_strength"]))
plt.colorbar()
plt.savefig(SAVE_PATH+"{0}_Bstrength.png".format(ds))
def print_ray(fname, startpoint, endpoint):
ds = yt.load(fname)
line = ds.ray(startpoint, endpoint)
labels = ["x", "y", "z"]
for f in ds.field_list:
labels.append(f[1])
print >> sys.stderr, "Columns: ", labels
print "#",
for l in labels:
print "{:<20}".format(l),
print
dlen = len(line[labels[0]].v)
for i in range(dlen):
for l in labels:
print "{:<20}".format(line[l].v[i]),
print
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 plot_dens(i,plane="z", velocity=False,grid=False,zmin ="",zmax="",magnetic=False, particle=False):
ds = yt.load("mhd_sphere_hdf5_chk_{}".format(str(i).zfill(4)))
physical_quantity="density"
slc = yt.SlicePlot(ds, plane,physical_quantity)#,center=(0.5,0.5,0.5))
slc.set_figure_size(5)
if grid: slc.annotate_grids()
if velocity: slc.annotate_velocity(normalize=True)
if magnetic: slc.annotate_magnetic_field()
slc.set_cmap("all","rainbow")
if zmin!="" and zmax!="": slc.set_zlim(physical_quantity,zmin,zmax)
if particle :
os.system("cp ../source/Simulation/SimulationMain/unitTest/SinkMomTest/utils/clean_sinks_evol.py .")
os.system("python clean_sinks_evol.py")
data =np.loadtxt("sinks_evol.dat_cleaned",skiprows=1)
pcl_indx_at_t = np.where(np.isclose(int(ds.current_time.in_cgs()),data[:,1]))[0]
print "Number of sink particles: " , len(pcl_indx_at_t)
pcl_pos_at_t = data[pcl_indx_at_t,2:5]
for pos in pcl_pos_at_t:
slc.annotate_marker(pos, coord_system='data',marker='.',plot_args={'color':'black','s':3})
slc.save(SAVE_PATH+"{0}_{1}.png".format(ds,physical_quantity))
def load_halo(pdf, hdf, nth_most_massive):
halo, center, left, right, radius = get_halo_bounds(nth_most_massive, pdf, hdf)
particles_midx = pdf + ".midx10"
halos_midx = hdf + ".midx7"
bbox = np.array([left, right]).T
pds = yt.load(pdf,
midx_filename=particles_midx,
bounding_box = bbox,
n_ref=64, over_refine_factor=2,
)
# the c code chokes on data types without these conversions
pds.domain_left_edge = pds.domain_left_edge.astype(np.float64)
pds.domain_right_edge = pds.domain_right_edge.astype(np.float64)
pds.domain_width = pds.domain_width.astype(np.float64)
pds.domain_center = pds.domain_center.astype(np.float64)
return pds, halo, center, radius
def test_add_particle_filter():
"""Test particle filters created via add_particle_filter
This accesses a deposition field using the particle filter, which was a
problem in previous versions on this dataset because there are chunks with
no stars in them.
"""
def stars(pfilter, data):
filter_field = (pfilter.filtered_type, "creation_time")
return (data.ds.current_time - data[filter_field]) > 0
add_particle_filter("stars", function=stars, filtered_type='all',
requires=["creation_time"])
ds = yt.load(iso_galaxy)
ds.add_particle_filter('stars')
ad = ds.all_data()
ad['deposit', 'stars_cic']
assert True
def setup(self):
if yt.__version__.startswith('3'):
self.ds = yt.load(self.dsname)
self.ad = self.ds.all_data()
self.field_name = "density"
else:
self.ds = load(self.dsname)
self.ad = self.ds.h.all_data()
self.field_name = "Density"
# Warmup hdd
self.ad[self.field_name]
if yt.__version__.startswith('3'):
mi, ma = self.ad.quantities['Extrema'](self.field_name)
self.tf = yt.ColorTransferFunction((np.log10(mi)+1, np.log10(ma)))
else:
mi, ma = self.ad.quantities['Extrema'](self.field_name)[0]
self.tf = ColorTransferFunction((np.log10(mi)+1, np.log10(ma)))
self.tf.add_layers(5, w=0.02, colormap="spectral")
self.c = [0.5, 0.5, 0.5]
self.L = [0.5, 0.2, 0.7]
self.W = 1.0
self.Npixels = 512
def test_from_yt(tmpdir):
from yt import load
ds = load(os.path.join(DATA, 'DD0010', 'moving7_0010'))
def _dust_density(field, data):
return data["density"].in_units('g/cm**3') * 0.01
ds.add_field(("gas", "dust_density"), function=_dust_density, units='g/cm**3')
amr = AMRGrid.from_yt(ds, quantity_mapping={'density': ('gas', 'dust_density')})
m = Model()
m.set_amr_grid(amr)
m.add_density_grid(amr['density'], get_realistic_test_dust())
s = m.add_point_source()
s.luminosity = 1000
s.temperature = 1000
m.set_n_initial_iterations(3)
m.set_n_photons(initial=1e5, imaging=0)
m.set_propagation_check_frequency(1)
m.set_copy_input(False)
input_file = tmpdir.join('test.rtin').strpath
output_file = tmpdir.join('test.rtout').strpath
m.write(input_file)
m.run(output_file)
def halofind(path="./",output="ascii/",rockfile='rockstar.cfg',rockout="rockstar/"):
"""
Perfom a number of initialization steps to run rockstar halo finder on
a ramses simulation. Specifically, conver the dark matter particles from
binary files to ascii lists, and generate a script that can be used to
run rockstar on each snapshot
Force resolution is computed assuming unigird.
Keep track of units throughout.
"""
import yt
import mypython as mp
import os
#create rock output folder
if not os.path.exists(rockout):
os.mkdir(rockout)
#find how many snapshots are there and process them
for out in os.listdir(path):
#check if valid output
if "output_" in out:
#grap snap number
snap=int(out.split("output_")[1])
print "Processing snap ", snap
#process in ascii
mp.ramses.ramses2ascii.ramses2ascii(snap,path=path,output=output)
#create rockstar output
rockfilename='output_'+format(snap, '05')+rockfile
if os.path.isfile(rockfilename):
print "rockstar config file exist..."
else:
#get sim properties
snapname=path+'output_'+format(snap, '05')+'/info_'+format(snap, '05')+'.txt'
sim = yt.load(snapname)
ad=sim.all_data()
#write config file
fle = open(rockout+rockfilename, 'w')
fle.write('#Configuration file for rockstar\n')
fle.write('FILE_FORMAT = "ASCII"\n')
#mass particle Msun/h
mpart=sim.mass_unit.v/2e33*ad['io', 'particle_mass'][0].v*sim.hubble_constant
fle.write('PARTICLE_MASS = '+str(mpart)+'\n')
#cosmology
aexp=1./(1.+sim.current_redshift)
fle.write('SCALE_NOW = '+str(aexp)+'\n')
fle.write('h0 = '+str(sim.hubble_constant)+'\n')
fle.write('Ol = '+str(sim.omega_lambda)+'\n')
fle.write('Om = '+str(sim.omega_matter)+'\n')
#box size
boxsize=(sim.domain_width.in_units('Mpccm/h'))[0].v
fle.write('BOX_SIZE = '+str(boxsize)+'\n')
fle.write('FORCE_RES = '+str(boxsize/sim.domain_dimensions[0]/2.)+'\n')
#utility
fle.write('OUTBASE = '+path+'output_'+format(snap, '05')+'/\n')
fle.write('MIN_HALO_OUTPUT_SIZE = 100\n')
fle.close()
exit()
def doit(plotfile, fname):
ds = yt.load(plotfile)
cm = "gist_rainbow"
if fname == "vz":
field = ('gas', 'velocity_z')
use_log = False
vals = [-1.e7, -5.e6, 5.e6, 1.e7]
sigma = 5.e5
cm = "coolwarm"
elif fname == "magvel":
field = ('gas', 'velocity_magnitude')
use_log = False
vals = [1.e6, 2.e6, 4.e6, 8.e6, 1.6e7]
sigma = 2.e5
elif fname == "radvel":
field = ('boxlib', 'radial_velocity')
use_log = False
vals = [-1.e7, -5.e6, -2.5e6, 2.5e6, 5.e6, 1.e7]
sigma = 2.e5
cm = "coolwarm"
dd = ds.all_data()
mi = min(vals)
ma = max(vals)
if use_log:
mi, ma = np.log10(mi), np.log10(ma)
# Instantiate the ColorTransferfunction.
tf = vr.ColorTransferFunction((mi, ma))
# Set up the camera parameters: center, looking direction, width, resolution
#c = np.array([0.0, 0.0, 0.0])
c = (ds.domain_right_edge + ds.domain_left_edge) / 2.0
L = np.array([1.0, 1.0, 1.0])
L = np.array([1.0, 1.0, 1.2])
W = 2.0 * ds.domain_width
N = 720
north = [0.0, 0.0, 1.0]
for v in vals:
if (use_log):
tf.sample_colormap(math.log10(v), sigma**2,
colormap=cm) #, alpha=0.2)
else:
tf.sample_colormap(v, sigma**2, colormap=cm) #, alpha=0.2)
# alternate attempt
ds.periodicity = (True, True, True)
# Create a camera object
cam = vr.Camera(c,
L,
W,
N,
transfer_function=tf,
ds=ds,
no_ghost=False,
north_vector=north,
fields=[field],
log_fields=[use_log])
# make an image
im = cam.snapshot()
# add an axes triad -- note if we do this, we HAVE to do draw
# domain, otherwise the image is blank (likely a bug)
cam.draw_coordinate_vectors(im)
# add the domain box to the image:
nim = cam.draw_domain(im)
# increase the contrast -- for some reason, the enhance default
# to save_annotated doesn't do the trick (likely a bug)
max_val = im[:, :, :3].std() * 4.0
nim[:, :, :3] /= max_val
f = pylab.figure()
pylab.text(0.2,
0.15,
"{:.3g} s".format(float(ds.current_time.d)),
transform=f.transFigure,
color="white")
cam._render_figure = f
# save annotated -- this added the transfer function values,
# but this messes up our image size defined above
cam.save_annotated("{}_{}.png".format(os.path.normpath(plotfile), fname),
nim,
dpi=145,
clear_fig=False)
print(str(sys.argv[t])),
print('')
print('-------------------------------------------------------------------\n')
idx_start = args.idx_start
idx_end = args.idx_end
didx = args.didx
prefix = args.prefix
colormap = 'arbre'
field = 'Dens'
center_mode = 'c'
dpi = 150
yt.enable_parallelism()
ts = yt.load([
prefix + '/Data_%06d' % idx for idx in range(idx_start, idx_end + 1, didx)
])
for ds in ts.piter():
sz = yt.SlicePlot(ds, 'z', field, center=center_mode)
# sz.set_zlim( field, 0.0, 9.0 )
sz.set_log(field, False)
sz.set_cmap(field, colormap)
sz.annotate_timestamp(corner='upper_right',
time_format='t = {time:.4f} {units}')
# sz.annotate_grids( periodic=False )
sz.save(mpl_kwargs={"dpi": dpi})
mixFrac = None
p_mixFrac = np.zeros(local_nfiles)
p_time = np.zeros(local_nfiles)
comm.Barrier()
## Display the correct message if multiple files are read/rank or just one/rank
if local_start != local_end - 1:
print('Processor ', Pid, 'will read from index file', local_start, 'to',
local_end - 1)
else:
print("Processor ", Pid, "will only read index file", local_start)
files = comm.bcast(files, root=master)
for i in range(local_start, local_end):
print("Reading file ", files[i], ' with processor ', Pid)
mylog.setLevel(40)
ds = yt.load(path + files[i])
# Create a 4cm-radius sphere
center = ds.domain_left_edge
sp = ds.sphere(center, r_max)
profile = yt.create_profile(sp,
'radius', ['fld1', 'mix'],
n_bins=p_res,
units={
'radius': 'cm',
"fld1": "",
"mix": ""
},
logs={
'radius': False,
"fld1": False,
"mix": False
# L_MPC = int(Lbox/1000)
# data_dir = '/raid/bruno/data/'
# data_dir = '/data/groups/comp-astro/bruno/'
# data_dir = '/home/bruno/Desktop/ssd_0/data/'
data_dir = '/gpfs/alpine/csc434/scratch/bvilasen/'
input_dir = data_dir + f'cosmo_sims/ics/enzo/{n_points}_{L_MPC}Mpc/'
output_dir = data_dir + f'cosmo_sims/ics/enzo/{n_points}_{L_MPC}Mpc/ics_{n_boxes}_z100/'
print(f'Input Dir: {input_dir}')
print(f'Output Dir: {output_dir}')
create_directory(output_dir)
# Load Enzo Dataset
nSnap = 0
snapKey = '{0:03}'.format(nSnap)
inFileName = 'DD0{0}/data0{0}'.format(snapKey)
ds = yt.load(input_dir + inFileName)
data = ds.all_data()
h = ds.hubble_constant
current_z = np.float(ds.current_redshift)
current_a = 1. / (current_z + 1)
data_enzo = {'dm': {}, 'gas': {}}
data_enzo['current_a'] = current_a
data_enzo['current_z'] = current_z
# Domain decomposition
box_size = [Lbox, Lbox, Lbox]
grid_size = [n_points, n_points, n_points]
if n_boxes == 1: proc_grid = [1, 1, 1]
if n_boxes == 2: proc_grid = [2, 1, 1]
if n_boxes == 8: proc_grid = [2, 2, 2]
import yt
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import AxesGrid
fn = "IsolatedGalaxy/galaxy0030/galaxy0030"
ds = yt.load(fn) # load data
fig = plt.figure()
# See http://matplotlib.org/mpl_toolkits/axes_grid/api/axes_grid_api.html
# These choices of keyword arguments produce two colorbars, both drawn on the
# right hand side. This means there are only two colorbar axes, one for Density
# and another for temperature. In addition, axes labels will be drawn for all
# plots.
grid = AxesGrid(
fig,
(0.075, 0.075, 0.85, 0.85),
nrows_ncols=(2, 2),
axes_pad=1.0,
label_mode="all",
share_all=True,
cbar_location="right",
cbar_mode="edge",
cbar_size="5%",
cbar_pad="0%",
)
cuts = ["x", "y", "z", "z"]
fields = ["density", "density", "density", "temperature"]
for i, (direction, field) in enumerate(zip(cuts, fields)):
plt.clf()
plt.close('all')
if 1:
prefix = 'ca03' #something to note the simulations you used.
directory = '/scratch1/dcollins/Paper49_EBQU/ca03_turb_weak'
frame = 339
if 1:
prefix = 'ca02' #something to note the simulations you used.
directory = '/scratch1/dcollins/Paper49_EBQU/ca02_turb'
frame = 100
plot_directory = "%s/PlotsToTransfer" % os.environ['HOME']
#This conditional tests to see if region has been defined.
#if not, define it. This process is slow so don't repeat if you don't have to.
ds = yt.load('%s/DD%04d/data%04d' % (directory, frame, frame))
region = ds.all_data()
#set up bins.
Nbins = 64
density_min = region['density'].min()
density_max = region['density'].max()
density_bins = np.logspace(np.log10(density_min), np.log10(density_max), Nbins)
velocity_magnitude_min = region['velocity_magnitude'].min()
velocity_magnitude_max = region['velocity_magnitude'].max()
velocity_magnitude_bins = np.logspace(np.log10(velocity_magnitude_min),
np.log10(velocity_magnitude_max),
Nbins / 4)
bins = {'density': density_bins, 'velocity_magnitude': velocity_magnitude_bins}
#
import yt
# Load the dataset.
ds = yt.load("Enzo_64/DD0029/data0029")
# Create a 1 Mpc radius sphere, centered on the max density.
sp = ds.sphere("max", (1.0, "Mpc"))
# Use the total_quantity derived quantity to sum up the
# values of the mass and particle_mass fields
# within the sphere.
baryon_mass, particle_mass = sp.quantities.total_quantity([("gas", "mass"),
("all",
"particle_mass")])
print(
"Total mass in sphere is %0.3e Msun (gas = %0.3e Msun, particles = %0.3e Msun)"
% (
(baryon_mass + particle_mass).in_units("Msun"),
baryon_mass.in_units("Msun"),
particle_mass.in_units("Msun"),
))
from scipy.optimize import curve_fit
from scipy.signal import argrelextrema
from yt.units import cm
from sys import exit as sysexit
start_time = time.time()
# load dataset
data_root_path = "/rds/user/dc-bamb1/rds-dirac-dp131/dc-bamb1/GRChombo_data/KerrSF"
#data_sub_dir = "run0031_KNL_l0_m0_a0_Al0_mu0.4_M1_correct_Ylm"
#data_sub_dir = "run0063_KNL_l0_m0_a0_Al0_mu1_M1_new_rho_more_levels"
data_sub_dir = "run0022_KNL_l0_m0_a0_Al0_mu1_M1_correct_Ylm"
number = 1460
dataset_path = data_root_path + "/" + data_sub_dir + "/KerrSFp_{:06d}.3d.hdf5".format(
number)
ds = yt.load(dataset_path)
print("loaded data from ", dataset_path)
print("time = ", time.time() - start_time)
# set centre
center = [512.0, 512.0, 0]
# set up parameters
dt = 0.25
t = number * dt
a = 0
M = 1
mu = 1
r_plus = M * (1 + np.sqrt(1 - a**2))
z_position = 0.001 # s position of slice
R_plus = 0.25 * r_plus # R_outer = r_+ / 4
import yt
# Load the dataset.
ds = yt.load("Enzo_64/DD0043/data0043")
# Make a density projection centered on the 'm'aximum density location
# with a width of 10 Mpc..
p = yt.ProjectionPlot(ds, "y", "density", center="m", width=(10, "Mpc"))
# Modify the projection
# The argument specifies the region along the line of sight
# for which particles will be gathered.
# 1.0 signifies the entire domain in the line of sight
# p.annotate_particles(1.0)
# but in this case we only go 10 Mpc in depth
p.annotate_particles((10, "Mpc"))
# Save the image.
# Optionally, give a string as an argument
# to name files with a keyword.
p.save()
lax.xaxis.set_ticklabels([r""])
for tick in lax.xaxis.get_ticklines():
tick.set_visible(False)
for ticklabel in lax.xaxis.get_majorticklabels():
ticklabel.set_visible(True)
ticklabel.set_horizontalalignment("left")
ticklabel.set_verticalalignment("top")
ticklabel.set_fontsize(10)
for ticklabel in lax.yaxis.get_majorticklabels():
ticklabel.set_fontsize(lfontsize)
if __name__ == "__main__":
halo_dir = "halo_2170858"
# halo_dir = "halo_2171203"
ds = yt.load(os.path.join(halo_dir, "bh_clump_distance_edge_1e-03.h5"))
ds_i = yt.load(
os.path.join(halo_dir, "bh_clump_distance_edge_1e-03_inner_0.25.h5"))
fontsize = 12
my_fig = GridFigure(1,
1,
top_buffer=0.12,
bottom_buffer=0.13,
left_buffer=0.14,
right_buffer=0.04,
horizontal_buffer=0.05,
vertical_buffer=0.1,
figsize=(6, 4.5))
my_axes = my_fig[0]
my_axes.set_yscale("log")
yt.enable_parallelism() # ============================================================================== # Spherical Horizon finder # # Assumes (Obviously) spherical symmetry !! Gives bad results if not ! # Input = Black hole center # # For valid calculation of Mass one has to slice along x-Axis !! # Output :-Black Hole mass (with Error estimate) # -Black Hole radius (with Error estimate) # ============================================================================== #Loading dataset loc = '/scratch2/kclough/BosonStars/ScalarField_001*' ds = yt.load(loc) #initial_location centerXYZ = ds[0].domain_right_edge/2 center = 512.0 - 30.0 #float(centerXYZ[0]) BHcenter = [center, float(centerXYZ[1]), float(centerXYZ[2])] max_radius = 10.0 # ================================= # INPUT # # ================================= # ================================ # ================================ time_data = [] CycleData = []
Voxelizing elements
We load our grid of positions and velocities from a file
rather than generating them ourselves.
We also use the 'periodic' options in PyPSI.elementBlocksFromGrid() and PyPSI.voxelize(),
along with 'box', to indicate that PSI should create appropriate ghost elements and wrap
the domain.
'''
i = 0
subgridsize = 64
while os.path.isdir("DD%04d" % i):
## using yt to load enzo data
ds = yt.load("DD%04d/data%04d" % (i, i))
## use a covering grid to include all data
all_data_level_0 = ds.covering_grid(level = 0, \
left_edge = [0, 0, 0], dims = ds.domain_dimensions)
ids = np.array(all_data_level_0['all', 'particle_index']).astype(int, \
order='C')
#print ids
pos = np.zeros([int(np.prod(ds.domain_dimensions)), 3])
pos[:,0] = all_data_level_0['all', 'particle_position_x']
pos[:,1] = all_data_level_0['all', 'particle_position_y']
pos[:,2] = all_data_level_0['all', 'particle_position_z']
pos = pos.reshape(np.append(ds.domain_dimensions, 3)).astype(np.float64, \
order='C')
pos = pos[:64,:64,:64,:]
vel = np.zeros([int(np.prod(ds.domain_dimensions)), 3])
def doit(plotfile):
ds = yt.load(plotfile)
ds.periodicity = (True, True, True)
field = ('gas', 'velocity_z')
ds._get_field_info(field).take_log = False
sc = Scene()
# add a volume: select a sphere
#center = (0, 0, 0)
#R = (5.e8, 'cm')
#dd = ds.sphere(center, R)
vol = VolumeSource(ds, field=field)
vol.use_ghost_zones = True
sc.add_source(vol)
# transfer function
vals = [-1.e7, -5.e6, 5.e6, 1.e7]
sigma = 5.e5
tf = yt.ColorTransferFunction((min(vals), max(vals)))
tf.clear()
cm = "coolwarm"
for v in vals:
tf.sample_colormap(v, sigma**2, colormap=cm) #, alpha=0.2)
sc.get_source(0).transfer_function = tf
cam = sc.add_camera(ds, lens_type="perspective")
cam.resolution = (1920, 1080)
center = 0.5*(ds.domain_left_edge + ds.domain_right_edge)
cam.position = [2.5*ds.domain_right_edge[0],
2.5*ds.domain_right_edge[1],
center[2]+0.25*ds.domain_right_edge[2]]
# look toward the center -- we are dealing with an octant
normal = (center - cam.position)
normal /= np.sqrt(normal.dot(normal))
cam.switch_orientation(normal_vector=normal,
north_vector=[0., 0., 1.])
cam.set_width(ds.domain_width)
sc.camera = cam
#sc.annotate_axes(alpha=0.05)
#sc.annotate_domain(ds, color=np.array([0.05, 0.05, 0.05, 0.05]))
#sc.annotate_grids(ds, alpha=0.05)
sc.render()
sc.save("{}_radvel".format(plotfile), sigma_clip=4.0)
sc.save_annotated("{}_radvel_annotated.png".format(plotfile),
sigma_clip=4.0,
text_annotate=[[(0.05, 0.05),
"t = {}".format(ds.current_time.d),
dict(horizontalalignment="left")],
[(0.5,0.95),
"Maestro simulation of convection in a mixed H/He XRB",
dict(color="y", fontsize="24",
horizontalalignment="center")]])
'''
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'))
image = proj_array * units['density_unit'].in_units('g/cm**3')
print('Made projection of file', fn, 'on rank', rank)
del proj
del proj_array
gc.collect()
def get_field(ds, search_field):
field = None
field_short_name = None
for f in ds.field_list + ds.derived_field_list:
if f[1] == search_field:
field_short_name = f[1]
field = f
return field, field_short_name
if not field:
print('Field {} not present.'.format(search_field))
return None, None
if __name__ == "__main__":
ds = yt.load(args.infile)
# Add Urca fields
ushell_fields = UrcaShellFields()
ushell_fields.setup(ds)
if args.list_fields:
if args.field:
field, field_short_name = get_field(ds, args.field)
assert (field)
print_field_stats(ds, field)
else:
for f in ds.field_list + ds.derived_field_list:
print_field_stats(ds, f)
exit()
def test_deposit_amr():
ds = load(ISOGAL)
for g in ds.index.grids:
gpm = g["all", "particle_mass"].sum()
dpm = g["deposit", "all_mass"].sum()
assert_allclose_units(gpm, dpm)
def test_gizmo_64():
ds = yt.load(g64)
assert isinstance(ds, GizmoDataset)
for test in sph_answer(ds, 'snap_N64L16_135', 524288, fields):
test_gizmo_64.__name__ = test.description
yield test
def surface_density_plot(ax,
snapnum,
filename="keplerian_ring",
density_limits=None,
vlim=None):
"""
Make the surface density plot (via yt).
Also returns the max and minimum values for the density so these can
be passed to the next call, as well as vlim which are the colourmap
max/min.
"""
unit_base = {
"length": (1.0, "cm"),
"velocity": (1.0, "cm/s"),
"mass": (1.0, "g")
}
filename = "{}_{:04d}.hdf5".format(filename, snapnum)
try:
snap = yt.load(filename, unit_base=unit_base, over_refine_factor=2)
except yt.utilities.exceptions.YTOutputNotIdentified:
# Probably the file isn't here because we supplied a too high snapshot
# number. Just return what we're given.
return density_limits, vlim
projection_plot = yt.ProjectionPlot(snap,
"z", ("gas", "cell_mass"),
width=5.5)
max_density = snap.all_data()[("gas", "cell_mass")].max()
min_density = snap.all_data()[("gas", "cell_mass")].min()
new_density_limits = (min_density, max_density)
if density_limits is None:
density_limits = new_density_limits
projection_plot.set_zlim("cell_mass", *density_limits)
data = get_yt_actual_data(projection_plot, ("gas", "cell_mass"))
# Becuase of the way plotting works, we also need a max/min for the colourmap.
new_vlim = (data[0].min(), data[0].max())
if vlim is None:
vlim = new_vlim
ax.imshow(data[0], cmap=data[1], vmin=vlim[0], vmax=vlim[1])
metadata = get_metadata(filename)
period = metadata["period"]
ax.text(
20,
80,
"Snapshot = {:04d}\nRotations = {:1.2f}".format(
snapnum,
float(snap.current_time) / period),
color="white",
)
# We now want to remove all of the ticklabels.
for axis in ["x", "y"]:
ax.tick_params(
axis=axis,
which="both",
bottom="off",
top="off",
left="off",
right="off",
labelleft="off",
labelbottom="off",
)
return density_limits, vlim
def __init__(self,
snapshot_path,
yield_table=None,
yield_model_Z=0.01,
yield_model_FIRE_Z_scaling=True):
self.work_dir = snapshot_path.split("output/snapshot_")[0]
self.snapshot_name = snapshot_path.split("output/")[-1]
if not (os.path.isfile(self.work_dir + 'output/snapshot_000.hdf5')):
print("Cannot find initial snapshot file at " + self.work_dir +
'output/snapshot_000.hdf5')
raise RuntimeError
self.age_bins = yield_model.get_bins(
config_file=self.work_dir + "/gizmo.out",
binfile=self.work_dir + "/age_bins.txt")
self.yield_model_Z = None
self.yield_model_FIRE_Z_scaling = None
self.yield_model_age_fraction = None
if yield_table is None:
self.generate_yield_table(yield_model_Z,
yield_model_FIRE_Z_scaling)
else:
self.yield_table = yield_table
# load datasets and define fields
self.ds0 = yt.load(self.work_dir + 'output/snapshot_000.hdf5')
self.data0 = self.ds0.all_data()
yield_model.generate_metal_fields(self.ds0,
_agebins=self.age_bins,
_yields=self.yield_table,
age_is_fraction=True)
self.metals = np.unique([
x[1] for x in self.ds0.field_list
if ((x[0] == 'PartType0') and ('Metal' in x[1]))
])
self.initial_abundance = np.zeros(15)
for i in np.arange(np.size(self.initial_abundance)):
z = self.data0[('PartType0', 'Metallicity_%02i' % (i))]
self.initial_abundance[i] = np.average(
z).value # should all be identical
self._logH = np.log10(self.ds0.hubble_constant)
self._littleh = 1.0
# load actual dataset
self.ds = yt.load(snapshot_path)
self.data = self.ds.all_data()
yield_model.generate_metal_fields(self.ds,
_agebins=self.age_bins,
_yields=self.yield_table,
age_is_fraction=True)
yield_model._generate_star_metal_fields(self.ds,
_agebins=self.age_bins,
_yields=self.yield_table,
age_is_fraction=True)
return
#!/usr/bin/env python
import matplotlib.pyplot as plt
import yt
from yt import derived_field
import yt.visualization.volume_rendering.api as vr
from yt.visualization.volume_rendering.transfer_function_helper import TransferFunctionHelper
from yt.visualization.volume_rendering.api import Scene, VolumeSource, Camera, ColorTransferFunction
import numpy as np
# Open Dataset
ds = yt.load('wd_512_rhoc4-5_plt64636')
# Hack: because rendering likes log fields ...
## create positive_radial_velocity and negative_radial_velocity fields.
@derived_field(name="pos_radial_velocity", units="cm/s")
def _pos_radial_velocity(field, data):
return np.maximum(data[('boxlib', 'radial_velocity')], 1.0e-99)
@derived_field(name="neg_radial_velocity", units="cm/s")
def _neg_radial_velocity(field, data):
return np.maximum(-data[('boxlib', 'radial_velocity')], 1.0e-99)
# Create lineout along x axis through the center
c = ds.domain_center
ax = 0 # take a line cut along the x axis
# cutting through the y0,z0 such that we hit the center
def calculate_mass_in_sphere(dd):
data_sub_dir = dd.name
a = dd.a
r_plus = M * (1 + math.sqrt(1 - a**2))
min_R = r_plus / 4
start_time = time.time()
# load dataset time series
if (new_rho == False):
# derived fields
@derived_field(name="rho_E_eff", units="")
def _rho_E_eff(field, data):
r_BL = (data["spherical_radius"] /
cm) * (1 + r_plus * cm / (4 * data["spherical_radius"]))**2
Sigma2 = r_BL**2 + (data["z"] * a * M / (r_BL * cm))
Delta = r_BL**2 + (a * M)**2 - 2 * M * r_BL
A = (r_BL**2 + (a * M)**2)**2 - (
(a * M)**2) * Delta * (data["x"]**2 + data["y"]**2) / (
(cm**2) * r_BL)
alpha = pow(Delta * Sigma2 / A, 0.5)
beta = 2 * a * (M**2) * r_BL / A
return (data["rho"] * alpha - beta * data["S_azimuth"]) * pow(
data["chi"], -3)
elif new_rho:
# derived fields
@derived_field(name="rho_E_eff", units="")
def _rho_E_eff(field, data):
return data["rho"]
dataset_path = data_root_path + "/" + data_sub_dir + "/KerrSFp_*.3d.hdf5"
ds = yt.load(dataset_path) # this loads a dataset time series
print("loaded data from ", dataset_path)
print("time = ", time.time() - start_time)
N = len(ds)
ds0 = ds[0] # get the first dataset
# set centre
center = [512.0, 512.0, 0]
L = 512.0
data_storage = {}
# iterate through datasets (forcing each to go to a different processor)
for sto, dsi in ds.piter(storage=data_storage):
time_0 = time.time()
# store time
current_time = dsi.current_time
output = [current_time]
# make sphere
sphere = dsi.sphere(center, max_R) - dsi.sphere(center, min_R)
volume = sphere.sum("cell_volume")
if half_box:
volume = 2 * volume
# calculate energy inside sphere
meanE = sphere.mean("rho_E_eff", weight="cell_volume")
E = volume * meanE
output.append(E)
# store output
sto.result = output
sto.result_id = str(dsi)
dt = 1.25
i = int(current_time / dt)
print("done {:d} of {:d} in {:.1f} s".format(i + 1, N,
time.time() - time_0),
flush=True)
if yt.is_root():
# make data directory if it does not already exist
makedirs(home_path + output_dir, exist_ok=True)
# output to file
dd.filename = "l={:d}_m={:d}_a={:s}_mass_in_r={:d}_true.csv".format(
dd.l, dd.m, str(dd.a), max_radius)
output_path = home_path + output_dir + "/" + dd.filename
# output header to file
f = open(output_path, "w+")
f.write("# t mass in r<=" + str(max_radius) + " #\n")
# output data
for key in sorted(data_storage.keys()):
data = data_storage[key]
f.write("{:.3f} ".format(data[0]))
f.write("{:.2f}\n".format(data[1]))
f.close()
print("saved data to file " + str(output_path))
def load_snapshot_enzo_yt(nSnap,
inDir,
cooling=False,
metals=False,
hydro=True,
dm=True):
import yt
snapKey = '{0:03}'.format(nSnap)
inFileName = 'DD0{0}/data0{0}'.format(snapKey)
# outFileName = 'cosmo_ICs_256_double.h5'
ds = yt.load(inDir + inFileName)
data = ds.all_data()
h = ds.hubble_constant
current_z = ds.current_redshift
current_a = 1. / (current_z + 1)
print(current_a)
print(current_z)
print(h)
if hydro:
data_grid = ds.covering_grid(level=0,
left_edge=ds.domain_left_edge,
dims=ds.domain_dimensions)
gas_dens = data_grid[
('gas', 'density')].in_units('msun/kpc**3') * current_a**3 / h**2
gas_temp = data_grid[('gas', 'temperature')]
gas_vel_x = data_grid[('gas', 'velocity_x')].in_units('km/s')
gas_vel_y = data_grid[('gas', 'velocity_y')].in_units('km/s')
gas_vel_z = data_grid[('gas', 'velocity_z')].in_units('km/s')
if cooling:
H_dens = data_grid[
('gas', 'H_density')].in_units('msun/kpc**3') * current_a**3 / h**2
H_0_dens = data_grid[('gas', 'H_p0_density'
)].in_units('msun/kpc**3') * current_a**3 / h**2
H_1_dens = data_grid[('gas', 'H_p1_density'
)].in_units('msun/kpc**3') * current_a**3 / h**2
He_dens = data_grid[(
'gas', 'He_density')].in_units('msun/kpc**3') * current_a**3 / h**2
He_0_dens = data_grid[('gas', 'He_p0_density'
)].in_units('msun/kpc**3') * current_a**3 / h**2
He_1_dens = data_grid[('gas', 'He_p1_density'
)].in_units('msun/kpc**3') * current_a**3 / h**2
He_2_dens = data_grid[('gas', 'He_p2_density'
)].in_units('msun/kpc**3') * current_a**3 / h**2
electron_dens = data_grid[(
'gas', 'El_density')].in_units('msun/kpc**3') * current_a**3 / h**2
if metals:
metal_dens = data_grid[('gas', 'metal_density')].in_units(
'msun/kpc**3') * current_a**3 / h**2
if dm:
p_mass = data[('all', 'particle_mass')].in_units('msun') * h
p_pos_x = data[('all',
'particle_position_x')].in_units('kpc') / current_a * h
p_pos_y = data[('all',
'particle_position_y')].in_units('kpc') / current_a * h
p_pos_z = data[('all',
'particle_position_z')].in_units('kpc') / current_a * h
p_vel_x = data[('all', 'particle_velocity_x')].in_units('km/s')
p_vel_y = data[('all', 'particle_velocity_y')].in_units('km/s')
p_vel_z = data[('all', 'particle_velocity_z')].in_units('km/s')
data_dic = {'dm': {}, 'gas': {}}
# data_dic['omega_l'] = ds.omega_lambda
# data_dic['omega_m'] = ds.omega_matter
data_dic['current_a'] = current_a
data_dic['current_z'] = current_z
if dm:
data_dic['dm']['mass'] = p_mass.v
data_dic['dm']['pos_x'] = p_pos_x.v
data_dic['dm']['pos_y'] = p_pos_y.v
data_dic['dm']['pos_z'] = p_pos_z.v
data_dic['dm']['vel_x'] = p_vel_x.v
data_dic['dm']['vel_y'] = p_vel_y.v
data_dic['dm']['vel_z'] = p_vel_z.v
if hydro:
data_dic['gas']['density'] = gas_dens.v
data_dic['gas']['temperature'] = gas_temp.v
data_dic['gas']['vel_x'] = gas_vel_x.v
data_dic['gas']['vel_y'] = gas_vel_y.v
data_dic['gas']['vel_z'] = gas_vel_z.v
if cooling:
data_dic['gas']['H_dens'] = H_0_dens
data_dic['gas']['HI_dens'] = H_1_dens
data_dic['gas']['He_dens'] = He_0_dens
data_dic['gas']['HeI_dens'] = He_1_dens
data_dic['gas']['HeII_dens'] = He_2_dens
data_dic['gas']['electron_dens'] = electron_dens
if metals:
data_dic['gas']['metal_dens'] = metal_dens
return data_dic
import re
import yt
yt.funcs.mylog.setLevel(0)
import numpy as np
import scipy.constants as scc
sys.path.insert(1, '../../../../warpx/Regression/Checksum/')
import checksumAPI
filename = sys.argv[1]
# Parse test name
averaged = True if re.search('averaged', filename) else False
current_correction = True if re.search('current_correction',
filename) else False
ds = yt.load(filename)
all_data = ds.covering_grid(level=0,
left_edge=ds.domain_left_edge,
dims=ds.domain_dimensions)
Ex = all_data['boxlib', 'Ex'].squeeze().v
Ey = all_data['boxlib', 'Ey'].squeeze().v
Ez = all_data['boxlib', 'Ez'].squeeze().v
if (averaged):
# energyE_ref was calculated with Galilean PSATD method (v_galilean = (0,0,0.99498743710662))
energyE_ref = 460315.9845556079
tolerance_rel = 1e-4
elif (current_correction):
# energyE_ref was calculated with standard PSATD method (v_galilean = (0.,0.,0.)):
energyE_ref = 75333.81851879464
def doit(plotfile):
ds = yt.load(plotfile)
ds.periodicity = (True, True, True)
field = ('gas', 'velocity_z')
ds._get_field_info(field).take_log = False
sc = Scene()
# add a volume: select a sphere
#center = (0, 0, 0)
#R = (5.e8, 'cm')
#dd = ds.sphere(center, R)
vol = VolumeSource(ds, field=field)
vol.use_ghost_zones = True
sc.add_source(vol)
# transfer function
vals = [-1.e7, -5.e6, 5.e6, 1.e7]
sigma = 5.e5
tf = yt.ColorTransferFunction((min(vals), max(vals)))
tf.clear()
cm = "coolwarm"
for v in vals:
tf.sample_colormap(v, sigma**2, colormap=cm) #, alpha=0.2)
sc.get_source(0).transfer_function = tf
cam = sc.add_camera(ds, lens_type="perspective")
cam.resolution = (1920, 1080)
center = 0.5 * (ds.domain_left_edge + ds.domain_right_edge)
cam.position = [
2.5 * ds.domain_right_edge[0], 2.5 * ds.domain_right_edge[1],
center[2] + 0.25 * ds.domain_right_edge[2]
]
# look toward the center -- we are dealing with an octant
normal = (center - cam.position)
normal /= np.sqrt(normal.dot(normal))
cam.switch_orientation(normal_vector=normal, north_vector=[0., 0., 1.])
cam.set_width(ds.domain_width)
sc.camera = cam
#sc.annotate_axes(alpha=0.05)
#sc.annotate_domain(ds, color=np.array([0.05, 0.05, 0.05, 0.05]))
#sc.annotate_grids(ds, alpha=0.05)
sc.render()
sc.save("{}_radvel".format(plotfile), sigma_clip=4.0)
sc.save_annotated(
"{}_radvel_annotated.png".format(plotfile),
sigma_clip=4.0,
text_annotate=[[(0.05, 0.05), "t = {}".format(ds.current_time.d),
dict(horizontalalignment="left")],
[(0.5, 0.95),
"Maestro simulation of convection in a mixed H/He XRB",
dict(color="y",
fontsize="24",
horizontalalignment="center")]])
import matplotlib.pyplot as plt
from scipy.constants import c, e, m_e, epsilon_0
import numpy as np
import yt
yt.funcs.mylog.setLevel(50)
# this will be the name of the plot file
fn = sys.argv[1]
# Parameters of the plasma
ux = 0.01
n0 = 1.e25
wp = (n0 * e**2 / (m_e * epsilon_0))**.5
# Load the dataset
ds = yt.load(fn)
t = ds.current_time.to_ndarray().mean() # in order to extract a single scalar
data = ds.covering_grid(0, ds.domain_left_edge, ds.domain_dimensions)
it = int(fn[-5:])
# Machine precision of the simulation
if ds.index._dtype == "float32":
dfloat = 0.7e-4 # single: ok, just somewhat larger than 1.e-6 (6 digits)
else:
dfloat = 2.0e-14 # double: ok, just a little larger than 1.e-15 (15 digits)
# Check the Jx field, which oscillates at wp
j_predicted = -n0 * e * c * ux * np.cos(wp * t *
(it - 0.5) / it) # j at half-timestep
jx = data['jx'].to_ndarray()
def get_dataset(enzo_output):
# open the enzo output
ds = yt.load(enzo_output)
return ds
madir = "/data/manda/jcibanezm/StratBox/RealProd/2pc/1pc/NoSG_Evol/"
#snp_nr = 2411
#for qq in range(42):
for qq in range(1):
snp_nr = (qq+offset)*10
plt_file = "Strat_Box_1pc_NoSG_Particles_hdf5_plt_cnt_%.4i" %snp_nr
par_file = "Strat_Box_1pc_NoSG_Particles_hdf5_part_%.4i" %snp_nr
my_plt_file = madir + plt_file
my_part_file = madir + par_file
# Load the data file
pf = yt.load(my_plt_file, particle_filename=my_part_file)
print ("=============================================================================")
print ("Running the Get cloud list routine.")
print ("Reading file: %s" %pf)
print ("")
# Define the cadence of the data files.
current_time10 = int(pf.current_time.in_units("Myr").value * 10 + 2 )
snapshot10= current_time10 - sg_init_time*10 + 8
snapshot = snapshot10
snapshot_str = str(snapshot)
print("My Snapshot name is %s" %snapshot_str)
'pre6': r'C$_6$ cm$^{-3}$'
}
names = [name for name in field_label.keys() if 'pre' in name]
for name in names:
field_label[name + '_scaled'] = field_label[name]
def _scale(field, data):
new_field = field.name
ftype, fname = new_field
original_fname = fname.split('_')[0]
return 1e13 * data[('all', original_fname)]
ds = yt.load(home_root + 'projects/moltres/tests/'
'twod_axi_coupled/2x_refined_from_base_with_gammas.e',
step=-1)
for i in range(1, 7):
fname = 'pre%d_scaled' % i
ds.add_field(('all', fname), function=_scale, take_log=False)
actual_domain_widths, actual_center = actual_center_and_widths(ds)
for field in ds.field_info.keys():
field_type, field_name = field
if field_type != 'all':
continue
if 'scaled' not in field_name and 'temp' not in field_name and 'group' not in field_name:
continue
slc = yt.SlicePlot(ds, 'z', field, origin='native', center=actual_center)
slc.set_log(field, False)
slc.set_width((actual_domain_widths[0], actual_domain_widths[1]))
import yt
import matplotlib.pyplot as plt
ds = yt.load("GasSloshing/sloshing_nomag2_hdf5_plt_cnt_0150")
# Get the first sphere
sp0 = ds.sphere(ds.domain_center, (500., "kpc"))
# Compute the bulk velocity from the cells in this sphere
bulk_vel = sp0.quantities.bulk_velocity()
# Get the second sphere
sp1 = ds.sphere(ds.domain_center, (500., "kpc"))
# Set the bulk velocity field parameter
sp1.set_field_parameter("bulk_velocity", bulk_vel)
# Radial profile without correction
rp0 = yt.create_profile(sp0,
'radius',
'radial_velocity',
units={'radius': 'kpc'},
logs={'radius': False})
# Radial profile with correction for bulk velocity
rp1 = yt.create_profile(sp1,
'radius',
'radial_velocity',
units={'radius': 'kpc'},
#yt.enable_parallelism()
import pickle
#Project B-field in a plane - derive these
def _Bxy(field, data): #define a new field in x-y plane
return np.sqrt(data["X-magnfield"]**2+data["Y-magnfield"]**2)
def _Byz(field, data): #define a new field in y-z plane
return np.sqrt(data["Y-magnfield"]**2+data["Z-magnfield"]**2)
def _Bxz(field, data): #define a new field in x-z plane
return np.sqrt(data["X-magnfield"]**2+data["Z-magnfield"]**2)
fn="data.0510.3d.hdf5"
ds=yt.load(fn) #load dataset
ds.add_field(("gas", "Bxy"), function=_Bxy, units="G", force_override=True) #add the new field
ds.add_field(("gas", "Byz"), function=_Byz, units="G", force_override=True)
ds.add_field(("gas", "Bxz"), function=_Bxz, units="G", force_override=True)
def Crop3DBfield(lvl): #make a 3D covering grid of all 3 of the B-field components, a function of the AMRgrid level (resolution)
dims=64*2**lvl
crop_data_lvl= ds.smoothed_covering_grid(level=lvl, left_edge=[-4.0e17,-4.0e17,-4.0e17], dims=[dims,dims,dims])
Xmag3d=crop_data_lvl["X-magnfield"]
Ymag3d=crop_data_lvl["Y-magnfield"]
Zmag3d=crop_data_lvl["Z-magnfield"]
return Xmag3d,Ymag3d,Zmag3d
def Crop3Ddensity(lvl): #make a 3D covering grid of the density
dims=64*2**lvl
import numpy as np
import os
import sys
import yt
from yt.utilities.logger import \
ytLogger
ytLogger.setLevel(40)
from yt.extensions.p3bh import *
if __name__ == "__main__":
black_holes = {}
sim_fn = os.path.abspath(sys.argv[1])
es = yt.load(sim_fn)
data_dir = os.path.dirname(sim_fn)
for i, fn in enumerate(es.data["filename"]):
if isinstance(fn, bytes):
fn = fn.decode("utf")
fn = os.path.basename(fn)
dsfn = os.path.join(data_dir, "black_holes_bh", "%s.h5" % fn)
if not os.path.exists(dsfn):
continue
ds = yt.load(dsfn)
if not ds.field_list:
continue
pids = ds.r["particle_index"].d.astype(int)
mbh = ds.arr(np.zeros(pids.size), "Msun")
mstar = ds.r["particle_mass"].to("Msun")
import numpy as np
import matplotlib.pyplot as plt
pc_to_cm = 3.0857E18 # 1pc in cm
plot_range = [0, 10, 20, 30, 40, 50, 60, 70, 80]
num_plots = len(plot_range)
colormap = plt.cm.gist_ncar
plt.gca().set_color_cycle([colormap(i) for i in np.linspace(0,0.9, num_plots)])
file_names = []
for i in plot_range:
file_names.append("../../DD" + `i`.zfill(4) + "/" + "cloud_collision_" + `i`.zfill(4))
labels = []
for file in file_names:
ds = yt.load(file)
sphere = ds.sphere("c", (75, "pc"))
plot = yt.ProfilePlot(sphere, "radius", ["temperature"],
weight_field="cell_mass")
profile = plot.profiles[0]
plt.loglog(profile.x/pc_to_cm, profile["temperature"])
labels.append(r"%0.2f Myr" % ds.current_time.value)
plt.xlabel(r"Radius $(\mathrm{pc})$")
plt.ylabel(r"Temperature $(\mathrm{K})$")
plt.xlim(1,75)
plt.legend(labels, loc="upper left", frameon=False, prop={'size':10})
plt.savefig("temperature_series")
import numpy as np
import yt
from yt.data_objects.level_sets.api import *
ds = yt.load("IsolatedGalaxy/galaxy0030/galaxy0030")
data_source = ds.disk([0.5, 0.5, 0.5], [0., 0., 1.], (8, 'kpc'), (1, 'kpc'))
# the field to be used for contouring
field = ("gas", "density")
# This is the multiplicative interval between contours.
step = 2.0
# Now we set some sane min/max values between which we want to find contours.
# This is how we tell the clump finder what to look for -- it won't look for
# contours connected below or above these threshold values.
c_min = 10**np.floor(np.log10(data_source[field]).min())
c_max = 10**np.floor(np.log10(data_source[field]).max() + 1)
# Now find get our 'base' clump -- this one just covers the whole domain.
master_clump = Clump(data_source, field)
# Add a "validator" to weed out clumps with less than 20 cells.
# As many validators can be added as you want.
master_clump.add_validator("min_cells", 20)
# Calculate center of mass for all clumps.
master_clump.add_info_item("center_of_mass")