warp-lang 1.0.0b2__py3-none-manylinux2014_x86_64.whl → 1.0.0b6__py3-none-manylinux2014_x86_64.whl

warp/utils.py

@@ -6,11 +6,10 @@
 # license agreement from NVIDIA CORPORATION is strictly prohibited.
 
 import cProfile
-import math
 import sys
 import timeit
 import warnings
-from typing import Any, Tuple, Union
+from typing import Any
 
 import numpy as np
 
@@ -31,157 +30,9 @@ def warn(message, category=None, stacklevel=1):
         warnings.warn(message, category, stacklevel + 1)  # Increment stacklevel by 1 since we are in a wrapper
 
 
-def length(a):
-    return np.linalg.norm(a)
-
-
-def length_sq(a):
-    return np.dot(a, a)
-
-
-def cross(a, b):
-    return np.array((a[1] * b[2] - a[2] * b[1], a[2] * b[0] - a[0] * b[2], a[0] * b[1] - a[1] * b[0]), dtype=np.float32)
-
-
-# NumPy has no normalize() method..
-def normalize(v):
-    norm = np.linalg.norm(v)
-    if norm == 0.0:
-        return v
-    return v / norm
-
-
-def skew(v):
-    return np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
-
-
-# math utils
-# def quat(i, j, k, w):
-#     return np.array([i, j, k, w])
-
-
-def quat_identity():
-    return np.array((0.0, 0.0, 0.0, 1.0))
-
-
-def quat_inverse(q):
-    return np.array((-q[0], -q[1], -q[2], q[3]))
-
-
-def quat_from_axis_angle(axis, angle):
-    v = normalize(np.array(axis))
-
-    half = angle * 0.5
-    w = math.cos(half)
-
-    sin_theta_over_two = math.sin(half)
-    v *= sin_theta_over_two
-
-    return np.array((v[0], v[1], v[2], w))
-
-
-def quat_to_axis_angle(quat):
-    w2 = quat[3] * quat[3]
-    if w2 > 1 - 1e-7:
-        return np.zeros(3), 0.0
-
-    angle = 2 * np.arccos(quat[3])
-    xyz = quat[:3] / np.sqrt(1 - w2)
-    return xyz, angle
-
-
-# quat_rotate a vector
-def quat_rotate(q, x):
-    x = np.array(x)
-    axis = np.array((q[0], q[1], q[2]))
-    return x * (2.0 * q[3] * q[3] - 1.0) + np.cross(axis, x) * q[3] * 2.0 + axis * np.dot(axis, x) * 2.0
-
-
-# multiply two quats
-def quat_multiply(a, b):
-    return np.array(
-        (
-            a[3] * b[0] + b[3] * a[0] + a[1] * b[2] - b[1] * a[2],
-            a[3] * b[1] + b[3] * a[1] + a[2] * b[0] - b[2] * a[0],
-            a[3] * b[2] + b[3] * a[2] + a[0] * b[1] - b[0] * a[1],
-            a[3] * b[3] - a[0] * b[0] - a[1] * b[1] - a[2] * b[2],
-        )
-    )
-
-
-# convert to mat33
-def quat_to_matrix(q):
-    c1 = quat_rotate(q, np.array((1.0, 0.0, 0.0)))
-    c2 = quat_rotate(q, np.array((0.0, 1.0, 0.0)))
-    c3 = quat_rotate(q, np.array((0.0, 0.0, 1.0)))
-
-    return np.array([c1, c2, c3]).T
-
-
-def quat_rpy(roll, pitch, yaw):
-    cy = math.cos(yaw * 0.5)
-    sy = math.sin(yaw * 0.5)
-    cr = math.cos(roll * 0.5)
-    sr = math.sin(roll * 0.5)
-    cp = math.cos(pitch * 0.5)
-    sp = math.sin(pitch * 0.5)
-
-    w = cy * cr * cp + sy * sr * sp
-    x = cy * sr * cp - sy * cr * sp
-    y = cy * cr * sp + sy * sr * cp
-    z = sy * cr * cp - cy * sr * sp
-
-    return (x, y, z, w)
-
-
-def quat_from_matrix(m):
-    tr = m[0, 0] + m[1, 1] + m[2, 2]
-    h = 0.0
-
-    if tr >= 0.0:
-        h = math.sqrt(tr + 1.0)
-        w = 0.5 * h
-        h = 0.5 / h
-
-        x = (m[2, 1] - m[1, 2]) * h
-        y = (m[0, 2] - m[2, 0]) * h
-        z = (m[1, 0] - m[0, 1]) * h
-
-    else:
-        i = 0
-        if m[1, 1] > m[0, 0]:
-            i = 1
-        if m[2, 2] > m[i, i]:
-            i = 2
-
-        if i == 0:
-            h = math.sqrt((m[0, 0] - (m[1, 1] + m[2, 2])) + 1.0)
-            x = 0.5 * h
-            h = 0.5 / h
-
-            y = (m[0, 1] + m[1, 0]) * h
-            z = (m[2, 0] + m[0, 2]) * h
-            w = (m[2, 1] - m[1, 2]) * h
-
-        elif i == 1:
-            h = math.sqrt((m[1, 1] - (m[2, 2] + m[0, 0])) + 1.0)
-            y = 0.5 * h
-            h = 0.5 / h
-
-            z = (m[1, 2] + m[2, 1]) * h
-            x = (m[0, 1] + m[1, 0]) * h
-            w = (m[0, 2] - m[2, 0]) * h
-
-        elif i == 2:
-            h = math.sqrt((m[2, 2] - (m[0, 0] + m[1, 1])) + 1.0)
-            z = 0.5 * h
-            h = 0.5 / h
-
-            x = (m[2, 0] + m[0, 2]) * h
-            y = (m[1, 2] + m[2, 1]) * h
-            w = (m[1, 0] - m[0, 1]) * h
-
-    return normalize(np.array([x, y, z, w]))
+# expand a 7-vec to a tuple of arrays
+def transform_expand(t):
+    return wp.transform(np.array(t[0:3]), np.array(t[3:7]))
 
 
 @wp.func
@@ -197,210 +48,6 @@ def quat_between_vectors(a: wp.vec3, b: wp.vec3) -> wp.quat:
     return wp.normalize(q)
 
 
-# rigid body transform
-
-
-# def transform(x, r):
-#     return (np.array(x), np.array(r))
-
-
-def transform_identity():
-    return wp.transform(np.array((0.0, 0.0, 0.0)), quat_identity())
-
-
-# se(3) -> SE(3), Park & Lynch pg. 105, screw in [w, v] normalized form
-def transform_exp(s, angle):
-    w = np.array(s[0:3])
-    v = np.array(s[3:6])
-
-    if length(w) < 1.0:
-        r = quat_identity()
-    else:
-        r = quat_from_axis_angle(w, angle)
-
-    t = v * angle + (1.0 - math.cos(angle)) * np.cross(w, v) + (angle - math.sin(angle)) * np.cross(w, np.cross(w, v))
-
-    return (t, r)
-
-
-def transform_inverse(t):
-    q_inv = quat_inverse(t.q)
-    return wp.transform(-quat_rotate(q_inv, t.p), q_inv)
-
-
-def transform_vector(t, v):
-    return quat_rotate(t.q, v)
-
-
-def transform_point(t, p):
-    return np.array(t.p) + quat_rotate(t.q, p)
-
-
-def transform_multiply(t, u):
-    return wp.transform(quat_rotate(t.q, u.p) + t.p, quat_multiply(t.q, u.q))
-
-
-# flatten an array of transforms (p,q) format to a 7-vector
-def transform_flatten(t):
-    return np.array([*t.p, *t.q])
-
-
-# expand a 7-vec to a tuple of arrays
-def transform_expand(t):
-    return wp.transform(np.array(t[0:3]), np.array(t[3:7]))
-
-
-# convert array of transforms to a array of 7-vecs
-def transform_flatten_list(xforms):
-    exp = lambda t: transform_flatten(t)
-    return list(map(exp, xforms))
-
-
-def transform_expand_list(xforms):
-    exp = lambda t: transform_expand(t)
-    return list(map(exp, xforms))
-
-
-def transform_inertia(m, I, p, q):
-    """
-    Transforms the inertia tensor described by the given mass and 3x3 inertia
-    matrix to a new frame described by the given position and orientation.
-    """
-    R = quat_to_matrix(q)
-
-    # Steiner's theorem
-    return R @ I @ R.T + m * (np.dot(p, p) * np.eye(3) - np.outer(p, p))
-
-
-# spatial operators
-
-
-# AdT
-def spatial_adjoint(t):
-    R = quat_to_matrix(t.q)
-    w = skew(t.p)
-
-    A = np.zeros((6, 6))
-    A[0:3, 0:3] = R
-    A[3:6, 0:3] = np.dot(w, R)
-    A[3:6, 3:6] = R
-
-    return A
-
-
-# (AdT)^-T
-def spatial_adjoint_dual(t):
-    R = quat_to_matrix(t.q)
-    w = skew(t.p)
-
-    A = np.zeros((6, 6))
-    A[0:3, 0:3] = R
-    A[0:3, 3:6] = np.dot(w, R)
-    A[3:6, 3:6] = R
-
-    return A
-
-
-# AdT*s
-def transform_twist(t_ab, s_b):
-    return np.dot(spatial_adjoint(t_ab), s_b)
-
-
-# AdT^{-T}*s
-def transform_wrench(t_ab, f_b):
-    return np.dot(spatial_adjoint_dual(t_ab), f_b)
-
-
-# transform spatial inertia (6x6) in b frame to a frame
-def transform_spatial_inertia(t_ab, I_b):
-    t_ba = transform_inverse(t_ab)
-
-    # todo: write specialized method
-    I_a = np.dot(np.dot(spatial_adjoint(t_ba).T, I_b), spatial_adjoint(t_ba))
-    return I_a
-
-
-def translate_twist(p_ab, s_b):
-    w = s_b[0:3]
-    v = np.cross(p_ab, s_b[0:3]) + s_b[3:6]
-
-    return np.array((*w, *v))
-
-
-def translate_wrench(p_ab, s_b):
-    w = s_b[0:3] + np.cross(p_ab, s_b[3:6])
-    v = s_b[3:6]
-
-    return np.array((*w, *v))
-
-
-# def spatial_vector(v=(0.0, 0.0, 0.0, 0.0, 0.0, 0.0)):
-#     return np.array(v)
-
-
-# ad_V pg. 289 L&P, pg. 25 Featherstone
-def spatial_cross(a, b):
-    w = np.cross(a[0:3], b[0:3])
-    v = np.cross(a[3:6], b[0:3]) + np.cross(a[0:3], b[3:6])
-
-    return np.array((*w, *v))
-
-
-# ad_V^T pg. 290 L&P,  pg. 25 Featurestone, note this does not includes the sign flip in the definition
-def spatial_cross_dual(a, b):
-    w = np.cross(a[0:3], b[0:3]) + np.cross(a[3:6], b[3:6])
-    v = np.cross(a[0:3], b[3:6])
-
-    return np.array((*w, *v))
-
-
-def spatial_dot(a, b):
-    return np.dot(a, b)
-
-
-def spatial_outer(a, b):
-    return np.outer(a, b)
-
-
-# def spatial_matrix():
-#     return np.zeros((6, 6))
-
-
-def spatial_matrix_from_inertia(I, m):
-    G = spatial_matrix()
-
-    G[0:3, 0:3] = I
-    G[3, 3] = m
-    G[4, 4] = m
-    G[5, 5] = m
-
-    return G
-
-
-# solves x = I^(-1)b
-def spatial_solve(I, b):
-    return np.dot(np.linalg.inv(I), b)
-
-
-# helper to retrive body angular velocity from a twist v_s in se(3)
-def get_body_angular_velocity(v_s):
-    return v_s[0:3]
-
-
-# helper to compute velocity of a point p on a body given it's spatial twist v_s
-def get_body_linear_velocity(v_s, p):
-    dpdt = v_s[3:6] + np.cross(v_s[0:3], p)
-    return dpdt
-
-
-# helper to build a body twist given the angular and linear velocity of
-# the center of mass specified in the world frame, returns the body
-# twist with respect to the origin (v_s)
-def get_body_twist(w_m, v_m, p_m):
-    lin = v_m + np.cross(p_m, w_m)
-    return (*w_m, *lin)
-
-
 def array_scan(in_array, out_array, inclusive=True):
     if in_array.device != out_array.device:
         raise RuntimeError("Array storage devices do not match")
@@ -411,6 +58,9 @@ def array_scan(in_array, out_array, inclusive=True):
     if in_array.dtype != out_array.dtype:
         raise RuntimeError("Array data types do not match")
 
+    if in_array.size == 0:
+        return
+
     from warp.context import runtime
 
     if in_array.device.is_cpu:
@@ -433,6 +83,9 @@ def radix_sort_pairs(keys, values, count: int):
     if keys.device != values.device:
         raise RuntimeError("Array storage devices do not match")
 
+    if count == 0:
+        return
+
     if keys.size < 2 * count or values.size < 2 * count:
         raise RuntimeError("Array storage must be large enough to contain 2*count elements")
 
@@ -469,14 +122,19 @@ def runlength_encode(values, run_values, run_lengths, run_count=None, value_coun
     # User can provide a device output array for storing the number of runs
     # For convenience, if no such array is provided, number of runs is returned on host
     if run_count is None:
-        host_return = True
+        if value_count == 0:
+            return 0
         run_count = wp.empty(shape=(1,), dtype=int, device=values.device)
+        host_return = True
     else:
-        host_return = False
         if run_count.device != values.device:
-            raise RuntimeError("run_count storage devices does not match other arrays")
+            raise RuntimeError("run_count storage device does not match other arrays")
         if run_count.dtype != wp.int32:
             raise RuntimeError("run_count array must be of type int32")
+        if value_count == 0:
+            run_count.zero_()
+            return 0
+        host_return = False
 
     from warp.context import runtime
 
@@ -532,6 +190,12 @@ def array_sum(values, out=None, value_count=None, axis=None):
         if out.shape != output_shape:
             raise RuntimeError(f"out array should have shape {output_shape}")
 
+    if value_count == 0:
+        out.zero_()
+        if axis is None and host_return:
+            return out.numpy()[0]
+        return out
+
     from warp.context import runtime
 
     if values.device.is_cpu:
@@ -578,7 +242,7 @@ def array_inner(a, b, out=None, count=None, axis=None):
         raise RuntimeError("Array storage sizes do not match")
 
     if a.device != b.device:
-        raise RuntimeError("Array storage sizes do not match")
+        raise RuntimeError("Array storage devices do not match")
 
     if a.dtype != b.dtype:
         raise RuntimeError("Array data types do not match")
@@ -615,6 +279,12 @@ def array_inner(a, b, out=None, count=None, axis=None):
         if out.shape != output_shape:
             raise RuntimeError(f"out array should have shape {output_shape}")
 
+    if count == 0:
+        if axis is None and host_return:
+            return 0.0
+        out.zero_()
+        return out
+
     from warp.context import runtime
 
     if a.device.is_cpu:
@@ -662,29 +332,16 @@ def array_inner(a, b, out=None, count=None, axis=None):
             return out
 
 
-_copy_kernel_cache = dict()
+@wp.kernel
+def _array_cast_kernel(
+    dest: Any,
+    src: Any,
+):
+    i = wp.tid()
+    dest[i] = dest.dtype(src[i])
 
 
 def array_cast(in_array, out_array, count=None):
-    def make_copy_kernel(dest_dtype, src_dtype):
-        import re
-
-        import warp.context
-
-        def copy_kernel(
-            dest: Any,
-            src: Any,
-        ):
-            dest[wp.tid()] = dest_dtype(src[wp.tid()])
-
-        module = wp.get_module(copy_kernel.__module__)
-        key = f"{copy_kernel.__name__}_{warp.context.type_str(src_dtype)}_{warp.context.type_str(dest_dtype)}"
-        key = re.sub("[^0-9a-zA-Z_]+", "", key)
-
-        if key not in _copy_kernel_cache:
-            _copy_kernel_cache[key] = wp.Kernel(func=copy_kernel, key=key, module=module)
-        return _copy_kernel_cache[key]
-
     if in_array.device != out_array.device:
         raise RuntimeError("Array storage devices do not match")
 
@@ -739,8 +396,7 @@ def array_cast(in_array, out_array, count=None):
         # Same data type, can simply copy
         wp.copy(dest=out_array, src=in_array, count=count)
     else:
-        copy_kernel = make_copy_kernel(src_dtype=in_array.dtype, dest_dtype=out_array.dtype)
-        wp.launch(kernel=copy_kernel, dim=dim, inputs=[out_array, in_array], device=out_array.device)
+        wp.launch(kernel=_array_cast_kernel, dim=dim, inputs=[out_array, in_array], device=out_array.device)
 
 
 # code snippet for invoking cProfile
@@ -816,11 +472,8 @@ class MeshAdjacency:
 
         self.edges[key] = edge
 
-    def opposite_vertex(self, edge):
-        pass
-
 
-def mem_report():
+def mem_report(): #pragma: no cover
     def _mem_report(tensors, mem_type):
         """Print the selected tensors of type
         There are two major storage types in our major concern:
@@ -872,12 +525,6 @@ def mem_report():
     print("=" * LEN)
 
 
-def lame_parameters(E, nu):
-    l = (E * nu) / ((1.0 + nu) * (1.0 - 2.0 * nu))
-    mu = E / (2.0 * (1.0 + nu))
-
-    return (l, mu)
-
 
 class ScopedDevice:
     def __init__(self, device):
@@ -1019,3 +666,17 @@ class ScopedTimer:
                 print("{}{} took {:.2f} ms".format(indent, self.name, self.elapsed))
 
             ScopedTimer.indent -= 1
+
+
+# helper kernels for adj_matmul
+@wp.kernel
+def add_kernel_2d(x: wp.array2d(dtype=Any), acc: wp.array2d(dtype=Any), beta: Any):
+    i, j = wp.tid()
+
+    x[i,j] = x[i,j] + beta * acc[i,j]
+
+@wp.kernel
+def add_kernel_3d(x: wp.array3d(dtype=Any), acc: wp.array3d(dtype=Any), beta: Any):
+    i, j, k = wp.tid()
+
+    x[i,j,k] = x[i,j,k] + beta * acc[i,j,k]