def create_field_info(self):
super(ARTDataset, self).create_field_info()
if "wspecies" in self.parameters:
# We create dark_matter and stars unions.
ptr = self.particle_types_raw
pu = ParticleUnion("darkmatter", list(ptr[:-1]))
self.add_particle_union(pu)
pu = ParticleUnion("stars", list(ptr[-1:]))
self.add_particle_union(pu)
def update_data(self, data):
"""
Update the stream data with a new data dict. If fields already exist,
they will be replaced, but if they do not, they will be added. Fields
already in the stream but not part of the data dict will be left
alone.
"""
particle_types = set_particle_types(data[0])
self.stream_handler.particle_types.update(particle_types)
self.ds._find_particle_types()
for i, grid in enumerate(self.grids):
field_units, gdata, number_of_particles = process_data(data[i])
self.stream_handler.particle_count[i] = number_of_particles
self.stream_handler.field_units.update(field_units)
for field in gdata:
if field in grid.field_data:
grid.field_data.pop(field, None)
self.stream_handler.fields[grid.id][field] = gdata[field]
self._reset_particle_count()
# We only want to create a superset of fields here.
for field in self.ds.field_list:
if field[0] == "all":
self.ds.field_list.remove(field)
self._detect_output_fields()
self.ds.create_field_info()
mylog.debug("Creating Particle Union 'all'")
pu = ParticleUnion("all", list(self.ds.particle_types_raw))
self.ds.add_particle_union(pu)
self.ds.particle_types = tuple(set(self.ds.particle_types))
def _setup_gas_alias(self):
"Alias the grid type to gas by making a particle union."
if "grid" in self.particle_types and "gas" not in self.particle_types:
pu = ParticleUnion("gas", ["grid"])
self.add_particle_union(pu)
# We have to alias this because particle unions only
# cover the field_list.
self.field_info.alias(("gas", "cell_volume"), ("grid", "cell_volume"))
def create_field_info(self):
super(ARTIODataset, self).create_field_info()
# only make the particle union if there are multiple DM species.
# If there are multiple, "N-BODY_0" will be the first species. If there
# are not multiple, they will be all under "N-BODY"
if "N-BODY_0" in self.particle_types_raw:
dm_labels = [label for label in self.particle_types_raw
if "N-BODY" in label]
# Use the N-BODY label for the union to be consistent with the
# previous single mass N-BODY case, where this label was used for
# all N-BODY particles by default
pu = ParticleUnion("N-BODY", dm_labels)
self.add_particle_union(pu)
def create_field_info(self):
self.field_dependencies = {}
self.derived_field_list = []
self.filtered_particle_types = []
self.field_info = self._field_info_class(self, self.field_list)
self.coordinates.setup_fields(self.field_info)
self.field_info.setup_fluid_fields()
for ptype in self.particle_types:
self.field_info.setup_particle_fields(ptype)
if "all" not in self.particle_types:
mylog.debug("Creating Particle Union 'all'")
pu = ParticleUnion("all", list(self.particle_types_raw))
self.add_particle_union(pu)
deps, unloaded = self.field_info.check_derived_fields()
self.field_dependencies.update(deps)
mylog.info("Loading field plugins.")
self.field_info.load_all_plugins()
def create_field_info(self):
super(ARTDataset, self).create_field_info()
ptr = self.particle_types_raw
pu = ParticleUnion("darkmatter", list(ptr))
self.add_particle_union(pu)
pass
def yt_octree_generate(fname, field_add):
print('[grid_construction]: bbox_lim = ', cfg.par.bbox_lim)
bbox = [[-2. * cfg.par.bbox_lim, 2. * cfg.par.bbox_lim],
[-2. * cfg.par.bbox_lim, 2. * cfg.par.bbox_lim],
[-2. * cfg.par.bbox_lim, 2. * cfg.par.bbox_lim]]
# load the DS and add pd fields; no need to put in stellar ages yet
# as this will happen downstream in zoom
pf = field_add(fname, bounding_box=bbox)
# zoom if necessary
# if cfg.par.zoom == True:
pf = octree_zoom_bbox_filter(fname, pf, bbox, field_add)
pf.index
ad = pf.all_data()
# ---------------------------------------------------------------
# PLOTTING DIAGNOSTIC PROJECTION PLOTS
# ---------------------------------------------------------------
# proj_plots(pf)
from yt.data_objects.particle_unions import ParticleUnion
pu = ParticleUnion("all", list(pf.particle_types_raw))
saved = pf.index.oct_handler.save_octree()
always = AlwaysSelector(None)
ir1 = pf.index.oct_handler.ires(always) # refinement levels
fc1 = pf.index.oct_handler.fcoords(always) # coordinates in code_length
fw1 = pf.index.oct_handler.fwidth(always) # width of cell in code_length
# convert fc1 and fw1 to YTArrays
fc1 = pf.arr(fc1, 'code_length')
fw1 = pf.arr(fw1, 'code_length')
print('----------------------------')
print('yt Octree Construction Stats')
print('----------------------------')
print(' n_ref = ', pf.index.oct_handler.n_ref)
print(' max_level = ', pf.index.oct_handler.max_level)
print(' nocts = ', pf.index.oct_handler.nocts)
print('----------------------------')
gridstats(ir1, fc1, fw1)
# ==================================
refined = saved['octree'].astype('bool')
try:
refined.reshape(len(ir1))
except:
refinements = 2**(3 * cfg.par.oref)
refined2 = []
for r in refined:
if r == 1:
refined2.append(True)
if r == 0:
refined2.append(np.zeros(refinements).astype('bool'))
refined = np.hstack(refined2)
# smooth the data on to the octree
volume = np.zeros(len(refined))
if cfg.par.CONSTANT_DUST_GRID == False:
# crash the code if the parameter choice for dust grid type isn't in
# the hard coded valid list below
dust_grid_type_list = ['dtm', 'rr', 'manual']
try:
dust_choice = dust_grid_type_list.index(cfg.par.dust_grid_type)
except ValueError as e:
print(
'You chose a dust_choice that isnt a valid selection within the list: '
+ dust_grid_type_list + '....crashing now!')
sys.exit()
if cfg.par.dust_grid_type == 'dtm':
dust_smoothed_dtm = dtm_grid(pf, refined)
dust_smoothed = dust_smoothed_dtm
if cfg.par.dust_grid_type == 'rr':
dust_smoothed_remy_ruyer = remy_ruyer(pf, refined)
dust_smoothed = dust_smoothed_remy_ruyer
if cfg.par.dust_grid_type == 'manual':
dust_smoothed_manual = manual(pf, refined)
dust_smoothed = dust_smoothed_manual
else:
print('cfg.par.CONSTANT_DUST_GRID=True')
print('setting constant dust grid to 4.e-22')
dust_smoothed = np.zeros(len(refined)) + 4.e-23
# return refined,dust_smoothed,xmin,xmax,ymin,ymax,zmin,zmax,boost
return refined, dust_smoothed, fc1, fw1, pf, ad
def run_rockstar(sim_dir, snap_base='*', particle_types=["darkmatter"],
recognized_star_types=["stars", "specie5"], multi_mass=False,
force_res=None, initial_metric_scaling=1, out_dir="rockstar_output",
num_readers=1):
'''
Run Rockstar for a set of snapshots in the given simulation directory
Parameters
----------
sim_dir : str
Simulation directory to be analyzed
snap_base : str
Snapshot files to be analyzed. You can use wild cards in SNAP_BASE
to select a set of snapshots that match a given pattern
particle_types: str list/array
The particle types to use for halo finding, as named in yt fields or
derived fields. This can be filtered or inherent types.
If particles with different masses are found for the given particle_types
only the particles with the lowest mass will be used, unless the
multi_mass option is set to True
recognized_star_types: str list/array
Particle types to be recognized as stars. All types in particle_types
that are not included in this list will be treated as dark matter particles
multi_mass : bool
If True all the particle species in particle_types will be passed to
Rockstar even if they have different masses
force_res : float
Force resolution in units of Mpc/h. Halos whose centers are closer than
FORCE_RES are usually noise, and are subject to stricter removal tests
than other halos. If no value is provided, this parameter is automatically
set to the width of the smallest grid element in the simulation from the
last data snapshot (i.e. the one where time has evolved the longest)
initial_metric_scaling : float
The position element of the fof distance metric is divided by this
parameter, set to 1 by default. If the initial_metric_scaling=0.1 the
position element will have 10 times more weight than the velocity element,
biasing the metric towards position information more so than velocity
information. That was found to be needed for hydro-ART simulations with
with 10's of parsecs resolution
out_dir : str
Directory where the Rockstar output will be placed
num_readers : int
If the number of files per snapshot is more than one, setting
num_readers > 1 can help I/O performance
Returns
-------
None
'''
if MPI.COMM_WORLD.Get_rank()==0:
print
print 'Beginning Rockstar analysis in ', sim_dir
print 'Placing output in ', out_dir
# Generate data series
snap_base = sim_dir+ '/' + snap_base + '*'
ts = yt.DatasetSeries(snap_base)
if len(ts) > 1:
first_snap = ts[0]
last_snap = ts[-1]
if MPI.COMM_WORLD.Get_rank()==0:
print 'first snapshot: ', first_snap.basename
print 'last snapshot: ', last_snap.basename
else:
last_snap = ts[0]
if MPI.COMM_WORLD.Get_rank()==0:
print 'snapshot: ', last_snap.basename
# Create particle union if needed
punion=None
if len(particle_types) > 1:
punion = ParticleUnion("rockstar_analysis_particles", particle_types)
last_snap.add_particle_union(punion)
particle_type = "rockstar_analysis_particles"
else:
particle_type = particle_types[0]
ad = last_snap.all_data()
particle_masses = yt.np.unique(ad[(particle_type, 'particle_mass')].in_units('Msun/h'))
total_particles = ad.quantities.total_quantity((particle_type, "particle_ones"))
if MPI.COMM_WORLD.Get_rank()==0:
print 'Using particle(s) ', particle_types, ' with mass(es) ', particle_masses
print
# Get the particle_type field of the star particles, and the min mass of the DM particles
dm_masses = []
star_types = []
for type in particle_types:
if type in recognized_star_types:
star_types = yt.np.concatenate((star_types, yt.np.unique(ad[(type, 'particle_type')])))
else:
dm_masses = \
yt.np.concatenate((dm_masses,yt.np.unique(ad[(type, 'particle_mass')].in_units('Msun/h'))))
dm_masses = last_snap.arr(dm_masses.value, 'Msun/h')
dm_min_mass = dm_masses.min()
# In not multi mass select only the DM particles with the min mass
if not multi_mass and len(particle_masses) > 1:
@yt.particle_filter('hires_particles', filtered_type=particle_type, requires=['particle_mass'])
def _hires_filter(pfilter, data):
return data['particle_mass'] == dm_min_mass
def setup_ds(ds):
if punion: ds.add_particle_union(punion)
ds.add_particle_filter('hires_particles')
particle_type = 'hires_particles'
else:
def setup_ds(ds):
if punion: ds.add_particle_union(punion)
# Remove last/first_snap from memory and cache
if len(ts) > 1 : del first_snap
del last_snap
_cached_datasets.clear()
# Pass the setup_function to the time series
ts._setup_function=setup_ds
# Run rockstar
rh = RockstarHaloFinder(ts,
num_readers=num_readers,
num_writers=None,
outbase=out_dir,
particle_type=particle_type,
star_types=star_types,
multi_mass=multi_mass,
force_res=force_res,
initial_metric_scaling=initial_metric_scaling,
particle_mass=dm_min_mass,
total_particles=total_particles)
MPI.COMM_WORLD.barrier()
rh.run()
MPI.COMM_WORLD.barrier()
if MPI.COMM_WORLD.Get_rank()==0:
print 'Successfully finished Rockstar analysis in ', sim_dir
print
