def test_nodes_from_selection(self):
a = self.arbor
sel = a.get_yt_selection(above=[("mass", 1e10, "Msun")])
sel_mass = sel["halos", "mass"].to("Msun")
sel_mass.sort()
nodes = list(a.get_nodes_from_selection(sel))
node_mass = a.arr([node["mass"] for node in nodes]).to("Msun")
node_mass.sort()
assert_array_rel_equal(node_mass, sel_mass, decimals=5)
def test_yt_all_data(self):
a = self.arbor
ds = a.ytds
assert isinstance(ds, YTreeDataset)
ad = ds.all_data()
for field, units in zip(["mass", "redshift"], ["Msun", ""]):
yt_data = ad["halos", field].to(units)
yt_data.sort()
ytree_data = uconcatenate([t["forest", field] for t in a])
ytree_data.sort()
assert_array_rel_equal(yt_data, ytree_data, decimals=5)
def test_conditionals(self):
a = self.arbor
ds = a.ytds
sel = a.get_yt_selection(conditionals=['obj["halos", "mass"] > 1e10'])
sel_mass = sel["halos", "mass"].to("Msun")
assert (sel_mass >= 1e10).all()
sel_mass.sort()
ad = ds.all_data()
ad_mass = ad["halos", "mass"].to("Msun")
yt_mass = ad_mass[ad_mass >= 1e10]
yt_mass.sort()
assert_array_rel_equal(sel_mass, yt_mass, decimals=5)
def test_equal_no_units(self):
a = self.arbor
ds = a.ytds
sel = a.get_yt_selection(equal=[("mmp?", 1)])
sel_mmp = sel["halos", "mmp?"].to("")
assert (sel_mmp == 1).all()
sel_mmp.sort()
ad = ds.all_data()
ad_mmp = ad["halos", "mmp?"].to("")
yt_mmp = ad_mmp[ad_mmp == 1]
yt_mmp.sort()
assert_array_rel_equal(sel_mmp, yt_mmp, decimals=5)
def test_above_no_units(self):
a = self.arbor
ds = a.ytds
mt = 1e10
sel = a.get_yt_selection(above=[("mass", mt)])
sel_mass = sel["halos", "mass"].to("Msun")
assert (sel_mass >= mt).all()
sel_mass.sort()
ad = ds.all_data()
ad_mass = ad["halos", "mass"].to("Msun")
yt_mass = ad_mass[ad_mass >= mt]
yt_mass.sort()
assert_array_rel_equal(sel_mass, yt_mass, decimals=5)
def test_conditionals_sphere(self):
a = self.arbor
ds = a.ytds
sp = ds.sphere(0.5*ds.domain_center, (20, "Mpc/h"))
sel = a.get_yt_selection(conditionals=['obj["halos", "mass"] > 1e10'],
data_source=sp)
sel_mass = sel["halos", "mass"].to("Msun")
assert (sel_mass >= 1e10).all()
sel_mass.sort()
sp_mass = sp["halos", "mass"].to("Msun")
yt_mass = sp_mass[sp_mass >= 1e10]
yt_mass.sort()
assert_array_rel_equal(sel_mass, yt_mass, decimals=5)
def test_equal_sphere(self):
a = self.arbor
ds = a.ytds
sp = ds.sphere(0.5*ds.domain_center, (20, "Mpc/h"))
sel = a.get_yt_selection(equal=[("mmp?", 1, "")],
data_source=sp)
sel_mmp = sel["halos", "mmp?"].to("")
assert (sel_mmp == 1).all()
sel_mmp.sort()
sp_mmp = sp["halos", "mmp?"].to("")
yt_mmp = sp_mmp[sp_mmp == 1]
yt_mmp.sort()
assert_array_rel_equal(sel_mmp, yt_mmp, decimals=5)
def test_about_no_units(self):
a = self.arbor
ds = a.ytds
mt = 1e10
within = 0.5
sel = a.get_yt_selection(about=[("mass", mt, within)])
sel_mass = sel["halos", "mass"].to("Msun")
assert ((sel_mass >= (1-within)*mt) & (sel_mass <= (1+within)*mt)).all()
sel_mass.sort()
ad = ds.all_data()
ad_mass = ad["halos", "mass"].to("Msun")
yt_mass = ad_mass[(ad_mass >= (1-within)*mt) & (ad_mass <= (1+within)*mt)]
yt_mass.sort()
assert_array_rel_equal(sel_mass, yt_mass, decimals=5)
def test_below_sphere(self):
a = self.arbor
ds = a.ytds
sp = ds.sphere(0.5*ds.domain_center, (20, "Mpc/h"))
mt = 1e10
sel = a.get_yt_selection(below=[("mass", mt, "Msun")],
data_source=sp)
sel_mass = sel["halos", "mass"].to("Msun")
assert (sel_mass <= mt).all()
sel_mass.sort()
sp_mass = sp["halos", "mass"].to("Msun")
yt_mass = sp_mass[sp_mass <= mt]
yt_mass.sort()
assert_array_rel_equal(sel_mass, yt_mass, decimals=5)
def test_about_sphere(self):
a = self.arbor
ds = a.ytds
sp = ds.sphere(0.5*ds.domain_center, (20, "Mpc/h"))
mt = 1e10
within = 0.5
sel = a.get_yt_selection(about=[("mass", mt, within)],
data_source=sp)
sel_mass = sel["halos", "mass"].to("Msun")
assert ((sel_mass >= (1-within)*mt) & (sel_mass <= (1+within)*mt)).all()
sel_mass.sort()
sp_mass = sp["halos", "mass"].to("Msun")
yt_mass = sp_mass[(sp_mass >= (1-within)*mt) & (sp_mass <= (1+within)*mt)]
yt_mass.sort()
assert_array_rel_equal(sel_mass, yt_mass, decimals=5)
def test_nodes_from_sphere(self):
a = self.arbor
ds = a.ytds
sp = ds.sphere(0.5*ds.domain_center, (20, "Mpc/h"))
nodes = list(a.get_nodes_from_selection(sp))
sp_mass = sp["halos", "mass"].to("Msun")
sp_mass.sort()
node_mass = a.arr([node["mass"] for node in nodes])
node_mass.sort()
assert_array_rel_equal(node_mass, sp_mass, decimals=5)
node_pos = a.arr([node["position"] for node in nodes])
r = a.quan(sp.radius.to("unitary"))
c = a.arr(sp.center.to("unitary"))
node_r = np.sqrt(((node_pos - c)**2).sum(axis=1))
assert (node_r <= r).all()
def test_multiple_criteria_2(self):
a = self.arbor
ds = a.ytds
sel = a.get_yt_selection(above=[("mass", 1e10, "Msun")],
below=[("redshift", 0.5)])
sel_mass = sel["halos", "mass"].to("Msun")
assert (sel_mass >= 1e10).all()
sel_redshift = sel["halos", "redshift"]
assert (sel_redshift <= 0.5).all()
sel_mass.sort()
ad = ds.all_data()
ad_mass = ad["halos", "mass"].to("Msun")
ad_redshift = ad["halos", "redshift"]
yt_mass = ad_mass[(ad_mass >= 1e10) & (ad_redshift <= 0.5)]
yt_mass.sort()
assert_array_rel_equal(sel_mass, yt_mass, decimals=5)
def test_yt_sphere(self):
a = self.arbor
ds = a.ytds
sp = ds.sphere(0.5*ds.domain_center, (20, "Mpc/h"))
ytree_pos = uconcatenate([t["forest", "position"] for t in a])
ytree_mass = uconcatenate([t["forest", "mass"] for t in a])
r = a.quan(sp.radius.to("unitary"))
c = a.arr(sp.center.to("unitary"))
ytree_r = np.sqrt(((ytree_pos - c)**2).sum(axis=1))
in_sphere = ytree_r <= r
ytree_sp_r = ytree_r[in_sphere].to("unitary")
ytree_sp_r.sort()
sp_r = sp["halos", "particle_radius"].to("unitary")
sp_r.sort()
assert_array_rel_equal(ytree_sp_r, sp_r, decimals=5)
sp_mass = sp["halos", "mass"].to("Msun")
sp_mass.sort()
ytree_sp_mass = ytree_mass[in_sphere].to("Msun")
ytree_sp_mass.sort()
assert_array_rel_equal(ytree_sp_mass, sp_mass, decimals=5)