def test_numpy_wrappers():
a1 = YTArray([1, 2, 3], 'cm')
a2 = YTArray([2, 3, 4, 5, 6], 'cm')
a3 = YTArray([7, 8, 9, 10, 11], 'cm')
catenate_answer = [1, 2, 3, 2, 3, 4, 5, 6]
intersect_answer = [2, 3]
union_answer = [1, 2, 3, 4, 5, 6]
vstack_answer = [[2, 3, 4, 5, 6], [7, 8, 9, 10, 11]]
vstack_answer_last_axis = [[2, 7], [3, 8], [4, 9], [5, 10], [6, 11]]
assert_array_equal(YTArray(catenate_answer, 'cm'), uconcatenate((a1, a2)))
assert_array_equal(catenate_answer, np.concatenate((a1, a2)))
assert_array_equal(YTArray(intersect_answer, 'cm'), uintersect1d(a1, a2))
assert_array_equal(intersect_answer, np.intersect1d(a1, a2))
assert_array_equal(YTArray(union_answer, 'cm'), uunion1d(a1, a2))
assert_array_equal(union_answer, np.union1d(a1, a2))
assert_array_equal(YTArray(catenate_answer, 'cm'), uhstack([a1, a2]))
assert_array_equal(catenate_answer, np.hstack([a1, a2]))
assert_array_equal(YTArray(vstack_answer, 'cm'), uvstack([a2, a3]))
assert_array_equal(vstack_answer, np.vstack([a2, a3]))
assert_array_equal(YTArray(vstack_answer, 'cm'), ustack([a2, a3]))
assert_array_equal(vstack_answer, np.stack([a2, a3]))
assert_array_equal(YTArray(vstack_answer_last_axis, 'cm'),
ustack([a2, a3], axis=-1))
assert_array_equal(vstack_answer_last_axis, np.stack([a2, a3], axis=-1))
def _magnetic_field_toroidal(field, data):
normal = data.get_field_parameter("normal")
Bfields = ustack([
data[ftype, "relative_magnetic_field_x"],
data[ftype, "relative_magnetic_field_y"],
data[ftype, "relative_magnetic_field_z"],
])
phi = data["index", "spherical_phi"]
return get_sph_phi_component(Bfields, phi, normal)
def _magnetic_field_poloidal(field, data):
normal = data.get_field_parameter("normal")
Bfields = ustack([
data[ftype, 'relative_magnetic_field_x'],
data[ftype, 'relative_magnetic_field_y'],
data[ftype, 'relative_magnetic_field_z']
])
theta = data["index", 'spherical_theta']
phi = data["index", 'spherical_phi']
return get_sph_theta_component(Bfields, theta, phi, normal)