WIP: API refactor

newton/__init__.py

@@ -106,11 +106,12 @@
 # ==================================================================================
 # submodule APIs
 # ==================================================================================
-from . import geometry, ik, selection, sensors, solvers, utils, viewer  # noqa: E402
+from . import geometry, ik, math, selection, sensors, solvers, utils, viewer  # noqa: E402
 
 __all__ += [
     "geometry",
     "ik",
+    "math",
     "selection",
     "sensors",
     "solvers",

newton/_src/core/types.py

@@ -165,6 +165,12 @@ def to_vec3(self) -> wp.vec3:
         """
         return wp.vec3(*self.to_vector())
 
+    def quat_between_axes(self, other: Axis) -> wp.quat:
+        """
+        Return the quaternion between two axes.
+        """
+        return wp.quat_between_vectors(self.to_vec3(), other.to_vec3())
+
 
 AxisType = Axis | Literal["X", "Y", "Z"] | Literal[0, 1, 2] | int | str
 """Type that can be used to represent an axis, including the enum, string, and integer representations."""

newton/_src/geometry/__init__.py

@@ -32,7 +32,7 @@
 )
 from .flags import ParticleFlags, ShapeFlags
 from .inertia import compute_shape_inertia, compute_sphere_inertia, transform_inertia
-from .terrain_generator import generate_terrain_grid, heightfield_to_mesh
+from .terrain_generator import create_mesh_heightfield, create_mesh_terrain
 from .types import (
     MESH_MAXHULLVERT,
     SDF,
@@ -67,8 +67,8 @@
     "compute_shape_inertia",
     "compute_shape_radius",
     "compute_sphere_inertia",
-    "generate_terrain_grid",
-    "heightfield_to_mesh",
+    "create_mesh_heightfield",
+    "create_mesh_terrain",
     "test_group_pair",
     "test_world_and_group_pair",
     "transform_inertia",

newton/_src/geometry/terrain_generator.py

@@ -444,7 +444,7 @@ def _heightfield_terrain(
         center_y = size[1] / 2
 
     # Use heightfield_to_mesh to convert heightfield to mesh
-    vertices, indices = heightfield_to_mesh(
+    vertices, indices = create_mesh_heightfield(
         heightfield=heightfield,
         extent_x=size[0],
         extent_y=size[1],
@@ -456,7 +456,7 @@ def _heightfield_terrain(
     return vertices, indices
 
 
-def generate_terrain_grid(
+def create_mesh_terrain(
     grid_size: tuple[int, int] = (4, 4),
     block_size: tuple[float, float] = (5.0, 5.0),
     terrain_types: list[str] | str | object | None = None,
@@ -602,7 +602,7 @@ def _to_newton_mesh(vertices: np.ndarray, indices: np.ndarray) -> tuple[np.ndarr
     return vertices.astype(np.float32), indices.astype(np.int32)
 
 
-def heightfield_to_mesh(
+def create_mesh_heightfield(
     heightfield: np.ndarray,
     extent_x: float,
     extent_y: float,

newton/_src/math/__init__.py

@@ -0,0 +1,230 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 The Newton Developers
+# 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
+
+
+@wp.func
+def boltzmann(a: float, b: float, alpha: float):
+    """
+    Compute the Boltzmann-weighted average of two values.
+
+    This function returns a smooth interpolation between `a` and `b` using a Boltzmann (softmax-like) weighting,
+    controlled by the parameter `alpha`. As `alpha` increases, the result approaches `max(a, b)`;
+    as `alpha` decreases, the result approaches the mean of `a` and `b`.
+
+    Args:
+        a (float): The first value.
+        b (float): The second value.
+        alpha (float): The sharpness parameter. Higher values make the function more "max-like".
+
+    Returns:
+        float: The Boltzmann-weighted average of `a` and `b`.
+    """
+    e1 = wp.exp(alpha * a)
+    e2 = wp.exp(alpha * b)
+    return (a * e1 + b * e2) / (e1 + e2)
+
+
+@wp.func
+def smooth_max(a: float, b: float, eps: float):
+    """
+    Compute a smooth approximation of the maximum of two values.
+
+    This function returns a value close to `max(a, b)`, but is differentiable everywhere.
+    The `eps` parameter controls the smoothness: larger values make the transition smoother.
+
+    Args:
+        a (float): The first value.
+        b (float): The second value.
+        eps (float): Smoothing parameter (should be small and positive).
+
+    Returns:
+        float: A smooth approximation of `max(a, b)`.
+    """
+    d = a - b
+    return 0.5 * (a + b + wp.sqrt(d * d + eps))
+
+
+@wp.func
+def smooth_min(a: float, b: float, eps: float):
+    """
+    Compute a smooth approximation of the minimum of two values.
+
+    This function returns a value close to `min(a, b)`, but is differentiable everywhere.
+    The `eps` parameter controls the smoothness: larger values make the transition smoother.
+
+    Args:
+        a (float): The first value.
+        b (float): The second value.
+        eps (float): Smoothing parameter (should be small and positive).
+
+    Returns:
+        float: A smooth approximation of `min(a, b)`.
+    """
+    d = a - b
+    return 0.5 * (a + b - wp.sqrt(d * d + eps))
+
+
+@wp.func
+def leaky_max(a: float, b: float):
+    """
+    Compute a numerically stable, differentiable approximation of `max(a, b)`.
+
+    This is equivalent to `smooth_max(a, b, 1e-5)`.
+
+    Args:
+        a (float): The first value.
+        b (float): The second value.
+
+    Returns:
+        float: A smooth, "leaky" maximum of `a` and `b`.
+    """
+    return smooth_max(a, b, 1e-5)
+
+
+@wp.func
+def leaky_min(a: float, b: float):
+    """
+    Compute a numerically stable, differentiable approximation of `min(a, b)`.
+
+    This is equivalent to `smooth_min(a, b, 1e-5)`.
+
+    Args:
+        a (float): The first value.
+        b (float): The second value.
+
+    Returns:
+        float: A smooth, "leaky" minimum of `a` and `b`.
+    """
+    return smooth_min(a, b, 1e-5)
+
+
+@wp.func
+def vec_min(a: wp.vec3, b: wp.vec3):
+    """
+    Compute the elementwise minimum of two 3D vectors.
+
+    Args:
+        a (wp.vec3): The first vector.
+        b (wp.vec3): The second vector.
+
+    Returns:
+        wp.vec3: The elementwise minimum.
+    """
+    return wp.vec3(wp.min(a[0], b[0]), wp.min(a[1], b[1]), wp.min(a[2], b[2]))
+
+
+@wp.func
+def vec_max(a: wp.vec3, b: wp.vec3):
+    """
+    Compute the elementwise maximum of two 3D vectors.
+
+    Args:
+        a (wp.vec3): The first vector.
+        b (wp.vec3): The second vector.
+
+    Returns:
+        wp.vec3: The elementwise maximum.
+    """
+    return wp.vec3(wp.max(a[0], b[0]), wp.max(a[1], b[1]), wp.max(a[2], b[2]))
+
+
+@wp.func
+def vec_leaky_min(a: wp.vec3, b: wp.vec3):
+    """
+    Compute the elementwise "leaky" minimum of two 3D vectors.
+
+    This uses `leaky_min` for each component.
+
+    Args:
+        a (wp.vec3): The first vector.
+        b (wp.vec3): The second vector.
+
+    Returns:
+        wp.vec3: The elementwise leaky minimum.
+    """
+    return wp.vec3(leaky_min(a[0], b[0]), leaky_min(a[1], b[1]), leaky_min(a[2], b[2]))
+
+
+@wp.func
+def vec_leaky_max(a: wp.vec3, b: wp.vec3):
+    """
+    Compute the elementwise "leaky" maximum of two 3D vectors.
+
+    This uses `leaky_max` for each component.
+
+    Args:
+        a (wp.vec3): The first vector.
+        b (wp.vec3): The second vector.
+
+    Returns:
+        wp.vec3: The elementwise leaky maximum.
+    """
+    return wp.vec3(leaky_max(a[0], b[0]), leaky_max(a[1], b[1]), leaky_max(a[2], b[2]))
+
+
+@wp.func
+def vec_abs(a: wp.vec3):
+    """
+    Compute the elementwise absolute value of a 3D vector.
+
+    Args:
+        a (wp.vec3): The input vector.
+
+    Returns:
+        wp.vec3: The elementwise absolute value.
+    """
+    return wp.vec3(wp.abs(a[0]), wp.abs(a[1]), wp.abs(a[2]))
+
+
+@wp.func
+def vec_allclose(a: Any, b: Any, rtol: float = 1e-5, atol: float = 1e-8) -> bool:
+    """
+    Check if two Warp vectors are all close.
+    """
+    for i in range(wp.static(len(a))):
+        if wp.abs(a[i] - b[i]) > atol + rtol * wp.abs(b[i]):
+            return False
+    return True
+
+
+@wp.func
+def vec_inside_limits(a: Any, lower: Any, upper: Any) -> bool:
+    """
+    Check if a Warp vector is inside the given limits.
+    """
+    for i in range(wp.static(len(a))):
+        if a[i] < lower[i] or a[i] > upper[i]:
+            return False
+    return True
+
+
+__all__ = [
+    "boltzmann",
+    "leaky_max",
+    "leaky_min",
+    "smooth_max",
+    "smooth_min",
+    "vec_abs",
+    "vec_allclose",
+    "vec_inside_limits",
+    "vec_leaky_max",
+    "vec_leaky_min",
+    "vec_max",
+    "vec_min",
+]