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)