def test_magnetic_fields():
ddims = (16, 16, 16)
data1 = {
"magnetic_field_x": (np.random.random(size=ddims), "T"),
"magnetic_field_y": (np.random.random(size=ddims), "T"),
"magnetic_field_z": (np.random.random(size=ddims), "T"),
}
data2 = {}
for field in data1:
data2[field] = (data1[field][0] * 1.0e4, "gauss")
ds1 = load_uniform_grid(data1, ddims, unit_system="cgs")
ds2 = load_uniform_grid(data2, ddims, unit_system="mks")
# For this test dataset, code units are cgs units
ds3 = load_uniform_grid(data2, ddims, unit_system="code")
ds1.index
ds2.index
ds3.index
dd1 = ds1.all_data()
dd2 = ds2.all_data()
dd3 = ds3.all_data()
assert ds1.fields.gas.magnetic_field_strength.units == "G"
assert ds1.fields.gas.magnetic_field_poloidal.units == "G"
assert ds1.fields.gas.magnetic_field_toroidal.units == "G"
assert ds2.fields.gas.magnetic_field_strength.units == "T"
assert ds2.fields.gas.magnetic_field_poloidal.units == "T"
assert ds2.fields.gas.magnetic_field_toroidal.units == "T"
assert ds3.fields.gas.magnetic_field_strength.units == "code_magnetic"
assert ds3.fields.gas.magnetic_field_poloidal.units == "code_magnetic"
assert ds3.fields.gas.magnetic_field_toroidal.units == "code_magnetic"
emag1 = (dd1["magnetic_field_x"]**2 + dd1["magnetic_field_y"]**2 +
dd1["magnetic_field_z"]**2) / (8.0 * np.pi)
emag1.convert_to_units("dyne/cm**2")
emag2 = (dd2["magnetic_field_x"]**2 + dd2["magnetic_field_y"]**2 +
dd2["magnetic_field_z"]**2) / (2.0 * mu_0)
emag2.convert_to_units("Pa")
emag3 = (dd3["magnetic_field_x"]**2 + dd3["magnetic_field_y"]**2 +
dd3["magnetic_field_z"]**2) / (2.0 * mu_0)
emag3.convert_to_units("code_pressure")
assert_almost_equal(emag1, dd1["magnetic_energy"])
assert_almost_equal(emag2, dd2["magnetic_energy"])
assert_almost_equal(emag3, dd3["magnetic_energy"])
assert str(emag1.units) == str(dd1["magnetic_energy"].units)
assert str(emag2.units) == str(dd2["magnetic_energy"].units)
assert str(emag3.units) == str(dd3["magnetic_energy"].units)
assert_almost_equal(emag1.in_cgs(), emag2.in_cgs())
assert_almost_equal(emag1.in_cgs(), emag3.in_cgs())
def test_ppv_nothermalbroad():
np.random.seed(seed=0x4D3D3D3)
dims = (16, 16, 128)
v_shift = 1.0e6 * u.cm / u.s
sigma_v = 2.0e6 * u.cm / u.s
data = {
"density": (np.ones(dims), "g/cm**3"),
"velocity_x": (np.zeros(dims), "cm/s"),
"velocity_y": (np.zeros(dims), "cm/s"),
"velocity_z": (
np.random.normal(loc=v_shift.v, scale=sigma_v.v, size=dims),
"cm/s",
),
}
ds = load_uniform_grid(data, dims)
cube = PPVCube(
ds,
"z",
("stream", "density"),
(-100.0, 100.0, 128, "km/s"),
dims=16,
thermal_broad=False,
)
dv = cube.dv
v_noth = np.sqrt(2) * (sigma_v).in_units("km/s")
a = cube.data.mean(axis=(0, 1)).v
b = (dv * np.exp(-(((cube.vmid + v_shift) / v_noth)**2)) /
(np.sqrt(np.pi) * v_noth))
assert_allclose_units(a, b, atol=5.0e-3)
def test_stream_non_cartesian_particles(loader):
eps = 1e-6
r, theta, phi = np.mgrid[0.0:1.0 - eps:64j, 0.0:np.pi - eps:64j,
0.0:2.0 * np.pi - eps:64j]
np.random.seed(0x4D3D3D3)
ind = np.random.randint(0, 64 * 64 * 64, size=1000)
particle_position_r = r.ravel()[ind]
particle_position_theta = theta.ravel()[ind]
particle_position_phi = phi.ravel()[ind]
ds = load_uniform_grid(
{
"density": r,
"temperature": phi,
"entropy": phi,
"particle_position_r": particle_position_r,
"particle_position_theta": particle_position_theta,
"particle_position_phi": particle_position_phi,
},
(64, 64, 64),
bbox=np.array([[0.0, 1.0], [0.0, np.pi], [0.0, 2.0 * np.pi]]),
geometry="spherical",
)
dd = ds.all_data()
assert_equal(dd["all", "particle_position_r"].v, particle_position_r)
assert_equal(dd["all", "particle_position_phi"].v, particle_position_phi)
assert_equal(dd["all", "particle_position_theta"].v,
particle_position_theta)
def test_on_off_compare():
# fake density field that varies in the x-direction only
den = np.arange(32**3) / 32**2 + 1
den = den.reshape(32, 32, 32)
den = np.array(den, dtype=np.float64)
data = dict(density=(den, "g/cm**3"))
bbox = np.array([[-1.5, 1.5], [-1.5, 1.5], [-1.5, 1.5]])
ds = load_uniform_grid(data,
den.shape,
length_unit="Mpc",
bbox=bbox,
nprocs=64)
sl_on = SlicePlot(ds, "z", [("gas", "density")])
L = [0, 0, 1]
north_vector = [0, 1, 0]
sl_off = OffAxisSlicePlot(ds,
L, ("gas", "density"),
center=[0, 0, 0],
north_vector=north_vector)
assert_array_almost_equal(sl_on.frb[("gas", "density")],
sl_off.frb[("gas", "density")])
sl_on.set_buff_size((800, 400))
sl_on._recreate_frb()
sl_off.set_buff_size((800, 400))
sl_off._recreate_frb()
assert_array_almost_equal(sl_on.frb[("gas", "density")],
sl_off.frb[("gas", "density")])
def test_ds_arr_invariance_under_projection_plot(tmp_path):
data_array = np.random.random((10, 10, 10))
bbox = np.array([[-100, 100], [-100, 100], [-100, 100]])
data = {("gas", "density"): (data_array, "g*cm**(-3)")}
ds = load_uniform_grid(data,
data_array.shape,
length_unit="kpc",
bbox=bbox)
start_source = np.array((0, 0, -0.5))
end_source = np.array((0, 0, 0.5))
start = ds.arr(start_source, "unitary")
end = ds.arr(end_source, "unitary")
start_i = start.copy()
end_i = end.copy()
p = ProjectionPlot(ds, 0, "number_density")
p.annotate_line(start, end)
p.save(tmp_path)
# for lack of a unyt.testing.assert_unit_array_equal function
np.testing.assert_array_equal(start_i, start)
assert start_i.units == start.units
np.testing.assert_array_equal(end_i, end)
assert end_i.units == end.units
def test_exclude_nan():
test_array = np.nan * np.ones((10, 10, 10))
test_array[1, 1, :] = 1
data = dict(density=test_array)
ds = load_uniform_grid(data, test_array.shape, length_unit="cm", nprocs=1)
ad = ds.all_data()
no_nan_ds = ad.exclude_nan(("gas", "density"))
assert_equal(no_nan_ds[("gas", "density")], np.array(np.ones(10)))
def fake_random_ds(
ndims,
peak_value=1.0,
fields=("density", "velocity_x", "velocity_y", "velocity_z"),
units=("g/cm**3", "cm/s", "cm/s", "cm/s"),
particle_fields=None,
particle_field_units=None,
negative=False,
nprocs=1,
particles=0,
length_unit=1.0,
unit_system="cgs",
bbox=None,
):
from yt.loaders import load_uniform_grid
prng = RandomState(0x4D3D3D3)
if not is_sequence(ndims):
ndims = [ndims, ndims, ndims]
else:
assert len(ndims) == 3
if not is_sequence(negative):
negative = [negative for f in fields]
assert len(fields) == len(negative)
offsets = []
for n in negative:
if n:
offsets.append(0.5)
else:
offsets.append(0.0)
data = {}
for field, offset, u in zip(fields, offsets, units):
v = (prng.random_sample(ndims) - offset) * peak_value
if field[0] == "all":
v = v.ravel()
data[field] = (v, u)
if particles:
if particle_fields is not None:
for field, unit in zip(particle_fields, particle_field_units):
if field in ("particle_position", "particle_velocity"):
data["io", field] = (prng.random_sample((int(particles), 3)), unit)
else:
data["io", field] = (prng.random_sample(size=int(particles)), unit)
else:
for f in (f"particle_position_{ax}" for ax in "xyz"):
data["io", f] = (prng.random_sample(size=particles), "code_length")
for f in (f"particle_velocity_{ax}" for ax in "xyz"):
data["io", f] = (prng.random_sample(size=particles) - 0.5, "cm/s")
data["io", "particle_mass"] = (prng.random_sample(particles), "g")
ug = load_uniform_grid(
data,
ndims,
length_unit=length_unit,
nprocs=nprocs,
unit_system=unit_system,
bbox=bbox,
)
return ug
def test_magnetic_code_units():
sqrt4pi = np.sqrt(4.0 * np.pi)
ddims = (16,) * 3
data = {"density": (np.random.uniform(size=ddims), "g/cm**3")}
ds1 = load_uniform_grid(
data, ddims, magnetic_unit=(sqrt4pi, "gauss"), unit_system="cgs"
)
assert_allclose(ds1.magnetic_unit.value, sqrt4pi)
assert str(ds1.magnetic_unit.units) == "G"
mucu = ds1.magnetic_unit.to("code_magnetic")
assert_allclose(mucu.value, 1.0)
assert str(mucu.units) == "code_magnetic"
ds2 = load_uniform_grid(data, ddims, magnetic_unit=(1.0, "T"), unit_system="cgs")
assert_allclose(ds2.magnetic_unit.value, 10000.0)
assert str(ds2.magnetic_unit.units) == "G"
mucu = ds2.magnetic_unit.to("code_magnetic")
assert_allclose(mucu.value, 1.0)
assert str(mucu.units) == "code_magnetic"
ds3 = load_uniform_grid(data, ddims, magnetic_unit=(1.0, "T"), unit_system="mks")
assert_allclose(ds3.magnetic_unit.value, 1.0)
assert str(ds3.magnetic_unit.units) == "T"
mucu = ds3.magnetic_unit.to("code_magnetic")
assert_allclose(mucu.value, 1.0)
assert str(mucu.units) == "code_magnetic"
ds4 = load_uniform_grid(
data, ddims, magnetic_unit=(1.0, "gauss"), unit_system="mks"
)
assert_allclose(ds4.magnetic_unit.value, 1.0e-4)
assert str(ds4.magnetic_unit.units) == "T"
mucu = ds4.magnetic_unit.to("code_magnetic")
assert_allclose(mucu.value, 1.0)
assert str(mucu.units) == "code_magnetic"
def test_equal():
test_array = np.ones((10, 10, 10))
test_array[1, 1, :] = 2.0
test_array[2, 1, :] = 3.0
data = dict(density=test_array)
ds = load_uniform_grid(data, test_array.shape, length_unit="cm", nprocs=1)
ad = ds.all_data()
no_ones = ad.exclude_equal(("gas", "density"), 1.0)
assert np.all(no_ones[("gas", "density")] != 1.0)
only_ones = ad.include_equal(("gas", "density"), 1.0)
assert np.all(only_ones[("gas", "density")] == 1.0)
def test_nan_data():
data = np.random.random((16, 16, 16)) - 0.5
data[:9, :9, :9] = np.nan
data = {"density": data}
ds = load_uniform_grid(data, [16, 16, 16])
plot = SlicePlot(ds, "z", ("gas", "density"))
with tempfile.NamedTemporaryFile(suffix="png") as f:
plot.save(f.name)
def test_symlog_extremely_small_vals():
# check that the plot can be constructed without crashing
# see https://github.com/yt-project/yt/issues/3858
shape = (64, 64, 1)
arr = np.full(shape, 5.0e-324)
arr[0, 0] = -1e12
arr[1, 1] = 200
d = {"scalar": arr}
ds = load_uniform_grid(d, shape)
p = SlicePlot(ds, "z", ("stream", "scalar"))
p["stream", "scalar"]
def test_ppv():
np.random.seed(seed=0x4D3D3D3)
dims = (8, 8, 128)
v_shift = 1.0e7 * u.cm / u.s
sigma_v = 2.0e7 * u.cm / u.s
T_0 = 1.0e8 * u.Kelvin
data = {
"density": (np.ones(dims), "g/cm**3"),
"temperature": (T_0.v * np.ones(dims), "K"),
"velocity_x": (np.zeros(dims), "cm/s"),
"velocity_y": (np.zeros(dims), "cm/s"),
"velocity_z": (
np.random.normal(loc=v_shift.v, scale=sigma_v.v, size=dims),
"cm/s",
),
}
ds = load_uniform_grid(data, dims)
cube = PPVCube(
ds,
"z",
("stream", "density"),
(-300.0, 300.0, 1024, "km/s"),
dims=8,
thermal_broad=True,
)
dv = cube.dv
v_th = np.sqrt(2.0 * kboltz * T_0 / (56.0 * mh) +
2.0 * sigma_v**2).in_units("km/s")
a = cube.data.mean(axis=(0, 1)).v
b = dv * np.exp(-((
(cube.vmid + v_shift) / v_th)**2)) / (np.sqrt(np.pi) * v_th)
assert_allclose_units(a, b, 1.0e-2)
E_0 = 6.8 * u.keV
cube.transform_spectral_axis(E_0.v, str(E_0.units))
dE = -cube.dv
delta_E = E_0 * v_th.in_cgs() / clight
E_shift = E_0 * (1.0 + v_shift / clight)
c = (dE * np.exp(-(((cube.vmid - E_shift) / delta_E)**2)) /
(np.sqrt(np.pi) * delta_E))
assert_allclose_units(a, c, 1.0e-2)
def export_dataset(self, fields=None, nprocs=1):
r"""Export a set of pixelized fields to an in-memory dataset that can be
analyzed as any other in yt. Unit information and other parameters (e.g.,
geometry, current_time, etc.) will be taken from the parent dataset.
Parameters
----------
fields : list of strings, optional
These fields will be pixelized and output. If "None", the keys of the
FRB will be used.
nprocs: integer, optional
If greater than 1, will create this number of subarrays out of data
Examples
--------
>>> import yt
>>> ds = yt.load("GasSloshing/sloshing_nomag2_hdf5_plt_cnt_0150")
>>> slc = ds.slice(2, 0.0)
>>> frb = slc.to_frb((500.0, "kpc"), 500)
>>> ds2 = frb.export_dataset(
... fields=[("gas", "density"), ("gas", "temperature")], nprocs=32
... )
"""
nx, ny = self.buff_size
data = {}
if fields is None:
fields = list(self.keys())
for field in fields:
arr = self[field]
data[field] = (arr.d.T.reshape(nx, ny, 1), str(arr.units))
bounds = [b.in_units("code_length").v for b in self.bounds]
bbox = np.array([[bounds[0], bounds[1]], [bounds[2], bounds[3]],
[0.0, 1.0]])
return load_uniform_grid(
data,
[nx, ny, 1],
length_unit=self.ds.length_unit,
bbox=bbox,
sim_time=self.ds.current_time.in_units("s").v,
mass_unit=self.ds.mass_unit,
time_unit=self.ds.time_unit,
velocity_unit=self.ds.velocity_unit,
magnetic_unit=self.ds.magnetic_unit,
periodicity=(False, False, False),
geometry=self.ds.geometry,
nprocs=nprocs,
)
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_particles_outside_domain():
np.random.seed(0x4D3D3D3)
posx_arr = np.random.uniform(low=-1.6, high=1.5, size=1000)
posy_arr = np.random.uniform(low=-1.5, high=1.5, size=1000)
posz_arr = np.random.uniform(low=-1.5, high=1.5, size=1000)
dens_arr = np.random.random((16, 16, 16))
data = dict(
density=dens_arr,
particle_position_x=posx_arr,
particle_position_y=posy_arr,
particle_position_z=posz_arr,
)
bbox = np.array([[-1.5, 1.5], [-1.5, 1.5], [-1.5, 1.5]])
ds = load_uniform_grid(data, (16, 16, 16), bbox=bbox, nprocs=4)
wh = (posx_arr < bbox[0, 0]).nonzero()[0]
assert wh.size == 1000 - ds.particle_type_counts["io"]
ad = ds.all_data()
assert ds.particle_type_counts["io"] == ad["particle_position_x"].size
def test_symlog_min_zero():
# see https://github.com/yt-project/yt/issues/3791
shape = (32, 16, 1)
a = np.linspace(0, 1, 16)
b = np.ones((32, 16))
c = np.reshape(a * b, shape)
data = {("gas", "density"): c}
ds = load_uniform_grid(
data,
shape,
bbox=np.array([[0.0, 5.0], [0, 1], [-0.1, +0.1]]),
)
p = SlicePlot(ds, "z", "density")
im_arr = p["gas", "density"].image.get_array()
# check that no data value was mapped to a NaN (log(0))
assert np.all(~np.isnan(im_arr))
# 0 should be mapped to itself since we expect a symlog norm
assert np.min(im_arr) == 0.0
def test_inside_outside():
test_array = np.ones((10, 10, 10))
test_array[1, 1, :] = 2.0
test_array[2, 1, :] = 3.0
data = dict(density=test_array)
ds = load_uniform_grid(data, test_array.shape, length_unit="cm", nprocs=1)
ad = ds.all_data()
only_ones_and_twos = ad.include_inside(("gas", "density"), 0.9, 2.1)
assert np.all(only_ones_and_twos[("gas", "density")] != 3.0)
assert len(only_ones_and_twos[("gas", "density")]) == 990
only_ones_and_twos = ad.exclude_outside(("gas", "density"), 0.9, 2.1)
assert len(only_ones_and_twos[("gas", "density")]) == 990
assert np.all(only_ones_and_twos[("gas", "density")] != 3.0)
only_threes = ad.include_outside(("gas", "density"), 0.9, 2.1)
assert np.all(only_threes[("gas", "density")] == 3)
assert len(only_threes[("gas", "density")]) == 10
only_threes = ad.include_outside(("gas", "density"), 0.9, 2.1)
assert np.all(only_threes[("gas", "density")] == 3)
assert len(only_threes[("gas", "density")]) == 10
# Repeat, but convert units to g/m**3
only_ones_and_twos = ad.include_inside(("gas", "density"), 0.9e6, 2.1e6, "g/m**3")
assert np.all(only_ones_and_twos[("gas", "density")] != 3.0)
assert len(only_ones_and_twos[("gas", "density")]) == 990
only_ones_and_twos = ad.exclude_outside(("gas", "density"), 0.9e6, 2.1e6, "g/m**3")
assert len(only_ones_and_twos[("gas", "density")]) == 990
assert np.all(only_ones_and_twos[("gas", "density")] != 3.0)
only_threes = ad.include_outside(("gas", "density"), 0.9e6, 2.1e6, "g/m**3")
assert np.all(only_threes[("gas", "density")] == 3)
assert len(only_threes[("gas", "density")]) == 10
only_threes = ad.include_outside(("gas", "density"), 0.9e6, 2.1e6, "g/m**3")
assert np.all(only_threes[("gas", "density")] == 3)
assert len(only_threes[("gas", "density")]) == 10
def setup_cluster():
R = 1000.
r_c = 100.
rho_c = 1.673e-26
beta = 1.
T0 = 4.
nx, ny, nz = 16, 16, 16
c = 0.17
a_c = 30.
a = 200.
v0 = 300. * cm_per_km
ddims = (nx, ny, nz)
x, y, z = np.mgrid[-R:R:nx * 1j, -R:R:ny * 1j, -R:R:nz * 1j]
r = np.sqrt(x**2 + y**2 + z**2)
dens = np.zeros(ddims)
dens = rho_c * (1. + (r / r_c)**2)**(-1.5 * beta)
temp = T0 * K_per_keV / (1. + r / a) * (c + r / a_c) / (1. + r / a_c)
velz = v0 * temp / (T0 * K_per_keV)
data = {}
data["density"] = (dens, "g/cm**3")
data["temperature"] = (temp, "K")
data["velocity_x"] = (np.zeros(ddims), "cm/s")
data["velocity_y"] = (np.zeros(ddims), "cm/s")
data["velocity_z"] = (velz, "cm/s")
L = 2 * R * cm_per_kpc
bbox = np.array([[-0.5, 0.5], [-0.5, 0.5], [-0.5, 0.5]]) * L
ds = load_uniform_grid(data, ddims, length_unit='cm', bbox=bbox)
ds.index
return ds
def test_magnetic_fields():
ddims = (16, 16, 16)
data1 = {
"magnetic_field_x": (np.random.random(size=ddims), "T"),
"magnetic_field_y": (np.random.random(size=ddims), "T"),
"magnetic_field_z": (np.random.random(size=ddims), "T"),
}
data2 = {}
for field in data1:
data2[field] = (data1[field][0] * 1.0e4, "gauss")
ds0 = load_uniform_grid(data1, ddims, unit_system="cgs")
ds1 = load_uniform_grid(data1,
ddims,
magnetic_unit=(1.0, "T"),
unit_system="cgs")
ds2 = load_uniform_grid(data2, ddims, unit_system="mks")
# For this test dataset, code units are cgs units
ds3 = load_uniform_grid(data2, ddims, unit_system="code")
# For this test dataset, code units are SI units
ds4 = load_uniform_grid(data1,
ddims,
magnetic_unit=(1.0, "T"),
unit_system="code")
ds0.index
ds1.index
ds2.index
ds3.index
ds4.index
dd0 = ds0.all_data()
dd1 = ds1.all_data()
dd2 = ds2.all_data()
dd3 = ds3.all_data()
dd4 = ds4.all_data()
assert ds0.fields.gas.magnetic_field_strength.units == "G"
assert ds1.fields.gas.magnetic_field_strength.units == "G"
assert ds1.fields.gas.magnetic_field_poloidal.units == "G"
assert ds1.fields.gas.magnetic_field_toroidal.units == "G"
assert ds2.fields.gas.magnetic_field_strength.units == "T"
assert ds2.fields.gas.magnetic_field_poloidal.units == "T"
assert ds2.fields.gas.magnetic_field_toroidal.units == "T"
assert ds3.fields.gas.magnetic_field_strength.units == "code_magnetic"
assert ds3.fields.gas.magnetic_field_poloidal.units == "code_magnetic"
assert ds3.fields.gas.magnetic_field_toroidal.units == "code_magnetic"
assert ds4.fields.gas.magnetic_field_strength.units == "code_magnetic"
assert ds4.fields.gas.magnetic_field_poloidal.units == "code_magnetic"
assert ds4.fields.gas.magnetic_field_toroidal.units == "code_magnetic"
emag0 = (dd0[("gas", "magnetic_field_x")]**2 +
dd0[("gas", "magnetic_field_y")]**2 +
dd0[("gas", "magnetic_field_z")]**2) / (8.0 * np.pi)
emag0.convert_to_units("dyne/cm**2")
emag1 = (dd1[("gas", "magnetic_field_x")]**2 +
dd1[("gas", "magnetic_field_y")]**2 +
dd1[("gas", "magnetic_field_z")]**2) / (8.0 * np.pi)
emag1.convert_to_units("dyne/cm**2")
emag2 = (dd2[("gas", "magnetic_field_x")]**2 +
dd2[("gas", "magnetic_field_y")]**2 +
dd2[("gas", "magnetic_field_z")]**2) / (2.0 * mu_0)
emag2.convert_to_units("Pa")
emag3 = (dd3[("gas", "magnetic_field_x")]**2 +
dd3[("gas", "magnetic_field_y")]**2 +
dd3[("gas", "magnetic_field_z")]**2) / (8.0 * np.pi)
emag3.convert_to_units("code_pressure")
emag4 = (dd4[("gas", "magnetic_field_x")]**2 +
dd4[("gas", "magnetic_field_y")]**2 +
dd4[("gas", "magnetic_field_z")]**2) / (2.0 * mu_0)
emag4.convert_to_units("code_pressure")
# note that "magnetic_energy_density" and "magnetic_pressure" are aliased
assert_almost_equal(emag0, dd0[("gas", "magnetic_energy_density")])
assert_almost_equal(emag1, dd1[("gas", "magnetic_energy_density")])
assert_almost_equal(emag2, dd2[("gas", "magnetic_energy_density")])
assert_almost_equal(emag3, dd3[("gas", "magnetic_energy_density")])
assert_almost_equal(emag4, dd4[("gas", "magnetic_energy_density")])
assert str(emag0.units) == str(dd0[("gas",
"magnetic_energy_density")].units)
assert str(emag1.units) == str(dd1[("gas",
"magnetic_energy_density")].units)
assert str(emag2.units) == str(dd2[("gas",
"magnetic_energy_density")].units)
assert str(emag3.units) == str(dd3[("gas",
"magnetic_energy_density")].units)
assert str(emag4.units) == str(dd4[("gas",
"magnetic_energy_density")].units)
assert_almost_equal(emag1.in_cgs(), emag0.in_cgs())
assert_almost_equal(emag2.in_cgs(), emag0.in_cgs())
assert_almost_equal(emag1.in_cgs(), emag2.in_cgs())
assert_almost_equal(emag1.in_cgs(), emag3.in_cgs())
assert_almost_equal(emag1.in_cgs(), emag4.in_cgs())
def fake_vr_orientation_test_ds(N=96, scale=1):
"""
create a toy dataset that puts a sphere at (0,0,0), a single cube
on +x, two cubes on +y, and three cubes on +z in a domain from
[-1*scale,1*scale]**3. The lower planes
(x = -1*scale, y = -1*scale, z = -1*scale) are also given non-zero
values.
This dataset allows you to easily explore orientations and
handiness in VR and other renderings
Parameters
----------
N : integer
The number of cells along each direction
scale : float
A spatial scale, the domain boundaries will be multiplied by scale to
test datasets that have spatial different scales (e.g. data in CGS units)
"""
from yt.loaders import load_uniform_grid
xmin = ymin = zmin = -1.0 * scale
xmax = ymax = zmax = 1.0 * scale
dcoord = (xmax - xmin) / N
arr = np.zeros((N, N, N), dtype=np.float64)
arr[:, :, :] = 1.0e-4
bbox = np.array([[xmin, xmax], [ymin, ymax], [zmin, zmax]])
# coordinates -- in the notation data[i, j, k]
x = (np.arange(N) + 0.5) * dcoord + xmin
y = (np.arange(N) + 0.5) * dcoord + ymin
z = (np.arange(N) + 0.5) * dcoord + zmin
x3d, y3d, z3d = np.meshgrid(x, y, z, indexing="ij")
# sphere at the origin
c = np.array([0.5 * (xmin + xmax), 0.5 * (ymin + ymax), 0.5 * (zmin + zmax)])
r = np.sqrt((x3d - c[0]) ** 2 + (y3d - c[1]) ** 2 + (z3d - c[2]) ** 2)
arr[r < 0.05] = 1.0
arr[abs(x3d - xmin) < 2 * dcoord] = 0.3
arr[abs(y3d - ymin) < 2 * dcoord] = 0.3
arr[abs(z3d - zmin) < 2 * dcoord] = 0.3
# single cube on +x
xc = 0.75 * scale
dx = 0.05 * scale
idx = np.logical_and(
np.logical_and(x3d > xc - dx, x3d < xc + dx),
np.logical_and(
np.logical_and(y3d > -dx, y3d < dx), np.logical_and(z3d > -dx, z3d < dx)
),
)
arr[idx] = 1.0
# two cubes on +y
dy = 0.05 * scale
for yc in [0.65 * scale, 0.85 * scale]:
idx = np.logical_and(
np.logical_and(y3d > yc - dy, y3d < yc + dy),
np.logical_and(
np.logical_and(x3d > -dy, x3d < dy), np.logical_and(z3d > -dy, z3d < dy)
),
)
arr[idx] = 0.8
# three cubes on +z
dz = 0.05 * scale
for zc in [0.5 * scale, 0.7 * scale, 0.9 * scale]:
idx = np.logical_and(
np.logical_and(z3d > zc - dz, z3d < zc + dz),
np.logical_and(
np.logical_and(x3d > -dz, x3d < dz), np.logical_and(y3d > -dz, y3d < dz)
),
)
arr[idx] = 0.6
data = dict(density=(arr, "g/cm**3"))
ds = load_uniform_grid(data, arr.shape, bbox=bbox)
return ds
def test_clump_finding():
n_c = 8
n_p = 1
dims = (n_c, n_c, n_c)
density = np.ones(dims)
high_rho = 10.0
# 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.0 / 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)
def _total_volume(clump):
total_vol = clump.data.quantities.total_quantity(["cell_volume"
]).in_units("cm**3")
return "Cell Volume: %6e cm**3.", total_vol
add_clump_info("total_volume", _total_volume)
master_clump.add_info_item("total_volume")
find_clumps(master_clump, 0.5, 2.0 * high_rho, 10.0)
# there should be two children
assert_equal(len(master_clump.children), 2)
leaf_clumps = master_clump.leaves
for l in leaf_clumps:
keys = l.info.keys()
assert "total_cells" in keys
assert "cell_mass" in keys
assert "max_grid_level" in keys
assert "total_volume" in keys
# 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)
# clean up global registry to avoid polluting other tests
del clump_info_registry["total_volume"]
def create_dataset(self,
domain_dimensions,
box_size,
left_edge=None,
**kwargs):
"""
Create an in-memory, uniformly gridded dataset in 3D using yt by
placing the clusters into a box. When adding multiple clusters,
per-volume quantities from each cluster such as density and
pressure are added, whereas per-mass quantites such as temperature
and velocity are mass-weighted.
Parameters
----------
domain_dimensions : 3-tuple of ints
The number of cells on a side for the domain.
box_size : float
The size of the box in kpc.
left_edge : array_like, optional
The minimum coordinate of the box in all three dimensions,
in kpc. Default: None, which means the left edge will
be [0, 0, 0].
"""
from yt.loaders import load_uniform_grid
from scipy.interpolate import InterpolatedUnivariateSpline
if left_edge is None:
left_edge = np.zeros(3)
left_edge = np.array(left_edge)
bbox = [[left_edge[0], left_edge[0] + box_size],
[left_edge[1], left_edge[1] + box_size],
[left_edge[2], left_edge[2] + box_size]]
x, y, z = np.mgrid[bbox[0][0]:bbox[0][1]:domain_dimensions[0] * 1j,
bbox[1][0]:bbox[1][1]:domain_dimensions[1] * 1j,
bbox[2][0]:bbox[2][1]:domain_dimensions[2] * 1j, ]
fields1 = [
"density", "pressure", "dark_matter_density"
"stellar_density", "gravitational_potential"
]
fields2 = ["temperature"]
fields3 = ["velocity_x", "velocity_y", "velocity_z"]
units = {
"density": "Msun/kpc**3",
"pressure": "Msun/kpc/Myr**2",
"dark_matter_density": "Msun/kpc**3",
"stellar_density": "Msun/kpc**3",
"temperature": "K",
"gravitational_potential": "kpc**2/Myr**2",
"velocity_x": "kpc/Myr",
"velocity_y": "kpc/Myr",
"velocity_z": "kpc/Myr",
"magnetic_field_strength": "G"
}
fields = fields1 + fields2
data = {}
for i, profile in enumerate(self.profiles):
p = ClusterModel.from_h5_file(profile)
xx = x - self.center.d[i][0]
yy = y - self.center.d[i][1]
zz = z - self.center.d[i][2]
rr = np.sqrt(xx * xx + yy * yy + zz * zz)
fd = InterpolatedUnivariateSpline(p["radius"].d, p["density"].d)
for field in fields:
if field not in p:
continue
if field not in data:
data[field] = (np.zeros(domain_dimensions), units[field])
f = InterpolatedUnivariateSpline(p["radius"].d, p[field].d)
if field in fields1:
data[field][0] += f(rr)
elif field in fields2:
data[field][0] += f(rr) * fd(rr)
for field in fields3:
data[field][0] += self.velocity.d[i][0] * fd(rr)
if "density" in data:
for field in fields2 + fields3:
data[field][0] /= data["density"][0]
return load_uniform_grid(data,
domain_dimensions,
length_unit="kpc",
bbox=bbox,
mass_unit="Msun",
time_unit="Myr",
**kwargs)
def test_stream_particles():
num_particles = 100000
domain_dims = (64, 64, 64)
dens = np.random.random(domain_dims)
x = np.random.uniform(size=num_particles)
y = np.random.uniform(size=num_particles)
z = np.random.uniform(size=num_particles)
m = np.ones(num_particles)
# Field operators and cell flagging methods
fo = []
fo.append(ic.TopHatSphere(0.1, [0.2, 0.3, 0.4], {"density": 2.0}))
fo.append(ic.TopHatSphere(0.05, [0.7, 0.4, 0.75], {"density": 20.0}))
# Add particles
fields1 = {
"density": dens,
"particle_position_x": x,
"particle_position_y": y,
"particle_position_z": z,
"particle_mass": m,
}
fields2 = fields1.copy()
ug1 = load_uniform_grid(fields1, domain_dims, 1.0)
ug2 = load_uniform_grid(fields2, domain_dims, 1.0, nprocs=8)
# Check to make sure the number of particles is the same
number_of_particles1 = np.sum(
[grid.NumberOfParticles for grid in ug1.index.grids])
number_of_particles2 = np.sum(
[grid.NumberOfParticles for grid in ug2.index.grids])
assert_equal(number_of_particles1, num_particles)
assert_equal(number_of_particles1, number_of_particles2)
for grid in ug2.index.grids:
tot_parts = grid["io", "particle_position_x"].size
tot_all_parts = grid["all", "particle_position_x"].size
assert tot_parts == grid.NumberOfParticles
assert tot_all_parts == grid.NumberOfParticles
# Check to make sure the fields have been defined correctly
for ptype in ("all", "io"):
assert (ug1._get_field_info(
ptype, "particle_position_x").sampling_type == "particle")
assert (ug1._get_field_info(
ptype, "particle_position_y").sampling_type == "particle")
assert (ug1._get_field_info(
ptype, "particle_position_z").sampling_type == "particle")
assert ug1._get_field_info(ptype,
"particle_mass").sampling_type == "particle"
assert not ug1._get_field_info("gas",
"density").sampling_type == "particle"
for ptype in ("all", "io"):
assert (ug2._get_field_info(
ptype, "particle_position_x").sampling_type == "particle")
assert (ug2._get_field_info(
ptype, "particle_position_y").sampling_type == "particle")
assert (ug2._get_field_info(
ptype, "particle_position_z").sampling_type == "particle")
assert ug2._get_field_info(ptype,
"particle_mass").sampling_type == "particle"
assert not ug2._get_field_info("gas",
"density").sampling_type == "particle"
# Now perform similar checks, but with multiple particle types
num_dm_particles = 30000
xd = np.random.uniform(size=num_dm_particles)
yd = np.random.uniform(size=num_dm_particles)
zd = np.random.uniform(size=num_dm_particles)
md = np.ones(num_dm_particles)
num_star_particles = 20000
xs = np.random.uniform(size=num_star_particles)
ys = np.random.uniform(size=num_star_particles)
zs = np.random.uniform(size=num_star_particles)
ms = 2.0 * np.ones(num_star_particles)
dens = np.random.random(domain_dims)
fields3 = {
"density": dens,
("dm", "particle_position_x"): xd,
("dm", "particle_position_y"): yd,
("dm", "particle_position_z"): zd,
("dm", "particle_mass"): md,
("star", "particle_position_x"): xs,
("star", "particle_position_y"): ys,
("star", "particle_position_z"): zs,
("star", "particle_mass"): ms,
}
fields4 = fields3.copy()
ug3 = load_uniform_grid(fields3, domain_dims, 1.0)
ug4 = load_uniform_grid(fields4, domain_dims, 1.0, nprocs=8)
# Check to make sure the number of particles is the same
number_of_particles3 = np.sum(
[grid.NumberOfParticles for grid in ug3.index.grids])
number_of_particles4 = np.sum(
[grid.NumberOfParticles for grid in ug4.index.grids])
assert_equal(number_of_particles3, num_dm_particles + num_star_particles)
assert_equal(number_of_particles3, number_of_particles4)
for grid in ug4.index.grids:
tot_parts = grid["dm", "particle_position_x"].size
tot_parts += grid["star", "particle_position_x"].size
tot_all_parts = grid["all", "particle_position_x"].size
assert tot_parts == grid.NumberOfParticles
assert tot_all_parts == grid.NumberOfParticles
# Check to make sure the fields have been defined correctly
for ptype in ("dm", "star"):
assert (ug3._get_field_info(
ptype, "particle_position_x").sampling_type == "particle")
assert (ug3._get_field_info(
ptype, "particle_position_y").sampling_type == "particle")
assert (ug3._get_field_info(
ptype, "particle_position_z").sampling_type == "particle")
assert ug3._get_field_info(ptype,
"particle_mass").sampling_type == "particle"
assert (ug4._get_field_info(
ptype, "particle_position_x").sampling_type == "particle")
assert (ug4._get_field_info(
ptype, "particle_position_y").sampling_type == "particle")
assert (ug4._get_field_info(
ptype, "particle_position_z").sampling_type == "particle")
assert ug4._get_field_info(ptype,
"particle_mass").sampling_type == "particle"
def test_particle_generator():
# First generate our dataset
domain_dims = (32, 32, 32)
dens = np.zeros(domain_dims) + 0.1
temp = 4.0 * np.ones(domain_dims)
fields = {
"density": (dens, "code_mass/code_length**3"),
"temperature": (temp, "K")
}
ds = load_uniform_grid(fields, domain_dims, 1.0)
# Now generate particles from density
field_list = [
("io", "particle_position_x"),
("io", "particle_position_y"),
("io", "particle_position_z"),
("io", "particle_index"),
("io", "particle_gas_density"),
]
num_particles = 10000
field_dict = {("gas", "density"): ("io", "particle_gas_density")}
sphere = ds.sphere(ds.domain_center, 0.45)
particles1 = WithDensityParticleGenerator(ds, sphere, num_particles,
field_list)
particles1.assign_indices()
particles1.map_grid_fields_to_particles(field_dict)
# Test to make sure we ended up with the right number of particles per grid
particles1.apply_to_stream()
particles_per_grid1 = [grid.NumberOfParticles for grid in ds.index.grids]
assert_equal(particles_per_grid1, particles1.NumberOfParticles)
particles_per_grid1 = [
len(grid["particle_position_x"]) for grid in ds.index.grids
]
assert_equal(particles_per_grid1, particles1.NumberOfParticles)
tags = uconcatenate([grid["particle_index"] for grid in ds.index.grids])
assert np.unique(tags).size == num_particles
del tags
# Set up a lattice of particles
pdims = np.array([32, 32, 32])
def new_indices():
# We just add new indices onto the existing ones
return np.arange(np.product(pdims)) + num_particles
le = np.array([0.25, 0.25, 0.25])
re = np.array([0.75, 0.75, 0.75])
particles2 = LatticeParticleGenerator(ds, pdims, le, re, field_list)
particles2.assign_indices(function=new_indices)
particles2.map_grid_fields_to_particles(field_dict)
# Test lattice positions
xpos = np.unique(particles2["io", "particle_position_x"])
ypos = np.unique(particles2["io", "particle_position_y"])
zpos = np.unique(particles2["io", "particle_position_z"])
xpred = np.linspace(le[0], re[0], num=pdims[0], endpoint=True)
ypred = np.linspace(le[1], re[1], num=pdims[1], endpoint=True)
zpred = np.linspace(le[2], re[2], num=pdims[2], endpoint=True)
assert_almost_equal(xpos, xpred)
assert_almost_equal(ypos, ypred)
assert_almost_equal(zpos, zpred)
del xpos, ypos, zpos
del xpred, ypred, zpred
# Test the number of particles again
particles2.apply_to_stream()
particles_per_grid2 = [grid.NumberOfParticles for grid in ds.index.grids]
assert_equal(particles_per_grid2,
particles1.NumberOfParticles + particles2.NumberOfParticles)
[grid.field_data.clear() for grid in ds.index.grids]
particles_per_grid2 = [
len(grid["particle_position_x"]) for grid in ds.index.grids
]
assert_equal(particles_per_grid2,
particles1.NumberOfParticles + particles2.NumberOfParticles)
# Test the uniqueness of tags
tags = np.concatenate([grid["particle_index"] for grid in ds.index.grids])
tags.sort()
assert_equal(tags, np.arange(np.product(pdims) + num_particles))
del tags
# Now dump all of these particle fields out into a dict
pdata = {}
dd = ds.all_data()
for field in field_list:
pdata[field] = dd[field]
# Test the "from-list" generator and particle field overwrite
num_particles3 = num_particles + np.product(pdims)
particles3 = FromListParticleGenerator(ds, num_particles3, pdata)
particles3.apply_to_stream(overwrite=True)
# Test the number of particles again
particles_per_grid3 = [grid.NumberOfParticles for grid in ds.index.grids]
assert_equal(particles_per_grid3,
particles1.NumberOfParticles + particles2.NumberOfParticles)
particles_per_grid2 = [
len(grid["particle_position_z"]) for grid in ds.index.grids
]
assert_equal(particles_per_grid3,
particles1.NumberOfParticles + particles2.NumberOfParticles)
assert_equal(particles_per_grid2, particles_per_grid3)
# Test adding in particles with a different particle type
num_star_particles = 20000
pdata2 = {
("star", "particle_position_x"):
np.random.uniform(size=num_star_particles),
("star", "particle_position_y"):
np.random.uniform(size=num_star_particles),
("star", "particle_position_z"):
np.random.uniform(size=num_star_particles),
}
particles4 = FromListParticleGenerator(ds,
num_star_particles,
pdata2,
ptype="star")
particles4.apply_to_stream()
dd = ds.all_data()
assert dd["star", "particle_position_x"].size == num_star_particles
assert dd["io", "particle_position_x"].size == num_particles3
assert dd[
"all",
"particle_position_x"].size == num_star_particles + num_particles3
del pdata
del pdata2
del ds
del particles1
del particles2
del particles4
del fields
del dens
del temp
