openscvx 0.3.2.dev143__tar.gz → 0.3.2.dev157__tar.gz

{openscvx-0.3.2.dev143/openscvx.egg-info → openscvx-0.3.2.dev157}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.3.2.dev143
+Version: 0.3.2.dev157
 Summary: A general Python-based successive convexification implementation which uses a JAX backend.
 Author-email: Chris Hayner and Griffin Norris <haynec@uw.edu>
 License: Apache Software License
@@ -32,9 +32,12 @@ Requires-Dist: cvxpygen; extra == "cvxpygen"
 Requires-Dist: qocogen; extra == "cvxpygen"
 Provides-Extra: stl
 Requires-Dist: stljax; extra == "stl"
+Provides-Extra: lie
+Requires-Dist: jaxlie; extra == "lie"
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
 Requires-Dist: scipy; extra == "test"
+Requires-Dist: jaxlie; extra == "test"
 Dynamic: license-file
 
 <a id="readme-top"></a>

openscvx-0.3.2.dev157/examples/arm/three_link_arm.py

@@ -0,0 +1,253 @@
+"""3-link robotic arm with Product of Exponentials forward kinematics.
+
+This example demonstrates trajectory optimization for a 3-link spatial arm
+using Lie algebra operations for forward kinematics. The problem includes:
+
+- 3 revolute joints with angle and velocity states
+- Product of Exponentials (PoE) forward kinematics using SE3Exp
+- End-effector position tracking objective
+- Joint torque control inputs
+
+The PoE formula computes forward kinematics as:
+    T_ee(q) = exp(ξ₁q₁) @ exp(ξ₂q₂) @ exp(ξ₃q₃) @ T_home
+
+where ξᵢ are the screw axes and T_home is the end-effector pose at q=0.
+
+Requires jaxlie: pip install openscvx[lie]
+"""
+
+import os
+import sys
+
+import numpy as np
+
+# Add grandparent directory to path to import examples.plotting
+current_dir = os.path.dirname(os.path.abspath(__file__))
+grandparent_dir = os.path.dirname(os.path.dirname(current_dir))
+sys.path.append(grandparent_dir)
+
+import openscvx as ox
+from openscvx import Problem
+from openscvx.plotting import plot_scp_iterations
+
+# =============================================================================
+# Robot Parameters
+# =============================================================================
+
+# Link lengths (meters)
+L1 = 0.4  # Base to shoulder
+L2 = 0.3  # Shoulder to elbow
+L3 = 0.2  # Elbow to end-effector
+
+# Joint inertias (simplified, kg*m^2)
+inertia = np.array([0.05, 0.03, 0.01])
+
+# Number of discretization nodes
+n = 3
+total_time = 2.0
+
+# =============================================================================
+# Screw Axes for Product of Exponentials
+# =============================================================================
+# Each screw axis ξ = [v; ω] where:
+#   ω = unit rotation axis
+#   v = -ω × q (q = point on the axis)
+#
+# Robot configuration at home (q=0):
+#   - Joint 1: z-rotation at origin
+#   - Joint 2: y-rotation at height L1
+#   - Joint 3: y-rotation at height L1, reach L2 in x
+
+# Screw axes as rows of a 3x6 matrix
+screw_axes = np.array(
+    [
+        [0.0, 0.0, 0.0, 0.0, 0.0, 1.0],  # Joint 1: z-rotation at origin
+        [-L1, 0.0, 0.0, 0.0, 1.0, 0.0],  # Joint 2: y-rotation at [0, 0, L1]
+        [-L1, 0.0, L2, 0.0, 1.0, 0.0],  # Joint 3: y-rotation at [L2, 0, L1]
+    ]
+)
+
+# Home configuration: end-effector at [L2+L3, 0, L1] with identity rotation
+T_home = np.array(
+    [
+        [1.0, 0.0, 0.0, L2 + L3],
+        [0.0, 1.0, 0.0, 0.0],
+        [0.0, 0.0, 1.0, L1],
+        [0.0, 0.0, 0.0, 1.0],
+    ]
+)
+
+# =============================================================================
+# States
+# =============================================================================
+
+# Joint angles (3,)
+angle = ox.State("angle", shape=(3,))
+angle.max = np.array([np.pi, np.pi / 2, np.pi])
+angle.min = np.array([-np.pi, -np.pi / 2, -np.pi])
+angle.initial = np.array([0.0, 0.0, 0.0])
+angle.final = [("free", 0.0), ("free", 0.0), ("free", 0.0)]
+
+# Joint velocities (3,)
+velocity = ox.State("velocity", shape=(3,))
+velocity.max = np.array([5.0, 5.0, 5.0])
+velocity.min = np.array([-5.0, -5.0, -5.0])
+velocity.initial = np.array([0.0, 0.0, 0.0])
+velocity.final = np.array([0.0, 0.0, 0.0])
+
+states = [angle, velocity]
+
+# =============================================================================
+# Controls
+# =============================================================================
+
+# Joint torques (3,)
+torque = ox.Control("torque", shape=(3,))
+torque.max = np.array([10.0, 5.0, 2.0])
+torque.min = np.array([-10.0, -5.0, -2.0])
+
+controls = [torque]
+
+# =============================================================================
+# Forward Kinematics using Product of Exponentials
+# =============================================================================
+# T_ee(q) = exp(ξ₁q₁) @ exp(ξ₂q₂) @ exp(ξ₃q₃) @ T_home
+
+# Symbolic screw axis expressions scaled by joint angles
+xi_1 = ox.Constant(screw_axes[0])
+xi_2 = ox.Constant(screw_axes[1])
+xi_3 = ox.Constant(screw_axes[2])
+
+twist_1 = xi_1 * angle[0]
+twist_2 = xi_2 * angle[1]
+twist_3 = xi_3 * angle[2]
+
+# Exponential maps for each joint
+T_01 = ox.lie.SE3Exp(twist_1)  # 4x4 transform
+T_12 = ox.lie.SE3Exp(twist_2)
+T_23 = ox.lie.SE3Exp(twist_3)
+
+# Chain the transforms: T_ee = T_01 @ T_12 @ T_23 @ T_home
+T_0_home = ox.Constant(T_home)
+T_ee = T_01 @ T_12 @ T_23 @ T_0_home
+
+# Extract end-effector position from homogeneous transform
+p_ee = ox.Concat(T_ee[0, 3], T_ee[1, 3], T_ee[2, 3])
+
+# =============================================================================
+# Dynamics (simplified second-order)
+# =============================================================================
+# Using simplified dynamics: I * qdd = tau
+#
+# Note: Full manipulator dynamics M(q)q̈ + C(q,q̇)q̇ + G(q) = τ are not needed
+# here. This example demonstrates the Lie algebra functionality (SE3Exp for
+# Product of Exponentials FK), which is independent of the dynamics model.
+
+I_inv = ox.Constant(1.0 / inertia)
+
+dynamics = {
+    "angle": velocity,
+    "velocity": I_inv * torque,  # Element-wise: qdd_i = tau_i / I_i
+}
+
+# =============================================================================
+# Constraints
+# =============================================================================
+
+# Target end-effector position
+target = ox.Parameter("target", shape=(3,), value=np.array([0.3, 0.3, 0.5]))
+
+# Box constraints
+constraints = []
+for state in states:
+    constraints.extend(
+        [
+            ox.ctcs(state <= state.max),
+            ox.ctcs(state.min <= state),
+        ]
+    )
+
+# End-effector target constraint at final node (commented out for debugging)
+ee_tolerance = 0.01  # 1cm tolerance
+ee_target_constraint = (ox.linalg.Norm(p_ee - target, ord=2) <= ee_tolerance).at([n - 1])
+constraints.append(ee_target_constraint)
+
+# =============================================================================
+# Initial Guesses
+# =============================================================================
+
+np.random.seed(42)  # For reproducibility
+# Terminal angles that reach the target (from workspace analysis)
+q_terminal = np.deg2rad([47.8, -38.6, 62.4])
+angle.guess = np.linspace(angle.initial, q_terminal, n)
+velocity.guess = np.zeros((n, 3))
+torque.guess = np.zeros((n, 3))
+
+# =============================================================================
+# Problem Setup
+# =============================================================================
+
+time = ox.Time(
+    initial=0.0,
+    final=total_time,
+    min=0.0,
+    max=total_time,
+)
+
+problem = Problem(
+    dynamics=dynamics,
+    states=states,
+    controls=controls,
+    time=time,
+    constraints=constraints,
+    N=n,
+)
+
+# Solver settings
+problem.settings.prp.dt = 0.01
+problem.settings.scp.w_tr = 1e1
+problem.settings.scp.lam_cost = 1e2
+problem.settings.scp.lam_vc = 1e1
+problem.settings.scp.lam_vb = 1e2
+
+if __name__ == "__main__":
+    print("3-Link Arm Trajectory Optimization with Product of Exponentials FK")
+    print("=" * 60)
+    print(f"Link lengths: L1={L1}m, L2={L2}m, L3={L3}m")
+    print(f"Home EE position: [{L2 + L3:.2f}, 0.00, {L1:.2f}]")
+    print(f"Target position: {target.value}")
+    print()
+
+    problem.initialize()
+    results = problem.solve()
+    results = problem.post_process()
+
+    # Extract final joint angles
+    final_q = results.trajectory["angle"][-1]
+
+    print()
+    print("Results:")
+    print(
+        f"Final joint angles [deg]: q1={np.rad2deg(final_q[0]):.1f}, "
+        f"q2={np.rad2deg(final_q[1]):.1f}, q3={np.rad2deg(final_q[2]):.1f}"
+    )
+
+    # Compute final EE position (using jaxlie for verification)
+    import jaxlie
+
+    def compute_ee_position(q_vals):
+        T1 = jaxlie.SE3.exp(screw_axes[0] * q_vals[0]).as_matrix()
+        T2 = jaxlie.SE3.exp(screw_axes[1] * q_vals[1]).as_matrix()
+        T3 = jaxlie.SE3.exp(screw_axes[2] * q_vals[2]).as_matrix()
+        T_final = T1 @ T2 @ T3 @ T_home
+        return T_final[:3, 3]
+
+    plot_scp_iterations(results).show()
+
+    tgt = target.value
+    final_ee = compute_ee_position(final_q)
+    error = np.linalg.norm(final_ee - tgt)
+
+    print(f"Final EE position: [{final_ee[0]:.3f}, {final_ee[1]:.3f}, {final_ee[2]:.3f}]")
+    print(f"Target position:   [{tgt[0]:.3f}, {tgt[1]:.3f}, {tgt[2]:.3f}]")
+    print(f"Position error:    {error:.4f} m")

{openscvx-0.3.2.dev143 → openscvx-0.3.2.dev157}/openscvx/__init__.py

@@ -5,6 +5,7 @@ os.environ["EQX_ON_ERROR"] = "nan"
 
 # Cache management
 # Core symbolic expressions - flat namespace for most common functions
+import openscvx.symbolic.expr.lie as lie
 import openscvx.symbolic.expr.linalg as linalg
 import openscvx.symbolic.expr.spatial as spatial
 import openscvx.symbolic.expr.stl as stl
@@ -106,6 +107,7 @@ __all__ = [
     "stl",
     "spatial",
     "linalg",
+    "lie",
     # Expert mode types
     "ByofSpec",
 ]

{openscvx-0.3.2.dev143 → openscvx-0.3.2.dev157}/openscvx/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.3.2.dev143'
-__version_tuple__ = version_tuple = (0, 3, 2, 'dev143')
+__version__ = version = '0.3.2.dev157'
+__version_tuple__ = version_tuple = (0, 3, 2, 'dev157')
 
-__commit_id__ = commit_id = 'gfab08d1d7'
+__commit_id__ = commit_id = 'ga9bc61751'

{openscvx-0.3.2.dev143 → openscvx-0.3.2.dev157}/openscvx/symbolic/expr/__init__.py

@@ -54,6 +54,13 @@ Module Organization:
         Direction Cosine Matrix), `SSMP` (4×4 skew-symmetric matrix for quaternion
         dynamics), and `SSM` (3×3 skew-symmetric matrix for cross products).
 
+    Lie Algebra Operations (lie.py):
+        Lie algebra operations for rigid body dynamics. Built-in operators include
+        `AdjointDual` (coadjoint for Coriolis/centrifugal forces) and `Adjoint`
+        (Lie bracket). jaxlie-backed operators (requires `pip install openscvx[lie]`)
+        include `SO3Exp`, `SO3Log`, `SE3Exp`, and `SE3Log` for exponential and
+        logarithm maps on rotation and rigid transformation groups.
+
     Constraint Specifications (constraint.py):
         `NodalConstraint` for enforcing constraints at discrete nodes and `CTCS` for
         continuous-time constraint satisfaction.
@@ -93,6 +100,18 @@ from .expr import (
     traverse,
 )
 
+# Lie algebra operations
+from .lie import (
+    Adjoint,
+    AdjointDual,
+    SE3Adjoint,
+    SE3AdjointDual,
+    SE3Exp,
+    SE3Log,
+    SO3Exp,
+    SO3Log,
+)
+
 # Linear algebra operations
 from .linalg import Diag, Norm, Sum, Transpose
 
@@ -183,6 +202,15 @@ __all__ = [
     "QDCM",
     "SSMP",
     "SSM",
+    # Lie algebra operations
+    "AdjointDual",
+    "Adjoint",
+    "SE3Adjoint",
+    "SE3AdjointDual",
+    "SO3Exp",
+    "SO3Log",
+    "SE3Exp",
+    "SE3Log",
     # Specialized constraints
     "NodalConstraint",
     "CrossNodeConstraint",

openscvx-0.3.2.dev157/openscvx/symbolic/expr/lie/__init__.py

@@ -0,0 +1,87 @@
+"""Lie algebra operations for rigid body dynamics.
+
+This module provides symbolic expression nodes for Lie algebra operations
+commonly used in 6-DOF rigid body dynamics, robotics, and geometric mechanics.
+These operations enable elegant formulations of Newton-Euler dynamics using
+spatial vectors (twists and wrenches).
+
+The module provides two tiers of functionality:
+
+**Built-in operators** work out of the box and include adjoint/coadjoint
+operators for dynamics (_e.g._ ``Adjoint``, ``AdjointDual``) and frame transformations
+(_e.g._ ``SE3Adjoint``, ``SE3AdjointDual``).
+
+**jaxlie-backed operators** require ``pip install openscvx[lie]`` and provide
+exponential/logarithm maps for SO(3) and SE(3) groups (_e.g._ ``SO3Exp``, ``SO3Log``,
+``SE3Exp``, ``SE3Log``).
+
+Conventions:
+    - Twist (spatial velocity): ξ = [v; ω] where v ∈ ℝ³ is linear velocity
+      and ω ∈ ℝ³ is angular velocity (both in body frame)
+    - Wrench (spatial force): F = [f; τ] where f ∈ ℝ³ is force and τ ∈ ℝ³
+      is torque (both in body frame)
+
+Note:
+    The twist convention [v; ω] (linear first, angular second) matches jaxlie's
+    SE3 tangent parameterization, so no reordering is needed during lowering.
+
+Example:
+    Newton-Euler dynamics for a rigid body using the coadjoint operator::
+
+        import openscvx as ox
+
+        twist = ox.State("twist", shape=(6,))
+        M = ox.Parameter("M", shape=(6, 6), value=spatial_inertia)
+        wrench = ox.Control("wrench", shape=(6,))
+
+        momentum = M @ twist
+        bias_force = ox.lie.AdjointDual(twist, momentum)
+        twist_dot = M_inv @ (wrench - bias_force)
+
+    Product of Exponentials forward kinematics (requires jaxlie)::
+
+        screw_axis = ox.Constant(np.array([0, 0, 0, 0, 0, 1]))
+        theta = ox.State("theta", shape=(1,))
+        T_joint = ox.lie.SE3Exp(screw_axis * theta)  # 4×4 matrix
+
+References:
+    - Murray, Li, Sastry: "A Mathematical Introduction to Robotic Manipulation"
+    - Featherstone: "Rigid Body Dynamics Algorithms"
+    - Sola et al.: "A micro Lie theory for state estimation in robotics"
+"""
+
+# Core operators - no dependencies
+from .adjoint import Adjoint, AdjointDual, SE3Adjoint, SE3AdjointDual
+
+# jaxlie-backed operators - optional dependency
+try:
+    from .se3 import SE3Exp, SE3Log
+    from .so3 import SO3Exp, SO3Log
+
+    _JAXLIE_AVAILABLE = True
+except ImportError:
+    _JAXLIE_AVAILABLE = False
+
+    def _make_stub(name: str):
+        """Create a stub class that raises ImportError on instantiation."""
+
+        def __init__(self, *args, **kwargs):
+            raise ImportError(f"{name} requires jaxlie. Install with: pip install openscvx[lie]")
+
+        return type(name, (), {"__init__": __init__})
+
+    SO3Exp = _make_stub("SO3Exp")
+    SO3Log = _make_stub("SO3Log")
+    SE3Exp = _make_stub("SE3Exp")
+    SE3Log = _make_stub("SE3Log")
+
+__all__ = [
+    "AdjointDual",
+    "Adjoint",
+    "SE3Adjoint",
+    "SE3AdjointDual",
+    "SO3Exp",
+    "SO3Log",
+    "SE3Exp",
+    "SE3Log",
+]