def test_decomposition_3d():
array = np.ones((33, 35, 37))
bbox = np.array([[0., 1.0], [-1.5, 1.5], [1.0, 2.5]])
ledge, redge, shapes, slices = dec.decompose_array(array.shape,
np.array([3, 2, 2]), bbox)
data = [array[slice] for slice in slices]
assert_array_equal(data[0].shape, np.array([11, 17, 18]))
gold_le = np.array(
[[0.00000, -1.50000, 1.00000], [0.00000, -1.50000, 1.72973],
[0.00000, -0.04286, 1.00000], [0.00000, -0.04286, 1.72973],
[0.33333, -1.50000, 1.00000], [0.33333, -1.50000, 1.72973],
[0.33333, -0.04286, 1.00000], [0.33333, -0.04286, 1.72973],
[0.66667, -1.50000, 1.00000], [0.66667, -1.50000, 1.72973],
[0.66667, -0.04286, 1.00000], [0.66667, -0.04286, 1.72973]]
)
assert_almost_equal(ledge, gold_le, 5)
gold_re = np.array(
[[0.33333, -0.04286, 1.72973], [0.33333, -0.04286, 2.50000],
[0.33333, 1.50000, 1.72973], [0.33333, 1.50000, 2.50000],
[0.66667, -0.04286, 1.72973], [0.66667, -0.04286, 2.50000],
[0.66667, 1.50000, 1.72973], [0.66667, 1.50000, 2.50000],
[1.00000, -0.04286, 1.72973], [1.00000, -0.04286, 2.50000],
[1.00000, 1.50000, 1.72973], [1.00000, 1.50000, 2.50000]]
)
assert_almost_equal(redge, gold_re, 5)
def test_morton_index():
ds = fake_amr_ds()
mi = ds.r["index", "morton_index"]
mi2 = mi.view("uint64")
assert_equal(np.unique(mi2).size, mi2.size)
a1 = np.argsort(mi)
a2 = np.argsort(mi2)
assert_array_equal(a1, a2)
def test_clone_cut_region():
fields = ("density", "temperature")
units = ("g/cm**3", "K")
ds = fake_random_ds(64, nprocs=4, fields=fields, units=units)
dd = ds.all_data()
reg1 = dd.cut_region(["obj['temperature'] > 0.5", "obj['density'] < 0.75"])
reg2 = reg1.clone()
assert_array_equal(reg1["density"], reg2["density"])
def test_old_profile_data():
tmpdir = tempfile.mkdtemp()
curdir = os.getcwd()
os.chdir(tmpdir)
ds = data_dir_load(enzotiny)
ad = ds.all_data()
profile_1d = create_profile(
ad,
("gas", "density"),
("gas", "temperature"),
weight_field=("gas", "cell_mass"),
)
fn = "DD0046_Profile1D.h5"
full_fn = os.path.join(ytdata_dir, fn)
prof_1d_ds = data_dir_load(full_fn)
compare_unit_attributes(ds, prof_1d_ds)
assert isinstance(prof_1d_ds, YTProfileDataset)
for field in profile_1d.standard_deviation:
assert_array_equal(
profile_1d.standard_deviation[field],
prof_1d_ds.profile.standard_deviation["data", field[1]],
)
p1 = ProfilePlot(
prof_1d_ds.data,
("gas", "density"),
("gas", "temperature"),
weight_field=("gas", "cell_mass"),
)
p1.save()
yield YTDataFieldTest(full_fn, ("gas", "temperature"), geometric=False)
yield YTDataFieldTest(full_fn, ("index", "x"), geometric=False)
yield YTDataFieldTest(full_fn, ("gas", "density"), geometric=False)
fn = "DD0046_Profile2D.h5"
full_fn = os.path.join(ytdata_dir, fn)
prof_2d_ds = data_dir_load(full_fn)
compare_unit_attributes(ds, prof_2d_ds)
assert isinstance(prof_2d_ds, YTProfileDataset)
p2 = PhasePlot(
prof_2d_ds.data,
("gas", "density"),
("gas", "temperature"),
("gas", "cell_mass"),
weight_field=None,
)
p2.save()
yield YTDataFieldTest(full_fn, ("gas", "density"), geometric=False)
yield YTDataFieldTest(full_fn, ("index", "x"), geometric=False)
yield YTDataFieldTest(full_fn, ("gas", "temperature"), geometric=False)
yield YTDataFieldTest(full_fn, ("index", "y"), geometric=False)
yield YTDataFieldTest(full_fn, ("gas", "cell_mass"), geometric=False)
os.chdir(curdir)
shutil.rmtree(tmpdir)
def test_nuclei_density_fields():
ds = data_dir_load(ecp)
ad = ds.all_data()
assert_array_equal(ad["H_nuclei_density"],
(ad["H_p0_number_density"] + ad["H_p1_number_density"]))
assert_array_equal(
ad["He_nuclei_density"],
(ad["He_p0_number_density"] + ad["He_p1_number_density"] +
ad["He_p2_number_density"]))
def _assert_05_075cm(self):
assert_array_equal(
[self.slc.xlim, self.slc.ylim, self.slc.width],
[
(YTQuantity(0.25, "cm"), YTQuantity(0.75, "cm")),
(YTQuantity(0.125, "cm"), YTQuantity(0.875, "cm")),
(YTQuantity(0.5, "cm"), YTQuantity(0.75, "cm")),
],
)
def test_profile_data():
tmpdir = make_tempdir()
curdir = os.getcwd()
os.chdir(tmpdir)
ds = data_dir_load(enzotiny)
ad = ds.all_data()
profile_1d = create_profile(ad,
"density",
"temperature",
weight_field="cell_mass")
fn = profile_1d.save_as_dataset()
full_fn = os.path.join(tmpdir, fn)
prof_1d_ds = load(full_fn)
compare_unit_attributes(ds, prof_1d_ds)
assert isinstance(prof_1d_ds, YTProfileDataset)
for field in profile_1d.standard_deviation:
assert_array_equal(
profile_1d.standard_deviation[field],
prof_1d_ds.profile.standard_deviation['data', field[1]])
p1 = ProfilePlot(prof_1d_ds.data,
"density",
"temperature",
weight_field="cell_mass")
p1.save()
yield YTDataFieldTest(full_fn, "temperature", geometric=False)
yield YTDataFieldTest(full_fn, "x", geometric=False)
yield YTDataFieldTest(full_fn, "density", geometric=False)
profile_2d = create_profile(ad, ["density", "temperature"],
"cell_mass",
weight_field=None,
n_bins=(128, 128))
fn = profile_2d.save_as_dataset()
full_fn = os.path.join(tmpdir, fn)
prof_2d_ds = load(full_fn)
compare_unit_attributes(ds, prof_2d_ds)
assert isinstance(prof_2d_ds, YTProfileDataset)
p2 = PhasePlot(prof_2d_ds.data,
"density",
"temperature",
"cell_mass",
weight_field=None)
p2.save()
yield YTDataFieldTest(full_fn, "density", geometric=False)
yield YTDataFieldTest(full_fn, "x", geometric=False)
yield YTDataFieldTest(full_fn, "temperature", geometric=False)
yield YTDataFieldTest(full_fn, "y", geometric=False)
yield YTDataFieldTest(full_fn, "cell_mass", geometric=False)
os.chdir(curdir)
if tmpdir != '.':
shutil.rmtree(tmpdir)
def test_2d_out():
ds = data_dir_load(twoD)
field_list = [('Hydrogen1+', 'particle_charge'),
('Hydrogen1+', 'particle_mass'),
('Hydrogen1+', 'particle_momentum_x'),
('Hydrogen1+', 'particle_momentum_y'),
('Hydrogen1+', 'particle_momentum_z'),
('Hydrogen1+', 'particle_positionCoarse_x'),
('Hydrogen1+', 'particle_positionCoarse_y'),
('Hydrogen1+', 'particle_positionCoarse_z'),
('Hydrogen1+', 'particle_positionOffset_x'),
('Hydrogen1+', 'particle_positionOffset_y'),
('Hydrogen1+', 'particle_positionOffset_z'),
('Hydrogen1+', 'particle_weighting'),
('all', 'particle_charge'), ('all', 'particle_mass'),
('all', 'particle_momentum_x'),
('all', 'particle_momentum_y'),
('all', 'particle_momentum_z'),
('all', 'particle_positionCoarse_x'),
('all', 'particle_positionCoarse_y'),
('all', 'particle_positionCoarse_z'),
('all', 'particle_positionOffset_x'),
('all', 'particle_positionOffset_y'),
('all', 'particle_positionOffset_z'),
('all', 'particle_weighting'),
('electrons', 'particle_charge'),
('electrons', 'particle_mass'),
('electrons', 'particle_momentum_x'),
('electrons', 'particle_momentum_y'),
('electrons', 'particle_momentum_z'),
('electrons', 'particle_positionCoarse_x'),
('electrons', 'particle_positionCoarse_y'),
('electrons', 'particle_positionCoarse_z'),
('electrons', 'particle_positionOffset_x'),
('electrons', 'particle_positionOffset_y'),
('electrons', 'particle_positionOffset_z'),
('electrons', 'particle_weighting'), ('openPMD', 'B_x'),
('openPMD', 'B_y'), ('openPMD', 'B_z'), ('openPMD', 'E_x'),
('openPMD', 'E_y'), ('openPMD', 'E_z'), ('openPMD', 'J_x'),
('openPMD', 'J_y'), ('openPMD', 'J_z'), ('openPMD', 'rho')]
domain_dimensions = [51, 201, 1] * np.ones_like(ds.domain_dimensions)
domain_width = [3.06e-05, 2.01e-05, 1e+0] * np.ones_like(
ds.domain_left_edge)
assert isinstance(ds, OpenPMDDataset)
assert_equal(str(ds), "data00000100.h5")
assert_equal(ds.dimensionality, 2)
assert_equal(ds.particle_types_raw, ('Hydrogen1+', 'electrons'))
assert "all" in ds.particle_unions
assert_array_equal(ds.field_list, field_list)
assert_array_equal(ds.domain_dimensions, domain_dimensions)
assert_almost_equal(ds.current_time,
3.29025596712e-14 * np.ones_like(ds.current_time))
assert_almost_equal(ds.domain_right_edge - ds.domain_left_edge,
domain_width)
def test_box_creation():
# test that creating a region with left and right edge
# with units works
ds = fake_random_ds(32, length_unit=2)
reg = ds.box([0, 0, 0] * cm, [2, 2, 2] * cm)
dens_units = reg[("gas", "density")]
reg = ds.box([0, 0, 0], [1, 1, 1])
dens_no_units = reg[("gas", "density")]
assert_array_equal(dens_units, dens_no_units)
def compare_light_ray_solutions(lr1, lr2):
assert len(lr1.light_ray_solution) == len(lr2.light_ray_solution)
if len(lr1.light_ray_solution) == 0:
return
for s1, s2 in zip(lr1.light_ray_solution, lr2.light_ray_solution):
for field in s1:
if field in ["next", "previous"]:
continue
if isinstance(s1[field], np.ndarray):
assert_array_equal(s1[field], s2[field])
else:
assert s1[field] == s2[field]
def test_out002():
ds = data_dir_load(out_s002)
field_list = [('all', 'conv_indicator'), ('all', 'conv_marker'),
('all', 'convected'), ('all', 'diffused'),
('connect1', 'conv_indicator'), ('connect1', 'conv_marker'),
('connect1', 'convected'), ('connect1', 'diffused'),
('connect2', 'conv_indicator'), ('connect2', 'conv_marker'),
('connect2', 'convected'), ('connect2', 'diffused')]
assert_equal(str(ds), "out.e-s002")
assert_equal(ds.dimensionality, 3)
assert_equal(ds.current_time, 2.0)
assert_array_equal(ds.field_list, field_list)
def test_get_morton_indices():
from yt.utilities.lib.geometry_utils import get_morton_indices, get_morton_indices_unravel
INDEX_MAX_64 = np.uint64(2097151)
li = np.arange(6, dtype=np.uint64).reshape((2, 3))
mi_ans = np.array([10, 229], dtype=np.uint64)
mi_out = get_morton_indices(li)
mi_out2 = get_morton_indices_unravel(li[:, 0], li[:, 1], li[:, 2])
assert_array_equal(mi_out, mi_ans)
assert_array_equal(mi_out2, mi_ans)
li[0, :] = INDEX_MAX_64 * np.ones(3, dtype=np.uint64)
assert_raises(ValueError, get_morton_indices, li)
assert_raises(ValueError, get_morton_indices_unravel, li[:, 0], li[:, 1],
li[:, 2])
def test_knn_direct(seed=1):
from yt.utilities.lib.geometry_utils import knn_direct
np.random.seed(seed)
k = 64
N = 1e5
idx = np.arange(N, dtype=np.uint64)
rad = np.arange(N, dtype=np.float64)
pos = np.vstack(3 * [rad**2 / 3.0]).T
sort_shf = np.arange(N, dtype=np.uint64)
for i in range(20):
np.random.shuffle(sort_shf)
sort_ans = np.argsort(sort_shf)[:k]
sort_out = knn_direct(pos[sort_shf, :], k, sort_ans[0], idx)
assert_array_equal(sort_out, sort_ans)
def test_halo_masses():
ds = data_dir_load(g298)
ad = ds.all_data()
for ptype in ["Group", "Subhalo"]:
nhalos = ds.index.particle_count[ptype]
mass = ds.arr(np.zeros(nhalos), "code_mass")
for i in range(nhalos):
halo = ds.halo(ptype, i)
mass[i] = halo.mass
# Check that masses from halo containers are the same
# as the array of all masses. This will test getting
# scalar fields for halos correctly.
assert_array_equal(ad[ptype, "particle_mass"], mass)
def test_get_vertex_centered_data():
ds = fake_random_ds(16)
g = ds.index.grids[0]
vec_list = g.get_vertex_centered_data([('gas', 'density')], no_ghost=True)
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter('always')
vec_str = g.get_vertex_centered_data('density', no_ghost=True)
assert len(w) == 1
assert issubclass(w[-1].category, DeprecationWarning)
assert 'requires list of fields' in str(w[-1].message)
vec_tuple = g.get_vertex_centered_data(('gas', 'density'), no_ghost=True)
assert_array_equal(vec_list[('gas', 'density')], vec_str)
assert_array_equal(vec_list[('gas', 'density')], vec_tuple)
def test_arbitrary_field_parameters():
def _test_field(field, data):
par = data.get_field_parameter('test_parameter')
return par * data['all', 'particle_mass']
ds = fake_random_ds(64, nprocs=8, particles=16**2)
ds.add_field(('all', 'test_field'), units='g',
function=_test_field, sampling_type='particle',
validators=[ValidateParameter('test_parameter')])
agrid = ds.arbitrary_grid([0.4, 0.4, 0.4], [0.99, 0.99, 0.99],
dims=[32, 32, 32])
agrid.set_field_parameter('test_parameter', 2)
assert_array_equal(2*agrid['all', 'particle_mass'], agrid['all', 'test_field'])
def test_plot_2d():
# Cartesian
ds = fake_random_ds((32, 32, 1), fields=('temperature', ), units=('K', ))
slc = SlicePlot(ds,
"z", ["temperature"],
width=(0.2, "unitary"),
center=[0.4, 0.3, 0.5])
slc2 = plot_2d(ds,
"temperature",
width=(0.2, "unitary"),
center=[0.4, 0.3])
slc3 = plot_2d(ds,
"temperature",
width=(0.2, "unitary"),
center=ds.arr([0.4, 0.3], "cm"))
assert_array_equal(slc.frb['temperature'], slc2.frb['temperature'])
assert_array_equal(slc.frb['temperature'], slc3.frb['temperature'])
# Cylindrical
ds = data_dir_load(WD)
slc = SlicePlot(ds, "theta", ["density"], width=(30000.0, "km"))
slc2 = plot_2d(ds, "density", width=(30000.0, "km"))
assert_array_equal(slc.frb['density'], slc2.frb['density'])
# Spherical
ds = data_dir_load(blast_wave)
slc = SlicePlot(ds, "phi", ["density"], width=(1, "unitary"))
slc2 = plot_2d(ds, "density", width=(1, "unitary"))
assert_array_equal(slc.frb['density'], slc2.frb['density'])
def test_out():
ds = data_dir_load(out)
field_list = [('all', 'conv_indicator'), ('all', 'conv_marker'),
('all', 'convected'), ('all', 'diffused'),
('connect1', 'conv_indicator'), ('connect1', 'conv_marker'),
('connect1', 'convected'), ('connect1', 'diffused'),
('connect2', 'conv_indicator'), ('connect2', 'conv_marker'),
('connect2', 'convected'), ('connect2', 'diffused')]
assert_equal(str(ds), "out.e")
assert_equal(ds.dimensionality, 3)
assert_equal(ds.current_time, 0.0)
assert_array_equal(ds.parameters['nod_names'], ['convected', 'diffused'])
assert_equal(ds.parameters['num_meshes'], 2)
assert_array_equal(ds.field_list, field_list)
def test_linear_interpolator_1d():
random_data = np.random.random(64)
fv = {'x': np.mgrid[0.0:1.0:64j]}
# evenly spaced bins
ufi = lin.UnilinearFieldInterpolator(random_data, (0.0, 1.0), "x", True)
assert_array_equal(ufi(fv), random_data)
# randomly spaced bins
size = 64
shift = (1. / size) * np.random.random(size) - (0.5 / size)
fv["x"] += shift
ufi = lin.UnilinearFieldInterpolator(random_data,
np.linspace(0.0, 1.0, size) + shift,
"x", True)
assert_array_almost_equal(ufi(fv), random_data, 15)
def test_refine_by():
ds = setup_fake_refby()
dd = ds.all_data()
# This checks that we always refine_by 1 in dimensions 2 and 3
dims = ds.domain_dimensions * ds.refine_by**ds.max_level
for i in range(1, 3):
# Check the refine_by == 1
ncoords = np.unique(dd.icoords[:, i]).size
assert_equal(ncoords, dims[i])
for g in ds.index.grids:
dims = ds.domain_dimensions * ds.refine_by**g.Level
# Now we can check converting back to the reference space
v = ((g.icoords + 1) / dims.astype("f8")).max(axis=0)
v *= ds.domain_width
assert_array_equal(v, g.RightEdge.d)
def test_non_square_frb():
tmpdir = tempfile.mkdtemp()
curdir = os.getcwd()
os.chdir(tmpdir)
# construct an arbitrary dataset
arr = np.arange(8.0 * 9.0 * 10.0).reshape((8, 9, 10))
data = dict(density=(arr, "g/cm**3"))
bbox = np.array([[-4, 4.0], [-4.5, 4.5], [-5.0, 5]])
ds = load_uniform_grid(data,
arr.shape,
length_unit="Mpc",
bbox=bbox,
periodicity=(False, False, False))
# make a slice
slc = ds.slice(axis="z", coord=ds.quan(0.0, "code_length"))
# make a frb and save it to disk
center = (ds.quan(0.0, "code_length"), ds.quan(0.0, "code_length"))
xax, yax = ds.coordinates.x_axis[slc.axis], ds.coordinates.y_axis[slc.axis]
res = [ds.domain_dimensions[xax], ds.domain_dimensions[yax]] # = [8,9]
width = ds.domain_right_edge[xax] - ds.domain_left_edge[
xax] # = 8 code_length
height = ds.domain_right_edge[yax] - ds.domain_left_edge[
yax] # = 9 code_length
frb = slc.to_frb(width=width, height=height, resolution=res, center=center)
fname = "test_frb_roundtrip.h5"
frb.save_as_dataset(fname, fields=[("gas", "density")])
expected_vals = arr[:, :, 5].T
print(
"\nConfirmation that initial frb results are expected:",
(expected_vals == frb[("gas", "density")].v).all(),
"\n",
)
# yt-reload:
reloaded_ds = load(fname)
assert_array_equal(frb[("gas", "density")].shape,
reloaded_ds.data[("gas", "density")].shape)
assert_array_equal(frb[("gas", "density")],
reloaded_ds.data[("gas", "density")])
os.chdir(curdir)
if tmpdir != ".":
shutil.rmtree(tmpdir)
def test_grid_datacontainer_data():
tmpdir = make_tempdir()
curdir = os.getcwd()
os.chdir(tmpdir)
ds = data_dir_load(enzotiny)
cg = ds.covering_grid(level=0, left_edge=[0.25] * 3, dims=[16] * 3)
fn = cg.save_as_dataset(fields=[
("gas", "density"),
("all", "particle_mass"),
("all", "particle_position"),
])
full_fn = os.path.join(tmpdir, fn)
cg_ds = load(full_fn)
compare_unit_attributes(ds, cg_ds)
assert isinstance(cg_ds, YTGridDataset)
assert (
cg["all",
"particle_position"].shape == cg_ds.r["all",
"particle_position"].shape)
yield YTDataFieldTest(full_fn, ("grid", "density"))
yield YTDataFieldTest(full_fn, ("all", "particle_mass"))
ag = ds.arbitrary_grid(left_edge=[0.25] * 3,
right_edge=[0.75] * 3,
dims=[16] * 3)
fn = ag.save_as_dataset(fields=[("gas", "density"), ("all",
"particle_mass")])
full_fn = os.path.join(tmpdir, fn)
ag_ds = load(full_fn)
compare_unit_attributes(ds, ag_ds)
assert isinstance(ag_ds, YTGridDataset)
yield YTDataFieldTest(full_fn, ("grid", "density"))
yield YTDataFieldTest(full_fn, ("all", "particle_mass"))
my_proj = ds.proj(("gas", "density"), "x", weight_field=("gas", "density"))
frb = my_proj.to_frb(1.0, (800, 800))
fn = frb.save_as_dataset(fields=[("gas", "density")])
frb_ds = load(fn)
assert_array_equal(frb[("gas", "density")],
frb_ds.data[("gas", "density")])
compare_unit_attributes(ds, frb_ds)
assert isinstance(frb_ds, YTGridDataset)
yield YTDataFieldTest(full_fn, ("grid", "density"), geometric=False)
os.chdir(curdir)
if tmpdir != ".":
shutil.rmtree(tmpdir)
def test_clump_finding():
n_c = 8
n_p = 1
dims = (n_c, n_c, n_c)
density = np.ones(dims)
high_rho = 10.
# add a couple disconnected density enhancements
density[2, 2, 2] = high_rho
density[6, 6, 6] = high_rho
# put a particle at the center of one of them
dx = 1. / n_c
px = 2.5 * dx * np.ones(n_p)
data = {
"density": density,
"particle_mass": np.ones(n_p),
"particle_position_x": px,
"particle_position_y": px,
"particle_position_z": px
}
ds = load_uniform_grid(data, dims)
ad = ds.all_data()
master_clump = Clump(ad, ("gas", "density"))
master_clump.add_validator("min_cells", 1)
find_clumps(master_clump, 0.5, 2. * high_rho, 10.)
# there should be two children
assert_equal(len(master_clump.children), 2)
leaf_clumps = get_lowest_clumps(master_clump)
# two leaf clumps
assert_equal(len(leaf_clumps), 2)
# check some clump fields
assert_equal(master_clump.children[0]["density"][0].size, 1)
assert_equal(master_clump.children[0]["density"][0], ad["density"].max())
assert_equal(master_clump.children[0]["particle_mass"].size, 1)
assert_array_equal(master_clump.children[0]["particle_mass"],
ad["particle_mass"])
assert_equal(master_clump.children[1]["density"][0].size, 1)
assert_equal(master_clump.children[1]["density"][0], ad["density"].max())
assert_equal(master_clump.children[1]["particle_mass"].size, 0)
def test_inout_bitarray():
# Check that we can do it for bitarrays that are funny-shaped
for i in range(7):
# Check we can feed in an array
arr_in = (np.random.random(32**3 + i) > 0.5)
b = ba.bitarray(arr = arr_in)
if i > 0:
assert_equal(b.ibuf.size, (32**3)/8.0 + 1)
arr_out = b.as_bool_array()
assert_equal(arr_in, arr_out)
# Let's check we can do it without feeding it at first
b = ba.bitarray(size = arr_in.size)
b.set_from_array(arr_in)
arr_out = b.as_bool_array()
assert_equal(arr_in, arr_out)
# Try a big array
arr_in = (np.random.random(32**3 + i) > 0.5)
b = ba.bitarray(arr = arr_in)
arr_out = b.as_bool_array()
assert_equal(arr_in, arr_out)
# Let's check we can do something interesting.
arr_in1 = (np.random.random(32**3) > 0.5)
arr_in2 = (np.random.random(32**3) > 0.5)
b1 = ba.bitarray(arr = arr_in1)
b2 = ba.bitarray(arr = arr_in2)
b3 = ba.bitarray(arr = (arr_in1 & arr_in2))
assert_equal((b1.ibuf & b2.ibuf), b3.ibuf)
b = ba.bitarray(10)
for i in range(10):
b.set_value(i, 2) # 2 should evaluate to True
arr = b.as_bool_array()
assert_equal(arr[:i+1].all(), True)
assert_equal(arr[i+1:].any(), False)
for i in range(10):
b.set_value(i, 0)
arr = b.as_bool_array()
assert_equal(arr.any(), False)
b.set_value(7, 1)
arr = b.as_bool_array()
assert_array_equal(arr, [0, 0, 0, 0, 0, 0, 0, 1, 0, 0])
b.set_value(2, 1)
arr = b.as_bool_array()
assert_array_equal(arr, [0, 0, 1, 0, 0, 0, 0, 1, 0, 0])
def test_clump_tree_save():
tmpdir = tempfile.mkdtemp()
curdir = os.getcwd()
os.chdir(tmpdir)
ds = data_dir_load(i30)
data_source = ds.disk([0.5, 0.5, 0.5], [0.0, 0.0, 1.0], (8, "kpc"),
(1, "kpc"))
field = ("gas", "density")
step = 2.0
c_min = 10**np.floor(np.log10(data_source[field]).min())
c_max = 10**np.floor(np.log10(data_source[field]).max() + 1)
master_clump = Clump(data_source, field)
master_clump.add_info_item("center_of_mass")
master_clump.add_validator("min_cells", 20)
find_clumps(master_clump, c_min, c_max, step)
leaf_clumps = master_clump.leaves
fn = master_clump.save_as_dataset(fields=[
("gas", "density"),
("index", "x"),
("index", "y"),
("index", "z"),
("all", "particle_mass"),
])
ds2 = load(fn)
# compare clumps in the tree
t1 = [c for c in master_clump]
t2 = [c for c in ds2.tree]
mt1 = ds.arr([c.info["cell_mass"][1] for c in t1])
mt2 = ds2.arr([c["clump", "cell_mass"] for c in t2])
it1 = np.array(np.argsort(mt1).astype(int))
it2 = np.array(np.argsort(mt2).astype(int))
assert_array_equal(mt1[it1], mt2[it2])
for i1, i2 in zip(it1, it2):
ct1 = t1[i1]
ct2 = t2[i2]
assert_array_equal(ct1["gas", "density"], ct2["grid", "density"])
assert_array_equal(ct1["all", "particle_mass"], ct2["all",
"particle_mass"])
# compare leaf clumps
c1 = [c for c in leaf_clumps]
c2 = [c for c in ds2.leaves]
mc1 = ds.arr([c.info["cell_mass"][1] for c in c1])
mc2 = ds2.arr([c["clump", "cell_mass"] for c in c2])
ic1 = np.array(np.argsort(mc1).astype(int))
ic2 = np.array(np.argsort(mc2).astype(int))
assert_array_equal(mc1[ic1], mc2[ic2])
os.chdir(curdir)
shutil.rmtree(tmpdir)
def test_halofinder_ptype(self):
ds = data_dir_load(enzotiny)
ds.add_particle_filter("dm")
for method in ["fof", "hop"]:
hc = HaloCatalog(data_ds=ds, finder_method=method,
output_dir="hc1",
finder_kwargs={"dm_only": True})
hc.create()
hc = HaloCatalog(data_ds=ds, finder_method=method,
output_dir="hc2",
finder_kwargs={"dm_only": False, "ptype": "dm"})
hc.create()
ds1 = load("hc1/hc1.0.h5")
ds2 = load("hc2/hc2.0.h5")
assert_array_equal(ds1.r["particle_mass"], ds2.r["particle_mass"])
def test_decomposition_2d():
array = np.ones((7, 5, 1))
bbox = np.array([[-0.7, 0.0], [1.5, 2.0], [0.0, 0.7]])
ledge, redge, shapes, slices = dec.decompose_array(array.shape,
np.array([2, 3, 1]),
bbox)
data = [array[slice] for slice in slices]
assert_array_equal(data[1].shape, np.array([3, 2, 1]))
gold_le = np.array([[-0.7, 1.5, 0.0], [-0.7, 1.6, 0.0], [-0.7, 1.8, 0.0],
[-0.4, 1.5, 0.0], [-0.4, 1.6, 0.0], [-0.4, 1.8, 0.0]])
assert_almost_equal(ledge, gold_le, 8)
gold_re = np.array([[-0.4, 1.6, 0.7], [-0.4, 1.8, 0.7], [-0.4, 2.0, 0.7],
[0.0, 1.6, 0.7], [0.0, 1.8, 0.7], [0.0, 2.0, 0.7]])
assert_almost_equal(redge, gold_re, 8)
def test_bitmask_pool():
bmp = BitmaskPool()
assert_equal(len(bmp), 0)
bmp.append(100)
assert_equal(len(bmp), 1)
assert_equal(bmp[0].size, 100)
bmp.append(200)
assert_equal(len(bmp), 2)
assert_equal(bmp[0].size, 100)
assert_equal(bmp[1].size, 200)
assert_equal(sum(_.size for _ in bmp.to_arrays()), 300)
arrs = bmp.to_arrays()
assert_equal(arrs[0].size, 100)
assert_equal(arrs[1].size, 200)
arrs[0][:] = 1
arrs = bmp.to_arrays()
assert_array_equal(arrs[0], 1)
def test_arbitrary_field_parameters():
def _test_field(field, data):
par = data.get_field_parameter("test_parameter")
return par * data["all", "particle_mass"]
ds = fake_random_ds(64, nprocs=8, particles=16 ** 2)
ds.add_field(
("all", "test_field"),
units="g",
function=_test_field,
sampling_type="particle",
validators=[ValidateParameter("test_parameter")],
)
agrid = ds.arbitrary_grid([0.4, 0.4, 0.4], [0.99, 0.99, 0.99], dims=[32, 32, 32])
agrid.set_field_parameter("test_parameter", 2)
assert_array_equal(2 * agrid["all", "particle_mass"], agrid["all", "test_field"])
def test_no_particles_out():
ds = data_dir_load(noParticles)
field_list = [('openPMD', 'E_x'), ('openPMD', 'E_y'), ('openPMD', 'E_z'),
('openPMD', 'rho')]
domain_dimensions = [51, 201, 1] * np.ones_like(ds.domain_dimensions)
domain_width = [3.06e-05, 2.01e-05, 1e+0] * np.ones_like(
ds.domain_left_edge)
assert isinstance(ds, OpenPMDDataset)
assert_equal(str(ds), "data00000400.h5")
assert_equal(ds.dimensionality, 2)
assert_equal(ds.particle_types_raw, ('io', ))
assert "all" not in ds.particle_unions
assert_array_equal(ds.field_list, field_list)
assert_array_equal(ds.domain_dimensions, domain_dimensions)
assert_almost_equal(ds.current_time,
1.3161023868481013e-13 * np.ones_like(ds.current_time))
assert_almost_equal(ds.domain_right_edge - ds.domain_left_edge,
domain_width)
def test_no_fields_out():
ds = data_dir_load(noFields)
particle_types = ("all", "io", "nbody")
no_fields_pfields = sorted(particle_fields + ["particle_id"])
field_list = list(product(particle_types, no_fields_pfields))
domain_dimensions = [1, 1, 1] * np.ones_like(ds.domain_dimensions)
domain_width = [1, 1, 1] * np.ones_like(ds.domain_left_edge)
assert isinstance(ds, OpenPMDDataset)
assert_equal(str(ds), "data00000400.h5")
assert_equal(ds.dimensionality, 3)
assert_equal(ds.particle_types_raw, ("io", ))
assert "all" in ds.particle_unions
assert_array_equal(ds.field_list, field_list)
assert_array_equal(ds.domain_dimensions, domain_dimensions)
assert_almost_equal(ds.current_time,
1.3161023868481013e-13 * np.ones_like(ds.current_time))
assert_almost_equal(ds.domain_right_edge - ds.domain_left_edge,
domain_width)
def test_psize_2d():
procs = dec.get_psize(np.array([5, 1, 7]), 6)
assert_array_equal(procs, np.array([3, 1, 2]))
procs = dec.get_psize(np.array([1, 7, 5]), 6)
assert_array_equal(procs, np.array([1, 2, 3]))
procs = dec.get_psize(np.array([7, 5, 1]), 6)
assert_array_equal(procs, np.array([2, 3, 1]))
def test_decomposition_2d():
array = np.ones((7, 5, 1))
bbox = np.array([[-0.7, 0.0], [1.5, 2.0], [0.0, 0.7]])
ledge, redge, shapes, slices = dec.decompose_array(array.shape,
np.array([2, 3, 1]), bbox)
data = [array[slice] for slice in slices]
assert_array_equal(data[1].shape, np.array([3, 2, 1]))
gold_le = np.array([
[-0.7, 1.5, 0.0], [-0.7, 1.6, 0.0],
[-0.7, 1.8, 0.0], [-0.4, 1.5, 0.0],
[-0.4, 1.6, 0.0], [-0.4, 1.8, 0.0]
])
assert_almost_equal(ledge, gold_le, 8)
gold_re = np.array(
[[-0.4, 1.6, 0.7], [-0.4, 1.8, 0.7],
[-0.4, 2.0, 0.7], [0.0, 1.6, 0.7],
[0.0, 1.8, 0.7], [0.0, 2.0, 0.7]]
)
assert_almost_equal(redge, gold_re, 8)
def test_psize_3d():
procs = dec.get_psize(np.array([33, 35, 37]), 12)
assert_array_equal(procs, np.array([3, 2, 2]))
def _assert_05_075cm(self):
assert_array_equal([self.slc.xlim, self.slc.ylim, self.slc.width],
[(YTQuantity(0.25, 'cm'), YTQuantity(0.75, 'cm')),
(YTQuantity(0.125, 'cm'), YTQuantity(0.875, 'cm')),
(YTQuantity(0.5, 'cm'), YTQuantity(0.75, 'cm'))])
