warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

warp/_src/math.py

@@ -0,0 +1,416 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import warp as wp
+
+"""
+Vector norm functions
+"""
+
+_wp_module_name_ = "warp.math"
+
+__all__ = [
+    "norm_huber",
+    "norm_l1",
+    "norm_l2",
+    "norm_pseudo_huber",
+    "smooth_normalize",
+    "transform_compose",
+    "transform_decompose",
+    "transform_from_matrix",
+    "transform_to_matrix",
+]
+
+
+@wp.func
+def norm_l1(v: Any):
+    """
+    Computes the L1 norm of a vector v.
+
+    .. math:: \\|v\\|_1 = \\sum_i |v_i|
+
+    Args:
+        v (Vector[Any,Float]): The vector to compute the L1 norm of.
+
+    Returns:
+        float: The L1 norm of the vector.
+    """
+    n = float(0.0)
+    for i in range(len(v)):
+        n += wp.abs(v[i])
+    return n
+
+
+@wp.func
+def norm_l2(v: Any):
+    """
+    Computes the L2 norm of a vector v.
+
+    .. math:: \\|v\\|_2 = \\sqrt{\\sum_i v_i^2}
+
+    Args:
+        v (Vector[Any,Float]): The vector to compute the L2 norm of.
+
+    Returns:
+        float: The L2 norm of the vector.
+    """
+    return wp.length(v)
+
+
+@wp.func
+def norm_huber(v: Any, delta: float = 1.0):
+    """
+    Computes the Huber norm of a vector v with a given delta.
+
+    .. math::
+        H(v) = \\begin{cases} \\frac{1}{2} \\|v\\|^2 & \\text{if } \\|v\\| \\leq \\delta \\\\ \\delta(\\|v\\| - \\frac{1}{2}\\delta) & \\text{otherwise} \\end{cases}
+
+    .. image:: /img/norm_huber.svg
+        :align: center
+
+    Args:
+        v (Vector[Any,Float]): The vector to compute the Huber norm of.
+        delta (float): The threshold value, defaults to 1.0.
+
+    Returns:
+        float: The Huber norm of the vector.
+    """
+    a = wp.dot(v, v)
+    if a <= delta * delta:
+        return 0.5 * a
+    return delta * (wp.sqrt(a) - 0.5 * delta)
+
+
+@wp.func
+def norm_pseudo_huber(v: Any, delta: float = 1.0):
+    """
+    Computes the "pseudo" Huber norm of a vector v with a given delta.
+
+    .. math::
+        H^\\prime(v) = \\delta \\sqrt{1 + \\frac{\\|v\\|^2}{\\delta^2}}
+
+    .. image:: /img/norm_pseudo_huber.svg
+        :align: center
+
+    Args:
+        v (Vector[Any,Float]): The vector to compute the Huber norm of.
+        delta (float): The threshold value, defaults to 1.0.
+
+    Returns:
+        float: The Huber norm of the vector.
+    """
+    a = wp.dot(v, v)
+    return delta * wp.sqrt(1.0 + a / (delta * delta))
+
+
+@wp.func
+def smooth_normalize(v: Any, delta: float = 1.0):
+    """
+    Normalizes a vector using the pseudo-Huber norm.
+
+    See :func:`norm_pseudo_huber`.
+
+    .. math::
+        \\frac{v}{H^\\prime(v)}
+
+    Args:
+        v (Vector[Any,Float]): The vector to normalize.
+        delta (float): The threshold value, defaults to 1.0.
+
+    Returns:
+        Vector[Any,Float]: The normalized vector.
+    """
+    return v / norm_pseudo_huber(v, delta)
+
+
+def create_transform_from_matrix_func(dtype):
+    mat44 = wp._src.types.matrix((4, 4), dtype)
+    vec3 = wp._src.types.vector(3, dtype)
+    transform = wp._src.types.transformation(dtype)
+
+    def transform_from_matrix(mat: mat44) -> transform:
+        """
+        Construct a transformation from a 4x4 matrix.
+
+        .. math::
+            M = \\begin{bmatrix}
+            R_{00} & R_{01} & R_{02} & p_x \\\\
+            R_{10} & R_{11} & R_{12} & p_y \\\\
+            R_{20} & R_{21} & R_{22} & p_z \\\\
+            0 & 0 & 0 & 1
+            \\end{bmatrix}
+
+        Where:
+
+        * :math:`R` is the 3x3 rotation matrix created from the orientation quaternion of the input transform.
+        * :math:`p` is the 3D position vector :math:`[p_x, p_y, p_z]` of the input transform.
+
+        Args:
+            mat (Matrix[4, 4, Float]): Matrix to convert.
+
+        Returns:
+            Transformation[Float]: The transformation.
+        """
+        p = vec3(mat[0][3], mat[1][3], mat[2][3])
+        q = wp.quat_from_matrix(mat)
+        return transform(p, q)
+
+    return transform_from_matrix
+
+
+transform_from_matrix = wp.func(
+    create_transform_from_matrix_func(wp.float32),
+    name="transform_from_matrix",
+)
+wp.func(
+    create_transform_from_matrix_func(wp.float16),
+    name="transform_from_matrix",
+)
+wp.func(
+    create_transform_from_matrix_func(wp.float64),
+    name="transform_from_matrix",
+)
+
+
+def create_transform_to_matrix_func(dtype):
+    mat44 = wp._src.types.matrix((4, 4), dtype)
+    transform = wp._src.types.transformation(dtype)
+
+    def transform_to_matrix(xform: transform) -> mat44:
+        """
+        Convert a transformation to a 4x4 matrix.
+
+        .. math::
+            M = \\begin{bmatrix}
+            R_{00} & R_{01} & R_{02} & p_x \\\\
+            R_{10} & R_{11} & R_{12} & p_y \\\\
+            R_{20} & R_{21} & R_{22} & p_z \\\\
+            0 & 0 & 0 & 1
+            \\end{bmatrix}
+
+        Where:
+
+        * :math:`R` is the 3x3 rotation matrix created from the orientation quaternion of the input transform.
+        * :math:`p` is the 3D position vector :math:`[p_x, p_y, p_z]` of the input transform.
+
+        Args:
+            xform (Transformation[Float]): Transformation to convert.
+
+        Returns:
+            Matrix[4, 4, Float]: The matrix.
+        """
+        p = wp.transform_get_translation(xform)
+        q = wp.transform_get_rotation(xform)
+        rot = wp.quat_to_matrix(q)
+        # fmt: off
+        return mat44(
+            rot[0][0], rot[0][1], rot[0][2], p[0],
+            rot[1][0], rot[1][1], rot[1][2], p[1],
+            rot[2][0], rot[2][1], rot[2][2], p[2],
+            dtype(0.0), dtype(0.0), dtype(0.0), dtype(1.0),
+        )
+        # fmt: on
+
+    return transform_to_matrix
+
+
+transform_to_matrix = wp.func(
+    create_transform_to_matrix_func(wp.float32),
+    name="transform_to_matrix",
+)
+wp.func(
+    create_transform_to_matrix_func(wp.float16),
+    name="transform_to_matrix",
+)
+wp.func(
+    create_transform_to_matrix_func(wp.float64),
+    name="transform_to_matrix",
+)
+
+
+def create_transform_compose_func(dtype):
+    mat44 = wp._src.types.matrix((4, 4), dtype)
+    quat = wp._src.types.quaternion(dtype)
+    vec3 = wp._src.types.vector(3, dtype)
+
+    def transform_compose(position: vec3, rotation: quat, scale: vec3):
+        """
+        Compose a 4x4 transformation matrix from a 3D position, quaternion orientation, and 3D scale.
+
+        .. math::
+            M = \\begin{bmatrix}
+            s_x R_{00} & s_y R_{01} & s_z R_{02} & p_x \\\\
+            s_x R_{10} & s_y R_{11} & s_z R_{12} & p_y \\\\
+            s_x R_{20} & s_y R_{21} & s_z R_{22} & p_z \\\\
+            0 & 0 & 0 & 1
+            \\end{bmatrix}
+
+        Where:
+
+        * :math:`R` is the 3x3 rotation matrix created from the orientation quaternion of the input transform.
+        * :math:`p` is the 3D position vector :math:`[p_x, p_y, p_z]` of the input transform.
+        * :math:`s` is the 3D scale vector :math:`[s_x, s_y, s_z]` of the input transform.
+
+        Args:
+            position (Vector[3, Float]): The 3D position vector.
+            rotation (Quaternion[Float]): The quaternion orientation.
+            scale (Vector[3, Float]): The 3D scale vector.
+
+        Returns:
+            Matrix[4, 4, Float]: The transformation matrix.
+        """
+        R = wp.quat_to_matrix(rotation)
+        # fmt: off
+        return mat44(
+            scale[0] * R[0,0], scale[1] * R[0,1], scale[2] * R[0,2], position[0],
+            scale[0] * R[1,0], scale[1] * R[1,1], scale[2] * R[1,2], position[1],
+            scale[0] * R[2,0], scale[1] * R[2,1], scale[2] * R[2,2], position[2],
+            dtype(0.0), dtype(0.0), dtype(0.0), dtype(1.0),
+        )
+        # fmt: on
+
+    return transform_compose
+
+
+transform_compose = wp.func(
+    create_transform_compose_func(wp.float32),
+    name="transform_compose",
+)
+wp.func(
+    create_transform_compose_func(wp.float16),
+    name="transform_compose",
+)
+wp.func(
+    create_transform_compose_func(wp.float64),
+    name="transform_compose",
+)
+
+
+def create_transform_decompose_func(dtype):
+    mat44 = wp._src.types.matrix((4, 4), dtype)
+    vec3 = wp._src.types.vector(3, dtype)
+    mat33 = wp._src.types.matrix((3, 3), dtype)
+    zero = dtype(0.0)
+
+    def transform_decompose(m: mat44):
+        """
+        Decompose a 4x4 transformation matrix into 3D position, quaternion orientation, and 3D scale.
+
+        .. math::
+            M = \\begin{bmatrix}
+            s_x R_{00} & s_y R_{01} & s_z R_{02} & p_x \\\\
+            s_x R_{10} & s_y R_{11} & s_z R_{12} & p_y \\\\
+            s_x R_{20} & s_y R_{21} & s_z R_{22} & p_z \\\\
+            0 & 0 & 0 & 1
+            \\end{bmatrix}
+
+        Where:
+
+        * :math:`R` is the 3x3 rotation matrix created from the orientation quaternion of the input transform.
+        * :math:`p` is the 3D position vector :math:`[p_x, p_y, p_z]` of the input transform.
+        * :math:`s` is the 3D scale vector :math:`[s_x, s_y, s_z]` of the input transform.
+
+        Args:
+            m (Matrix[4, 4, Float]): The matrix to decompose.
+
+        Returns:
+            Tuple[Vector[3, Float], Quaternion[Float], Vector[3, Float]]: A tuple containing the position vector, quaternion orientation, and scale vector.
+        """
+        # extract position
+        position = vec3(m[0, 3], m[1, 3], m[2, 3])
+        # extract rotation matrix components
+        r00, r01, r02 = m[0, 0], m[0, 1], m[0, 2]
+        r10, r11, r12 = m[1, 0], m[1, 1], m[1, 2]
+        r20, r21, r22 = m[2, 0], m[2, 1], m[2, 2]
+        # get scale magnitudes
+        sx = wp.sqrt(r00 * r00 + r10 * r10 + r20 * r20)
+        sy = wp.sqrt(r01 * r01 + r11 * r11 + r21 * r21)
+        sz = wp.sqrt(r02 * r02 + r12 * r12 + r22 * r22)
+        # normalize rotation matrix components
+        if sx != zero:
+            r00 /= sx
+            r10 /= sx
+            r20 /= sx
+        if sy != zero:
+            r01 /= sy
+            r11 /= sy
+            r21 /= sy
+        if sz != zero:
+            r02 /= sz
+            r12 /= sz
+            r22 /= sz
+        # extract rotation (quaternion)
+        rotation = wp.quat_from_matrix(mat33(r00, r01, r02, r10, r11, r12, r20, r21, r22))
+        # extract scale
+        scale = vec3(sx, sy, sz)
+        return position, rotation, scale
+
+    return transform_decompose
+
+
+transform_decompose = wp.func(
+    create_transform_decompose_func(wp.float32),
+    name="transform_decompose",
+)
+wp.func(
+    create_transform_decompose_func(wp.float16),
+    name="transform_decompose",
+)
+wp.func(
+    create_transform_decompose_func(wp.float64),
+    name="transform_decompose",
+)
+
+
+# register API functions so they appear in the documentation
+
+wp._src.context.register_api_function(
+    norm_l1,
+    group="Vector Math",
+)
+wp._src.context.register_api_function(
+    norm_l2,
+    group="Vector Math",
+)
+wp._src.context.register_api_function(
+    norm_huber,
+    group="Vector Math",
+)
+wp._src.context.register_api_function(
+    norm_pseudo_huber,
+    group="Vector Math",
+)
+wp._src.context.register_api_function(
+    smooth_normalize,
+    group="Vector Math",
+)
+wp._src.context.register_api_function(
+    transform_from_matrix,
+    group="Transformations",
+)
+wp._src.context.register_api_function(
+    transform_to_matrix,
+    group="Transformations",
+)
+wp._src.context.register_api_function(
+    transform_compose,
+    group="Transformations",
+)
+wp._src.context.register_api_function(
+    transform_decompose,
+    group="Transformations",
+)

warp/_src/optim/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

warp/_src/optim/adam.py

@@ -0,0 +1,165 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import warp as wp
+
+_wp_module_name_ = "warp.optim.adam"
+
+
+@wp.kernel
+def adam_step_kernel_vec3(
+    g: wp.array(dtype=wp.vec3),
+    m: wp.array(dtype=wp.vec3),
+    v: wp.array(dtype=wp.vec3),
+    lr: float,
+    beta1: float,
+    beta2: float,
+    t: float,
+    eps: float,
+    params: wp.array(dtype=wp.vec3),
+):
+    i = wp.tid()
+    m[i] = beta1 * m[i] + (1.0 - beta1) * g[i]
+    v[i] = beta2 * v[i] + (1.0 - beta2) * wp.cw_mul(g[i], g[i])
+    mhat = m[i] / (1.0 - wp.pow(beta1, (t + 1.0)))
+    vhat = v[i] / (1.0 - wp.pow(beta2, (t + 1.0)))
+    sqrt_vhat = wp.vec3(wp.sqrt(vhat[0]), wp.sqrt(vhat[1]), wp.sqrt(vhat[2]))
+    eps_vec3 = wp.vec3(eps, eps, eps)
+    params[i] = params[i] - lr * wp.cw_div(mhat, (sqrt_vhat + eps_vec3))
+
+
+@wp.kernel
+def adam_step_kernel_float(
+    g: wp.array(dtype=float),
+    m: wp.array(dtype=float),
+    v: wp.array(dtype=float),
+    lr: float,
+    beta1: float,
+    beta2: float,
+    t: float,
+    eps: float,
+    params: wp.array(dtype=float),
+):
+    i = wp.tid()
+    m[i] = beta1 * m[i] + (1.0 - beta1) * g[i]
+    v[i] = beta2 * v[i] + (1.0 - beta2) * g[i] * g[i]
+    mhat = m[i] / (1.0 - wp.pow(beta1, (t + 1.0)))
+    vhat = v[i] / (1.0 - wp.pow(beta2, (t + 1.0)))
+    params[i] = params[i] - lr * mhat / (wp.sqrt(vhat) + eps)
+
+
+@wp.kernel
+def adam_step_kernel_half(
+    g: wp.array(dtype=wp.float16),
+    m: wp.array(dtype=float),
+    v: wp.array(dtype=float),
+    lr: float,
+    beta1: float,
+    beta2: float,
+    t: float,
+    eps: float,
+    params: wp.array(dtype=wp.float16),
+):
+    i = wp.tid()
+    m[i] = beta1 * m[i] + (1.0 - beta1) * float(g[i])
+    v[i] = beta2 * v[i] + (1.0 - beta2) * float(g[i]) * float(g[i])
+    mhat = m[i] / (1.0 - wp.pow(beta1, (t + 1.0)))
+    vhat = v[i] / (1.0 - wp.pow(beta2, (t + 1.0)))
+    params[i] = params[i] - wp.float16(lr * mhat / (wp.sqrt(vhat) + eps))
+
+
+class Adam:
+    """An implementation of the Adam Optimizer
+    It is designed to mimic Pytorch's version.
+    https://pytorch.org/docs/stable/generated/torch.optim.Adam.html#torch.optim.Adam
+    """
+
+    def __init__(self, params=None, lr=0.001, betas=(0.9, 0.999), eps=1e-08):
+        self.m = []  # first moment
+        self.v = []  # second moment
+        self.set_params(params)
+        self.lr = lr
+        self.beta1 = betas[0]
+        self.beta2 = betas[1]
+        self.eps = eps
+        self.t = 0
+
+    def set_params(self, params):
+        self.params = params
+        if params is not None and isinstance(params, list) and len(params) > 0:
+            if len(self.m) != len(params):
+                self.m = [None] * len(params)  # reset first moment
+            if len(self.v) != len(params):
+                self.v = [None] * len(params)  # reset second moment
+            for i in range(len(params)):
+                param = params[i]
+
+                if param.dtype == wp.vec3:
+                    dtype = wp.vec3
+                elif param.dtype == wp.float32:
+                    dtype = wp.float32
+                elif param.dtype == wp.float16:
+                    dtype = wp.float32  # we always use fp32 for moments, even if params are fp16
+                else:
+                    raise RuntimeError(f"Unsupported dtype for Warp Adam optimizer: {param.dtype}")
+
+                if self.m[i] is None or self.m[i].shape != param.shape or self.m[i].dtype != param.dtype:
+                    self.m[i] = wp.zeros(shape=param.shape, dtype=dtype, device=param.device)
+                if self.v[i] is None or self.v[i].shape != param.shape or self.v[i].dtype != param.dtype:
+                    self.v[i] = wp.zeros(shape=param.shape, dtype=dtype, device=param.device)
+
+    def reset_internal_state(self):
+        for m_i in self.m:
+            m_i.zero_()
+        for v_i in self.v:
+            v_i.zero_()
+        self.t = 0
+
+    def step(self, grad):
+        assert self.params is not None
+        for i in range(len(self.params)):
+            Adam.step_detail(
+                grad[i], self.m[i], self.v[i], self.lr, self.beta1, self.beta2, self.t, self.eps, self.params[i]
+            )
+        self.t = self.t + 1
+
+    @staticmethod
+    def step_detail(g, m, v, lr, beta1, beta2, t, eps, params):
+        assert params.dtype == g.dtype
+        assert params.shape == g.shape
+        kernel_inputs = [g, m, v, lr, beta1, beta2, t, eps, params]
+        if params.dtype == wp._src.types.float32:
+            wp.launch(
+                kernel=adam_step_kernel_float,
+                dim=len(params),
+                inputs=kernel_inputs,
+                device=params.device,
+            )
+        elif params.dtype == wp._src.types.float16:
+            wp.launch(
+                kernel=adam_step_kernel_half,
+                dim=len(params),
+                inputs=kernel_inputs,
+                device=params.device,
+            )
+        elif params.dtype == wp._src.types.vec3:
+            wp.launch(
+                kernel=adam_step_kernel_vec3,
+                dim=len(params),
+                inputs=kernel_inputs,
+                device=params.device,
+            )
+        else:
+            raise RuntimeError("Params data type not supported in Adam step kernels.")