Spaces:
Sleeping
Sleeping
| # Copyright (c) 2022 NVIDIA CORPORATION. All rights reserved. | |
| # NVIDIA CORPORATION and its licensors retain all intellectual property | |
| # and proprietary rights in and to this software, related documentation | |
| # and any modifications thereto. Any use, reproduction, disclosure or | |
| # distribution of this software and related documentation without an express | |
| # license agreement from NVIDIA CORPORATION is strictly prohibited. | |
| import unittest | |
| import numpy as np | |
| import warp as wp | |
| from warp.tests.unittest_utils import * | |
| wp.init() | |
| def kernel_1d(a: wp.array(dtype=int, ndim=1)): | |
| i = wp.tid() | |
| wp.expect_eq(a[i], wp.tid()) | |
| a[i] = a[i] * 2 | |
| wp.atomic_add(a, i, 1) | |
| wp.expect_eq(a[i], wp.tid() * 2 + 1) | |
| def test_1d(test, device): | |
| dim_x = 4 | |
| a = np.arange(0, dim_x, dtype=np.int32) | |
| arr = wp.array(a, device=device) | |
| test.assertEqual(arr.shape, a.shape) | |
| test.assertEqual(arr.size, a.size) | |
| test.assertEqual(arr.ndim, a.ndim) | |
| with CheckOutput(test): | |
| wp.launch(kernel_1d, dim=arr.size, inputs=[arr], device=device) | |
| def kernel_2d(a: wp.array(dtype=int, ndim=2), m: int, n: int): | |
| i = wp.tid() // n | |
| j = wp.tid() % n | |
| wp.expect_eq(a[i, j], wp.tid()) | |
| wp.expect_eq(a[i][j], wp.tid()) | |
| a[i, j] = a[i, j] * 2 | |
| wp.atomic_add(a, i, j, 1) | |
| wp.expect_eq(a[i, j], wp.tid() * 2 + 1) | |
| def test_2d(test, device): | |
| dim_x = 4 | |
| dim_y = 2 | |
| a = np.arange(0, dim_x * dim_y, dtype=np.int32) | |
| a = a.reshape(dim_x, dim_y) | |
| arr = wp.array(a, device=device) | |
| test.assertEqual(arr.shape, a.shape) | |
| test.assertEqual(arr.size, a.size) | |
| test.assertEqual(arr.ndim, a.ndim) | |
| with CheckOutput(test): | |
| wp.launch(kernel_2d, dim=arr.size, inputs=[arr, dim_x, dim_y], device=device) | |
| def kernel_3d(a: wp.array(dtype=int, ndim=3), m: int, n: int, o: int): | |
| i = wp.tid() // (n * o) | |
| j = wp.tid() % (n * o) // o | |
| k = wp.tid() % o | |
| wp.expect_eq(a[i, j, k], wp.tid()) | |
| wp.expect_eq(a[i][j][k], wp.tid()) | |
| a[i, j, k] = a[i, j, k] * 2 | |
| a[i][j][k] = a[i][j][k] * 2 | |
| wp.atomic_add(a, i, j, k, 1) | |
| wp.expect_eq(a[i, j, k], wp.tid() * 4 + 1) | |
| def test_3d(test, device): | |
| dim_x = 8 | |
| dim_y = 4 | |
| dim_z = 2 | |
| a = np.arange(0, dim_x * dim_y * dim_z, dtype=np.int32) | |
| a = a.reshape(dim_x, dim_y, dim_z) | |
| arr = wp.array(a, device=device) | |
| test.assertEqual(arr.shape, a.shape) | |
| test.assertEqual(arr.size, a.size) | |
| test.assertEqual(arr.ndim, a.ndim) | |
| with CheckOutput(test): | |
| wp.launch(kernel_3d, dim=arr.size, inputs=[arr, dim_x, dim_y, dim_z], device=device) | |
| def kernel_4d(a: wp.array(dtype=int, ndim=4), m: int, n: int, o: int, p: int): | |
| i = wp.tid() // (n * o * p) | |
| j = wp.tid() % (n * o * p) // (o * p) | |
| k = wp.tid() % (o * p) / p | |
| l = wp.tid() % p | |
| wp.expect_eq(a[i, j, k, l], wp.tid()) | |
| wp.expect_eq(a[i][j][k][l], wp.tid()) | |
| def test_4d(test, device): | |
| dim_x = 16 | |
| dim_y = 8 | |
| dim_z = 4 | |
| dim_w = 2 | |
| a = np.arange(0, dim_x * dim_y * dim_z * dim_w, dtype=np.int32) | |
| a = a.reshape(dim_x, dim_y, dim_z, dim_w) | |
| arr = wp.array(a, device=device) | |
| test.assertEqual(arr.shape, a.shape) | |
| test.assertEqual(arr.size, a.size) | |
| test.assertEqual(arr.ndim, a.ndim) | |
| with CheckOutput(test): | |
| wp.launch(kernel_4d, dim=arr.size, inputs=[arr, dim_x, dim_y, dim_z, dim_w], device=device) | |
| def kernel_4d_transposed(a: wp.array(dtype=int, ndim=4), m: int, n: int, o: int, p: int): | |
| i = wp.tid() // (n * o * p) | |
| j = wp.tid() % (n * o * p) // (o * p) | |
| k = wp.tid() % (o * p) / p | |
| l = wp.tid() % p | |
| wp.expect_eq(a[l, k, j, i], wp.tid()) | |
| wp.expect_eq(a[l][k][j][i], wp.tid()) | |
| def test_4d_transposed(test, device): | |
| dim_x = 16 | |
| dim_y = 8 | |
| dim_z = 4 | |
| dim_w = 2 | |
| a = np.arange(0, dim_x * dim_y * dim_z * dim_w, dtype=np.int32) | |
| a = a.reshape(dim_x, dim_y, dim_z, dim_w) | |
| arr = wp.array(a, device=device) | |
| # Transpose the array manually, as using the wp.array() constructor with arr.T would make it contiguous first | |
| a_T = a.T | |
| arr_T = wp.array( | |
| dtype=arr.dtype, | |
| shape=a_T.shape, | |
| strides=a_T.__array_interface__["strides"], | |
| capacity=arr.capacity, | |
| ptr=arr.ptr, | |
| owner=False, | |
| requires_grad=arr.requires_grad, | |
| device=device, | |
| ) | |
| test.assertFalse(arr_T.is_contiguous) | |
| test.assertEqual(arr_T.shape, a_T.shape) | |
| test.assertEqual(arr_T.strides, a_T.__array_interface__["strides"]) | |
| test.assertEqual(arr_T.size, a_T.size) | |
| test.assertEqual(arr_T.ndim, a_T.ndim) | |
| with CheckOutput(test): | |
| wp.launch(kernel_4d_transposed, dim=arr_T.size, inputs=[arr_T, dim_x, dim_y, dim_z, dim_w], device=device) | |
| def lower_bound_kernel(values: wp.array(dtype=float), arr: wp.array(dtype=float), indices: wp.array(dtype=int)): | |
| tid = wp.tid() | |
| indices[tid] = wp.lower_bound(arr, values[tid]) | |
| def test_lower_bound(test, device): | |
| arr = wp.array(np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0], dtype=float), dtype=float, device=device) | |
| values = wp.array(np.array([-0.1, 0.0, 2.5, 4.0, 5.0, 5.5], dtype=float), dtype=float, device=device) | |
| indices = wp.zeros(6, dtype=int, device=device) | |
| wp.launch(kernel=lower_bound_kernel, dim=6, inputs=[values, arr, indices], device=device) | |
| test.assertTrue((np.array([0, 0, 3, 4, 5, 5]) == indices.numpy()).all()) | |
| def f1(arr: wp.array(dtype=float)): | |
| wp.expect_eq(arr.shape[0], 10) | |
| def f2(arr: wp.array2d(dtype=float)): | |
| wp.expect_eq(arr.shape[0], 10) | |
| wp.expect_eq(arr.shape[1], 20) | |
| slice = arr[0] | |
| wp.expect_eq(slice.shape[0], 20) | |
| def f3(arr: wp.array3d(dtype=float)): | |
| wp.expect_eq(arr.shape[0], 10) | |
| wp.expect_eq(arr.shape[1], 20) | |
| wp.expect_eq(arr.shape[2], 30) | |
| slice = arr[0, 0] | |
| wp.expect_eq(slice.shape[0], 30) | |
| def f4(arr: wp.array4d(dtype=float)): | |
| wp.expect_eq(arr.shape[0], 10) | |
| wp.expect_eq(arr.shape[1], 20) | |
| wp.expect_eq(arr.shape[2], 30) | |
| wp.expect_eq(arr.shape[3], 40) | |
| slice = arr[0, 0, 0] | |
| wp.expect_eq(slice.shape[0], 40) | |
| def test_shape(test, device): | |
| with CheckOutput(test): | |
| a1 = wp.zeros(dtype=float, shape=10, device=device) | |
| wp.launch(f1, dim=1, inputs=[a1], device=device) | |
| a2 = wp.zeros(dtype=float, shape=(10, 20), device=device) | |
| wp.launch(f2, dim=1, inputs=[a2], device=device) | |
| a3 = wp.zeros(dtype=float, shape=(10, 20, 30), device=device) | |
| wp.launch(f3, dim=1, inputs=[a3], device=device) | |
| a4 = wp.zeros(dtype=float, shape=(10, 20, 30, 40), device=device) | |
| wp.launch(f4, dim=1, inputs=[a4], device=device) | |
| def sum_array(arr: wp.array(dtype=float), loss: wp.array(dtype=float)): | |
| tid = wp.tid() | |
| wp.atomic_add(loss, 0, arr[tid]) | |
| def test_flatten(test, device): | |
| np_arr = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]], dtype=float) | |
| arr = wp.array(np_arr, dtype=float, shape=np_arr.shape, device=device, requires_grad=True) | |
| arr_flat = arr.flatten() | |
| arr_comp = wp.array(np_arr.flatten(), dtype=float, device=device) | |
| assert_array_equal(arr_flat, arr_comp) | |
| loss = wp.zeros(1, dtype=float, device=device, requires_grad=True) | |
| tape = wp.Tape() | |
| with tape: | |
| wp.launch(kernel=sum_array, dim=len(arr_flat), inputs=[arr_flat, loss], device=device) | |
| tape.backward(loss=loss) | |
| grad = tape.gradients[arr_flat] | |
| ones = wp.array( | |
| np.ones( | |
| (8,), | |
| dtype=float, | |
| ), | |
| dtype=float, | |
| device=device, | |
| ) | |
| assert_array_equal(grad, ones) | |
| test.assertEqual(loss.numpy()[0], 36) | |
| def test_reshape(test, device): | |
| np_arr = np.arange(6, dtype=float) | |
| arr = wp.array(np_arr, dtype=float, device=device, requires_grad=True) | |
| arr_reshaped = arr.reshape((3, 2)) | |
| arr_comp = wp.array(np_arr.reshape((3, 2)), dtype=float, device=device) | |
| assert_array_equal(arr_reshaped, arr_comp) | |
| arr_reshaped = arr_reshaped.reshape(6) | |
| assert_array_equal(arr_reshaped, arr) | |
| loss = wp.zeros(1, dtype=float, device=device, requires_grad=True) | |
| tape = wp.Tape() | |
| with tape: | |
| wp.launch(kernel=sum_array, dim=len(arr_reshaped), inputs=[arr_reshaped, loss], device=device) | |
| tape.backward(loss=loss) | |
| grad = tape.gradients[arr_reshaped] | |
| ones = wp.array( | |
| np.ones( | |
| (6,), | |
| dtype=float, | |
| ), | |
| dtype=float, | |
| device=device, | |
| ) | |
| assert_array_equal(grad, ones) | |
| test.assertEqual(loss.numpy()[0], 15) | |
| np_arr = np.arange(6, dtype=float) | |
| arr = wp.array(np_arr, dtype=float, device=device) | |
| arr_infer = arr.reshape((-1, 3)) | |
| arr_comp = wp.array(np_arr.reshape((-1, 3)), dtype=float, device=device) | |
| assert_array_equal(arr_infer, arr_comp) | |
| def compare_stepped_window_a(x: wp.array2d(dtype=float)): | |
| wp.expect_eq(x[0, 0], 1.0) | |
| wp.expect_eq(x[0, 1], 2.0) | |
| wp.expect_eq(x[1, 0], 9.0) | |
| wp.expect_eq(x[1, 1], 10.0) | |
| def compare_stepped_window_b(x: wp.array2d(dtype=float)): | |
| wp.expect_eq(x[0, 0], 3.0) | |
| wp.expect_eq(x[0, 1], 4.0) | |
| wp.expect_eq(x[1, 0], 7.0) | |
| wp.expect_eq(x[1, 1], 8.0) | |
| wp.expect_eq(x[2, 0], 11.0) | |
| wp.expect_eq(x[2, 1], 12.0) | |
| def test_slicing(test, device): | |
| np_arr = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]], dtype=float) | |
| arr = wp.array(np_arr, dtype=float, shape=np_arr.shape, device=device, requires_grad=True) | |
| slice_a = arr[1, :, :] # test indexing | |
| slice_b = arr[1:2, :, :] # test slicing | |
| slice_c = arr[-1, :, :] # test negative indexing | |
| slice_d = arr[-2:-1, :, :] # test negative slicing | |
| slice_e = arr[-1:3, :, :] # test mixed slicing | |
| slice_e2 = slice_e[0, 0, :] # test 2x slicing | |
| slice_f = arr[0:3:2, 0, :] # test step | |
| assert_array_equal(slice_a, wp.array(np_arr[1, :, :], dtype=float, device=device)) | |
| assert_array_equal(slice_b, wp.array(np_arr[1:2, :, :], dtype=float, device=device)) | |
| assert_array_equal(slice_c, wp.array(np_arr[-1, :, :], dtype=float, device=device)) | |
| assert_array_equal(slice_d, wp.array(np_arr[-2:-1, :, :], dtype=float, device=device)) | |
| assert_array_equal(slice_e, wp.array(np_arr[-1:3, :, :], dtype=float, device=device)) | |
| assert_array_equal(slice_e2, wp.array(np_arr[2, 0, :], dtype=float, device=device)) | |
| # wp does not support copying from/to non-contiguous arrays | |
| # stepped windows must read on the device the original array was created on | |
| wp.launch(kernel=compare_stepped_window_a, dim=1, inputs=[slice_f], device=device) | |
| slice_flat = slice_b.flatten() | |
| loss = wp.zeros(1, dtype=float, device=device, requires_grad=True) | |
| tape = wp.Tape() | |
| with tape: | |
| wp.launch(kernel=sum_array, dim=len(slice_flat), inputs=[slice_flat, loss], device=device) | |
| tape.backward(loss=loss) | |
| grad = tape.gradients[slice_flat] | |
| ones = wp.array( | |
| np.ones( | |
| (4,), | |
| dtype=float, | |
| ), | |
| dtype=float, | |
| device=device, | |
| ) | |
| assert_array_equal(grad, ones) | |
| test.assertEqual(loss.numpy()[0], 26) | |
| index_a = arr[1] | |
| index_b = arr[2, 1] | |
| index_c = arr[1, :] | |
| index_d = arr[:, 1] | |
| assert_array_equal(index_a, wp.array(np_arr[1], dtype=float, device=device)) | |
| assert_array_equal(index_b, wp.array(np_arr[2, 1], dtype=float, device=device)) | |
| assert_array_equal(index_c, wp.array(np_arr[1, :], dtype=float, device=device)) | |
| wp.launch(kernel=compare_stepped_window_b, dim=1, inputs=[index_d], device=device) | |
| np_arr = np.zeros(10, dtype=int) | |
| wp_arr = wp.array(np_arr, dtype=int, device=device) | |
| assert_array_equal(wp_arr[:5], wp.array(np_arr[:5], dtype=int, device=device)) | |
| assert_array_equal(wp_arr[1:5], wp.array(np_arr[1:5], dtype=int, device=device)) | |
| assert_array_equal(wp_arr[-9:-5:1], wp.array(np_arr[-9:-5:1], dtype=int, device=device)) | |
| assert_array_equal(wp_arr[:5,], wp.array(np_arr[:5], dtype=int, device=device)) # noqa: E231 | |
| def test_view(test, device): | |
| np_arr_a = np.arange(1, 10, 1, dtype=np.uint32) | |
| np_arr_b = np.arange(1, 10, 1, dtype=np.float32) | |
| np_arr_c = np.arange(1, 10, 1, dtype=np.uint16) | |
| np_arr_d = np.arange(1, 10, 1, dtype=np.float16) | |
| np_arr_e = np.ones((4, 4), dtype=np.float32) | |
| wp_arr_a = wp.array(np_arr_a, dtype=wp.uint32, device=device) | |
| wp_arr_b = wp.array(np_arr_b, dtype=wp.float32, device=device) | |
| wp_arr_c = wp.array(np_arr_a, dtype=wp.uint16, device=device) | |
| wp_arr_d = wp.array(np_arr_b, dtype=wp.float16, device=device) | |
| wp_arr_e = wp.array(np_arr_e, dtype=wp.vec4, device=device) | |
| wp_arr_f = wp.array(np_arr_e, dtype=wp.quat, device=device) | |
| assert np.array_equal(np_arr_a.view(dtype=np.float32), wp_arr_a.view(dtype=wp.float32).numpy()) | |
| assert np.array_equal(np_arr_b.view(dtype=np.uint32), wp_arr_b.view(dtype=wp.uint32).numpy()) | |
| assert np.array_equal(np_arr_c.view(dtype=np.float16), wp_arr_c.view(dtype=wp.float16).numpy()) | |
| assert np.array_equal(np_arr_d.view(dtype=np.uint16), wp_arr_d.view(dtype=wp.uint16).numpy()) | |
| assert_array_equal(wp_arr_e.view(dtype=wp.quat), wp_arr_f) | |
| def compare_2darrays(x: wp.array2d(dtype=float), y: wp.array2d(dtype=float), z: wp.array2d(dtype=int)): | |
| i, j = wp.tid() | |
| if x[i, j] == y[i, j]: | |
| z[i, j] = 1 | |
| def compare_3darrays(x: wp.array3d(dtype=float), y: wp.array3d(dtype=float), z: wp.array3d(dtype=int)): | |
| i, j, k = wp.tid() | |
| if x[i, j, k] == y[i, j, k]: | |
| z[i, j, k] = 1 | |
| def test_transpose(test, device): | |
| # test default transpose in non-square 2d case | |
| # wp does not support copying from/to non-contiguous arrays so check in kernel | |
| np_arr = np.array([[1, 2], [3, 4], [5, 6]], dtype=float) | |
| arr = wp.array(np_arr, dtype=float, device=device) | |
| arr_transpose = arr.transpose() | |
| arr_compare = wp.array(np_arr.transpose(), dtype=float, device=device) | |
| check = wp.zeros(shape=(2, 3), dtype=int, device=device) | |
| wp.launch(compare_2darrays, dim=(2, 3), inputs=[arr_transpose, arr_compare, check], device=device) | |
| assert np.array_equal(check.numpy(), np.ones((2, 3), dtype=int)) | |
| # test transpose in square 3d case | |
| # wp does not support copying from/to non-contiguous arrays so check in kernel | |
| np_arr = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]], dtype=float) | |
| arr = wp.array3d(np_arr, dtype=float, shape=np_arr.shape, device=device, requires_grad=True) | |
| arr_transpose = arr.transpose((0, 2, 1)) | |
| arr_compare = wp.array3d(np_arr.transpose((0, 2, 1)), dtype=float, device=device) | |
| check = wp.zeros(shape=(3, 2, 2), dtype=int, device=device) | |
| wp.launch(compare_3darrays, dim=(3, 2, 2), inputs=[arr_transpose, arr_compare, check], device=device) | |
| assert np.array_equal(check.numpy(), np.ones((3, 2, 2), dtype=int)) | |
| # test transpose in square 3d case without axes supplied | |
| arr_transpose = arr.transpose() | |
| arr_compare = wp.array3d(np_arr.transpose(), dtype=float, device=device) | |
| check = wp.zeros(shape=(2, 2, 3), dtype=int, device=device) | |
| wp.launch(compare_3darrays, dim=(2, 2, 3), inputs=[arr_transpose, arr_compare, check], device=device) | |
| assert np.array_equal(check.numpy(), np.ones((2, 2, 3), dtype=int)) | |
| # test transpose in 1d case (should be noop) | |
| np_arr = np.array([1, 2, 3], dtype=float) | |
| arr = wp.array(np_arr, dtype=float, device=device) | |
| assert np.array_equal(np_arr.transpose(), arr.transpose().numpy()) | |
| def test_fill_scalar(test, device): | |
| dim_x = 4 | |
| for nptype, wptype in wp.types.np_dtype_to_warp_type.items(): | |
| a1 = wp.zeros(dim_x, dtype=wptype, device=device) | |
| a2 = wp.zeros((dim_x, dim_x), dtype=wptype, device=device) | |
| a3 = wp.zeros((dim_x, dim_x, dim_x), dtype=wptype, device=device) | |
| a4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=wptype, device=device) | |
| assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype)) | |
| assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype)) | |
| assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype)) | |
| assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype)) | |
| # fill with int value | |
| fill_value = 42 | |
| a1.fill_(fill_value) | |
| a2.fill_(fill_value) | |
| a3.fill_(fill_value) | |
| a4.fill_(fill_value) | |
| assert_np_equal(a1.numpy(), np.full(a1.shape, fill_value, dtype=nptype)) | |
| assert_np_equal(a2.numpy(), np.full(a2.shape, fill_value, dtype=nptype)) | |
| assert_np_equal(a3.numpy(), np.full(a3.shape, fill_value, dtype=nptype)) | |
| assert_np_equal(a4.numpy(), np.full(a4.shape, fill_value, dtype=nptype)) | |
| a1.zero_() | |
| a2.zero_() | |
| a3.zero_() | |
| a4.zero_() | |
| assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype)) | |
| assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype)) | |
| assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype)) | |
| assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype)) | |
| if wptype in wp.types.float_types: | |
| # fill with float value | |
| fill_value = 13.37 | |
| a1.fill_(fill_value) | |
| a2.fill_(fill_value) | |
| a3.fill_(fill_value) | |
| a4.fill_(fill_value) | |
| assert_np_equal(a1.numpy(), np.full(a1.shape, fill_value, dtype=nptype)) | |
| assert_np_equal(a2.numpy(), np.full(a2.shape, fill_value, dtype=nptype)) | |
| assert_np_equal(a3.numpy(), np.full(a3.shape, fill_value, dtype=nptype)) | |
| assert_np_equal(a4.numpy(), np.full(a4.shape, fill_value, dtype=nptype)) | |
| # fill with Warp scalar value | |
| fill_value = wptype(17) | |
| a1.fill_(fill_value) | |
| a2.fill_(fill_value) | |
| a3.fill_(fill_value) | |
| a4.fill_(fill_value) | |
| assert_np_equal(a1.numpy(), np.full(a1.shape, fill_value.value, dtype=nptype)) | |
| assert_np_equal(a2.numpy(), np.full(a2.shape, fill_value.value, dtype=nptype)) | |
| assert_np_equal(a3.numpy(), np.full(a3.shape, fill_value.value, dtype=nptype)) | |
| assert_np_equal(a4.numpy(), np.full(a4.shape, fill_value.value, dtype=nptype)) | |
| def test_fill_vector(test, device): | |
| # test filling a vector array with scalar or vector values (vec_type, list, or numpy array) | |
| dim_x = 4 | |
| for nptype, wptype in wp.types.np_dtype_to_warp_type.items(): | |
| # vector types | |
| vector_types = [ | |
| wp.types.vector(2, wptype), | |
| wp.types.vector(3, wptype), | |
| wp.types.vector(4, wptype), | |
| wp.types.vector(5, wptype), | |
| ] | |
| for vec_type in vector_types: | |
| vec_len = vec_type._length_ | |
| a1 = wp.zeros(dim_x, dtype=vec_type, device=device) | |
| a2 = wp.zeros((dim_x, dim_x), dtype=vec_type, device=device) | |
| a3 = wp.zeros((dim_x, dim_x, dim_x), dtype=vec_type, device=device) | |
| a4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=vec_type, device=device) | |
| assert_np_equal(a1.numpy(), np.zeros((*a1.shape, vec_len), dtype=nptype)) | |
| assert_np_equal(a2.numpy(), np.zeros((*a2.shape, vec_len), dtype=nptype)) | |
| assert_np_equal(a3.numpy(), np.zeros((*a3.shape, vec_len), dtype=nptype)) | |
| assert_np_equal(a4.numpy(), np.zeros((*a4.shape, vec_len), dtype=nptype)) | |
| # fill with int scalar | |
| fill_value = 42 | |
| a1.fill_(fill_value) | |
| a2.fill_(fill_value) | |
| a3.fill_(fill_value) | |
| a4.fill_(fill_value) | |
| assert_np_equal(a1.numpy(), np.full((*a1.shape, vec_len), fill_value, dtype=nptype)) | |
| assert_np_equal(a2.numpy(), np.full((*a2.shape, vec_len), fill_value, dtype=nptype)) | |
| assert_np_equal(a3.numpy(), np.full((*a3.shape, vec_len), fill_value, dtype=nptype)) | |
| assert_np_equal(a4.numpy(), np.full((*a4.shape, vec_len), fill_value, dtype=nptype)) | |
| # test zeroing | |
| a1.zero_() | |
| a2.zero_() | |
| a3.zero_() | |
| a4.zero_() | |
| assert_np_equal(a1.numpy(), np.zeros((*a1.shape, vec_len), dtype=nptype)) | |
| assert_np_equal(a2.numpy(), np.zeros((*a2.shape, vec_len), dtype=nptype)) | |
| assert_np_equal(a3.numpy(), np.zeros((*a3.shape, vec_len), dtype=nptype)) | |
| assert_np_equal(a4.numpy(), np.zeros((*a4.shape, vec_len), dtype=nptype)) | |
| # vector values can be passed as a list, numpy array, or Warp vector instance | |
| fill_list = [17, 42, 99, 101, 127][:vec_len] | |
| fill_arr = np.array(fill_list, dtype=nptype) | |
| fill_vec = vec_type(fill_list) | |
| expected1 = np.tile(fill_arr, a1.size).reshape((*a1.shape, vec_len)) | |
| expected2 = np.tile(fill_arr, a2.size).reshape((*a2.shape, vec_len)) | |
| expected3 = np.tile(fill_arr, a3.size).reshape((*a3.shape, vec_len)) | |
| expected4 = np.tile(fill_arr, a4.size).reshape((*a4.shape, vec_len)) | |
| # fill with list of vector length | |
| a1.fill_(fill_list) | |
| a2.fill_(fill_list) | |
| a3.fill_(fill_list) | |
| a4.fill_(fill_list) | |
| assert_np_equal(a1.numpy(), expected1) | |
| assert_np_equal(a2.numpy(), expected2) | |
| assert_np_equal(a3.numpy(), expected3) | |
| assert_np_equal(a4.numpy(), expected4) | |
| # clear | |
| a1.zero_() | |
| a2.zero_() | |
| a3.zero_() | |
| a4.zero_() | |
| # fill with numpy array of vector length | |
| a1.fill_(fill_arr) | |
| a2.fill_(fill_arr) | |
| a3.fill_(fill_arr) | |
| a4.fill_(fill_arr) | |
| assert_np_equal(a1.numpy(), expected1) | |
| assert_np_equal(a2.numpy(), expected2) | |
| assert_np_equal(a3.numpy(), expected3) | |
| assert_np_equal(a4.numpy(), expected4) | |
| # clear | |
| a1.zero_() | |
| a2.zero_() | |
| a3.zero_() | |
| a4.zero_() | |
| # fill with vec instance | |
| a1.fill_(fill_vec) | |
| a2.fill_(fill_vec) | |
| a3.fill_(fill_vec) | |
| a4.fill_(fill_vec) | |
| assert_np_equal(a1.numpy(), expected1) | |
| assert_np_equal(a2.numpy(), expected2) | |
| assert_np_equal(a3.numpy(), expected3) | |
| assert_np_equal(a4.numpy(), expected4) | |
| if wptype in wp.types.float_types: | |
| # fill with float scalar | |
| fill_value = 13.37 | |
| a1.fill_(fill_value) | |
| a2.fill_(fill_value) | |
| a3.fill_(fill_value) | |
| a4.fill_(fill_value) | |
| assert_np_equal(a1.numpy(), np.full((*a1.shape, vec_len), fill_value, dtype=nptype)) | |
| assert_np_equal(a2.numpy(), np.full((*a2.shape, vec_len), fill_value, dtype=nptype)) | |
| assert_np_equal(a3.numpy(), np.full((*a3.shape, vec_len), fill_value, dtype=nptype)) | |
| assert_np_equal(a4.numpy(), np.full((*a4.shape, vec_len), fill_value, dtype=nptype)) | |
| # fill with float list of vector length | |
| fill_list = [-2.5, -1.25, 1.25, 2.5, 5.0][:vec_len] | |
| a1.fill_(fill_list) | |
| a2.fill_(fill_list) | |
| a3.fill_(fill_list) | |
| a4.fill_(fill_list) | |
| expected1 = np.tile(np.array(fill_list, dtype=nptype), a1.size).reshape((*a1.shape, vec_len)) | |
| expected2 = np.tile(np.array(fill_list, dtype=nptype), a2.size).reshape((*a2.shape, vec_len)) | |
| expected3 = np.tile(np.array(fill_list, dtype=nptype), a3.size).reshape((*a3.shape, vec_len)) | |
| expected4 = np.tile(np.array(fill_list, dtype=nptype), a4.size).reshape((*a4.shape, vec_len)) | |
| assert_np_equal(a1.numpy(), expected1) | |
| assert_np_equal(a2.numpy(), expected2) | |
| assert_np_equal(a3.numpy(), expected3) | |
| assert_np_equal(a4.numpy(), expected4) | |
| def test_fill_matrix(test, device): | |
| # test filling a matrix array with scalar or matrix values (mat_type, nested list, or 2d numpy array) | |
| dim_x = 4 | |
| for nptype, wptype in wp.types.np_dtype_to_warp_type.items(): | |
| # matrix types | |
| matrix_types = [ | |
| # square matrices | |
| wp.types.matrix((2, 2), wptype), | |
| wp.types.matrix((3, 3), wptype), | |
| wp.types.matrix((4, 4), wptype), | |
| wp.types.matrix((5, 5), wptype), | |
| # non-square matrices | |
| wp.types.matrix((2, 3), wptype), | |
| wp.types.matrix((3, 2), wptype), | |
| wp.types.matrix((3, 4), wptype), | |
| wp.types.matrix((4, 3), wptype), | |
| ] | |
| for mat_type in matrix_types: | |
| mat_len = mat_type._length_ | |
| mat_shape = mat_type._shape_ | |
| a1 = wp.zeros(dim_x, dtype=mat_type, device=device) | |
| a2 = wp.zeros((dim_x, dim_x), dtype=mat_type, device=device) | |
| a3 = wp.zeros((dim_x, dim_x, dim_x), dtype=mat_type, device=device) | |
| a4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=mat_type, device=device) | |
| assert_np_equal(a1.numpy(), np.zeros((*a1.shape, *mat_shape), dtype=nptype)) | |
| assert_np_equal(a2.numpy(), np.zeros((*a2.shape, *mat_shape), dtype=nptype)) | |
| assert_np_equal(a3.numpy(), np.zeros((*a3.shape, *mat_shape), dtype=nptype)) | |
| assert_np_equal(a4.numpy(), np.zeros((*a4.shape, *mat_shape), dtype=nptype)) | |
| # fill with scalar | |
| fill_value = 42 | |
| a1.fill_(fill_value) | |
| a2.fill_(fill_value) | |
| a3.fill_(fill_value) | |
| a4.fill_(fill_value) | |
| assert_np_equal(a1.numpy(), np.full((*a1.shape, *mat_shape), fill_value, dtype=nptype)) | |
| assert_np_equal(a2.numpy(), np.full((*a2.shape, *mat_shape), fill_value, dtype=nptype)) | |
| assert_np_equal(a3.numpy(), np.full((*a3.shape, *mat_shape), fill_value, dtype=nptype)) | |
| assert_np_equal(a4.numpy(), np.full((*a4.shape, *mat_shape), fill_value, dtype=nptype)) | |
| # test zeroing | |
| a1.zero_() | |
| a2.zero_() | |
| a3.zero_() | |
| a4.zero_() | |
| assert_np_equal(a1.numpy(), np.zeros((*a1.shape, *mat_shape), dtype=nptype)) | |
| assert_np_equal(a2.numpy(), np.zeros((*a2.shape, *mat_shape), dtype=nptype)) | |
| assert_np_equal(a3.numpy(), np.zeros((*a3.shape, *mat_shape), dtype=nptype)) | |
| assert_np_equal(a4.numpy(), np.zeros((*a4.shape, *mat_shape), dtype=nptype)) | |
| # matrix values can be passed as a 1d numpy array, 2d numpy array, flat list, nested list, or Warp matrix instance | |
| if wptype != wp.bool: | |
| fill_arr1 = np.arange(mat_len, dtype=nptype) | |
| else: | |
| fill_arr1 = np.ones(mat_len, dtype=nptype) | |
| fill_arr2 = fill_arr1.reshape(mat_shape) | |
| fill_list1 = list(fill_arr1) | |
| fill_list2 = [list(row) for row in fill_arr2] | |
| fill_mat = mat_type(fill_arr1) | |
| expected1 = np.tile(fill_arr1, a1.size).reshape((*a1.shape, *mat_shape)) | |
| expected2 = np.tile(fill_arr1, a2.size).reshape((*a2.shape, *mat_shape)) | |
| expected3 = np.tile(fill_arr1, a3.size).reshape((*a3.shape, *mat_shape)) | |
| expected4 = np.tile(fill_arr1, a4.size).reshape((*a4.shape, *mat_shape)) | |
| # fill with 1d numpy array | |
| a1.fill_(fill_arr1) | |
| a2.fill_(fill_arr1) | |
| a3.fill_(fill_arr1) | |
| a4.fill_(fill_arr1) | |
| assert_np_equal(a1.numpy(), expected1) | |
| assert_np_equal(a2.numpy(), expected2) | |
| assert_np_equal(a3.numpy(), expected3) | |
| assert_np_equal(a4.numpy(), expected4) | |
| # clear | |
| a1.zero_() | |
| a2.zero_() | |
| a3.zero_() | |
| a4.zero_() | |
| # fill with 2d numpy array | |
| a1.fill_(fill_arr2) | |
| a2.fill_(fill_arr2) | |
| a3.fill_(fill_arr2) | |
| a4.fill_(fill_arr2) | |
| assert_np_equal(a1.numpy(), expected1) | |
| assert_np_equal(a2.numpy(), expected2) | |
| assert_np_equal(a3.numpy(), expected3) | |
| assert_np_equal(a4.numpy(), expected4) | |
| # clear | |
| a1.zero_() | |
| a2.zero_() | |
| a3.zero_() | |
| a4.zero_() | |
| # fill with flat list | |
| a1.fill_(fill_list1) | |
| a2.fill_(fill_list1) | |
| a3.fill_(fill_list1) | |
| a4.fill_(fill_list1) | |
| assert_np_equal(a1.numpy(), expected1) | |
| assert_np_equal(a2.numpy(), expected2) | |
| assert_np_equal(a3.numpy(), expected3) | |
| assert_np_equal(a4.numpy(), expected4) | |
| # clear | |
| a1.zero_() | |
| a2.zero_() | |
| a3.zero_() | |
| a4.zero_() | |
| # fill with nested list | |
| a1.fill_(fill_list2) | |
| a2.fill_(fill_list2) | |
| a3.fill_(fill_list2) | |
| a4.fill_(fill_list2) | |
| assert_np_equal(a1.numpy(), expected1) | |
| assert_np_equal(a2.numpy(), expected2) | |
| assert_np_equal(a3.numpy(), expected3) | |
| assert_np_equal(a4.numpy(), expected4) | |
| # clear | |
| a1.zero_() | |
| a2.zero_() | |
| a3.zero_() | |
| a4.zero_() | |
| # fill with mat instance | |
| a1.fill_(fill_mat) | |
| a2.fill_(fill_mat) | |
| a3.fill_(fill_mat) | |
| a4.fill_(fill_mat) | |
| assert_np_equal(a1.numpy(), expected1) | |
| assert_np_equal(a2.numpy(), expected2) | |
| assert_np_equal(a3.numpy(), expected3) | |
| assert_np_equal(a4.numpy(), expected4) | |
| class FillStruct: | |
| # scalar members (make sure to test float16) | |
| i1: wp.int8 | |
| i2: wp.int16 | |
| i4: wp.int32 | |
| i8: wp.int64 | |
| f2: wp.float16 | |
| f4: wp.float32 | |
| f8: wp.float16 | |
| # vector members (make sure to test vectors of float16) | |
| v2: wp.types.vector(2, wp.int64) | |
| v3: wp.types.vector(3, wp.float32) | |
| v4: wp.types.vector(4, wp.float16) | |
| v5: wp.types.vector(5, wp.uint8) | |
| # matrix members (make sure to test matrices of float16) | |
| m2: wp.types.matrix((2, 2), wp.float64) | |
| m3: wp.types.matrix((3, 3), wp.int32) | |
| m4: wp.types.matrix((4, 4), wp.float16) | |
| m5: wp.types.matrix((5, 5), wp.int8) | |
| # arrays | |
| a1: wp.array(dtype=float) | |
| a2: wp.array2d(dtype=float) | |
| a3: wp.array3d(dtype=float) | |
| a4: wp.array4d(dtype=float) | |
| def test_fill_struct(test, device): | |
| dim_x = 4 | |
| nptype = FillStruct.numpy_dtype() | |
| a1 = wp.zeros(dim_x, dtype=FillStruct, device=device) | |
| a2 = wp.zeros((dim_x, dim_x), dtype=FillStruct, device=device) | |
| a3 = wp.zeros((dim_x, dim_x, dim_x), dtype=FillStruct, device=device) | |
| a4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=FillStruct, device=device) | |
| assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype)) | |
| assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype)) | |
| assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype)) | |
| assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype)) | |
| s = FillStruct() | |
| # fill with default struct value (should be all zeros) | |
| a1.fill_(s) | |
| a2.fill_(s) | |
| a3.fill_(s) | |
| a4.fill_(s) | |
| assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype)) | |
| assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype)) | |
| assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype)) | |
| assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype)) | |
| # scalars | |
| s.i1 = -17 | |
| s.i2 = 42 | |
| s.i4 = 99 | |
| s.i8 = 101 | |
| s.f2 = -1.25 | |
| s.f4 = 13.37 | |
| s.f8 = 0.125 | |
| # vectors | |
| s.v2 = [21, 22] | |
| s.v3 = [31, 32, 33] | |
| s.v4 = [41, 42, 43, 44] | |
| s.v5 = [51, 52, 53, 54, 55] | |
| # matrices | |
| s.m2 = [[61, 62]] * 2 | |
| s.m3 = [[71, 72, 73]] * 3 | |
| s.m4 = [[81, 82, 83, 84]] * 4 | |
| s.m5 = [[91, 92, 93, 94, 95]] * 5 | |
| # arrays | |
| s.a1 = wp.zeros((2,) * 1, dtype=float, device=device) | |
| s.a2 = wp.zeros((2,) * 2, dtype=float, device=device) | |
| s.a3 = wp.zeros((2,) * 3, dtype=float, device=device) | |
| s.a4 = wp.zeros((2,) * 4, dtype=float, device=device) | |
| # fill with custom struct value | |
| a1.fill_(s) | |
| a2.fill_(s) | |
| a3.fill_(s) | |
| a4.fill_(s) | |
| ns = s.numpy_value() | |
| expected1 = np.empty(a1.shape, dtype=nptype) | |
| expected2 = np.empty(a2.shape, dtype=nptype) | |
| expected3 = np.empty(a3.shape, dtype=nptype) | |
| expected4 = np.empty(a4.shape, dtype=nptype) | |
| expected1.fill(ns) | |
| expected2.fill(ns) | |
| expected3.fill(ns) | |
| expected4.fill(ns) | |
| assert_np_equal(a1.numpy(), expected1) | |
| assert_np_equal(a2.numpy(), expected2) | |
| assert_np_equal(a3.numpy(), expected3) | |
| assert_np_equal(a4.numpy(), expected4) | |
| # test clearing | |
| a1.zero_() | |
| a2.zero_() | |
| a3.zero_() | |
| a4.zero_() | |
| assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype)) | |
| assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype)) | |
| assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype)) | |
| assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype)) | |
| def test_fill_slices(test, device): | |
| # test fill_ and zero_ for non-contiguous arrays | |
| # Note: we don't need to test the whole range of dtypes (vectors, matrices, structs) here | |
| dim_x = 8 | |
| for nptype, wptype in wp.types.np_dtype_to_warp_type.items(): | |
| a1 = wp.zeros(dim_x, dtype=wptype, device=device) | |
| a2 = wp.zeros((dim_x, dim_x), dtype=wptype, device=device) | |
| a3 = wp.zeros((dim_x, dim_x, dim_x), dtype=wptype, device=device) | |
| a4 = wp.zeros((dim_x, dim_x, dim_x, dim_x), dtype=wptype, device=device) | |
| assert_np_equal(a1.numpy(), np.zeros(a1.shape, dtype=nptype)) | |
| assert_np_equal(a2.numpy(), np.zeros(a2.shape, dtype=nptype)) | |
| assert_np_equal(a3.numpy(), np.zeros(a3.shape, dtype=nptype)) | |
| assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype)) | |
| # partititon each array into even and odd slices | |
| a1a = a1[::2] | |
| a1b = a1[1::2] | |
| a2a = a2[::2] | |
| a2b = a2[1::2] | |
| a3a = a3[::2] | |
| a3b = a3[1::2] | |
| a4a = a4[::2] | |
| a4b = a4[1::2] | |
| # fill even slices | |
| fill_a = 17 | |
| a1a.fill_(fill_a) | |
| a2a.fill_(fill_a) | |
| a3a.fill_(fill_a) | |
| a4a.fill_(fill_a) | |
| # ensure filled slices are correct | |
| assert_np_equal(a1a.numpy(), np.full(a1a.shape, fill_a, dtype=nptype)) | |
| assert_np_equal(a2a.numpy(), np.full(a2a.shape, fill_a, dtype=nptype)) | |
| assert_np_equal(a3a.numpy(), np.full(a3a.shape, fill_a, dtype=nptype)) | |
| assert_np_equal(a4a.numpy(), np.full(a4a.shape, fill_a, dtype=nptype)) | |
| # ensure unfilled slices are unaffected | |
| assert_np_equal(a1b.numpy(), np.zeros(a1b.shape, dtype=nptype)) | |
| assert_np_equal(a2b.numpy(), np.zeros(a2b.shape, dtype=nptype)) | |
| assert_np_equal(a3b.numpy(), np.zeros(a3b.shape, dtype=nptype)) | |
| assert_np_equal(a4b.numpy(), np.zeros(a4b.shape, dtype=nptype)) | |
| # fill odd slices | |
| fill_b = 42 | |
| a1b.fill_(fill_b) | |
| a2b.fill_(fill_b) | |
| a3b.fill_(fill_b) | |
| a4b.fill_(fill_b) | |
| # ensure filled slices are correct | |
| assert_np_equal(a1b.numpy(), np.full(a1b.shape, fill_b, dtype=nptype)) | |
| assert_np_equal(a2b.numpy(), np.full(a2b.shape, fill_b, dtype=nptype)) | |
| assert_np_equal(a3b.numpy(), np.full(a3b.shape, fill_b, dtype=nptype)) | |
| assert_np_equal(a4b.numpy(), np.full(a4b.shape, fill_b, dtype=nptype)) | |
| # ensure unfilled slices are unaffected | |
| assert_np_equal(a1a.numpy(), np.full(a1a.shape, fill_a, dtype=nptype)) | |
| assert_np_equal(a2a.numpy(), np.full(a2a.shape, fill_a, dtype=nptype)) | |
| assert_np_equal(a3a.numpy(), np.full(a3a.shape, fill_a, dtype=nptype)) | |
| assert_np_equal(a4a.numpy(), np.full(a4a.shape, fill_a, dtype=nptype)) | |
| # clear even slices | |
| a1a.zero_() | |
| a2a.zero_() | |
| a3a.zero_() | |
| a4a.zero_() | |
| # ensure cleared slices are correct | |
| assert_np_equal(a1a.numpy(), np.zeros(a1a.shape, dtype=nptype)) | |
| assert_np_equal(a2a.numpy(), np.zeros(a2a.shape, dtype=nptype)) | |
| assert_np_equal(a3a.numpy(), np.zeros(a3a.shape, dtype=nptype)) | |
| assert_np_equal(a4a.numpy(), np.zeros(a4a.shape, dtype=nptype)) | |
| # ensure uncleared slices are unaffected | |
| assert_np_equal(a1b.numpy(), np.full(a1b.shape, fill_b, dtype=nptype)) | |
| assert_np_equal(a2b.numpy(), np.full(a2b.shape, fill_b, dtype=nptype)) | |
| assert_np_equal(a3b.numpy(), np.full(a3b.shape, fill_b, dtype=nptype)) | |
| assert_np_equal(a4b.numpy(), np.full(a4b.shape, fill_b, dtype=nptype)) | |
| # re-fill even slices | |
| a1a.fill_(fill_a) | |
| a2a.fill_(fill_a) | |
| a3a.fill_(fill_a) | |
| a4a.fill_(fill_a) | |
| # clear odd slices | |
| a1b.zero_() | |
| a2b.zero_() | |
| a3b.zero_() | |
| a4b.zero_() | |
| # ensure cleared slices are correct | |
| assert_np_equal(a1b.numpy(), np.zeros(a1b.shape, dtype=nptype)) | |
| assert_np_equal(a2b.numpy(), np.zeros(a2b.shape, dtype=nptype)) | |
| assert_np_equal(a3b.numpy(), np.zeros(a3b.shape, dtype=nptype)) | |
| assert_np_equal(a4b.numpy(), np.zeros(a4b.shape, dtype=nptype)) | |
| # ensure uncleared slices are unaffected | |
| assert_np_equal(a1a.numpy(), np.full(a1a.shape, fill_a, dtype=nptype)) | |
| assert_np_equal(a2a.numpy(), np.full(a2a.shape, fill_a, dtype=nptype)) | |
| assert_np_equal(a3a.numpy(), np.full(a3a.shape, fill_a, dtype=nptype)) | |
| assert_np_equal(a4a.numpy(), np.full(a4a.shape, fill_a, dtype=nptype)) | |
| def test_full_scalar(test, device): | |
| dim = 4 | |
| for ndim in range(1, 5): | |
| shape = (dim,) * ndim | |
| for nptype, wptype in wp.types.np_dtype_to_warp_type.items(): | |
| # fill with int value and specific dtype | |
| fill_value = 42 | |
| a = wp.full(shape, fill_value, dtype=wptype, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.dtype, wptype) | |
| test.assertEqual(na.shape, shape) | |
| test.assertEqual(na.dtype, nptype) | |
| assert_np_equal(na, np.full(shape, fill_value, dtype=nptype)) | |
| if wptype in wp.types.float_types: | |
| # fill with float value and specific dtype | |
| fill_value = 13.37 | |
| a = wp.full(shape, fill_value, dtype=wptype, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.dtype, wptype) | |
| test.assertEqual(na.shape, shape) | |
| test.assertEqual(na.dtype, nptype) | |
| assert_np_equal(na, np.full(shape, fill_value, dtype=nptype)) | |
| # fill with int value and automatically inferred dtype | |
| fill_value = 42 | |
| a = wp.full(shape, fill_value, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.dtype, wp.int32) | |
| test.assertEqual(na.shape, shape) | |
| test.assertEqual(na.dtype, np.int32) | |
| assert_np_equal(na, np.full(shape, fill_value, dtype=np.int32)) | |
| # fill with float value and automatically inferred dtype | |
| fill_value = 13.37 | |
| a = wp.full(shape, fill_value, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.dtype, wp.float32) | |
| test.assertEqual(na.shape, shape) | |
| test.assertEqual(na.dtype, np.float32) | |
| assert_np_equal(na, np.full(shape, fill_value, dtype=np.float32)) | |
| def test_full_vector(test, device): | |
| dim = 4 | |
| for ndim in range(1, 5): | |
| shape = (dim,) * ndim | |
| # full from scalar | |
| for veclen in [2, 3, 4, 5]: | |
| npshape = (*shape, veclen) | |
| for nptype, wptype in wp.types.np_dtype_to_warp_type.items(): | |
| vectype = wp.types.vector(veclen, wptype) | |
| # fill with scalar int value and specific dtype | |
| fill_value = 42 | |
| a = wp.full(shape, fill_value, dtype=vectype, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.dtype, vectype) | |
| test.assertEqual(na.shape, npshape) | |
| test.assertEqual(na.dtype, nptype) | |
| assert_np_equal(na, np.full(a.size * veclen, fill_value, dtype=nptype).reshape(npshape)) | |
| if wptype in wp.types.float_types: | |
| # fill with scalar float value and specific dtype | |
| fill_value = 13.37 | |
| a = wp.full(shape, fill_value, dtype=vectype, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.dtype, vectype) | |
| test.assertEqual(na.shape, npshape) | |
| test.assertEqual(na.dtype, nptype) | |
| assert_np_equal(na, np.full(a.size * veclen, fill_value, dtype=nptype).reshape(npshape)) | |
| # fill with vector value and specific dtype | |
| fill_vec = vectype(42) | |
| a = wp.full(shape, fill_vec, dtype=vectype, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.dtype, vectype) | |
| test.assertEqual(na.shape, npshape) | |
| test.assertEqual(na.dtype, nptype) | |
| assert_np_equal(na, np.full(a.size * veclen, 42, dtype=nptype).reshape(npshape)) | |
| # fill with vector value and automatically inferred dtype | |
| a = wp.full(shape, fill_vec, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.dtype, vectype) | |
| test.assertEqual(na.shape, npshape) | |
| test.assertEqual(na.dtype, nptype) | |
| assert_np_equal(na, np.full(a.size * veclen, 42, dtype=nptype).reshape(npshape)) | |
| fill_lists = [ | |
| [17, 42], | |
| [17, 42, 99], | |
| [17, 42, 99, 101], | |
| [17, 42, 99, 101, 127], | |
| ] | |
| # full from list and numpy array | |
| for fill_list in fill_lists: | |
| veclen = len(fill_list) | |
| npshape = (*shape, veclen) | |
| for nptype, wptype in wp.types.np_dtype_to_warp_type.items(): | |
| vectype = wp.types.vector(veclen, wptype) | |
| # fill with list and specific dtype | |
| a = wp.full(shape, fill_list, dtype=vectype, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.dtype, vectype) | |
| test.assertEqual(na.shape, npshape) | |
| test.assertEqual(na.dtype, nptype) | |
| expected = np.tile(np.array(fill_list, dtype=nptype), a.size).reshape(npshape) | |
| assert_np_equal(na, expected) | |
| fill_arr = np.array(fill_list, dtype=nptype) | |
| # fill with numpy array and specific dtype | |
| a = wp.full(shape, fill_arr, dtype=vectype, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.dtype, vectype) | |
| test.assertEqual(na.shape, npshape) | |
| test.assertEqual(na.dtype, nptype) | |
| assert_np_equal(na, expected) | |
| # fill with numpy array and automatically infer dtype | |
| a = wp.full(shape, fill_arr, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertTrue(wp.types.types_equal(a.dtype, vectype)) | |
| test.assertEqual(na.shape, npshape) | |
| test.assertEqual(na.dtype, nptype) | |
| assert_np_equal(na, expected) | |
| # fill with list and automatically infer dtype | |
| a = wp.full(shape, fill_list, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| # check that the inferred dtype is a vector | |
| # Note that we cannot guarantee the scalar type, because it depends on numpy and may vary by platform | |
| # (e.g. int64 on Linux and int32 on Windows). | |
| test.assertEqual(a.dtype._wp_generic_type_str_, "vec_t") | |
| test.assertEqual(a.dtype._length_, veclen) | |
| expected = np.tile(np.array(fill_list), a.size).reshape(npshape) | |
| assert_np_equal(na, expected) | |
| def test_full_matrix(test, device): | |
| dim = 4 | |
| for ndim in range(1, 5): | |
| shape = (dim,) * ndim | |
| for nptype, wptype in wp.types.np_dtype_to_warp_type.items(): | |
| matrix_types = [ | |
| # square matrices | |
| wp.types.matrix((2, 2), wptype), | |
| wp.types.matrix((3, 3), wptype), | |
| wp.types.matrix((4, 4), wptype), | |
| wp.types.matrix((5, 5), wptype), | |
| # non-square matrices | |
| wp.types.matrix((2, 3), wptype), | |
| wp.types.matrix((3, 2), wptype), | |
| wp.types.matrix((3, 4), wptype), | |
| wp.types.matrix((4, 3), wptype), | |
| ] | |
| for mattype in matrix_types: | |
| npshape = (*shape, *mattype._shape_) | |
| # fill with scalar int value and specific dtype | |
| fill_value = 42 | |
| a = wp.full(shape, fill_value, dtype=mattype, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.dtype, mattype) | |
| test.assertEqual(na.shape, npshape) | |
| test.assertEqual(na.dtype, nptype) | |
| assert_np_equal(na, np.full(a.size * mattype._length_, fill_value, dtype=nptype).reshape(npshape)) | |
| if wptype in wp.types.float_types: | |
| # fill with scalar float value and specific dtype | |
| fill_value = 13.37 | |
| a = wp.full(shape, fill_value, dtype=mattype, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.dtype, mattype) | |
| test.assertEqual(na.shape, npshape) | |
| test.assertEqual(na.dtype, nptype) | |
| assert_np_equal(na, np.full(a.size * mattype._length_, fill_value, dtype=nptype).reshape(npshape)) | |
| # fill with matrix value and specific dtype | |
| fill_mat = mattype(42) | |
| a = wp.full(shape, fill_mat, dtype=mattype, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.dtype, mattype) | |
| test.assertEqual(na.shape, npshape) | |
| test.assertEqual(na.dtype, nptype) | |
| assert_np_equal(na, np.full(a.size * mattype._length_, 42, dtype=nptype).reshape(npshape)) | |
| # fill with matrix value and automatically inferred dtype | |
| fill_mat = mattype(42) | |
| a = wp.full(shape, fill_mat, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.dtype, mattype) | |
| test.assertEqual(na.shape, npshape) | |
| test.assertEqual(na.dtype, nptype) | |
| assert_np_equal(na, np.full(a.size * mattype._length_, 42, dtype=nptype).reshape(npshape)) | |
| # fill with 1d numpy array and specific dtype | |
| if wptype != wp.bool: | |
| fill_arr1d = np.arange(mattype._length_, dtype=nptype) | |
| else: | |
| fill_arr1d = np.ones(mattype._length_, dtype=nptype) | |
| a = wp.full(shape, fill_arr1d, dtype=mattype, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.dtype, mattype) | |
| test.assertEqual(na.shape, npshape) | |
| test.assertEqual(na.dtype, nptype) | |
| expected = np.tile(fill_arr1d, a.size).reshape(npshape) | |
| assert_np_equal(na, expected) | |
| # fill with 2d numpy array and specific dtype | |
| fill_arr2d = fill_arr1d.reshape(mattype._shape_) | |
| a = wp.full(shape, fill_arr2d, dtype=mattype, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.dtype, mattype) | |
| test.assertEqual(na.shape, npshape) | |
| test.assertEqual(na.dtype, nptype) | |
| assert_np_equal(na, expected) | |
| # fill with 2d numpy array and automatically infer dtype | |
| a = wp.full(shape, fill_arr2d, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertTrue(wp.types.types_equal(a.dtype, mattype)) | |
| test.assertEqual(na.shape, npshape) | |
| test.assertEqual(na.dtype, nptype) | |
| assert_np_equal(na, expected) | |
| # fill with flat list and specific dtype | |
| fill_list1d = list(fill_arr1d) | |
| a = wp.full(shape, fill_list1d, dtype=mattype, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.dtype, mattype) | |
| test.assertEqual(na.shape, npshape) | |
| test.assertEqual(na.dtype, nptype) | |
| assert_np_equal(na, expected) | |
| # fill with nested list and specific dtype | |
| fill_list2d = [list(row) for row in fill_arr2d] | |
| a = wp.full(shape, fill_list2d, dtype=mattype, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.dtype, mattype) | |
| test.assertEqual(na.shape, npshape) | |
| test.assertEqual(na.dtype, nptype) | |
| assert_np_equal(na, expected) | |
| mat_lists = [ | |
| # square matrices | |
| [[1, 2], [3, 4]], | |
| [[1, 2, 3], [4, 5, 6], [7, 8, 9]], | |
| [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]], | |
| # non-square matrices | |
| [[1, 2, 3, 4], [5, 6, 7, 8]], | |
| [[1, 2], [3, 4], [5, 6], [7, 8]], | |
| ] | |
| # fill with nested lists and automatically infer dtype | |
| for fill_list in mat_lists: | |
| num_rows = len(fill_list) | |
| num_cols = len(fill_list[0]) | |
| npshape = (*shape, num_rows, num_cols) | |
| a = wp.full(shape, fill_list, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| # check that the inferred dtype is a correctly shaped matrix | |
| # Note that we cannot guarantee the scalar type, because it depends on numpy and may vary by platform | |
| # (e.g. int64 on Linux and int32 on Windows). | |
| test.assertEqual(a.dtype._wp_generic_type_str_, "mat_t") | |
| test.assertEqual(a.dtype._shape_, (num_rows, num_cols)) | |
| expected = np.tile(np.array(fill_list).flatten(), a.size).reshape(npshape) | |
| assert_np_equal(na, expected) | |
| def test_full_struct(test, device): | |
| dim = 4 | |
| for ndim in range(1, 5): | |
| shape = (dim,) * ndim | |
| s = FillStruct() | |
| # fill with default struct (should be zeros) | |
| a = wp.full(shape, s, dtype=FillStruct, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.dtype, FillStruct) | |
| test.assertEqual(na.shape, shape) | |
| test.assertEqual(na.dtype, FillStruct.numpy_dtype()) | |
| assert_np_equal(na, np.zeros(a.size, dtype=FillStruct.numpy_dtype())) | |
| # scalars | |
| s.i1 = -17 | |
| s.i2 = 42 | |
| s.i4 = 99 | |
| s.i8 = 101 | |
| s.f2 = -1.25 | |
| s.f4 = 13.37 | |
| s.f8 = 0.125 | |
| # vectors | |
| s.v2 = [21, 22] | |
| s.v3 = [31, 32, 33] | |
| s.v4 = [41, 42, 43, 44] | |
| s.v5 = [51, 52, 53, 54, 55] | |
| # matrices | |
| s.m2 = [[61, 62]] * 2 | |
| s.m3 = [[71, 72, 73]] * 3 | |
| s.m4 = [[81, 82, 83, 84]] * 4 | |
| s.m5 = [[91, 92, 93, 94, 95]] * 5 | |
| # arrays | |
| s.a1 = wp.zeros((2,) * 1, dtype=float, device=device) | |
| s.a2 = wp.zeros((2,) * 2, dtype=float, device=device) | |
| s.a3 = wp.zeros((2,) * 3, dtype=float, device=device) | |
| s.a4 = wp.zeros((2,) * 4, dtype=float, device=device) | |
| # fill with initialized struct and explicit dtype | |
| a = wp.full(shape, s, dtype=FillStruct, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.dtype, FillStruct) | |
| test.assertEqual(na.shape, shape) | |
| test.assertEqual(na.dtype, FillStruct.numpy_dtype()) | |
| expected = np.empty(shape, dtype=FillStruct.numpy_dtype()) | |
| expected.fill(s.numpy_value()) | |
| assert_np_equal(na, expected) | |
| # fill with initialized struct and automatically inferred dtype | |
| a = wp.full(shape, s, device=device) | |
| na = a.numpy() | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.dtype, FillStruct) | |
| test.assertEqual(na.shape, shape) | |
| test.assertEqual(na.dtype, FillStruct.numpy_dtype()) | |
| assert_np_equal(na, expected) | |
| def test_round_trip(test, device): | |
| rng = np.random.default_rng(123) | |
| dim_x = 4 | |
| for nptype, wptype in wp.types.np_dtype_to_warp_type.items(): | |
| a_np = rng.standard_normal(size=dim_x).astype(nptype) | |
| a = wp.array(a_np, device=device) | |
| test.assertEqual(a.dtype, wptype) | |
| assert_np_equal(a.numpy(), a_np) | |
| v_np = rng.standard_normal(size=(dim_x, 3)).astype(nptype) | |
| v = wp.array(v_np, dtype=wp.types.vector(3, wptype), device=device) | |
| assert_np_equal(v.numpy(), v_np) | |
| def test_empty_array(test, device): | |
| # Test whether common operations work with empty (zero-sized) arrays | |
| # without throwing exceptions. | |
| def test_empty_ops(ndim, nrows, ncols, wptype, nptype): | |
| shape = (0,) * ndim | |
| dtype_shape = () | |
| if wptype in wp.types.scalar_types: | |
| # scalar, vector, or matrix | |
| if ncols > 0: | |
| if nrows > 0: | |
| wptype = wp.types.matrix((nrows, ncols), wptype) | |
| else: | |
| wptype = wp.types.vector(ncols, wptype) | |
| dtype_shape = wptype._shape_ | |
| fill_value = wptype(42) | |
| else: | |
| # struct | |
| fill_value = wptype() | |
| # create a zero-sized array | |
| a = wp.empty(shape, dtype=wptype, device=device, requires_grad=True) | |
| test.assertEqual(a.ptr, None) | |
| test.assertEqual(a.size, 0) | |
| test.assertEqual(a.shape, shape) | |
| test.assertEqual(a.grad.ptr, None) | |
| test.assertEqual(a.grad.size, 0) | |
| test.assertEqual(a.grad.shape, shape) | |
| # all of these methods should succeed with zero-sized arrays | |
| a.zero_() | |
| a.fill_(fill_value) | |
| b = a.flatten() | |
| b = a.reshape((0,)) | |
| b = a.transpose() | |
| b = a.contiguous() | |
| b = wp.empty_like(a) | |
| b = wp.zeros_like(a) | |
| b = wp.full_like(a, fill_value) | |
| b = wp.clone(a) | |
| wp.copy(a, b) | |
| a.assign(b) | |
| na = a.numpy() | |
| test.assertEqual(na.size, 0) | |
| test.assertEqual(na.shape, (*shape, *dtype_shape)) | |
| test.assertEqual(na.dtype, nptype) | |
| test.assertEqual(a.list(), []) | |
| for ndim in range(1, 5): | |
| # test with scalars, vectors, and matrices | |
| for nptype, wptype in wp.types.np_dtype_to_warp_type.items(): | |
| # scalars | |
| test_empty_ops(ndim, 0, 0, wptype, nptype) | |
| for ncols in [2, 3, 4, 5]: | |
| # vectors | |
| test_empty_ops(ndim, 0, ncols, wptype, nptype) | |
| # square matrices | |
| test_empty_ops(ndim, ncols, ncols, wptype, nptype) | |
| # non-square matrices | |
| test_empty_ops(ndim, 2, 3, wptype, nptype) | |
| test_empty_ops(ndim, 3, 2, wptype, nptype) | |
| test_empty_ops(ndim, 3, 4, wptype, nptype) | |
| test_empty_ops(ndim, 4, 3, wptype, nptype) | |
| # test with structs | |
| test_empty_ops(ndim, 0, 0, FillStruct, FillStruct.numpy_dtype()) | |
| def test_empty_from_numpy(test, device): | |
| # Test whether wrapping an empty (zero-sized) numpy array works correctly | |
| def test_empty_from_data(ndim, nrows, ncols, wptype, nptype): | |
| shape = (0,) * ndim | |
| dtype_shape = () | |
| if ncols > 0: | |
| if nrows > 0: | |
| wptype = wp.types.matrix((nrows, ncols), wptype) | |
| else: | |
| wptype = wp.types.vector(ncols, wptype) | |
| dtype_shape = wptype._shape_ | |
| npshape = (*shape, *dtype_shape) | |
| na = np.empty(npshape, dtype=nptype) | |
| a = wp.array(na, dtype=wptype, device=device) | |
| test.assertEqual(a.size, 0) | |
| test.assertEqual(a.shape, shape) | |
| for ndim in range(1, 5): | |
| # test with scalars, vectors, and matrices | |
| for nptype, wptype in wp.types.np_dtype_to_warp_type.items(): | |
| # scalars | |
| test_empty_from_data(ndim, 0, 0, wptype, nptype) | |
| for ncols in [2, 3, 4, 5]: | |
| # vectors | |
| test_empty_from_data(ndim, 0, ncols, wptype, nptype) | |
| # square matrices | |
| test_empty_from_data(ndim, ncols, ncols, wptype, nptype) | |
| # non-square matrices | |
| test_empty_from_data(ndim, 2, 3, wptype, nptype) | |
| test_empty_from_data(ndim, 3, 2, wptype, nptype) | |
| test_empty_from_data(ndim, 3, 4, wptype, nptype) | |
| test_empty_from_data(ndim, 4, 3, wptype, nptype) | |
| def test_empty_from_list(test, device): | |
| # Test whether creating an array from an empty Python list works correctly | |
| def test_empty_from_data(nrows, ncols, wptype): | |
| if ncols > 0: | |
| if nrows > 0: | |
| wptype = wp.types.matrix((nrows, ncols), wptype) | |
| else: | |
| wptype = wp.types.vector(ncols, wptype) | |
| a = wp.array([], dtype=wptype, device=device) | |
| test.assertEqual(a.size, 0) | |
| test.assertEqual(a.shape, (0,)) | |
| # test with scalars, vectors, and matrices | |
| for wptype in wp.types.scalar_types: | |
| # scalars | |
| test_empty_from_data(0, 0, wptype) | |
| for ncols in [2, 3, 4, 5]: | |
| # vectors | |
| test_empty_from_data(0, ncols, wptype) | |
| # square matrices | |
| test_empty_from_data(ncols, ncols, wptype) | |
| # non-square matrices | |
| test_empty_from_data(2, 3, wptype) | |
| test_empty_from_data(3, 2, wptype) | |
| test_empty_from_data(3, 4, wptype) | |
| test_empty_from_data(4, 3, wptype) | |
| def test_to_list_scalar(test, device): | |
| dim = 3 | |
| fill_value = 42 | |
| for ndim in range(1, 5): | |
| shape = (dim,) * ndim | |
| for wptype in wp.types.scalar_types: | |
| a = wp.full(shape, fill_value, dtype=wptype, device=device) | |
| l = a.list() | |
| test.assertEqual(len(l), a.size) | |
| test.assertTrue(all(x == fill_value for x in l)) | |
| def test_to_list_vector(test, device): | |
| dim = 3 | |
| for ndim in range(1, 5): | |
| shape = (dim,) * ndim | |
| for veclen in [2, 3, 4, 5]: | |
| for wptype in wp.types.scalar_types: | |
| vectype = wp.types.vector(veclen, wptype) | |
| fill_value = vectype(42) | |
| a = wp.full(shape, fill_value, dtype=vectype, device=device) | |
| l = a.list() | |
| test.assertEqual(len(l), a.size) | |
| test.assertTrue(all(x == fill_value for x in l)) | |
| def test_to_list_matrix(test, device): | |
| dim = 3 | |
| for ndim in range(1, 5): | |
| shape = (dim,) * ndim | |
| for wptype in wp.types.scalar_types: | |
| matrix_types = [ | |
| # square matrices | |
| wp.types.matrix((2, 2), wptype), | |
| wp.types.matrix((3, 3), wptype), | |
| wp.types.matrix((4, 4), wptype), | |
| wp.types.matrix((5, 5), wptype), | |
| # non-square matrices | |
| wp.types.matrix((2, 3), wptype), | |
| wp.types.matrix((3, 2), wptype), | |
| wp.types.matrix((3, 4), wptype), | |
| wp.types.matrix((4, 3), wptype), | |
| ] | |
| for mattype in matrix_types: | |
| fill_value = mattype(42) | |
| a = wp.full(shape, fill_value, dtype=mattype, device=device) | |
| l = a.list() | |
| test.assertEqual(len(l), a.size) | |
| test.assertTrue(all(x == fill_value for x in l)) | |
| def test_to_list_struct(test, device): | |
| class Inner: | |
| h: wp.float16 | |
| v: wp.vec3 | |
| class ListStruct: | |
| i: int | |
| f: float | |
| h: wp.float16 | |
| vi: wp.vec2i | |
| vf: wp.vec3f | |
| vh: wp.vec4h | |
| mi: wp.types.matrix((2, 2), int) | |
| mf: wp.types.matrix((3, 3), float) | |
| mh: wp.types.matrix((4, 4), wp.float16) | |
| inner: Inner | |
| a1: wp.array(dtype=int) | |
| a2: wp.array2d(dtype=float) | |
| a3: wp.array3d(dtype=wp.float16) | |
| bool: wp.bool | |
| dim = 3 | |
| s = ListStruct() | |
| s.i = 42 | |
| s.f = 2.5 | |
| s.h = -1.25 | |
| s.vi = wp.vec2i(1, 2) | |
| s.vf = wp.vec3f(0.1, 0.2, 0.3) | |
| s.vh = wp.vec4h(1.0, 2.0, 3.0, 4.0) | |
| s.mi = [[1, 2], [3, 4]] | |
| s.mf = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] | |
| s.mh = [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]] | |
| s.inner = Inner() | |
| s.inner.h = 1.5 | |
| s.inner.v = [1, 2, 3] | |
| s.a1 = wp.empty(1, dtype=int, device=device) | |
| s.a2 = wp.empty((1, 1), dtype=float, device=device) | |
| s.a3 = wp.empty((1, 1, 1), dtype=wp.float16, device=device) | |
| s.bool = True | |
| for ndim in range(1, 5): | |
| shape = (dim,) * ndim | |
| a = wp.full(shape, s, dtype=ListStruct, device=device) | |
| l = a.list() | |
| for i in range(a.size): | |
| test.assertEqual(l[i].i, s.i) | |
| test.assertEqual(l[i].f, s.f) | |
| test.assertEqual(l[i].h, s.h) | |
| test.assertEqual(l[i].vi, s.vi) | |
| test.assertEqual(l[i].vf, s.vf) | |
| test.assertEqual(l[i].vh, s.vh) | |
| test.assertEqual(l[i].mi, s.mi) | |
| test.assertEqual(l[i].mf, s.mf) | |
| test.assertEqual(l[i].mh, s.mh) | |
| test.assertEqual(l[i].bool, s.bool) | |
| test.assertEqual(l[i].inner.h, s.inner.h) | |
| test.assertEqual(l[i].inner.v, s.inner.v) | |
| test.assertEqual(l[i].a1.dtype, s.a1.dtype) | |
| test.assertEqual(l[i].a1.ndim, s.a1.ndim) | |
| test.assertEqual(l[i].a2.dtype, s.a2.dtype) | |
| test.assertEqual(l[i].a2.ndim, s.a2.ndim) | |
| test.assertEqual(l[i].a3.dtype, s.a3.dtype) | |
| test.assertEqual(l[i].a3.ndim, s.a3.ndim) | |
| def kernel_array_to_bool(array_null: wp.array(dtype=float), array_valid: wp.array(dtype=float)): | |
| if not array_null: | |
| # always succeed | |
| wp.expect_eq(0, 0) | |
| else: | |
| # force failure | |
| wp.expect_eq(1, 2) | |
| if array_valid: | |
| # always succeed | |
| wp.expect_eq(0, 0) | |
| else: | |
| # force failure | |
| wp.expect_eq(1, 2) | |
| def test_array_to_bool(test, device): | |
| arr = wp.zeros(8, dtype=float, device=device) | |
| wp.launch(kernel_array_to_bool, dim=1, inputs=[None, arr], device=device) | |
| class InputStruct: | |
| param1: int | |
| param2: float | |
| param3: wp.vec3 | |
| param4: wp.array(dtype=float) | |
| class OutputStruct: | |
| param1: int | |
| param2: float | |
| param3: wp.vec3 | |
| def struct_array_kernel(inputs: wp.array(dtype=InputStruct), outputs: wp.array(dtype=OutputStruct)): | |
| tid = wp.tid() | |
| wp.expect_eq(inputs[tid].param1, tid) | |
| wp.expect_eq(inputs[tid].param2, float(tid * tid)) | |
| wp.expect_eq(inputs[tid].param3[0], 1.0) | |
| wp.expect_eq(inputs[tid].param3[1], 2.0) | |
| wp.expect_eq(inputs[tid].param3[2], 3.0) | |
| wp.expect_eq(inputs[tid].param4[0], 1.0) | |
| wp.expect_eq(inputs[tid].param4[1], 2.0) | |
| wp.expect_eq(inputs[tid].param4[2], 3.0) | |
| o = OutputStruct() | |
| o.param1 = inputs[tid].param1 | |
| o.param2 = inputs[tid].param2 | |
| o.param3 = inputs[tid].param3 | |
| outputs[tid] = o | |
| def test_array_of_structs(test, device): | |
| num_items = 10 | |
| l = [] | |
| for i in range(num_items): | |
| s = InputStruct() | |
| s.param1 = i | |
| s.param2 = float(i * i) | |
| s.param3 = wp.vec3(1.0, 2.0, 3.0) | |
| s.param4 = wp.array([1.0, 2.0, 3.0], dtype=float, device=device) | |
| l.append(s) | |
| # initialize array from list of structs | |
| inputs = wp.array(l, dtype=InputStruct, device=device) | |
| outputs = wp.zeros(num_items, dtype=OutputStruct, device=device) | |
| # pass to our compute kernel | |
| wp.launch(struct_array_kernel, dim=num_items, inputs=[inputs, outputs], device=device) | |
| out_numpy = outputs.numpy() | |
| out_list = outputs.list() | |
| out_cptr = outputs.to("cpu").cptr() | |
| for i in range(num_items): | |
| test.assertEqual(out_numpy[i][0], l[i].param1) | |
| test.assertEqual(out_numpy[i][1], l[i].param2) | |
| assert_np_equal(out_numpy[i][2], np.array(l[i].param3)) | |
| # test named slices of numpy structured array | |
| test.assertEqual(out_numpy["param1"][i], l[i].param1) | |
| test.assertEqual(out_numpy["param2"][i], l[i].param2) | |
| assert_np_equal(out_numpy["param3"][i], np.array(l[i].param3)) | |
| test.assertEqual(out_list[i].param1, l[i].param1) | |
| test.assertEqual(out_list[i].param2, l[i].param2) | |
| test.assertEqual(out_list[i].param3, l[i].param3) | |
| test.assertEqual(out_cptr[i].param1, l[i].param1) | |
| test.assertEqual(out_cptr[i].param2, l[i].param2) | |
| test.assertEqual(out_cptr[i].param3, l[i].param3) | |
| class GradStruct: | |
| param1: int | |
| param2: float | |
| param3: wp.vec3 | |
| def test_array_of_structs_grad_kernel(inputs: wp.array(dtype=GradStruct), loss: wp.array(dtype=float)): | |
| tid = wp.tid() | |
| wp.atomic_add(loss, 0, inputs[tid].param2 * 2.0) | |
| def test_array_of_structs_grad(test, device): | |
| num_items = 10 | |
| l = [] | |
| for i in range(num_items): | |
| g = GradStruct() | |
| g.param2 = float(i) | |
| l.append(g) | |
| a = wp.array(l, dtype=GradStruct, device=device, requires_grad=True) | |
| loss = wp.zeros(1, dtype=float, device=device, requires_grad=True) | |
| with wp.Tape() as tape: | |
| wp.launch(test_array_of_structs_grad_kernel, dim=num_items, inputs=[a, loss], device=device) | |
| tape.backward(loss) | |
| grads = a.grad.numpy() | |
| assert_np_equal(grads["param2"], np.full(num_items, 2.0, dtype=np.float32)) | |
| class NumpyStruct: | |
| x: int | |
| v: wp.vec3 | |
| def test_array_of_structs_from_numpy(test, device): | |
| num_items = 10 | |
| na = np.zeros(num_items, dtype=NumpyStruct.numpy_dtype()) | |
| na["x"] = 17 | |
| na["v"] = (1, 2, 3) | |
| a = wp.array(data=na, dtype=NumpyStruct, device=device) | |
| assert_np_equal(a.numpy(), na) | |
| def test_array_of_structs_roundtrip(test, device): | |
| num_items = 10 | |
| value = NumpyStruct() | |
| value.x = 17 | |
| value.v = wp.vec3(1.0, 2.0, 3.0) | |
| # create Warp structured array | |
| a = wp.full(num_items, value, device=device) | |
| # convert to NumPy structured array | |
| na = a.numpy() | |
| expected = np.zeros(num_items, dtype=NumpyStruct.numpy_dtype()) | |
| expected["x"] = value.x | |
| expected["v"] = value.v | |
| assert_np_equal(na, expected) | |
| # modify a field | |
| na["x"] = 42 | |
| # convert back to Warp array | |
| a = wp.from_numpy(na, NumpyStruct, device=device) | |
| expected["x"] = 42 | |
| assert_np_equal(a.numpy(), expected) | |
| def test_array_from_numpy(test, device): | |
| arr = np.array((1.0, 2.0, 3.0), dtype=float) | |
| result = wp.from_numpy(arr) | |
| expected = wp.array((1.0, 2.0, 3.0), dtype=wp.float32, shape=(3,)) | |
| assert_np_equal(result.numpy(), expected.numpy()) | |
| result = wp.from_numpy(arr, dtype=wp.vec3) | |
| expected = wp.array(((1.0, 2.0, 3.0),), dtype=wp.vec3, shape=(1,)) | |
| assert_np_equal(result.numpy(), expected.numpy()) | |
| # -------------------------------------------------------------------------- | |
| arr = np.array(((1.0, 2.0, 3.0), (4.0, 5.0, 6.0)), dtype=float) | |
| result = wp.from_numpy(arr) | |
| expected = wp.array(((1.0, 2.0, 3.0), (4.0, 5.0, 6.0)), dtype=wp.vec3, shape=(2,)) | |
| assert_np_equal(result.numpy(), expected.numpy()) | |
| result = wp.from_numpy(arr, dtype=wp.float32) | |
| expected = wp.array(((1.0, 2.0, 3.0), (4.0, 5.0, 6.0)), dtype=wp.float32, shape=(2, 3)) | |
| assert_np_equal(result.numpy(), expected.numpy()) | |
| result = wp.from_numpy(arr, dtype=wp.float32, shape=(6,)) | |
| expected = wp.array((1.0, 2.0, 3.0, 4.0, 5.0, 6.0), dtype=wp.float32, shape=(6,)) | |
| assert_np_equal(result.numpy(), expected.numpy()) | |
| # -------------------------------------------------------------------------- | |
| arr = np.array( | |
| ( | |
| ( | |
| (1.0, 2.0, 3.0, 4.0), | |
| (2.0, 3.0, 4.0, 5.0), | |
| (3.0, 4.0, 5.0, 6.0), | |
| (4.0, 5.0, 6.0, 7.0), | |
| ), | |
| ( | |
| (2.0, 3.0, 4.0, 5.0), | |
| (3.0, 4.0, 5.0, 6.0), | |
| (4.0, 5.0, 6.0, 7.0), | |
| (5.0, 6.0, 7.0, 8.0), | |
| ), | |
| ), | |
| dtype=float, | |
| ) | |
| result = wp.from_numpy(arr) | |
| expected = wp.array( | |
| ( | |
| ( | |
| (1.0, 2.0, 3.0, 4.0), | |
| (2.0, 3.0, 4.0, 5.0), | |
| (3.0, 4.0, 5.0, 6.0), | |
| (4.0, 5.0, 6.0, 7.0), | |
| ), | |
| ( | |
| (2.0, 3.0, 4.0, 5.0), | |
| (3.0, 4.0, 5.0, 6.0), | |
| (4.0, 5.0, 6.0, 7.0), | |
| (5.0, 6.0, 7.0, 8.0), | |
| ), | |
| ), | |
| dtype=wp.mat44, | |
| shape=(2,), | |
| ) | |
| assert_np_equal(result.numpy(), expected.numpy()) | |
| result = wp.from_numpy(arr, dtype=wp.float32) | |
| expected = wp.array( | |
| ( | |
| ( | |
| (1.0, 2.0, 3.0, 4.0), | |
| (2.0, 3.0, 4.0, 5.0), | |
| (3.0, 4.0, 5.0, 6.0), | |
| (4.0, 5.0, 6.0, 7.0), | |
| ), | |
| ( | |
| (2.0, 3.0, 4.0, 5.0), | |
| (3.0, 4.0, 5.0, 6.0), | |
| (4.0, 5.0, 6.0, 7.0), | |
| (5.0, 6.0, 7.0, 8.0), | |
| ), | |
| ), | |
| dtype=wp.float32, | |
| shape=(2, 4, 4), | |
| ) | |
| assert_np_equal(result.numpy(), expected.numpy()) | |
| result = wp.from_numpy(arr, dtype=wp.vec4) | |
| expected = wp.array( | |
| ( | |
| (1.0, 2.0, 3.0, 4.0), | |
| (2.0, 3.0, 4.0, 5.0), | |
| (3.0, 4.0, 5.0, 6.0), | |
| (4.0, 5.0, 6.0, 7.0), | |
| (2.0, 3.0, 4.0, 5.0), | |
| (3.0, 4.0, 5.0, 6.0), | |
| (4.0, 5.0, 6.0, 7.0), | |
| (5.0, 6.0, 7.0, 8.0), | |
| ), | |
| dtype=wp.vec4, | |
| shape=(8,), | |
| ) | |
| assert_np_equal(result.numpy(), expected.numpy()) | |
| result = wp.from_numpy(arr, dtype=wp.float32, shape=(32,)) | |
| expected = wp.array( | |
| ( | |
| 1.0, | |
| 2.0, | |
| 3.0, | |
| 4.0, | |
| 2.0, | |
| 3.0, | |
| 4.0, | |
| 5.0, | |
| 3.0, | |
| 4.0, | |
| 5.0, | |
| 6.0, | |
| 4.0, | |
| 5.0, | |
| 6.0, | |
| 7.0, | |
| 2.0, | |
| 3.0, | |
| 4.0, | |
| 5.0, | |
| 3.0, | |
| 4.0, | |
| 5.0, | |
| 6.0, | |
| 4.0, | |
| 5.0, | |
| 6.0, | |
| 7.0, | |
| 5.0, | |
| 6.0, | |
| 7.0, | |
| 8.0, | |
| ), | |
| dtype=wp.float32, | |
| shape=(32,), | |
| ) | |
| assert_np_equal(result.numpy(), expected.numpy()) | |
| devices = get_test_devices() | |
| class TestArray(unittest.TestCase): | |
| pass | |
| add_function_test(TestArray, "test_shape", test_shape, devices=devices) | |
| add_function_test(TestArray, "test_flatten", test_flatten, devices=devices) | |
| add_function_test(TestArray, "test_reshape", test_reshape, devices=devices) | |
| add_function_test(TestArray, "test_slicing", test_slicing, devices=devices) | |
| add_function_test(TestArray, "test_transpose", test_transpose, devices=devices) | |
| add_function_test(TestArray, "test_view", test_view, devices=devices) | |
| add_function_test(TestArray, "test_1d_array", test_1d, devices=devices) | |
| add_function_test(TestArray, "test_2d_array", test_2d, devices=devices) | |
| add_function_test(TestArray, "test_3d_array", test_3d, devices=devices) | |
| add_function_test(TestArray, "test_4d_array", test_4d, devices=devices) | |
| add_function_test(TestArray, "test_4d_array_transposed", test_4d_transposed, devices=devices) | |
| add_function_test(TestArray, "test_fill_scalar", test_fill_scalar, devices=devices) | |
| add_function_test(TestArray, "test_fill_vector", test_fill_vector, devices=devices) | |
| add_function_test(TestArray, "test_fill_matrix", test_fill_matrix, devices=devices) | |
| add_function_test(TestArray, "test_fill_struct", test_fill_struct, devices=devices) | |
| add_function_test(TestArray, "test_fill_slices", test_fill_slices, devices=devices) | |
| add_function_test(TestArray, "test_full_scalar", test_full_scalar, devices=devices) | |
| add_function_test(TestArray, "test_full_vector", test_full_vector, devices=devices) | |
| add_function_test(TestArray, "test_full_matrix", test_full_matrix, devices=devices) | |
| add_function_test(TestArray, "test_full_struct", test_full_struct, devices=devices) | |
| add_function_test(TestArray, "test_empty_array", test_empty_array, devices=devices) | |
| add_function_test(TestArray, "test_empty_from_numpy", test_empty_from_numpy, devices=devices) | |
| add_function_test(TestArray, "test_empty_from_list", test_empty_from_list, devices=devices) | |
| add_function_test(TestArray, "test_to_list_scalar", test_to_list_scalar, devices=devices) | |
| add_function_test(TestArray, "test_to_list_vector", test_to_list_vector, devices=devices) | |
| add_function_test(TestArray, "test_to_list_matrix", test_to_list_matrix, devices=devices) | |
| add_function_test(TestArray, "test_to_list_struct", test_to_list_struct, devices=devices) | |
| add_function_test(TestArray, "test_lower_bound", test_lower_bound, devices=devices) | |
| add_function_test(TestArray, "test_round_trip", test_round_trip, devices=devices) | |
| add_function_test(TestArray, "test_array_to_bool", test_array_to_bool, devices=devices) | |
| add_function_test(TestArray, "test_array_of_structs", test_array_of_structs, devices=devices) | |
| add_function_test(TestArray, "test_array_of_structs_grad", test_array_of_structs_grad, devices=devices) | |
| add_function_test(TestArray, "test_array_of_structs_from_numpy", test_array_of_structs_from_numpy, devices=devices) | |
| add_function_test(TestArray, "test_array_of_structs_roundtrip", test_array_of_structs_roundtrip, devices=devices) | |
| add_function_test(TestArray, "test_array_from_numpy", test_array_from_numpy, devices=devices) | |
| if __name__ == "__main__": | |
| wp.build.clear_kernel_cache() | |
| unittest.main(verbosity=2) | |