warp-lang 1.9.1__py3-none-manylinux_2_34_aarch64.whl → 1.10.0rc2__py3-none-manylinux_2_34_aarch64.whl

warp/tests/test_quat.py

@@ -18,7 +18,6 @@ import unittest
 import numpy as np
 
 import warp as wp
-import warp.sim
 from warp.tests.unittest_utils import *
 
 np_float_types = [np.float32, np.float64, np.float16]
@@ -26,6 +25,43 @@ np_float_types = [np.float32, np.float64, np.float16]
 kernel_cache = {}
 
 
+@wp.func
+def quat_from_euler(e: wp.vec3, i: int, j: int, k: int) -> wp.quat:
+    """
+    Convert Euler angles to a quaternion.
+
+    :math:`i, j, k` are the indices in :math:`[0, 1, 2]` of the axes in which the Euler angles are provided
+    (:math:`i \\neq j, j \\neq k`), e.g. (0, 1, 2) for Euler sequence XYZ.
+
+    Args:
+        e (vec3): The Euler angles (in radians)
+        i (int): The index of the first axis
+        j (int): The index of the second axis
+        k (int): The index of the third axis
+
+    Returns:
+        quat: The quaternion
+    """
+    # Half angles
+    half_e = e / 2.0
+
+    # Precompute sines and cosines of half angles
+    cr = wp.cos(half_e[i])
+    sr = wp.sin(half_e[i])
+    cp = wp.cos(half_e[j])
+    sp = wp.sin(half_e[j])
+    cy = wp.cos(half_e[k])
+    sy = wp.sin(half_e[k])
+
+    # Components of the quaternion based on the rotation sequence
+    return wp.quat(
+        (cy * sr * cp - sy * cr * sp),
+        (cy * cr * sp + sy * sr * cp),
+        (sy * cr * cp - cy * sr * sp),
+        (cy * cr * cp + sy * sr * sp),
+    )
+
+
 def getkernel(func, suffix=""):
     key = func.__name__ + "_" + suffix
     if key not in kernel_cache:
@@ -56,9 +92,9 @@ def test_constructors(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    vec3 = wp.types.vector(length=3, dtype=wptype)
-    quat = wp.types.quaternion(dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    vec3 = wp._src.types.vector(length=3, dtype=wptype)
+    quat = wp._src.types.quaternion(dtype=wptype)
 
     def check_component_constructor(
         input: wp.array(dtype=wptype),
@@ -130,17 +166,17 @@ def test_constructors(test, device, dtype, register_kernels=False):
 
 def test_casting_constructors(test, device, dtype, register_kernels=False):
     np_type = np.dtype(dtype)
-    wp_type = wp.types.np_dtype_to_warp_type[np_type]
-    quat = wp.types.quaternion(dtype=wp_type)
+    wp_type = wp._src.types.np_dtype_to_warp_type[np_type]
+    quat = wp._src.types.quaternion(dtype=wp_type)
 
     np16 = np.dtype(np.float16)
-    wp16 = wp.types.np_dtype_to_warp_type[np16]
+    wp16 = wp._src.types.np_dtype_to_warp_type[np16]
 
     np32 = np.dtype(np.float32)
-    wp32 = wp.types.np_dtype_to_warp_type[np32]
+    wp32 = wp._src.types.np_dtype_to_warp_type[np32]
 
     np64 = np.dtype(np.float64)
-    wp64 = wp.types.np_dtype_to_warp_type[np64]
+    wp64 = wp._src.types.np_dtype_to_warp_type[np64]
 
     def cast_float16(a: wp.array(dtype=wp_type, ndim=2), b: wp.array(dtype=wp16, ndim=2)):
         tid = wp.tid()
@@ -243,8 +279,8 @@ def test_inverse(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    quat = wp.types.quaternion(dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    quat = wp._src.types.quaternion(dtype=wptype)
 
     output_select_kernel = get_select_kernel(wptype)
 
@@ -299,8 +335,8 @@ def test_dotproduct(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    quat = wp.types.quaternion(dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    quat = wp._src.types.quaternion(dtype=wptype)
 
     def check_quat_dot(
         s: wp.array(dtype=quat),
@@ -351,8 +387,8 @@ def test_length(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-7,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    quat = wp.types.quaternion(dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    quat = wp._src.types.quaternion(dtype=wptype)
 
     def check_quat_length(
         q: wp.array(dtype=quat),
@@ -408,8 +444,8 @@ def test_normalize(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    quat = wp.types.quaternion(dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    quat = wp._src.types.quaternion(dtype=wptype)
 
     def check_normalize(
         q: wp.array(dtype=quat),
@@ -509,8 +545,8 @@ def test_addition(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    quat = wp.types.quaternion(dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    quat = wp._src.types.quaternion(dtype=wptype)
 
     def check_quat_add(
         q: wp.array(dtype=quat),
@@ -581,8 +617,8 @@ def test_subtraction(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    quat = wp.types.quaternion(dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    quat = wp._src.types.quaternion(dtype=wptype)
 
     def check_quat_sub(
         q: wp.array(dtype=quat),
@@ -654,8 +690,8 @@ def test_scalar_multiplication(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    quat = wp.types.quaternion(dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    quat = wp._src.types.quaternion(dtype=wptype)
 
     def check_quat_scalar_mul(
         s: wp.array(dtype=wptype),
@@ -752,8 +788,8 @@ def test_scalar_division(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    quat = wp.types.quaternion(dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    quat = wp._src.types.quaternion(dtype=wptype)
 
     def check_quat_scalar_div(
         s: wp.array(dtype=wptype),
@@ -825,8 +861,8 @@ def test_quat_multiplication(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    quat = wp.types.quaternion(dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    quat = wp._src.types.quaternion(dtype=wptype)
 
     def check_quat_mul(
         s: wp.array(dtype=quat),
@@ -929,8 +965,8 @@ def test_indexing(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    quat = wp.types.quaternion(dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    quat = wp._src.types.quaternion(dtype=wptype)
 
     def check_quat_indexing(
         q: wp.array(dtype=quat),
@@ -996,8 +1032,8 @@ def test_quat_lerp(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    quat = wp.types.quaternion(dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    quat = wp._src.types.quaternion(dtype=wptype)
 
     def check_quat_lerp(
         s: wp.array(dtype=quat),
@@ -1077,9 +1113,9 @@ def test_quat_rotate(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    quat = wp.types.quaternion(dtype=wptype)
-    vec3 = wp.types.vector(length=3, dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    quat = wp._src.types.quaternion(dtype=wptype)
+    vec3 = wp._src.types.vector(length=3, dtype=wptype)
 
     def check_quat_rotate(
         q: wp.array(dtype=quat),
@@ -1203,9 +1239,9 @@ def test_quat_to_matrix(test, device, dtype, register_kernels=False):
         np.float64: 1.0e-8,
     }.get(dtype, 0)
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    quat = wp.types.quaternion(dtype=wptype)
-    vec3 = wp.types.vector(length=3, dtype=wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    quat = wp._src.types.quaternion(dtype=wptype)
+    vec3 = wp._src.types.vector(length=3, dtype=wptype)
 
     def check_quat_to_matrix(
         q: wp.array(dtype=quat),
@@ -1321,9 +1357,9 @@ def test_slerp_grad(test, device, dtype, register_kernels=False):
     rng = np.random.default_rng(123)
     seed = 42
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    vec3 = wp.types.vector(3, wptype)
-    quat = wp.types.quaternion(wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    vec3 = wp._src.types.vector(3, wptype)
+    quat = wp._src.types.quaternion(wptype)
 
     def slerp_kernel(
         q0: wp.array(dtype=quat),
@@ -1480,14 +1516,13 @@ def test_slerp_grad(test, device, dtype, register_kernels=False):
 
 
 def test_quat_to_axis_angle_grad(test, device, dtype, register_kernels=False):
-    rng = np.random.default_rng(123)
     seed = 42
     num_rand = 50
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    vec3 = wp.types.vector(3, wptype)
-    vec4 = wp.types.vector(4, wptype)
-    quat = wp.types.quaternion(wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    vec3 = wp._src.types.vector(3, wptype)
+    vec4 = wp._src.types.vector(4, wptype)
+    quat = wp._src.types.quaternion(wptype)
 
     def quat_to_axis_angle_kernel(quats: wp.array(dtype=quat), loss: wp.array(dtype=wptype), coord_idx: int):
         tid = wp.tid()
@@ -1616,10 +1651,10 @@ def test_quat_rpy_grad(test, device, dtype, register_kernels=False):
     rng = np.random.default_rng(123)
     N = 3
 
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
 
-    vec3 = wp.types.vector(3, wptype)
-    quat = wp.types.quaternion(wptype)
+    vec3 = wp._src.types.vector(3, wptype)
+    quat = wp._src.types.quaternion(wptype)
 
     def rpy_to_quat_kernel(rpy_arr: wp.array(dtype=vec3), loss: wp.array(dtype=wptype), coord_idx: int):
         tid = wp.tid()
@@ -1708,10 +1743,10 @@ def test_quat_rpy_grad(test, device, dtype, register_kernels=False):
 
 
 def test_quat_from_matrix(test, device, dtype, register_kernels=False):
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-    mat33 = wp.types.matrix((3, 3), wptype)
-    mat44 = wp.types.matrix((4, 4), wptype)
-    quat = wp.types.quaternion(wptype)
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    mat33 = wp._src.types.matrix((3, 3), wptype)
+    mat44 = wp._src.types.matrix((4, 4), wptype)
+    quat = wp._src.types.quaternion(wptype)
 
     def quat_from_matrix(m: wp.array2d(dtype=wptype), loss: wp.array(dtype=wptype), idx: int):
         tid = wp.tid()
@@ -1867,7 +1902,7 @@ def test_quat_from_matrix(test, device, dtype, register_kernels=False):
 
 
 def test_quat_identity(test, device, dtype, register_kernels=False):
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
 
     def quat_identity_test(output: wp.array(dtype=wptype)):
         q = wp.quat_identity(dtype=wptype)
@@ -1909,7 +1944,7 @@ def test_quat_euler_conversion(test, device, dtype, register_kernels=False):
 
     rpy_arr = rng.uniform(low=-np.pi, high=np.pi, size=(N, 3))
 
-    quats_from_euler = [list(wp.sim.quat_from_euler(wp.vec3(*rpy), 0, 1, 2)) for rpy in rpy_arr]
+    quats_from_euler = [list(quat_from_euler(wp.vec3(*rpy), 0, 1, 2)) for rpy in rpy_arr]
     quats_from_rpy = [list(wp.quat_rpy(rpy[0], rpy[1], rpy[2])) for rpy in rpy_arr]
 
     assert_np_equal(np.array(quats_from_euler), np.array(quats_from_rpy), tol=1e-4)
@@ -1917,7 +1952,7 @@ def test_quat_euler_conversion(test, device, dtype, register_kernels=False):
 
 def test_anon_type_instance(test, device, dtype, register_kernels=False):
     rng = np.random.default_rng(123)
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
 
     def quat_create_test(input: wp.array(dtype=wptype), output: wp.array(dtype=wptype)):
         # component constructor:
@@ -1990,16 +2025,16 @@ def test_constructor_default():
 
 
 def test_py_arithmetic_ops(test, device, dtype):
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
 
     def make_quat(*args):
-        if wptype in wp.types.int_types:
+        if wptype in wp._src.types.int_types:
             # Cast to the correct integer type to simulate wrapping.
             return tuple(wptype._type_(x).value for x in args)
 
         return args
 
-    quat_cls = wp.types.quaternion(wptype)
+    quat_cls = wp._src.types.quaternion(wptype)
 
     v = quat_cls(1, -2, 3, -4)
     test.assertSequenceEqual(+v, make_quat(1, -2, 3, -4))

warp/tests/test_runlength_encode.py

@@ -19,8 +19,8 @@ from functools import partial
 import numpy as np
 
 import warp as wp
+from warp._src.utils import runlength_encode
 from warp.tests.unittest_utils import *
-from warp.utils import runlength_encode
 
 
 def test_runlength_encode_int(test, device, n):

warp/tests/test_scalar_ops.py

@@ -44,10 +44,10 @@ np_scalar_types = np_int_types + np_float_types
 
 
 def test_py_arithmetic_ops(test, device, dtype):
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
 
     def make_scalar(value):
-        if wptype in wp.types.int_types:
+        if wptype in wp._src.types.int_types:
             # Cast to the correct integer type to simulate wrapping.
             return wptype._type_(value).value
 
@@ -69,14 +69,7 @@ def test_py_arithmetic_ops(test, device, dtype):
 
 
 def test_py_math_ops(test, device, dtype):
-    wptype = wp.types.np_dtype_to_warp_type[np.dtype(dtype)]
-
-    def make_scalar(value):
-        if wptype in wp.types.int_types:
-            # Cast to the correct integer type to simulate wrapping.
-            return wptype._type_(value).value
-
-        return value
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
 
     a = wptype(1)
     test.assertAlmostEqual(wp.abs(a), 1)
@@ -86,6 +79,38 @@ def test_py_math_ops(test, device, dtype):
         test.assertAlmostEqual(wp.radians(a), 0.01745329251994329577, places=5)
 
 
+def test_py_array_ops(test, device, dtype):
+    wptype = wp._src.types.np_dtype_to_warp_type[np.dtype(dtype)]
+
+    tol = {
+        np.float16: 5.0e-3,
+        np.float32: 1.0e-6,
+        np.float64: 1.0e-8,
+    }.get(dtype, 0)
+
+    def make_array(*args):
+        if wptype in wp._src.types.int_types:
+            # Cast to the correct integer type to simulate wrapping.
+            return np.array(tuple(wptype._type_(x).value for x in args), dtype=dtype)
+
+        return np.array(args, dtype=dtype)
+
+    arr = wp.array((1, 2, 3), dtype=wptype)
+    s = wptype(5)
+    assert_np_equal((arr + s).numpy(), make_array(6, 7, 8), tol=tol)
+    assert_np_equal((s + arr).numpy(), make_array(6, 7, 8), tol=tol)
+    assert_np_equal((arr - s).numpy(), make_array(-4, -3, -2), tol=tol)
+    assert_np_equal((s - arr).numpy(), make_array(4, 3, 2), tol=tol)
+    assert_np_equal((arr * s).numpy(), make_array(5, 10, 15), tol=tol)
+    assert_np_equal((s * arr).numpy(), make_array(5, 10, 15), tol=tol)
+    assert_np_equal((arr % s).numpy(), make_array(1, 2, 3), tol=tol)
+    assert_np_equal((s % arr).numpy(), make_array(0, 1, 2), tol=tol)
+
+    if dtype in np_float_types:
+        assert_np_equal((arr / s).numpy(), make_array(0.2, 0.4, 0.6), tol=tol)
+        assert_np_equal((s / arr).numpy(), make_array(5.0, 2.5, 1.666666667), tol=tol)
+
+
 devices = get_test_devices()
 
 
@@ -98,6 +123,9 @@ for dtype in np_scalar_types:
         TestScalarOps, f"test_py_arithmetic_ops_{dtype.__name__}", test_py_arithmetic_ops, devices=None, dtype=dtype
     )
     add_function_test(TestScalarOps, f"test_py_math_ops_{dtype.__name__}", test_py_math_ops, devices=None, dtype=dtype)
+    add_function_test(
+        TestScalarOps, f"test_py_array_ops_{dtype.__name__}", test_py_array_ops, devices=None, dtype=dtype
+    )
 
 
 if __name__ == "__main__":

warp/tests/test_smoothstep.py

@@ -78,7 +78,7 @@ def test_smoothstep(test, device):
             a = wp.array([test_data.a], dtype=data_type, device=device, requires_grad=True)
             b = wp.array([test_data.b], dtype=data_type, device=device, requires_grad=True)
             t = wp.array([test_data.t], dtype=float, device=device, requires_grad=True)
-            out = wp.array([0] * wp.types.type_size(data_type), dtype=data_type, device=device, requires_grad=True)
+            out = wp.array([0] * wp._src.types.type_size(data_type), dtype=data_type, device=device, requires_grad=True)
 
             with wp.Tape() as tape:
                 wp.launch(kernel, dim=1, inputs=[a, b, t, out], device=device)