pyforge3d 1.0.0__py3-none-any.whl

forge3d/dynamics/rnea.py

@@ -0,0 +1,260 @@
+"""Recursive Newton-Euler Algorithm (RNEA) — inverse dynamics.
+
+Given a RigidBodyModel and joint state (q, qd, qdd), computes joint torques tau
+such that the equations of motion M(q)*qdd + h(q,qd) = tau are satisfied.
+
+Also provides:
+  - compute_mass_matrix   : O(n^2) CRBA via repeated RNEA calls
+  - forward_dynamics      : compute qdd from tau (uses CRBA + RNEA for bias)
+  - semi_implicit_euler   : one integration step
+
+Reference: Featherstone, Rigid Body Dynamics Algorithms (2008), Algorithm 5.1.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+from forge3d.dynamics.model import RigidBodyModel
+from forge3d.math.spatial import Xrot, crf, crm
+
+# ── helpers ───────────────────────────────────────────────────────────────────
+
+
+def _joint_transform(S_i: Any, q_i: float) -> Any:
+    """Spatial transform due to joint motion.
+
+    For revolute joint (S = [axis; 0]) rotating by q_i about joint axis.
+    For zero motion returns identity.
+    """
+    axis = S_i[:3]
+    n = np.linalg.norm(axis)
+    if n < 1e-10:
+        # Prismatic joint (not used in Phase 1 but handled gracefully)
+        return np.eye(6)
+    return Xrot(axis / n, q_i)
+
+
+def _gravity_spatial(gravity: Any) -> Any:
+    """Convert gravity 3-vector to spatial base acceleration (Featherstone trick).
+
+    Setting a_base = [0; -gravity] accounts for gravitational forces without
+    adding explicit external forces.  The sign flip makes gravity appear as an
+    inertial upward acceleration propagating through the tree.
+    """
+    g = np.asarray(gravity, dtype=float)
+    return np.array([0.0, 0.0, 0.0, -g[0], -g[1], -g[2]])
+
+
+# ── RNEA ─────────────────────────────────────────────────────────────────────
+
+
+def inverse_dynamics(
+    model: RigidBodyModel,
+    q: Any,
+    qd: Any,
+    qdd: Any,
+    gravity: Any | None = None,
+) -> Any:
+    """Recursive Newton-Euler: tau = ID(model, q, qd, qdd).
+
+    Parameters
+    ----------
+    model   : RigidBodyModel
+    q, qd, qdd : (n,) joint angles, velocities, accelerations
+    gravity : (3,) gravity vector (overrides model.gravity if given)
+
+    Returns
+    -------
+    tau : (n,) joint torques
+    """
+    q = np.asarray(q, dtype=float)
+    qd = np.asarray(qd, dtype=float)
+    qdd = np.asarray(qdd, dtype=float)
+    n = model.n_links
+    grav = np.asarray(gravity if gravity is not None else model.gravity, dtype=float)
+
+    # Storage
+    v = np.zeros((n, 6))  # spatial velocities
+    a = np.zeros((n, 6))  # spatial accelerations
+    f = np.zeros((n, 6))  # spatial forces
+    Xup = np.zeros((n, 6, 6))  # cumulative joint+tree transforms
+
+    # Base spatial acceleration (gravity trick)
+    a_base = _gravity_spatial(grav)
+
+    # ── Forward pass ──────────────────────────────────────────────────────────
+    for i in range(n):
+        X_J = _joint_transform(model.S[i], q[i])
+        Xup[i] = X_J @ model.X_tree[i]
+
+        if model.parent[i] == -1:
+            v[i] = model.S[i] * qd[i]
+            a[i] = Xup[i] @ a_base + model.S[i] * qdd[i] + crm(v[i]) @ model.S[i] * qd[i]
+        else:
+            p = model.parent[i]
+            v[i] = Xup[i] @ v[p] + model.S[i] * qd[i]
+            a[i] = Xup[i] @ a[p] + model.S[i] * qdd[i] + crm(v[i]) @ model.S[i] * qd[i]
+
+        f[i] = model.I_link[i] @ a[i] + crf(v[i]) @ model.I_link[i] @ v[i]
+
+    # ── Backward pass ─────────────────────────────────────────────────────────
+    tau = np.zeros(n)
+    for i in range(n - 1, -1, -1):
+        tau[i] = float(model.S[i] @ f[i])
+        p = model.parent[i]
+        if p != -1:
+            f[p] = f[p] + Xup[i].T @ f[i]
+
+    return tau
+
+
+# ── Mass matrix (CRBA via column-by-column RNEA) ──────────────────────────────
+
+
+def compute_mass_matrix(model: RigidBodyModel, q: Any) -> Any:
+    """n×n joint-space mass matrix M(q) via O(n^2) RNEA calls."""
+    q = np.asarray(q, dtype=float)
+    n = model.n_links
+    M = np.zeros((n, n))
+    zero_qd = np.zeros(n)
+    zero_grav = np.zeros(3)
+    for j in range(n):
+        qdd_j = np.zeros(n)
+        qdd_j[j] = 1.0
+        col = inverse_dynamics(model, q, zero_qd, qdd_j, gravity=zero_grav)
+        M[:, j] = col
+    return M
+
+
+# ── Forward dynamics ──────────────────────────────────────────────────────────
+
+
+def forward_dynamics(
+    model: RigidBodyModel,
+    q: Any,
+    qd: Any,
+    tau: Any,
+    gravity: Any | None = None,
+) -> Any:
+    """Compute joint accelerations qdd given torques tau.
+
+    qdd = M(q)^{-1} * (tau - h(q, qd))
+    where h = RNEA(q, qd, 0, gravity)  (bias: Coriolis + gravity)
+    """
+    q = np.asarray(q, dtype=float)
+    qd = np.asarray(qd, dtype=float)
+    tau = np.asarray(tau, dtype=float)
+    grav = gravity if gravity is not None else model.gravity
+
+    M = compute_mass_matrix(model, q)
+    h = inverse_dynamics(model, q, qd, np.zeros_like(q), gravity=grav)
+    return np.linalg.solve(M, tau - h)
+
+
+# ── Integrator ────────────────────────────────────────────────────────────────
+
+
+def semi_implicit_euler(
+    model: RigidBodyModel,
+    q: Any,
+    qd: Any,
+    tau: Any,
+    dt: float,
+    gravity: Any | None = None,
+) -> tuple[Any, Any]:
+    """Semi-implicit (symplectic) Euler integration step.
+
+    1. qdd = forward_dynamics(q, qd, tau)
+    2. qd_new = qd + dt * qdd   (velocity updated first)
+    3. q_new  = q  + dt * qd_new
+
+    Returns (q_new, qd_new) — no in-place mutation.
+    """
+    q = np.asarray(q, dtype=float)
+    qd = np.asarray(qd, dtype=float)
+    qdd = forward_dynamics(model, q, qd, tau, gravity=gravity)
+    qd_new = qd + dt * qdd
+    q_new = q + dt * qd_new
+    return q_new, qd_new
+
+
+# ── Energy utilities ──────────────────────────────────────────────────────────
+
+
+def kinetic_energy(model: RigidBodyModel, q: Any, qd: Any) -> float:
+    """T = 0.5 * qd^T * M(q) * qd."""
+    q = np.asarray(q, dtype=float)
+    qd = np.asarray(qd, dtype=float)
+    M = compute_mass_matrix(model, q)
+    return float(0.5 * qd @ M @ qd)
+
+
+def potential_energy(model: RigidBodyModel, q: Any) -> float:
+    """V = -sum_i m_i * g · r_i  (height above reference).
+
+    Uses gravity direction from model.gravity.
+    """
+    from forge3d.math.se3 import aa_to_rot, unskew
+
+    q = np.asarray(q, dtype=float)
+    n = model.n_links
+    grav = model.gravity
+
+    # Forward kinematics via SE(3).
+    # Convention: T_world[i] maps points FROM link i's joint frame TO world frame.
+    #   T_world[i] = T_world[parent] @ T_parent_to_link[i]
+    #   T_parent_to_link[i] = make_se3(R_tree @ R_joint, p_tree)
+    # where (Featherstone passive convention):
+    #   E = X_tree[:3,:3]   (E = R_tree.T)
+    #   R_tree = E.T         (active rotation of tree frame in parent)
+    #   p_tree = unskew(-(R_tree @ X_tree[3:,:3]))
+    #   R_joint = aa_to_rot(S[:3]/|S[:3]|, q_i)
+
+    R_world = [np.eye(3)] * n  # accumulated world-frame rotation of link i
+    p_world = [np.zeros(3)] * n  # world-frame position of link i joint origin
+
+    for i in range(n):
+        X_t = model.X_tree[i]
+        E = X_t[:3, :3]  # E = R_tree.T  (Featherstone passive)
+        R_tree = E.T  # active rotation
+        p_tree = unskew(-(R_tree @ X_t[3:, :3]))  # joint position in parent frame
+
+        axis = model.S[i][:3]
+        n_ax = float(np.linalg.norm(axis))
+        R_joint = aa_to_rot(axis / n_ax, q[i]) if n_ax > 1e-10 else np.eye(3)
+
+        if model.parent[i] == -1:
+            R_parent = np.eye(3)
+            p_parent = np.zeros(3)
+        else:
+            R_parent = R_world[model.parent[i]]
+            p_parent = p_world[model.parent[i]]
+
+        # World position of this link's joint origin
+        p_world[i] = p_parent + R_parent @ p_tree
+        # World orientation: parent @ tree-rotation @ joint-rotation
+        R_world[i] = R_parent @ R_tree @ R_joint
+
+    # Accumulate PE contribution from each link's CoM
+    V = 0.0
+    for i in range(n):
+        I_i = model.I_link[i]
+        mass = float(I_i[3, 3])
+        if mass < 1e-12:
+            continue
+        # CoM in link frame from spatial inertia: I[:3, 3:] = mass * skew(com)
+        skew_com = I_i[:3, 3:] / mass
+        com_local = np.array([-skew_com[1, 2], skew_com[0, 2], -skew_com[0, 1]])
+
+        com_world = p_world[i] + R_world[i] @ com_local
+        V += -mass * float(np.dot(grav, com_world))
+
+    return V
+
+
+def total_energy(model: RigidBodyModel, q: Any, qd: Any) -> float:
+    """Total mechanical energy E = T + V."""
+    return kinetic_energy(model, q, qd) + potential_energy(model, q)

forge3d/errors.py

@@ -0,0 +1,98 @@
+"""forge3d exception hierarchy and validation utilities."""
+
+from __future__ import annotations
+
+import warnings
+from typing import Any
+
+
+# ── Exception hierarchy ───────────────────────────────────────────────────────
+
+
+class Forge3dError(Exception):
+    """Base exception for all forge3d errors."""
+
+
+class PhysicsError(Forge3dError):
+    """Physics configuration or state error."""
+
+
+class ValidationError(Forge3dError, ValueError):
+    """Invalid argument passed to a forge3d API call.
+
+    The message always follows the pattern:
+        ``ClassName.method_name() — description, got value=<value>``
+    """
+
+
+class RenderError(Forge3dError):
+    """Renderer configuration or capability error."""
+
+
+class AssetError(Forge3dError):
+    """Asset loading failure (OBJ/texture not found, parse error, etc.)."""
+
+
+# ── Input validation helpers ───────────────────────────────────────────────────
+
+
+def require_positive(value: float, name: str, caller: str) -> float:
+    """Raise ValidationError if value <= 0."""
+    if value <= 0:
+        raise ValidationError(
+            f"{caller} — {name} must be positive, got {name}={value!r}"
+        )
+    return value
+
+
+def require_nonneg(value: float, name: str, caller: str) -> float:
+    """Raise ValidationError if value < 0."""
+    if value < 0:
+        raise ValidationError(
+            f"{caller} — {name} must be ≥ 0, got {name}={value!r}"
+        )
+    return value
+
+
+def require_range(value: float, lo: float, hi: float, name: str, caller: str) -> float:
+    """Raise ValidationError if value not in [lo, hi]."""
+    if not (lo <= value <= hi):
+        raise ValidationError(
+            f"{caller} — {name} must be in [{lo}, {hi}], got {name}={value!r}"
+        )
+    return value
+
+
+def require_sequence(value: Any, length: int, name: str, caller: str) -> Any:
+    """Raise ValidationError if value is not a sequence of the given length."""
+    try:
+        seq = tuple(value)
+    except TypeError:
+        raise ValidationError(
+            f"{caller} — {name} must be a {length}-element sequence, "
+            f"got {name}={value!r}"
+        ) from None
+    if len(seq) != length:
+        raise ValidationError(
+            f"{caller} — {name} must be a {length}-element sequence, "
+            f"got {len(seq)} elements"
+        )
+    return value
+
+
+def require_all_positive(value: Any, name: str, caller: str) -> None:
+    """Raise ValidationError if any component of a size/extents tuple is <= 0."""
+    for i, v in enumerate(value):
+        if v <= 0:
+            raise ValidationError(
+                f"{caller} — all {name} components must be positive, "
+                f"got {name}[{i}]={v!r}"
+            )
+
+
+# ── Deprecation helper ────────────────────────────────────────────────────────
+
+
+def deprecated(message: str, stacklevel: int = 2) -> None:
+    """Issue a DeprecationWarning with the given message."""
+    warnings.warn(message, DeprecationWarning, stacklevel=stacklevel + 1)

forge3d/events.py

@@ -0,0 +1,267 @@
+"""Collision event system for forge3d.
+
+Provides per-frame collision callbacks (begin / stay / end) and
+pair-specific collision handlers, similar to Unity's OnCollision* API
+and Pymunk's collision handler system.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Callable
+
+import numpy as np
+
+if TYPE_CHECKING:
+    pass
+
+# Type alias for callback functions
+CollisionCallback = Callable[["CollisionEvent"], None]
+
+
+@dataclass
+class CollisionEvent:
+    """Data describing a single collision contact.
+
+    Attributes:
+        body_a: First body involved in the collision.
+        body_b: Second body involved in the collision.
+        contact_point: Approximate world-frame contact point (3,).
+        normal: Contact normal pointing from body_b toward body_a (3,).
+        impulse: Estimated contact impulse magnitude (N·s).
+        relative_speed: Relative normal speed at contact (m/s, unsigned).
+    """
+
+    body_a: Any  # forge3d.facade.Body
+    body_b: Any  # forge3d.facade.Body
+    contact_point: np.ndarray
+    normal: np.ndarray
+    impulse: float
+    relative_speed: float
+
+
+@dataclass
+class CollisionHandler:
+    """Callbacks for a specific pair of bodies.
+
+    Attributes:
+        on_begin: Called once when contact first appears.
+        on_stay:  Called every step while contact persists.
+        on_end:   Called once when contact disappears.
+    """
+
+    on_begin: CollisionCallback | None = None
+    on_stay: CollisionCallback | None = None
+    on_end: CollisionCallback | None = None
+
+
+@dataclass
+class TriggerZone:
+    """Pure-data trigger zone — no physics body, no collision geometry.
+
+    Detects bodies entering / exiting a box region each step.
+    The zone does NOT generate physics contacts.
+
+    Attributes:
+        position: World-space centre (3,).
+        half_extents: Box half-extents (3,).
+    """
+
+    position: np.ndarray
+    half_extents: np.ndarray
+    _prev_inside: set[int] = field(default_factory=set)
+    _on_enter: list[Callable[[Any], None]] = field(default_factory=list)
+    _on_exit: list[Callable[[Any], None]] = field(default_factory=list)
+
+    def on_enter(self, fn: Callable[[Any], None]) -> Callable[[Any], None]:
+        """Register a callback (or use as decorator) for when a body enters this zone."""
+        self._on_enter.append(fn)
+        return fn
+
+    def on_exit(self, fn: Callable[[Any], None]) -> Callable[[Any], None]:
+        """Register a callback for when a body exits this zone."""
+        self._on_exit.append(fn)
+        return fn
+
+
+# Backward-compatible alias
+_TriggerZone = TriggerZone
+
+
+class EventDispatcher:
+    """Manages collision event state and dispatches callbacks.
+
+    Attached to ``PhysicsWorld`` and updated every ``step()``.
+    """
+
+    def __init__(self) -> None:
+        # Global listeners
+        self._on_begin: list[CollisionCallback] = []
+        self._on_stay: list[CollisionCallback] = []
+        self._on_end: list[CollisionCallback] = []
+
+        # Per-pair handlers — key = frozenset({id_a, id_b})
+        self._pair_handlers: dict[frozenset[int], CollisionHandler] = {}
+
+        # Contact state tracking — set of frozenset({id_a, id_b})
+        self._prev_contacts: set[frozenset[int]] = set()
+
+        # Body-id → Body facade handle (set by World)
+        self._bodies: dict[int, Any] = {}
+
+        # Trigger zones
+        self._triggers: list[_TriggerZone] = []
+
+    # ── Listener registration ─────────────────────────────────────────────────
+
+    def add_begin_listener(self, fn: CollisionCallback) -> None:
+        self._on_begin.append(fn)
+
+    def add_stay_listener(self, fn: CollisionCallback) -> None:
+        self._on_stay.append(fn)
+
+    def add_end_listener(self, fn: CollisionCallback) -> None:
+        self._on_end.append(fn)
+
+    def add_pair_handler(self, id_a: int, id_b: int) -> CollisionHandler:
+        key = frozenset({id_a, id_b})
+        if key not in self._pair_handlers:
+            self._pair_handlers[key] = CollisionHandler()
+        return self._pair_handlers[key]
+
+    def add_trigger_zone(
+        self,
+        position: np.ndarray,
+        half_extents: np.ndarray,
+    ) -> TriggerZone:
+        """Create and register a trigger zone (pure-data, no physics body)."""
+        zone = TriggerZone(
+            position=np.asarray(position, dtype=float),
+            half_extents=np.asarray(half_extents, dtype=float),
+        )
+        self._triggers.append(zone)
+        return zone
+
+    def ignore_pair(self, id_a: int, id_b: int) -> None:
+        """Ignore all events and collisions between two bodies (handled at broad-phase)."""
+        key = frozenset({id_a, id_b})
+        self._pair_handlers[key] = CollisionHandler()  # empty handler (no callbacks)
+
+    # ── Per-step dispatch ─────────────────────────────────────────────────────
+
+    def dispatch(self, contacts: list[Any]) -> None:
+        """Compare current contacts with previous; fire begin/stay/end callbacks.
+
+        Args:
+            contacts: List of contact data objects from the physics step.
+                Each contact has ``.body_id_a``, ``.body_id_b``,
+                ``.contact_point``, ``.normal``, ``.impulse``.
+        """
+        curr_contacts: set[frozenset[int]] = set()
+        contact_map: dict[frozenset[int], Any] = {}
+
+        for c in contacts:
+            key = frozenset({c.body_id_a, c.body_id_b})
+            curr_contacts.add(key)
+            contact_map[key] = c
+
+        began = curr_contacts - self._prev_contacts
+        stayed = curr_contacts & self._prev_contacts
+        ended = self._prev_contacts - curr_contacts
+
+        for key in began:
+            event = self._make_event(key, contact_map[key])
+            if event is None:
+                continue
+            for fn in self._on_begin:
+                fn(event)
+            handler = self._pair_handlers.get(key)
+            if handler and handler.on_begin:
+                handler.on_begin(event)
+
+        for key in stayed:
+            event = self._make_event(key, contact_map[key])
+            if event is None:
+                continue
+            for fn in self._on_stay:
+                fn(event)
+            handler = self._pair_handlers.get(key)
+            if handler and handler.on_stay:
+                handler.on_stay(event)
+
+        for key in ended:
+            # No contact data available for ended contacts; create minimal event
+            ids = list(key)
+            ba = self._bodies.get(ids[0])
+            bb = self._bodies.get(ids[1])
+            if ba is None or bb is None:
+                continue
+            event = CollisionEvent(
+                body_a=ba, body_b=bb,
+                contact_point=np.zeros(3),
+                normal=np.zeros(3),
+                impulse=0.0,
+                relative_speed=0.0,
+            )
+            for fn in self._on_end:
+                fn(event)
+            handler = self._pair_handlers.get(key)
+            if handler and handler.on_end:
+                handler.on_end(event)
+
+        self._prev_contacts = curr_contacts
+
+        # Trigger zones
+        self._dispatch_triggers()
+
+    def _make_event(self, key: frozenset[int], contact: Any) -> CollisionEvent | None:
+        ids = list(key)
+        ba = self._bodies.get(contact.body_id_a) or self._bodies.get(ids[0])
+        bb = self._bodies.get(contact.body_id_b) or self._bodies.get(ids[1])
+        if ba is None or bb is None:
+            return None
+        cp = np.asarray(contact.contact_point, dtype=float)
+        n = np.asarray(contact.normal, dtype=float)
+        imp = float(getattr(contact, "impulse", 0.0))
+        spd = float(getattr(contact, "relative_speed", 0.0))
+        return CollisionEvent(body_a=ba, body_b=bb, contact_point=cp,
+                               normal=n, impulse=imp, relative_speed=spd)
+
+    def _dispatch_triggers(self) -> None:
+        """Check which bodies are inside each trigger zone (AABB check)."""
+        for zone in self._triggers:
+            curr_inside: set[int] = set()
+            for bid, body in self._bodies.items():
+                try:
+                    bpos = body.position
+                    diff = np.abs(bpos - zone.position)
+                    if np.all(diff <= zone.half_extents):
+                        curr_inside.add(bid)
+                except Exception:
+                    pass
+
+            entered = curr_inside - zone._prev_inside
+            exited = zone._prev_inside - curr_inside
+
+            for bid in entered:
+                body = self._bodies.get(bid)
+                if body is not None:
+                    for fn in zone._on_enter:
+                        fn(body)
+
+            for bid in exited:
+                body = self._bodies.get(bid)
+                if body is not None:
+                    for fn in zone._on_exit:
+                        fn(body)
+
+            zone._prev_inside = curr_inside
+
+    def register_body(self, body_id: int, body: Any) -> None:
+        self._bodies[body_id] = body
+
+    def unregister_body(self, body_id: int) -> None:
+        self._bodies.pop(body_id, None)
+        # Remove from trigger zone tracking
+        for zone in self._triggers:
+            zone._prev_inside.discard(body_id)