def test_particle_plot_wf(self):
test_ds = fake_particle_ds()
for dim in range(3):
for weight_field in WEIGHT_FIELDS:
pplot_wf = ParticleProjectionPlot(
test_ds, dim, "particle_mass", weight_field=weight_field)
pplot_wf.save()
def test_center_squeeze():
# checks that the center is reshaped correctly
# create and test amr, random and particle data
check_single_ds(fake_amr_ds(fields=("Density", )))
check_single_ds(fake_random_ds(16, fields=("Density", )))
check_single_ds(fake_particle_ds(npart=100))
def setUpClass(cls):
test_ds = fake_particle_ds()
data_sources = [test_ds.region([0.5] * 3, [0.4] * 3, [0.6] * 3),
test_ds.all_data()]
particle_phases = []
for source in data_sources:
for x_field, y_field, z_fields in PHASE_FIELDS:
particle_phases.append(ParticlePhasePlot(source,
x_field,
y_field,
z_fields,
x_bins=16,
y_bins=16))
particle_phases.append(ParticlePhasePlot(source,
x_field,
y_field,
z_fields,
x_bins=16,
y_bins=16,
deposition='cic'))
pp = create_profile(source, [x_field, y_field], z_fields,
weight_field='particle_ones',
n_bins=[16, 16])
particle_phases.append(ParticlePhasePlot.from_profile(pp))
cls.particle_phases = particle_phases
cls.ds = test_ds
def test_derived_field(self):
# Test that derived field on filtered particles do not require
# their parent field to be created
ds = fake_particle_ds()
dd = ds.all_data()
@particle_filter(requires=["particle_mass"], filtered_type="io")
def massive(pfilter, data):
return data[(pfilter.filtered_type, "particle_mass")].to("code_mass") > 0.5
ds.add_particle_filter("massive")
def fun(field, data):
return data[field.name[0], "particle_mass"]
# Add the field to the massive particles
ds.add_field(
("massive", "test"),
function=fun,
sampling_type="particle",
units="code_mass",
)
expected_size = (dd["io", "particle_mass"].to("code_mass") > 0.5).sum()
fields_to_test = [f for f in ds.derived_field_list if f[0] == "massive"]
def test_this(fname):
data = dd[fname]
assert_equal(data.shape[0], expected_size)
for fname in fields_to_test:
test_this(fname)
def test_load_particles_with_data_source():
ds1 = fake_particle_ds()
# Load from dataset
ad = ds1.all_data()
fields = ["particle_mass"]
fields += [f"particle_position_{ax}" for ax in "xyz"]
data = {field: ad[field] for field in fields}
ds2 = load_particles(data, data_source=ad)
def in_cgs(quan):
return quan.in_cgs().v
# Test bbox is parsed correctly
for attr in ["domain_left_edge", "domain_right_edge"]:
assert np.allclose(in_cgs(getattr(ds1, attr)),
in_cgs(getattr(ds2, attr)))
# Test sim_time is parsed correctly
assert in_cgs(ds1.current_time) == in_cgs(ds2.current_time)
# Test code units are parsed correctly
def get_cu(ds, dim):
return ds.quan(1, "code_" + dim)
for dim in ["length", "mass", "time", "velocity", "magnetic"]:
assert in_cgs(get_cu(ds1, dim)) == in_cgs(get_cu(ds2, dim))
def test_write_out(self):
filename = "sphere.txt"
ds = fake_particle_ds()
sp = ds.sphere(ds.domain_center, 0.25)
sp.write_out(filename, fields=["cell_volume"])
with open(filename, "r") as file:
file_row_1 = file.readline()
file_row_2 = file.readline()
file_row_2 = np.array(file_row_2.split('\t'), dtype=np.float64)
sorted_keys = sorted(sp.field_data.keys())
keys = [str(k) for k in sorted_keys]
keys = "\t".join(["#"] + keys + ["\n"])
data = [sp.field_data[k][0] for k in sorted_keys]
assert_equal(keys, file_row_1)
assert_array_equal(data, file_row_2)
# Test for exception
with assert_raises(YTException) as ex:
sp.write_out(filename, fields=["particle_position_x"])
desired = (
"Field type ['all'] of the supplied field ['particle_position_x']"
" is in consistent with field type 'gas'.")
assert_equal(str(ex.exception)[:50], desired[:50])
def test_neighbor_value_search():
np.random.seed(0x4d3d3d3)
ds = fake_particle_ds(npart=16**3, over_refine_factor=2)
ds.periodicity = (True, True, True)
ds.index
fn, = add_nearest_neighbor_value_field("all", "particle_position",
"particle_mass", ds.field_info)
dd = ds.all_data()
# Set up our positions onto which the field will be deposited
index_pos = np.array([dd["index", ax]
for ax in 'xyz']) * dd["index", "x"].uq
particle_pos = dd["particle_position"]
values_in = dd["particle_mass"]
values_out = dd[fn]
for i in range(index_pos.shape[0]):
r2 = particle_pos[:, 0] * 0
r2 = r2 * r2
for j in range(3):
DR = (index_pos[i, j] - particle_pos[:, j])
DRo = DR.copy()
DR[DRo > ds.domain_width[j] / 2.0] -= ds.domain_width[j]
DR[DRo < -ds.domain_width[j] / 2.0] += ds.domain_width[j]
r2 += DR * DR
radius = np.sqrt(r2)
assert (values_in[np.argmin(radius)] == values_out[i])
def test_extract_isocontours(self):
# Test isocontour properties for AMRGridData
fields = ["density", "cell_mass"]
units = ["g/cm**3", "g"]
ds = fake_amr_ds(fields=fields, units=units, particles=16**3)
dd = ds.all_data()
q = dd.quantities["WeightedAverageQuantity"]
rho = q(("gas", "density"), weight=("gas", "cell_mass"))
dd.extract_isocontours(("gas", "density"), rho, "triangles.obj", True)
dd.calculate_isocontour_flux(
("gas", "density"),
rho,
("index", "x"),
("index", "y"),
("index", "z"),
("index", "dx"),
)
# Test error in case of ParticleData
ds = fake_particle_ds()
dd = ds.all_data()
q = dd.quantities["WeightedAverageQuantity"]
rho = q(("all", "particle_velocity_x"),
weight=("all", "particle_mass"))
with assert_raises(NotImplementedError):
dd.extract_isocontours("density", rho, sample_values="x")
def test_particle_phase_plot(self):
test_ds = fake_particle_ds()
data_sources = [test_ds.region([0.5] * 3, [0.4] * 3, [0.6] * 3),
test_ds.all_data()]
particle_phases = []
for source in data_sources:
for x_field, y_field, z_fields in PHASE_FIELDS:
particle_phases.append(ParticlePhasePlot(source,
x_field,
y_field,
z_fields,
x_bins=16,
y_bins=16))
particle_phases.append(ParticlePhasePlot(source,
x_field,
y_field,
z_fields,
x_bins=16,
y_bins=16,
deposition='cic'))
pp = create_profile(source, [x_field, y_field], z_fields,
weight_field='particle_ones',
n_bins=[16, 16])
particle_phases.append(ParticlePhasePlot.from_profile(pp))
particle_phases[0]._repr_html_()
for p in particle_phases:
for fname in TEST_FLNMS:
assert assert_fname(p.save(fname)[0])
def test_load_particles_with_data_source():
ds1 = fake_particle_ds()
# Load from dataset
ad = ds1.all_data()
fields = ['particle_mass']
fields += ['particle_position_{}'.format(ax) for ax in 'xyz']
data = {field: ad[field] for field in fields}
ds2 = load_particles(data, data_source=ad)
def in_cgs(quan):
return quan.in_cgs().v
# Test bbox is parsed correctly
for attr in ['domain_left_edge', 'domain_right_edge']:
assert np.allclose(
in_cgs(getattr(ds1, attr)),
in_cgs(getattr(ds2, attr))
)
# Test sim_time is parsed correctly
assert in_cgs(ds1.current_time) == in_cgs(ds2.current_time)
# Test code units are parsed correctly
def get_cu(ds, dim):
return ds.quan(1, 'code_' + dim)
for dim in ['length', 'mass', 'time', 'velocity', 'magnetic']:
assert in_cgs(get_cu(ds1, dim)) == in_cgs(get_cu(ds2, dim))
def test_derived_quantities_with_particle_types():
ds = fake_particle_ds()
@particle_filter(requires=["particle_position_x"], filtered_type='all')
def low_x(pfilter,data):
return data['particle_position_x'].in_units('code_length')<0.5
ds.add_particle_filter('low_x')
ad=ds.all_data()
for ptype in ['all','low_x']:
#Check bulk velocity
bulk_vx=(ad[(ptype,'particle_mass')]*ad[(ptype,'particle_velocity_x')]/ad[(ptype,'particle_mass')].sum()).sum()
assert_almost_equal(ad.quantities.bulk_velocity(use_gas=False,use_particles=True,particle_type=ptype)[0],bulk_vx,5)
#Check center of mass
com_x=(ad[(ptype,'particle_mass')]*ad[(ptype,'particle_position_x')]/ad[(ptype,'particle_mass')].sum()).sum()
assert_almost_equal(ad.quantities.center_of_mass(use_gas=False,use_particles=True,particle_type=ptype)[0],com_x,5)
#Check angular momentum vector
l_x=(ad[(ptype,'particle_specific_angular_momentum_x')]*ad[(ptype,'particle_mass')]/ad[(ptype,'particle_mass')].sum()).sum()
assert_almost_equal(ad.quantities.angular_momentum_vector(use_gas=False,use_particles=True,particle_type=ptype)[0],l_x,5)
#Check spin parameter values
assert_almost_equal(ad.quantities.spin_parameter(use_gas=False,use_particles=True),655.7311454765503)
assert_almost_equal(ad.quantities.spin_parameter(use_gas=False,use_particles=True,particle_type='low_x'),1309.164886405665)
def test_ptype():
n_particles = 100
fields = ['particle_position_x', 'particle_position_y', 'particle_position_z', 'particle_index',
'particle_dummy']
negative = [False, False, False, False, False]
units = ['cm', 'cm', 'cm', '1', '1']
# Setup filters on the 'particle_dummy' field, keeping only the first 50
@particle_filter(name='dummy', requires=["particle_dummy"])
def dummy(pfilter, data):
return data[(pfilter.filtered_type, "particle_dummy")] <= n_particles // 2
# Setup fake particle datasets with repeated ids. This should work because
# the ids are unique among `dummy_particles` so let's test this
data = {'particle_index': np.arange(n_particles) % (n_particles // 2),
'particle_dummy': np.arange(n_particles)}
all_ds = [fake_particle_ds(fields=fields, negative=negative, units=units,
npart=n_particles, data=data)]
for ds in all_ds:
ds.add_particle_filter('dummy')
ts = DatasetSeries(all_ds)
# Select all dummy particles
print(ts[0].derived_field_list)
ids = ts[0].all_data()['dummy', 'particle_index']
# Build trajectories
ts.particle_trajectories(ids, ptype='dummy')
def test_uniqueness():
n_particles = 2
n_steps = 2
ids = np.arange(n_particles, dtype=int) % (n_particles // 2)
data = {"particle_index": ids}
fields = [
"particle_position_x",
"particle_position_y",
"particle_position_z",
"particle_index",
]
negative = [False, False, False, False]
units = ["cm", "cm", "cm", "1"]
ts = DatasetSeries(
[
fake_particle_ds(
fields=fields,
negative=negative,
units=units,
npart=n_particles,
data=data,
)
for i in range(n_steps)
]
)
assert_raises(YTIllDefinedParticleData, ts.particle_trajectories, [0])
def test_derived_field(self):
# Test that derived field on filtered particles do not require
# their parent field to be created
ds = fake_particle_ds()
dd = ds.all_data()
@particle_filter(requires=['particle_mass'], filtered_type='io')
def massive(pfilter, data):
return data[(pfilter.filtered_type, 'particle_mass')].to('code_mass') > 0.5
ds.add_particle_filter('massive')
def fun(field, data):
return data[field.name[0], 'particle_mass']
# Add the field to the massive particles
ds.add_field(('massive', 'test'), function=fun,
sampling_type='particle', units='code_mass')
expected_size = (dd['io', 'particle_mass'].to('code_mass') > 0.5).sum()
fields_to_test = (f for f in ds.derived_field_list
if f[0] == 'massive')
def test_this(fname):
data = dd[fname]
assert_equal(data.shape[0], expected_size)
for fname in fields_to_test:
test_this(fname)
def test_add_field_string_aliasing():
ds = fake_random_ds(16)
def density_alias(field, data):
return data['density']
ds.add_field('density_alias',
sampling_type='cell',
function=density_alias,
units='g/cm**3')
ds.field_info['density_alias']
ds.field_info['gas', 'density_alias']
ds = fake_particle_ds()
def pmass_alias(field, data):
return data['particle_mass']
ds.add_field('particle_mass_alias',
function=pmass_alias,
sampling_type='particle',
units='g')
ds.field_info['particle_mass_alias']
ds.field_info['all', 'particle_mass_alias']
def test_pickleability():
# tests the pickleability of the selection objects.
def pickle_test(sel_obj):
assert_fields = sel_obj._get_state_attnames(
) # the attrs used in get/set state
new_sel_obj = pickle.loads(pickle.dumps(sel_obj))
for attr in assert_fields:
assert_equal(getattr(new_sel_obj, attr), getattr(sel_obj, attr))
# list of selection types and argument tuples for each selection type
c = np.array([0.5, 0.5, 0.5])
sargs = (
("point", (c, )),
("sphere", (c, 0.25)),
("box", (c - 0.3, c)),
("ellipsoid", (c, 0.3, 0.2, 0.1, c - 0.4, 0.2)),
("disk", (c, [1, 0, 0], 0.2, 0.2)),
("cutting", ([0.1, 0.2, -0.9], [0.5, 0.42, 0.6])),
("ortho_ray", ("z", c)),
("ray", (c, [0.1, 0.1, 0.1])),
)
# load fake data
ds = fake_particle_ds()
for sel_type, args in sargs:
sel = getattr(ds, sel_type)(*args) # instantiate this selection type
pickle_test(sel.selector) # make sure it (un)pickles
def test_neighbor_search():
np.random.seed(0x4d3d3d3)
ds = fake_particle_ds(npart=16**3)
ds.periodicity = (True, True, True)
ds.index
fn, = add_nearest_neighbor_field("all", "particle_position", ds)
dd = ds.all_data()
nearest_neighbors = dd[fn]
pos = dd["particle_position"]
all_neighbors = np.zeros_like(nearest_neighbors)
any_eq = np.zeros(pos.shape[0], dtype='bool')
min_in = np.zeros(pos.shape[0], dtype='int64')
for i in range(pos.shape[0]):
dd.set_field_parameter("center", pos[i, :])
#radius = dd["particle_radius"]
#radius.sort()
r2 = (pos[:, 0] * pos[:, 0]) * 0
for j in range(3):
DR = (pos[i, j] - pos[:, j])
DRo = DR.copy()
DR[DRo > ds.domain_width[j] / 2.0] -= ds.domain_width[j]
DR[DRo < -ds.domain_width[j] / 2.0] += ds.domain_width[j]
r2 += DR * DR
radius = np.sqrt(r2)
radius.sort()
assert (radius[0] == 0.0)
all_neighbors[i] = radius[63]
any_eq[i] = np.any(np.abs(radius - nearest_neighbors[i]) < 1e-7)
min_in[i] = np.argmin(np.abs(radius - nearest_neighbors[i]))
#if i == 34: raise RuntimeError
#dd.field_data.pop(("all", "particle_radius"))
assert_equal((min_in == 63).sum(), min_in.size)
assert_array_almost_equal(nearest_neighbors, all_neighbors)
def test_particle_plot_c(self):
test_ds = fake_particle_ds()
for center in CENTER_SPECS:
for dim in range(3):
pplot_c = ParticleProjectionPlot(
test_ds, dim, "particle_mass", center=center)
pplot_c.save()
def test_particle_plot_ds(self):
test_ds = fake_particle_ds()
ds_region = test_ds.region([0.5] * 3, [0.4] * 3, [0.6] * 3)
for dim in range(3):
pplot_ds = ParticleProjectionPlot(
test_ds, dim, "particle_mass", data_source=ds_region)
pplot_ds.save()
def test_add_field_string_aliasing():
ds = fake_random_ds(16)
def density_alias(field, data):
return data["density"]
ds.add_field("density_alias",
sampling_type="cell",
function=density_alias,
units="g/cm**3")
ds.field_info["density_alias"]
ds.field_info["gas", "density_alias"]
ds = fake_particle_ds()
def pmass_alias(field, data):
return data["particle_mass"]
ds.add_field("particle_mass_alias",
function=pmass_alias,
units="g",
sampling_type="particle")
ds.field_info["particle_mass_alias"]
ds.field_info["all", "particle_mass_alias"]
def test_update_data_particle():
npart = 100
ds = fake_particle_ds(npart=npart)
part_data = {"temperature": np.random.rand(npart)}
ds.index.update_data(part_data)
assert ("io", "temperature") in ds.field_list
dd = ds.all_data()
dd["temperature"]
def test_add_sph_fields():
ds = fake_particle_ds()
ds.index
assert set(ds.particle_types) == {'io', 'all', 'nbody'}
ds.add_sph_fields()
assert set(ds.particle_types) == {'io', 'all'}
assert ('io', 'smoothing_length') in ds.field_list
assert ('io', 'density') in ds.field_list
def test_add_sph_fields():
ds = fake_particle_ds()
ds.index
assert set(ds.particle_types) == {"io", "all", "nbody"}
ds.add_sph_fields()
assert set(ds.particle_types) == {"io", "all"}
assert ("io", "smoothing_length") in ds.field_list
assert ("io", "density") in ds.field_list
def test_firefly_JSON_object():
tmpdir = tempfile.mkdtemp()
ds = fake_particle_ds()
ad = ds.all_data()
reader = ad.create_firefly_object(
path_to_firefly = tmpdir,
velocity_units = 'cm/s',
coordinate_units = 'cm',
dataset_name='test')
reader.dumpToJSON()
def test_particle_plot_wf(self):
test_ds = fake_particle_ds()
for dim in range(3):
for weight_field in WEIGHT_FIELDS:
pplot_wf = ParticleProjectionPlot(
test_ds, dim, "particle_mass", weight_field=weight_field
)
with mock.patch(
"yt.visualization._mpl_imports.FigureCanvasAgg.print_figure"
):
pplot_wf.save()
def test_save_object(self):
ds = fake_particle_ds()
sp = ds.sphere(ds.domain_center, 0.25)
sp.save_object("my_sphere_1", filename="test_save_obj")
obj = shelve.open("test_save_obj", protocol=-1)
loaded_sphere = obj["my_sphere_1"][1]
obj.close()
assert_array_equal(loaded_sphere.center, sp.center)
assert_equal(loaded_sphere.radius, sp.radius)
for k in loaded_sphere._key_fields:
assert_array_equal(loaded_sphere[k], sp[k])
def test_particle_plot_ds(self):
test_ds = fake_particle_ds()
ds_region = test_ds.region([0.5] * 3, [0.4] * 3, [0.6] * 3)
for dim in range(3):
pplot_ds = ParticleProjectionPlot(test_ds,
dim, ("all", "particle_mass"),
data_source=ds_region)
with mock.patch(
"yt.visualization._mpl_imports.FigureCanvasAgg.print_figure"
):
pplot_ds.save()
def test_particle_plot_c(self):
test_ds = fake_particle_ds()
for center in CENTER_SPECS:
for dim in range(3):
pplot_c = ParticleProjectionPlot(test_ds,
dim, ("all", "particle_mass"),
center=center)
with mock.patch(
"yt.visualization._mpl_imports.FigureCanvasAgg.print_figure"
):
pplot_c.save()
def test_firefly_write_to_disk():
tmpdir = tempfile.mkdtemp()
ds = fake_particle_ds()
ad = ds.all_data()
reader = ad.create_firefly_object(
tmpdir,
velocity_units="cm/s",
coordinate_units="cm",
)
reader.dumpToJSON()
def test_creation_with_width(self):
test_ds = fake_particle_ds()
for width, (xlim, ylim, pwidth, _aun) in WIDTH_SPECS.items():
plot = ParticleProjectionPlot(test_ds, 0, "particle_mass", width=width)
xlim = [plot.ds.quan(el[0], el[1]) for el in xlim]
ylim = [plot.ds.quan(el[0], el[1]) for el in ylim]
pwidth = [plot.ds.quan(el[0], el[1]) for el in pwidth]
[assert_array_almost_equal(px, x, 14) for px, x in zip(plot.xlim, xlim)]
[assert_array_almost_equal(py, y, 14) for py, y in zip(plot.ylim, ylim)]
[assert_array_almost_equal(pw, w, 14) for pw, w in zip(plot.width, pwidth)]
def setUpClass(cls):
test_ds = fake_particle_ds()
ds_region = test_ds.region([0.5] * 3, [0.4] * 3, [0.6] * 3)
pplots = []
pplots_ds = []
pplots_c = []
pplots_wf = []
pplots_w = {}
pplots_m = []
for dim in range(3):
pplots.append(ParticleProjectionPlot(test_ds, dim,
"particle_mass"))
pplots_ds.append(ParticleProjectionPlot(test_ds, dim,
"particle_mass",
data_source=ds_region))
for center in CENTER_SPECS:
pplots_c.append(ParticleProjectionPlot(test_ds, dim,
"particle_mass",
center=center))
for width in WIDTH_SPECS:
pplots_w[width] = ParticleProjectionPlot(test_ds, dim,
'particle_mass',
width=width)
for wf in WEIGHT_FIELDS:
pplots_wf.append(ParticleProjectionPlot(test_ds, dim,
"particle_mass",
weight_field=wf))
cls.pplots = pplots
cls.pplots_ds = pplots_ds
cls.pplots_c = pplots_c
cls.pplots_wf = pplots_wf
cls.pplots_w = pplots_w
cls.pplots_m = pplots_m