def test_pipeline(self):
a = ytree.load(self.test_filename)
fn = a.save_arbor(trees=a[:8])
a = ytree.load(fn)
if "test_field" not in a.field_list:
a.add_analysis_field("test_field", default=-1, units="Msun")
a.add_analysis_field("test_always", dtype=int, default=0, units="")
ap = ytree.AnalysisPipeline(output_dir="analysis")
ap.add_recipe(my_recipe)
trees = list(a[:])
for tree in trees:
for node in tree["forest"]:
ap.process_target(node)
assert os.path.exists("analysis/my_analysis")
fn = a.save_arbor(trees=trees)
a2 = ytree.load(fn)
mf = a2["mass"] > 3e11
assert_array_equal(a2["test_field"][mf], 2 * a2["mass"][mf])
assert_array_equal(a2["test_field"][~mf], -np.ones((~mf).sum()))
assert_equal(a2["test_always"].sum(), a2.size)
for tree in a2:
mf = tree["forest", "mass"] > 3e11
assert_array_equal(tree["forest", "test_field"][mf],
2 * tree["forest", "mass"][mf])
assert_array_equal(tree["forest", "test_field"][~mf],
-np.ones((~mf).sum()))
def test_analysis_fields_int(self):
a = ytree.load(CT)
my_dtype = np.int64
a.add_analysis_field('bears', 'Msun/yr', dtype=my_dtype)
assert_equal(my_dtype, a['bears'].dtype)
fake_bears = np.zeros(a.size)
assert_array_equal(fake_bears, a['bears'])
all_trees = list(a[:])
my_tree = all_trees[0]
my_tree['bears'] = 9
fake_bears[0] = 9
assert_array_equal(fake_bears, a['bears'])
my_tree = all_trees[1]
assert_equal(my_dtype, my_tree['tree', 'bears'].dtype)
fake_tree_bears = np.zeros(my_tree.tree_size)
assert_array_equal(fake_tree_bears, my_tree['tree', 'bears'])
fake_tree_bears[72] = 99
my_halos = list(my_tree['tree'])
my_halos[72]['bears'] = 99
assert_array_equal(fake_tree_bears, my_tree['tree', 'bears'])
fn = a.save_arbor(trees=all_trees)
a2 = ytree.load(fn)
assert_array_equal(fake_bears, a2['bears'])
assert_array_equal(fake_tree_bears, a2[1]['tree', 'bears'])
assert_equal(a2['bears'].dtype, my_dtype)
assert_equal(a2[1]['tree', 'bears'].dtype, my_dtype)
def arbor(self):
if self._arbor is not None:
return self._arbor
a = ytree.load(CTG)
fn = a.save_arbor()
self._arbor = ytree.load(fn)
return self._arbor
def test_save_tree(self):
a = ytree.load(CT)
for t in [a[0], a[0]["tree"][1]]:
fn = t.save_tree()
a2 = ytree.load(fn)
fields = a2.field_list[:]
fields.remove("desc_uid")
compare_trees(t, a2[0], fields=fields)
def test_save_tree(self):
a = ytree.load(CT)
t = a[0]
fn = t.save_tree()
a2 = ytree.load(fn)
fields = a2.field_list[:]
fields.remove("desc_uid")
compare_trees(t, a2[0], fields=fields)
def test_save_tree_nonroot(self):
a = ytree.load(CT)
t = list(a[0]["tree"])[1]
fn = t.save_tree()
a2 = ytree.load(fn)
fields = a2.field_list[:]
fields.remove("desc_uid")
compare_trees(t, a2[0], fields=fields)
def test_save_non_roots(self):
a = ytree.load(CT)
my_trees = [list(a[0]["tree"])[1], list(a[1]["tree"])[1]]
fn = a.save_arbor(trees=my_trees)
a2 = ytree.load(fn)
fields = a2.field_list[:]
fields.remove("desc_uid")
for t1, t2 in zip(a2, my_trees):
compare_trees(t1, t2, fields=fields)
def test_significance(self):
"""
Test John Wise's significance.
"""
a = ytree.load(CTG)
a.add_analysis_field('significance', 'Msun*Myr')
my_trees = a[:]
for tree in my_trees:
get_significance(tree)
fn = a.save_arbor('sig_tree', trees=my_trees)
a2 = ytree.load(fn)
a2.set_selector('max_field_value', 'significance')
prog = a2[0]['prog']
def get_my_tree_rare():
"""
Halo with the most black holes from the normal region.
"""
### Get the tree with the most black holes.
ds = yt.load("halo_catalogs_nosub/RD0041/RD0041.0.h5")
a = ytree.load("rockstar_halos/trees/tree_0_0_0.dat")
# halo with highest relative growth
gsort = ds.r["relative_growth_max"].argsort()[::-1]
hid = ds.r["particle_identifier"][gsort][0].d
# # halo with second highest relative growth
# gsort = ds.r["relative_growth_max"].argsort()[::-1]
# hid = ds.r["particle_identifier"][gsort][1].d
t = a[a["Orig_halo_ID"] == hid][0]
uid = t["uid"]
data_dir = "halo_%d" % uid
ensure_dir(data_dir)
fn = t.save_tree(filename=os.path.join(data_dir, "tree_%d" % uid))
return fn
def test_1x_tree_farms(self):
tree_farm_filenames = \
glob.glob(os.path.join(test_data_dir,
"ytree_1x/arbors/tree_farm*"))
fields = ["redshift", "particle_mass"]
for filename in tree_farm_filenames:
basename = os.path.basename(filename)
arbor_glob = \
os.path.join(test_data_dir, "ytree_1x/arbors",
basename, "*.h5")
arbor_filename = glob.glob(arbor_glob)[0]
results_filename = \
os.path.join(test_data_dir, "ytree_1x/results",
"%s.h5" % basename)
a = ytree.load(arbor_filename)
basename = "%s.h5" % os.path.basename(filename)
generate_results_file(a, basename, fields)
compare_hdf5(basename,
results_filename,
compare_groups=False,
compare=assert_array_rel_equal,
decimals=15)
def get_rockstar_data(rstar_fn, halo_id):
"""
Use ytree to get all of the halo centroids, virial radii, and redshift info; store in a dict
"""
# load up dataset and get appropriate TreeNode
import ytree
a = ytree.load(rstar_fn)
t = a[a["Orig_halo_ID"] == halo_id][0]
redshift_arr = t['prog', 'redshift']
x_arr = t['prog', 'x'].in_units('unitary')
y_arr = t['prog', 'y'].in_units('unitary')
z_arr = t['prog', 'z'].in_units('unitary')
vx_arr = t['prog', 'vx'].in_units('km/s')
vy_arr = t['prog', 'vy'].in_units('km/s')
vz_arr = t['prog', 'vz'].in_units('km/s')
rvir_arr = t['prog', 'Rvir'].convert_to_units('kpc')
return {
'redshift_arr': redshift_arr,
'x_arr': x_arr,
'y_arr': y_arr,
'z_arr': z_arr,
'vx_arr': vx_arr,
'vy_arr': vy_arr,
'vz_arr': vz_arr,
'rvir_arr': rvir_arr
}
def test_tree_plot_custom_edge(self):
a = ytree.load(AHF, hubble_constant=0.7)
p = ytree.TreePlot(a[0],
dot_kwargs={'rankdir': "BT"},
node_function=my_node,
edge_function=my_edge)
p.save('tree_custom_edge.png')
def run():
input_fn, output_fn, selection, group = sys.argv[1:5]
njobs = tuple([int(arg) for arg in sys.argv[5:7]])
dynamic = tuple([bool(int(arg)) for arg in sys.argv[7:9]])
if len(sys.argv) > 9:
save_every = int(sys.argv[9])
else:
save_every = None
a = ytree.load(input_fn)
if "test_field" not in a.field_list:
a.add_analysis_field("test_field", default=-1, units="Msun")
if selection == "all":
trees = list(a[:8])
elif selection == "nonroot":
trees = get_tree_split(a)
else:
print(f"Bad selection: {selection}.")
sys.exit(1)
for node in ytree.parallel_nodes(trees,
group=group,
njobs=njobs,
dynamic=dynamic,
save_every=save_every,
filename=output_fn):
root = node.root
yt.mylog.info(
f"Doing {node.tree_id}/{root.tree_size} of {root._arbor_index}")
node["test_field"] = 2 * node["mass"]
def test_edge_function(self):
def my_func(desc, anc):
return {"color": "red"}
a = ytree.load(CT)
p = ytree.TreePlot(a[0], edge_function=my_func)
p.save()
def test_save_field_list(self):
a = ytree.load(CT)
my_trees = [list(a[0]["tree"])[1], list(a[1]["tree"])[1]]
fn = a.save_arbor(trees=my_trees, fields=["mass", "redshift"])
a2 = ytree.load(fn)
assert len(a2.field_list) == 4
assert sorted(["mass", "redshift", "uid", "desc_uid"]) == \
sorted(a2.field_list)
fields = a2.field_list[:]
fields.remove("desc_uid")
for t1, t2 in zip(a2, my_trees):
compare_trees(t1, t2, fields=fields)
def test_node_function(self):
def my_func(halo):
label = f"{halo['uid']}"
return {"label": label}
a = ytree.load(CT)
p = ytree.TreePlot(a[0], node_function=my_func)
p.save()
def test_halo_age(self):
"""
Test halo age example.
"""
a = ytree.load(CTG)
my_tree = a[0]
age = t50(my_tree).to('Gyr')
def test_tree_plot(self):
a = ytree.load(AHF, hubble_constant=0.7)
p = ytree.TreePlot(a[0],
dot_kwargs={
'rankdir': 'LR',
'size': '"12,4"'
})
p.save('tree.png')
def test_tree_plot_small(self):
a = ytree.load(AHF, hubble_constant=0.7)
p = ytree.TreePlot(a[0],
dot_kwargs={
'rankdir': 'LR',
'size': '"12,4"'
})
p.min_mass_ratio = 0.01
p.save('tree_small.png')
def test_halo_age(self):
"""
Test halo age example.
"""
a = ytree.load(CTG)
a.add_analysis_field("a50", "")
ap = ytree.AnalysisPipeline()
ap.add_operation(calc_a50)
trees = list(a[:])
for tree in trees:
ap.process_target(tree)
fn = a.save_arbor(filename="halo_age", trees=trees)
a2 = ytree.load(fn)
print(a2[0]["a50"])
def my_halo_filter(tree_file):
mass_string = "1E9"
all_criteria = [
'(tree["tree", "redshift"] > 12.0) *'
'(tree["tree", "redshift"] < 20.0) *'
'(tree["tree", "mass"] > ' + mass_string + ')',
'(tree["tree", "redshift"] > 11.0) *'
'(tree["tree", "redshift"] < 12.0) *'
'(tree["tree", "mass"] > ' + mass_string + ')',
'(tree["tree", "redshift"] > 10.0) *'
'(tree["tree", "redshift"] < 11.0) *'
'(tree["tree", "mass"] > ' + mass_string + ')',
'(tree["tree", "redshift"] > 9.0) *'
'(tree["tree", "redshift"] < 10.0) *'
'(tree["tree", "mass"] > ' + mass_string + ')',
'(tree["tree", "redshift"] > 8.0) *'
'(tree["tree", "redshift"] < 9.0) *'
'(tree["tree", "mass"] > ' + mass_string + ')',
'(tree["tree", "redshift"] > 7.0) *'
'(tree["tree", "redshift"] < 8.0) *'
'(tree["tree", "mass"] > ' + mass_string + ')',
'(tree["tree", "redshift"] > 6.0) *'
'(tree["tree", "redshift"] < 7.0) *'
'(tree["tree", "mass"] > ' + mass_string + ')',
'(tree["tree", "redshift"] > 5.0) *'
'(tree["tree", "redshift"] < 6.0) *'
'(tree["tree", "mass"] > ' + mass_string + ')'
]
fields = ['redshift', 'mass']
a = ytree.load(tree_file)
all_halos = get_root_halos(a,
all_criteria,
fields,
unique=True,
mutually_exclusive='sequential')
selection = a['mass'] > 1.0E8
all_neighbors = get_neighbors(a, all_halos, n=1, selection=selection)
outfiles, outplots = [], []
for z in np.arange(12, 4.5, -1):
outfiles.append(mass_string + "_above_z%i.dat" % (z))
outplots.append(mass_string + "_above_z%i.png" % (z))
for i in np.arange(len(all_criteria)):
print_halo_info(a,
all_halos[i],
header=all_criteria[i],
neighbors=all_neighbors[i],
outfile=outfiles[i])
plot_halos(a, all_halos[i], outfile=outplots[i])
return
def test_select_halos():
a = ytree.load(CT)
halos = a.select_halos('tree["tree", "Orig_halo_ID"] == 0')
hids = a.arr([h["Orig_halo_ID"] for h in halos])
assert (hids == 0).all()
halos = a.select_halos('tree["prog", "Orig_halo_ID"] == 0')
hids = a.arr([h["Orig_halo_ID"] for h in halos])
assert (hids == 0).all()
def test_alter_vector_field(self):
a = ytree.load(CT)
for i, ax in enumerate('xyz'):
a.add_analysis_field(f'thing_{ax}',
'',
default=i,
dtype=np.float64)
fn = a.save_arbor()
a2 = ytree.load(fn)
a2.add_vector_field('thing')
t = a2[0]
for i, ax in enumerate('xyz'):
t[f'thing_{ax}'] += 1
assert_array_equal(t['forest', f'thing_{ax}'], t['forest',
'thing'][:, i])
assert_array_equal(a2[f'thing_{ax}'], a2['thing'][:, i])
def test_significance(self):
"""
Test John Wise's significance.
"""
a = ytree.load(CTG)
a.add_analysis_field("significance", "Msun*Myr")
ap = ytree.AnalysisPipeline()
ap.add_operation(calc_significance)
trees = list(a[:])
for tree in trees:
ap.process_target(tree)
fn = a.save_arbor(filename="significance", trees=trees)
a2 = ytree.load(fn)
a2.set_selector("max_field_value", "significance")
prog = list(a2[0]["prog"])
print(prog)
def test_add_vector_field(self):
a = ytree.load(CT)
data = []
for i, ax in enumerate('xyz'):
a.add_analysis_field(f'thing_{ax}',
'',
default=i,
dtype=np.float64)
data.append(i * np.ones(a.size, dtype=np.float64))
data = np.rollaxis(np.vstack(data), 1)
fn = a.save_arbor()
a2 = ytree.load(fn)
a2.add_vector_field('thing')
assert_array_equal(a2['thing'], data)
data_mag = np.sqrt((data**2).sum(axis=1))
assert_array_equal(a2['thing_magnitude'], data_mag)
def test_tree_farm_descendents(self):
ts = yt.DatasetSeries(
os.path.join(test_data_dir, "fof_subfind/groups_02*/*.0.hdf5"))
my_tree = TreeFarm(ts, setup_function=setup_ds)
my_tree.set_selector("all")
my_tree.trace_descendents("Group",
filename="my_descendents/",
fields=["virial_radius"])
a = ytree.load("my_descendents/fof_subhalo_tab_020.0.h5")
assert isinstance(a, TreeFarmArbor)
save_and_compare(a)
def test_derived_fields_int(self):
a = ytree.load(CT)
my_dtype = np.int64
a.add_derived_field('uid_mod7', _uid_mod7, units='', dtype=my_dtype)
um7 = a['uid'] % 7
assert_array_equal(a['uid_mod7'], um7)
assert_equal(a['uid_mod7'].dtype, my_dtype)
my_tree = a[0]
um7 = my_tree['tree', 'uid'] % 7
assert_array_equal(my_tree['tree', 'uid_mod7'], um7)
assert_equal(my_tree['tree', 'uid_mod7'].dtype, my_dtype)
def test_non_defaults(self):
attrs = {
'size_field': 'virial_radius',
'size_log': False,
'min_mass': 1e14,
'min_mass_ratio': 0.1
}
a = ytree.load(CT)
for attr, val in attrs.items():
p = ytree.TreePlot(a[0])
setattr(p, attr, val)
p.save()
def test_select_halos_bad_input():
a = ytree.load(CT)
with assert_raises(ArborFieldNotFound):
list(
a.select_halos(
"(tree['forest', 'not_a_field'].to('Msun') > 1e13)"))
with assert_raises(ValueError):
list(a.select_halos("(tree['flock', 'mass'].to('Msun') > 1e13)"))
list(
a.select_halos(
"(tree['forest', 'mass'].to('Msun') > 1e13) & (tree['tree', 'redshift'] < 0.5)"
))
def test_analysis_fields(self):
a = ytree.load(CT)
a.add_analysis_field('bears', 'Msun/yr')
fake_bears = np.zeros(a.size)
assert_array_equal(fake_bears, a['bears'])
a[0]['bears'] = 9
fake_bears[0] = 9
assert_array_equal(fake_bears, a['bears'])
fake_tree_bears = np.zeros(a[1].tree_size)
assert_array_equal(
fake_tree_bears, a[1]['tree', 'bears'])
fake_tree_bears[72] = 99
a[1]['tree'][72]['bears'] = 99
assert_array_equal(fake_tree_bears, a[1]['tree', 'bears'])
fn = a.save_arbor()
a2 = ytree.load(fn)
assert_array_equal(fake_bears, a2['bears'])
assert_array_equal(fake_tree_bears, a2[1]['tree', 'bears'])