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

warp/fem/geometry/closest_point.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,85 +13,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any
+# isort: skip_file
 
-import warp as wp
-from warp.fem.types import Coords
+from warp._src.fem.geometry.closest_point import project_on_seg_at_origin as project_on_seg_at_origin
+from warp._src.fem.geometry.closest_point import project_on_tri_at_origin as project_on_tri_at_origin
+from warp._src.fem.geometry.closest_point import project_on_tet_at_origin as project_on_tet_at_origin
+from warp._src.fem.geometry.closest_point import project_on_box_at_origin as project_on_box_at_origin
 
+# TODO: Remove after cleaning up the public API.
 
-@wp.func
-def project_on_seg_at_origin(q: Any, seg: Any, len_sq: float):
-    s = wp.clamp(wp.dot(q, seg) / len_sq, 0.0, 1.0)
-    return wp.length_sq(q - s * seg), s
+from warp._src.fem.geometry import closest_point as _closest_point
 
 
-@wp.func
-def project_on_tri_at_origin(q: Any, e1: Any, e2: Any):
-    e1e1 = wp.dot(e1, e1)
-    e1e2 = wp.dot(e1, e2)
-    e2e2 = wp.dot(e2, e2)
+def __getattr__(name):
+    from warp._src.utils import get_deprecated_api
 
-    det = e1e1 * e2e2 - e1e2 * e1e2
-
-    if det > e1e1 * e2e2 * 1.0e-6:
-        e1p = wp.dot(e1, q)
-        e2p = wp.dot(e2, q)
-
-        s = (e2e2 * e1p - e1e2 * e2p) / det
-        t = (e1e1 * e2p - e1e2 * e1p) / det
-
-        if s >= 0.0 and t >= 0.0 and s + t <= 1.0:
-            # point inside triangle (distance can be non-zero in 3D case)
-            return wp.length_sq(q - s * e1 - t * e2), Coords(1.0 - s - t, s, t)
-
-    d1, s1 = project_on_seg_at_origin(q, e1, e1e1)
-    d2, s2 = project_on_seg_at_origin(q, e2, e2e2)
-    d12, s12 = project_on_seg_at_origin(q - e1, e2 - e1, wp.length_sq(e2 - e1))
-
-    if d1 <= d2:
-        if d1 <= d12:
-            return d1, Coords(1.0 - s1, s1, 0.0)
-    elif d2 <= d12:
-        return d2, Coords(1.0 - s2, 0.0, s2)
-
-    return d12, Coords(0.0, 1.0 - s12, s12)
-
-
-@wp.func
-def project_on_tet_at_origin(q: wp.vec3, e1: wp.vec3, e2: wp.vec3, e3: wp.vec3):
-    mat = wp.inverse(wp.matrix_from_cols(e1, e2, e3))
-    coords = mat * q
-
-    if wp.min(coords) >= 0.0 and coords[0] + coords[1] + coords[2] <= 1.0:
-        return 0.0, coords
-
-    # Not inside tet, compare closest point on each tri
-
-    d12, s12 = project_on_tri_at_origin(q, e1, e2)
-    d23, s23 = project_on_tri_at_origin(q, e2, e3)
-    d31, s31 = project_on_tri_at_origin(q, e3, e1)
-    d123, s123 = project_on_tri_at_origin(q - e1, e2 - e1, e3 - e1)
-
-    dmin = wp.min(wp.vec4(d12, d23, d31, d123))
-
-    if dmin == d12:
-        return dmin, Coords(s12[1], s12[2], 0.0)
-    elif dmin == d23:
-        return dmin, Coords(0.0, s23[1], s23[2])
-    elif dmin == d31:
-        return dmin, Coords(s31[2], 0.0, s31[1])
-    else:
-        return dmin, s123
-
-
-@wp.func
-def project_on_box_at_origin(coords: wp.vec3, sizes: wp.vec3):
-    proj_coords = wp.min(wp.max(coords, wp.vec3(0.0)), sizes)
-    return wp.length_sq(coords - proj_coords), wp.cw_div(proj_coords, sizes)
-
-
-@wp.func
-def project_on_box_at_origin(coords: wp.vec2, sizes: wp.vec2):
-    proj_coords = wp.min(wp.max(coords, wp.vec2(0.0)), sizes)
-    norm_coords = wp.cw_div(proj_coords, sizes)
-    return wp.length_sq(coords - proj_coords), Coords(norm_coords[0], norm_coords[1], 0.0)
+    return get_deprecated_api(_closest_point, "wp.fem.geometry", name)

warp/fem/linalg.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,374 +13,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any
+# isort: skip_file
 
-import warp as wp
-import warp.types
+from warp._src.fem.linalg import array_axpy as array_axpy
+from warp._src.fem.linalg import generalized_inner as generalized_inner
+from warp._src.fem.linalg import generalized_outer as generalized_outer
+from warp._src.fem.linalg import spherical_part as spherical_part
+from warp._src.fem.linalg import symmetric_part as symmetric_part
+from warp._src.fem.linalg import skew_part as skew_part
+from warp._src.fem.linalg import householder_make_hessenberg as householder_make_hessenberg
+from warp._src.fem.linalg import householder_qr_decomposition as householder_qr_decomposition
+from warp._src.fem.linalg import inverse_qr as inverse_qr
+from warp._src.fem.linalg import solve_triangular as solve_triangular
+from warp._src.fem.linalg import symmetric_eigenvalues_qr as symmetric_eigenvalues_qr
 
+# TODO: Remove after cleaning up the public API.
 
-@wp.func
-def generalized_outer(x: wp.vec(Any, wp.Scalar), y: wp.vec(Any, wp.Scalar)):
-    """Generalized outer product allowing for vector or scalar arguments"""
-    return wp.outer(x, y)
+from warp._src.fem import linalg as _linalg
 
 
-@wp.func
-def generalized_outer(x: wp.Scalar, y: wp.vec(Any, wp.Scalar)):
-    return x * y
+def __getattr__(name):
+    from warp._src.utils import get_deprecated_api
 
-
-@wp.func
-def generalized_outer(x: wp.vec(Any, wp.Scalar), y: wp.Scalar):
-    return x * y
-
-
-@wp.func
-def generalized_outer(x: wp.quatf, y: wp.vec(Any, wp.Scalar)):
-    return generalized_outer(wp.vec4(x[0], x[1], x[2], x[3]), y)
-
-
-@wp.func
-def generalized_inner(x: wp.vec(Any, wp.Scalar), y: wp.vec(Any, wp.Scalar)):
-    """Generalized inner product allowing for vector, tensor and scalar arguments"""
-    return wp.dot(x, y)
-
-
-@wp.func
-def generalized_inner(x: wp.Scalar, y: wp.Scalar):
-    return x * y
-
-
-@wp.func
-def generalized_inner(x: wp.mat((Any, Any), wp.Scalar), y: wp.vec(Any, wp.Scalar)):
-    return y @ x
-
-
-@wp.func
-def generalized_inner(x: wp.vec(Any, wp.Scalar), y: wp.mat((Any, Any), wp.Scalar)):
-    return y @ x
-
-
-@wp.func
-def basis_coefficient(val: wp.Scalar, i: int):
-    return val
-
-
-@wp.func
-def basis_coefficient(val: wp.mat((Any, Any), wp.Scalar), i: int):
-    cols = int(type(val[0]).length)
-    row = i // cols
-    col = i - row * cols
-    return val[row, col]
-
-
-@wp.func
-def basis_coefficient(val: Any, i: int):
-    return val[i]
-
-
-@wp.func
-def basis_coefficient(val: wp.vec(Any, wp.Scalar), i: int, j: int):
-    # treat as row vector
-    return val[j]
-
-
-@wp.func
-def basis_coefficient(val: wp.mat((Any, Any), wp.Scalar), i: int, j: int):
-    return val[i, j]
-
-
-@wp.func
-def symmetric_part(x: Any):
-    """Symmetric part of a square tensor"""
-    return 0.5 * (x + wp.transpose(x))
-
-
-@wp.func
-def spherical_part(x: wp.mat22):
-    """Spherical part of a square tensor"""
-    return 0.5 * wp.trace(x) * wp.identity(n=2, dtype=float)
-
-
-@wp.func
-def spherical_part(x: wp.mat33):
-    """Spherical part of a square tensor"""
-    return (wp.trace(x) / 3.0) * wp.identity(n=3, dtype=float)
-
-
-@wp.func
-def skew_part(x: wp.mat22):
-    """Skew part of a 2x2 tensor as corresponding rotation angle"""
-    return 0.5 * (x[1, 0] - x[0, 1])
-
-
-@wp.func
-def skew_part(x: wp.mat33):
-    """Skew part of a 3x3 tensor as the corresponding rotation vector"""
-    a = 0.5 * (x[2, 1] - x[1, 2])
-    b = 0.5 * (x[0, 2] - x[2, 0])
-    c = 0.5 * (x[1, 0] - x[0, 1])
-    return wp.vec3(a, b, c)
-
-
-@wp.func
-def householder_qr_decomposition(A: Any):
-    """
-    QR decomposition of a square matrix using Householder reflections
-
-    Returns Q and R such that Q R = A, Q orthonormal (such that QQ^T = Id), R upper triangular
-    """
-
-    x = type(A[0])()
-    Q = wp.identity(n=type(x).length, dtype=A.dtype)
-
-    zero = x.dtype(0.0)
-    two = x.dtype(2.0)
-
-    for i in range(type(x).length):
-        for k in range(type(x).length):
-            x[k] = wp.where(k < i, zero, A[k, i])
-
-        alpha = wp.length(x) * wp.sign(x[i])
-        x[i] += alpha
-        two_over_x_sq = wp.where(alpha == zero, zero, two / wp.length_sq(x))
-
-        A -= wp.outer(two_over_x_sq * x, x * A)
-        Q -= wp.outer(Q * x, two_over_x_sq * x)
-
-    return Q, A
-
-
-@wp.func
-def householder_make_hessenberg(A: Any):
-    """Transforms a square matrix to Hessenberg form (single lower diagonal) using Householder reflections
-
-    Returns:
-        Q and H such that Q H Q^T = A, Q orthonormal, H under Hessenberg form
-        If A is symmetric, H will be tridiagonal
-    """
-
-    x = type(A[0])()
-    Q = wp.identity(n=type(x).length, dtype=A.dtype)
-
-    zero = x.dtype(0.0)
-    two = x.dtype(2.0)
-
-    for i in range(1, type(x).length):
-        for k in range(type(x).length):
-            x[k] = wp.where(k < i, zero, A[k, i - 1])
-
-        alpha = wp.length(x) * wp.sign(x[i])
-        x[i] += alpha
-        two_over_x_sq = wp.where(alpha == zero, zero, two / wp.length_sq(x))
-
-        # apply on both sides
-        A -= wp.outer(two_over_x_sq * x, x * A)
-        A -= wp.outer(A * x, two_over_x_sq * x)
-        Q -= wp.outer(Q * x, two_over_x_sq * x)
-
-    return Q, A
-
-
-@wp.func
-def solve_triangular(R: Any, b: Any):
-    """Solves for R x = b where R is an upper triangular matrix
-
-    Returns x
-    """
-    zero = b.dtype(0)
-    x = type(b)(b.dtype(0))
-    for i in range(b.length, 0, -1):
-        j = i - 1
-        r = b[j] - wp.dot(R[j], x)
-        x[j] = wp.where(R[j, j] == zero, zero, r / R[j, j])
-
-    return x
-
-
-@wp.func
-def inverse_qr(A: Any):
-    # Computes a square matrix inverse using QR factorization
-
-    Q, R = householder_qr_decomposition(A)
-
-    A_inv = type(A)()
-    for i in range(type(A[0]).length):
-        A_inv[i] = solve_triangular(R, Q[i])  # ith column of Q^T
-
-    return wp.transpose(A_inv)
-
-
-@wp.func
-def _wilkinson_shift(a: Any, b: Any, c: Any, tol: Any):
-    # Wilkinson shift: estimate eigenvalue of 2x2 symmetric matrix [a, c, c, b]
-    d = (a - b) * type(tol)(0.5)
-    return b + d - wp.sign(d) * wp.sqrt(d * d + c * c)
-
-
-@wp.func
-def _givens_rotation(a: Any, b: Any):
-    # Givens rotation [[c -s], [s c]] such that sa+cb =0
-    zero = type(a)(0.0)
-    one = type(a)(1.0)
-
-    b2 = b * b
-    if b2 == zero:
-        # id rotation
-        return one, zero
-
-    scale = one / wp.sqrt(a * a + b2)
-    return a * scale, -b * scale
-
-
-@wp.func
-def tridiagonal_symmetric_eigenvalues_qr(D: Any, L: Any, Q: Any, tol: Any):
-    """
-    Computes the eigenvalues and eigen vectors of a symmetric tridiagonal matrix using the
-    Symmetric tridiagonal QR algorithm with implicit Wilkinson shift
-
-    Args:
-        D: Main diagonal of the matrix
-        L: Lower diagonal of the matrix, indexed such that L[i] = A[i+1, i]
-        Q: Initialization for the eigenvectors, useful if a pre-transformation has been applied, otherwise may be identity
-        tol: Tolerance for the diagonalization residual (Linf norm of off-diagonal over diagonal terms)
-
-    Returns a tuple (D: vector of eigenvalues, P: matrix with one eigenvector per row) such that A = P^T D P
-
-
-    Ref: Arbenz P, Numerical Methods for Solving Large Scale Eigenvalue Problems, Chapter 4 (QR algorithm, Mar 13, 2018)
-    """
-
-    two = D.dtype(2.0)
-
-    # so that we can use the type length in expressions
-    # this will prevent unrolling by warp, but should be ok for native code
-    m = int(0)
-    for _ in range(type(D).length):
-        m += 1
-
-    start = int(0)
-    y = D.dtype(0.0)  # moving buldge
-    x = D.dtype(0.0)  # coeff atop buldge
-
-    for _ in range(32 * m):  # failsafe, usually converges faster than that
-        # Iterate over all independent (deflated) blocks
-        end = int(-1)
-
-        for k in range(m - 1):
-            if k >= end:
-                # Check if new block is starting
-                if k == end or wp.abs(L[k]) <= tol * (wp.abs(D[k]) + wp.abs(D[k + 1])):
-                    continue
-
-                # Find end of block
-                start = k
-                end = start + 1
-                while end + 1 < m:
-                    if wp.abs(L[end]) <= tol * (wp.abs(D[end + 1]) + wp.abs(D[end])):
-                        break
-                    end += 1
-
-                # Wilkinson shift (an eigenvalue of the last 2x2 block)
-                shift = _wilkinson_shift(D[end - 1], D[end], L[end - 1], tol)
-
-                # start with eliminating lower diag of first column of shifted matrix
-                # (i.e. first step of explicit QR factorization)
-                # Then all further steps eliminate the buldge (second diag) of the non-shifted matrix
-                x = D[start] - shift
-                y = L[start]
-
-            c, s = _givens_rotation(x, y)
-
-            # Apply Givens rotation on both sides of tridiagonal matrix
-
-            # middle block
-            d = D[k] - D[k + 1]
-            z = (two * c * L[k] + d * s) * s
-            D[k] -= z
-            D[k + 1] += z
-            L[k] = d * c * s + (c * c - s * s) * L[k]
-
-            if k > start:
-                L[k - 1] = c * x - s * y
-
-            x = L[k]
-            y = -s * L[k + 1]  # new buldge
-            L[k + 1] *= c
-
-            # apply givens rotation on left of Q
-            # note: Q is transposed compared to usual impls, as Warp makes it easier to index rows
-            Qk0 = Q[k]
-            Qk1 = Q[k + 1]
-            Q[k] = c * Qk0 - s * Qk1
-            Q[k + 1] = c * Qk1 + s * Qk0
-
-        if end <= 0:
-            # We did nothing, so diagonalization must have been achieved
-            break
-
-    return D, Q
-
-
-@wp.func
-def symmetric_eigenvalues_qr(A: Any, tol: Any):
-    """
-    Computes the eigenvalues and eigen vectors of a square symmetric matrix A using the QR algorithm
-
-    Args:
-        A: square symmetric matrix
-        tol: Tolerance for the diagonalization residual (Linf norm of off-diagonal over diagonal terms)
-
-    Returns a tuple (D: vector of eigenvalues, P: matrix with one eigenvector per row) such that A = P^T D P
-    """
-
-    # Put A under Hessenberg form (tridiagonal)
-    Q, H = householder_make_hessenberg(A)
-
-    # tridiagonal storage for H
-    D = wp.get_diag(H)
-    L = type(D)(A.dtype(0.0))
-    for i in range(1, type(D).length):
-        L[i - 1] = H[i, i - 1]
-
-    Qt = wp.transpose(Q)
-    ev, P = tridiagonal_symmetric_eigenvalues_qr(D, L, Qt, tol)
-    return ev, P
-
-
-def array_axpy(x: wp.array, y: wp.array, alpha: float = 1.0, beta: float = 1.0):
-    """Performs y = alpha*x + beta*y"""
-
-    dtype = wp.types.type_scalar_type(y.dtype)
-
-    alpha = dtype(alpha)
-    beta = dtype(beta)
-
-    if x.shape != y.shape or x.device != y.device:
-        raise ValueError("x and y arrays must have the same shape and device")
-
-    # array_axpy requires a custom adjoint; unfortunately we cannot use `wp.func_grad`
-    # as generic functions are not supported yet. Instead we use a non-differentiable kernel
-    # and record a custom adjoint function on the tape.
-
-    # temporarily disable tape to avoid printing warning that kernel is not differentiable
-    (tape, wp.context.runtime.tape) = (wp.context.runtime.tape, None)
-    wp.launch(kernel=_array_axpy_kernel, dim=x.shape, device=x.device, inputs=[x, y, alpha, beta])
-    wp.context.runtime.tape = tape
-
-    if tape is not None and (x.requires_grad or y.requires_grad):
-
-        def backward_axpy():
-            # adj_x += adj_y * alpha
-            # adj_y = adj_y * beta
-            array_axpy(x=y.grad, y=x.grad, alpha=alpha, beta=1.0)
-            if beta != 1.0:
-                array_axpy(x=y.grad, y=y.grad, alpha=0.0, beta=beta)
-
-        tape.record_func(backward_axpy, arrays=[x, y])
-
-
-@wp.kernel(enable_backward=False)
-def _array_axpy_kernel(x: wp.array(dtype=Any), y: wp.array(dtype=Any), alpha: Any, beta: Any):
-    i = wp.tid()
-    y[i] = beta * y[i] + alpha * y.dtype(x[i])
+    return get_deprecated_api(_linalg, "wp.fem", name)