openscvx 0.3.2.dev184__tar.gz → 0.3.2.dev204__tar.gz

{openscvx-0.3.2.dev184/openscvx.egg-info → openscvx-0.3.2.dev204}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.3.2.dev184
+Version: 0.3.2.dev204
 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

{openscvx-0.3.2.dev184 → openscvx-0.3.2.dev204}/examples/drone/dr_vp.py

@@ -116,14 +116,13 @@ for center in gate_centers:
 ### End Gate Parameters ###
 
 n_subs = 10
-init_poses = []
 np.random.seed(0)
-for i in range(n_subs):
-    init_pose = np.array([100.0, -60.0, 20.0])
-    init_pose[:2] = init_pose[:2] + np.random.random(2) * 20.0
-    init_poses.append(init_pose)
-
-init_poses = init_poses
+init_poses = np.array(
+    [
+        [100.0 + np.random.random() * 20.0, -60.0 + np.random.random() * 20.0, 20.0]
+        for _ in range(n_subs)
+    ]
+)  # Shape: (n_subs, 3)
 
 
 # Define list of all states (needed for Problem and constraints)
@@ -142,9 +141,16 @@ constraints = []
 for state in states:
     constraints.extend([ox.ctcs(state <= state.max), ox.ctcs(state.min <= state)])
 
-# Add visibility constraints using symbolic expressions
-for pose in init_poses:
-    constraints.append(ox.ctcs(g_vp(pose, position, attitude) <= 0.0))
+# Add visibility constraints using Vmap for parallel evaluation
+# Single CTCS constraint with vectorized evaluation over all target poses
+visibility_constraint = ox.ctcs(
+    ox.Vmap(
+        lambda pose: g_vp(pose, position, attitude),
+        batch=init_poses,
+    )
+    <= 0.0
+)
+constraints.append(visibility_constraint)
 
 # Add gate constraints using symbolic expressions
 for node, cen in zip(gate_nodes, A_gate_cen):
@@ -207,12 +213,9 @@ for _ in range(n_gates + 1):
 # Modify attitude to point sensor at targets
 R_sb = R_sb  # Sensor to body frame
 b = R_sb @ np.array([0, 1, 0])
+mean_target = np.mean(init_poses, axis=0)  # Average target position
 for k in range(n):
-    kp = []
-    for pose in init_poses:
-        kp.append(pose)
-    kp = np.mean(kp, axis=0)
-    a = kp - position_bar[k]
+    a = mean_target - position_bar[k]
     # Determine the direction cosine matrix that aligns the z-axis of the sensor frame with the
     # relative position vector
     q_xyz = np.cross(b, a)

{openscvx-0.3.2.dev184 → openscvx-0.3.2.dev204}/examples/drone/dr_vp_nodal.py

@@ -116,14 +116,13 @@ for center in gate_centers:
 ### End Gate Parameters ###
 
 n_subs = 10
-init_poses = []
 np.random.seed(5)
-for i in range(n_subs):
-    init_pose = np.array([100.0, -70.0, 20.0])
-    init_pose[:2] = init_pose[:2] + np.random.random(2) * 20.0
-    init_poses.append(init_pose)
-
-init_poses = init_poses
+init_poses = np.array(
+    [
+        [100.0 + np.random.random() * 20.0, -70.0 + np.random.random() * 20.0, 20.0]
+        for _ in range(n_subs)
+    ]
+)  # Shape: (n_subs, 3)
 
 
 # Define list of all states (needed for Problem and constraints)
@@ -142,9 +141,16 @@ constraints = []
 for state in states:
     constraints.extend([ox.ctcs(state <= state.max), ox.ctcs(state.min <= state)])
 
-# Add visibility constraints using symbolic expressions
-for pose in init_poses:
-    constraints.append((g_vp(pose, position, attitude) <= 0.0))
+# Add visibility constraints using Vmap for parallel evaluation
+# Single nodal constraint with vectorized evaluation over all target poses
+visibility_constraint = (
+    ox.Vmap(
+        lambda pose: g_vp(pose, position, attitude),
+        batch=init_poses,
+    )
+    <= 0.0
+)
+constraints.append(visibility_constraint)
 
 # Add gate constraints using symbolic expressions
 for node, cen in zip(gate_nodes, A_gate_cen):
@@ -207,12 +213,9 @@ for _ in range(n_gates + 1):
 # Modify attitude to point sensor at targets
 R_sb = R_sb  # Sensor to body frame
 b = R_sb @ np.array([0, 1, 0])
+mean_target = np.mean(init_poses, axis=0)  # Average target position
 for k in range(n):
-    kp = []
-    for pose in init_poses:
-        kp.append(pose)
-    kp = np.mean(kp, axis=0)
-    a = kp - position_bar[k]
+    a = mean_target - position_bar[k]
     # Determine the direction cosine matrix that aligns the z-axis of the sensor frame with the
     # relative position vector
     q_xyz = np.cross(b, a)

{openscvx-0.3.2.dev184 → openscvx-0.3.2.dev204}/examples/drone/dr_vp_polytope.py

@@ -111,10 +111,7 @@ polytope_point = np.array(
         [107.47, -50.00, 22.85],
     ]
 )
-init_poses = []
-for point in polytope_point:
-    init_poses.append(point)
-init_poses = init_poses
+init_poses = polytope_point  # Shape: (20, 3)
 ### Gate Parameters ###
 n_gates = 10
 gate_centers = [
@@ -162,9 +159,16 @@ constraints = []
 for state in states:
     constraints.extend([ox.ctcs(state <= state.max), ox.ctcs(state.min <= state)])
 
-# Add visibility constraints for polytope points using symbolic expressions
-for pose in init_poses:
-    constraints.append(ox.ctcs(g_vp(pose, position, attitude) <= 0.0))
+# Add visibility constraints using Vmap for parallel evaluation
+# Single CTCS constraint with vectorized evaluation over all polytope points
+visibility_constraint = ox.ctcs(
+    ox.Vmap(
+        lambda pose: g_vp(pose, position, attitude),
+        batch=init_poses,
+    )
+    <= 0.0
+)
+constraints.append(visibility_constraint)
 
 # Add gate constraints using symbolic expressions
 for node, cen in zip(gate_nodes, A_gate_cen):
@@ -227,12 +231,9 @@ for _ in range(n_gates + 1):
 # Modify attitude to point sensor at targets
 R_sb = R_sb  # Sensor to body frame
 b = R_sb @ np.array([0, 1, 0])
+mean_target = np.mean(init_poses, axis=0)  # Average target position
 for k in range(n):
-    kp = []
-    for pose in init_poses:
-        kp.append(pose)
-    kp = np.mean(kp, axis=0)
-    a = kp - position_bar[k]
+    a = mean_target - position_bar[k]
     # Determine the direction cosine matrix that aligns the z-axis of the sensor frame with the
     # relative position vector
     q_xyz = np.cross(b, a)

openscvx-0.3.2.dev204/examples/drone/obstacle_avoidance_vmap.py

@@ -0,0 +1,238 @@
+"""Double integrator obstacle avoidance with Vmap for parallel constraint evaluation.
+
+This example demonstrates using ox.Vmap directly in constraints for data-parallel
+obstacle avoidance with numerous spherical obstacles.
+
+The approach:
+
+1. Stack all obstacle centers into a single array
+2. Use Vmap directly in the constraint to compute distances in parallel
+3. The vector-valued constraint is enforced element-wise
+
+Compare with:
+  - obstacle_avoidance.py (manual loop over 3 obstacles)
+"""
+
+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 examples.plotting_viser import (
+    create_animated_plotting_server,
+    create_scp_animated_plotting_server,
+)
+from openscvx import Problem
+
+n = 6  # Number of nodes
+total_time = 10.0  # Total time for the simulation
+
+# =============================================================================
+# State and Control Definitions
+# =============================================================================
+
+# 3D position
+position = ox.State("position", shape=(3,))
+position.max = np.array([15.0, 15.0, 15.0])
+position.min = np.array([-15.0, -15.0, 0.0])
+position.initial = np.array([-10.0, -10.0, 2.0])
+position.final = np.array([10.0, 10.0, 2.0])
+
+# 3D velocity
+velocity = ox.State("velocity", shape=(3,))
+velocity.max = np.array([10.0, 10.0, 10.0])
+velocity.min = np.array([-10.0, -10.0, -10.0])
+velocity.initial = np.array([0.0, 0.0, 0.0])
+velocity.final = [("free", 0.0), ("free", 0.0), ("free", 0.0)]
+
+# Control force
+force = ox.Control("force", shape=(3,))
+f_max = 20.0
+force.max = np.array([f_max, f_max, f_max])
+force.min = np.array([-f_max, -f_max, -f_max])
+
+# Physical parameters
+m = 1.0  # Mass
+g = -9.81  # Gravity
+
+# =============================================================================
+# Obstacle Configuration (3D grid)
+# =============================================================================
+
+obstacle_radius_min, obstacle_radius_max = 1.0, 2.5
+
+# Generate obstacle positions in a 3D grid pattern between start and goal
+np.random.seed(42)
+obstacle_centers = []
+
+# Tweak number of obstacles by changing the grid sizes
+n_rows = 4
+n_cols = 4
+n_lays = 4
+
+# Create a 3D field of obstacles: rows (x) x columns (y) x layers (z)
+for i in range(n_rows):
+    for j in range(n_cols):
+        for k in range(n_lays):
+            # Base grid position
+            x = -6.0 + i * 6.0
+            y = -7.5 + j * 5.0
+            z = 1.0 + k * 2.5  # Layers at z = 1.0, 3.5, 6.0
+            # Add some randomness
+            x += np.random.uniform(-1.0, 1.0)
+            y += np.random.uniform(-1.0, 1.0)
+            z += np.random.uniform(-0.5, 0.5)
+            obstacle_centers.append([x, y, z])
+
+n_obstacles = len(obstacle_centers)  # 36 obstacles
+
+obstacle_centers = np.array(obstacle_centers)  # Shape: (n_obstacles, 3)
+obstacle_radii = np.random.uniform(obstacle_radius_min, obstacle_radius_max, size=n_obstacles)
+
+print(f"Created {n_obstacles} obstacles")
+print(f"Obstacle centers shape: {obstacle_centers.shape}")
+
+# =============================================================================
+# Dynamics (simple double integrator)
+# =============================================================================
+
+dynamics = {
+    "position": velocity,
+    "velocity": (1.0 / m) * force + np.array([0.0, 0.0, g]),
+}
+
+# =============================================================================
+# Constraints
+# =============================================================================
+
+states = [position, velocity]
+controls = [force]
+
+constraints = []
+
+# Box constraints on states
+for state in states:
+    constraints.extend(
+        [
+            ox.ctcs(state <= state.max),
+            ox.ctcs(state.min <= state),
+        ]
+    )
+
+# Box constraints on controls
+constraints.extend(
+    [
+        force <= force.max,
+        force.min <= force,
+    ]
+)
+
+# =============================================================================
+# Obstacle Avoidance
+# =============================================================================
+# Two approaches are shown below for benchmarking. Comment/uncomment to toggle.
+
+# --- APPROACH 1: Vmap (single constraint, vectorized) ---
+# - Single CTCS constraint with vectorized evaluation
+# - Vmap computes distance to all obstacles in parallel
+
+obstacle_avoidance = ox.ctcs(
+    obstacle_radii
+    <= ox.Vmap(
+        lambda obs_center: ox.linalg.Norm(position - obs_center),
+        batch=obstacle_centers,
+    )
+)
+constraints.append(obstacle_avoidance)
+
+# --- APPROACH 2: Individual constraints (loop, no vmap) ---
+# - Creates n_obstacles separate CTCS constraints
+# - Each constraint is lowered and traced independently
+
+# for i in range(n_obstacles):
+#     obs_center = obstacle_centers[i]
+#     radius = obstacle_radii[i]
+#     obstacle_constraint = ox.ctcs(
+#         radius <= ox.linalg.Norm(position - obs_center)
+#     )
+#     constraints.append(obstacle_constraint)
+
+# =============================================================================
+# Initial Guesses
+# =============================================================================
+
+straight_line = np.linspace(position.initial, position.final, n)
+position.guess = straight_line
+velocity.guess = np.zeros((n, 3))
+force.guess = np.tile([0.0, 0.0, -m * g], (n, 1))  # Hover thrust
+
+# =============================================================================
+# Problem Setup
+# =============================================================================
+
+time = ox.Time(
+    initial=0.0,
+    final=("minimize", total_time),
+    min=0.0,
+    max=total_time,
+)
+
+problem = Problem(
+    dynamics=dynamics,
+    states=states,
+    controls=controls,
+    time=time,
+    constraints=constraints,
+    N=n,
+)
+
+# SCP settings
+problem.settings.prp.dt = 0.01
+problem.settings.scp.w_tr = 5e0
+problem.settings.scp.lam_cost = 1e0
+problem.settings.scp.lam_vc = 1e2
+problem.settings.scp.ep_tr = 1e-3
+problem.settings.scp.ep_vb = 1e-4
+problem.settings.scp.ep_vc = 1e-8
+problem.settings.scp.cost_drop = 5
+problem.settings.scp.cost_relax = 0.7
+problem.settings.scp.w_tr_adapt = 1.5
+
+# =============================================================================
+# Solve and Visualize
+# =============================================================================
+
+if __name__ == "__main__":
+    problem.initialize()
+    results = problem.solve()
+    results = problem.post_process()
+
+    # Store obstacle info for visualization
+    results.update(
+        {
+            "obstacles_centers": [c for c in obstacle_centers],
+            "obstacles_radii": [[1 / r, 1 / r, 1 / r] for r in obstacle_radii],
+            "obstacles_axes": [np.eye(3) for _ in range(n_obstacles)],
+        }
+    )
+
+    # Create viser visualization servers
+    traj_server = create_animated_plotting_server(
+        results,
+        thrust_key="force",
+        viewcone_scale=5.0,
+    )
+    scp_server = create_scp_animated_plotting_server(
+        results,
+        attitude_stride=3,
+        frame_duration_ms=200,
+    )
+
+    # Keep servers running
+    traj_server.sleep_forever()

{openscvx-0.3.2.dev184 → openscvx-0.3.2.dev204}/openscvx/__init__.py

@@ -15,6 +15,7 @@ from openscvx.symbolic.expr import (
     CTCS,
     Abs,
     Add,
+    Bilerp,
     Block,
     Concat,
     Constant,
@@ -32,6 +33,7 @@ from openscvx.symbolic.expr import (
     Index,
     Inequality,
     Leaf,
+    Linterp,
     Log,
     LogSumExp,
     MatMul,
@@ -51,6 +53,7 @@ from openscvx.symbolic.expr import (
     Sum,
     Tan,
     Variable,
+    Vmap,
     Vstack,
     ctcs,
 )
@@ -106,6 +109,8 @@ __all__ = [
     "Log",
     "LogSumExp",
     "Max",
+    "Linterp",
+    "Bilerp",
     # Constraints
     "Constraint",
     "Equality",
@@ -113,6 +118,8 @@ __all__ = [
     "NodalConstraint",
     "CTCS",
     "ctcs",
+    # Data parallelism
+    "Vmap",
     # Submodules
     "stl",
     "spatial",

{openscvx-0.3.2.dev184 → openscvx-0.3.2.dev204}/openscvx/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.3.2.dev184'
-__version_tuple__ = version_tuple = (0, 3, 2, 'dev184')
+__version__ = version = '0.3.2.dev204'
+__version_tuple__ = version_tuple = (0, 3, 2, 'dev204')
 
-__commit_id__ = commit_id = 'gb43481d6f'
+__commit_id__ = commit_id = 'g3b21ffc13'

{openscvx-0.3.2.dev184 → openscvx-0.3.2.dev204}/openscvx/symbolic/expr/__init__.py

@@ -118,9 +118,11 @@ from .linalg import Diag, Norm, Sum, Transpose
 # Mathematical functions
 from .math import (
     Abs,
+    Bilerp,
     Cos,
     Exp,
     Huber,
+    Linterp,
     Log,
     LogSumExp,
     Max,
@@ -144,6 +146,9 @@ from .stl import Or
 # Variable
 from .variable import Variable
 
+# Vmap
+from .vmap import Vmap, _Placeholder  # noqa: F401
+
 __all__ = [
     # Core base classes and fundamental operations
     "Expr",
@@ -191,6 +196,8 @@ __all__ = [
     "Log",
     "LogSumExp",
     "Max",
+    "Linterp",
+    "Bilerp",
     # Linear algebra operations
     "Transpose",
     "Stack",
@@ -219,4 +226,6 @@ __all__ = [
     "ctcs",
     # STL operations
     "Or",
+    # Vmap
+    "Vmap",
 ]