getting rocket landing example to work

Author: Transurgeon

cvxpy/sandbox/control_of_car.py

@@ -4,7 +4,7 @@
 import cvxpy as cp
 
 
-def solve_car_control(x_final, L=0.1, N=10, h=0.1, gamma=10):
+def solve_car_control(x_final, L=0.1, N=25, h=0.1, gamma=10):
     """
     Solve the nonlinear optimal control problem for car trajectory planning.
     
@@ -19,9 +19,9 @@ def solve_car_control(x_final, L=0.1, N=10, h=0.1, gamma=10):
     - x_opt: optimal states (N+1 x 3)
     - u_opt: optimal controls (N x 2)
     """
-    
-    # Variables
-    # States: x[k] = [p1(k), p2(k), theta(k)]
+    # Add random seed for reproducibility
+    np.random.seed(78)
+
     x = [cp.Variable(3) for _ in range(N+1)]
     # Controls: u[k] = [s(k), phi(k)]
     u = [cp.Variable(2) for _ in range(N)]
@@ -78,10 +78,6 @@ def plot_trajectory(x_opt, u_opt, L, h, title="Car Trajectory"):
     Plot the car trajectory with orientation indicators.
     """
     fig, ax = plt.subplots(1, 1, figsize=(8, 8))
-    
-    # Convert x_opt to numpy array for easier indexing
-    x_opt = np.array([xi.value for xi in x_opt])
-    u_opt = np.array([ui.value for ui in u_opt])
     # Plot trajectory
     ax.plot(x_opt[:, 0], x_opt[:, 1], 'b-', linewidth=2, label='Trajectory')
 
@@ -150,29 +146,28 @@ def plot_trajectory(x_opt, u_opt, L, h, title="Car Trajectory"):
                 f"p2={x_opt[-1].value[1]:.3f}, "
                 f"theta={x_opt[-1].value[2]:.3f}"
             )
-            
+            x_opt = np.array([xi.value for xi in x_opt])
+            u_opt = np.array([ui.value for ui in u_opt])
             # Plot the trajectory
             fig, ax = plot_trajectory(x_opt, u_opt, L=0.1, h=0.1, title=description)
             plt.show()
         else:
             print("Optimization failed!")
-    """
-    # Additional analysis: plot control inputs
-    if x_opt is not None and u_opt is not None:
-        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 6))
-        
-        time_steps = np.arange(len(u_opt)) * 0.1  # h = 0.1
-        
-        ax1.plot(time_steps, u_opt[:, 0], 'b-', linewidth=2)
-        ax1.set_ylabel('Speed s(t)')
-        ax1.set_xlabel('Time')
-        ax1.grid(True, alpha=0.3)
-        
-        ax2.plot(time_steps, u_opt[:, 1], 'r-', linewidth=2)
-        ax2.set_ylabel('Steering angle φ(t)')
-        ax2.set_xlabel('Time')
-        ax2.grid(True, alpha=0.3)
-        
-        plt.tight_layout()
-        plt.show()
-    """
+        # Additional analysis: plot control inputs
+        if x_opt is not None and u_opt is not None:
+            fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 6))
+
+            time_steps = np.arange(len(u_opt)) * 0.1  # h = 0.1
+
+            ax1.plot(time_steps, u_opt[:, 0], 'b-', linewidth=2)
+            ax1.set_ylabel('Speed s(t)')
+            ax1.set_xlabel('Time')
+            ax1.grid(True, alpha=0.3)
+            
+            ax2.plot(time_steps, u_opt[:, 1], 'r-', linewidth=2)
+            ax2.set_ylabel('Steering angle φ(t)')
+            ax2.set_xlabel('Time')
+            ax2.grid(True, alpha=0.3)
+            
+            plt.tight_layout()
+            plt.show()

cvxpy/sandbox/control_of_car_matrix.py

@@ -0,0 +1,186 @@
+import cvxpy as cp
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+def solve_car_control(x_final, L=0.1, N=10, h=0.1, gamma=10):
+    """
+    Solve the nonlinear optimal control problem for car trajectory planning.
+    
+    Parameters:
+    - x_final: tuple (p1, p2, theta) for final position and orientation
+    - L: wheelbase length
+    - N: number of time steps
+    - h: time step size
+    - gamma: weight for control smoothness term
+    
+    Returns:
+    - x_opt: optimal states (N+1 x 3)
+    - u_opt: optimal controls (N x 2)
+    """
+    # Add random seed for reproducibility
+    np.random.seed(42)
+    # Variables
+    # States: x[k] = [p1(k), p2(k), theta(k)]
+    x = cp.Variable((N+1, 3))
+    # Controls: u[k] = [s(k), phi(k)]
+    u = cp.Variable((N, 2))
+    
+    # Initial state (starting at origin with zero orientation)
+    x_init = np.array([0, 0, 0])
+    
+    # Objective function
+    objective = 0
+    
+    # Sum of squared control inputs
+    for k in range(N):
+        objective += cp.sum_squares(u[k, :])
+    
+    # Control smoothness term
+    for k in range(N-1):
+        objective += gamma * cp.sum_squares(u[k+1, :] - u[k, :])
+    
+    # Constraints
+    constraints = []
+    
+    # Initial state constraint
+    constraints.append(x[0, :] == x_init)
+    
+    # Dynamics constraints
+    # Note: We're pretending cp.sin, cp.cos, and cp.tan exist as atoms
+    for k in range(N):
+        # x[k+1] = f(x[k], u[k])
+        # where f(x, u) = x + h * [u[0]*cos(x[2]), u[0]*sin(x[2]), u[0]*tan(u[1])/L]
+        
+        # Position dynamics
+        constraints.append(x[k+1, 0] == x[k, 0] + h * u[k, 0] * cp.cos(x[k, 2]))
+        constraints.append(x[k+1, 1] == x[k, 1] + h * u[k, 0] * cp.sin(x[k, 2]))
+        
+        # Orientation dynamics
+        constraints.append(x[k+1, 2] == x[k, 2] + h * (u[k, 0] / L) * cp.tan(u[k, 1]))
+    
+    # Final state constraint
+    constraints.append(x[N, :] == x_final)
+    
+    # Steering angle limits (optional but realistic)
+    # Assuming max steering angle of 45 degrees
+    max_steering = np.pi / 4
+    constraints.append(u[:, 1] >= -max_steering)
+    constraints.append(u[:, 1] <= max_steering)
+    
+    # Create and solve the problem
+    problem = cp.Problem(cp.Minimize(objective), constraints)
+    
+    # Solve using an appropriate solver
+    # For nonlinear problems, we might need special solver options
+    problem.solve(solver=cp.SCS, verbose=True)
+    
+    # Extract solution
+    x_opt = x.value
+    u_opt = u.value
+    
+    return x_opt, u_opt
+
+
+def plot_trajectory(x_opt, u_opt, L, h, title="Car Trajectory"):
+    """
+    Plot the car trajectory with orientation indicators.
+    """
+    fig, ax = plt.subplots(1, 1, figsize=(8, 8))
+    
+    # Plot trajectory
+    ax.plot(x_opt[:, 0], x_opt[:, 1], 'b-', linewidth=2, label='Trajectory')
+    
+    # Plot car position and orientation at several time steps
+    car_length = L
+    car_width = L * 0.6
+    
+    # Select time steps to show car outline (every 5th step)
+    steps_to_show = range(0, len(x_opt), 5)
+    
+    for k in steps_to_show:
+        p1, p2, theta = x_opt[k]
+        
+        # Car outline (simplified rectangle)
+        # Front of car
+        front_x = p1 + (car_length/2) * np.cos(theta)
+        front_y = p2 + (car_length/2) * np.sin(theta)
+        
+        # Rear of car
+        rear_x = p1 - (car_length/2) * np.cos(theta)
+        rear_y = p2 - (car_length/2) * np.sin(theta)
+        
+        # Draw car as a line with orientation
+        ax.plot([rear_x, front_x], [rear_y, front_y], 'k-', linewidth=3, alpha=0.5)
+        
+        # Draw steering angle indicator if not at final position
+        if k < len(u_opt):
+            phi = u_opt[k, 1]
+            # Steering direction from front of car
+            steer_x = front_x + (car_length/3) * np.cos(theta + phi)
+            steer_y = front_y + (car_length/3) * np.sin(theta + phi)
+            ax.plot([front_x, steer_x], [front_y, steer_y], 'r-', linewidth=2, alpha=0.7)
+    
+    # Mark start and end points
+    ax.plot(x_opt[0, 0], x_opt[0, 1], 'go', markersize=10, label='Start')
+    ax.plot(x_opt[-1, 0], x_opt[-1, 1], 'ro', markersize=10, label='Goal')
+    
+    ax.set_xlabel('p1')
+    ax.set_ylabel('p2')
+    ax.set_title(title)
+    ax.legend()
+    ax.grid(True, alpha=0.3)
+    ax.axis('equal')
+    
+    return fig, ax
+
+
+# Example usage
+if __name__ == "__main__":
+    # Test cases from the figure
+    test_cases = [
+        ((0, 1, 0), "Move forward to (0, 1)"),
+        #((0, 1, np.pi/2), "Move to (0, 1) and turn 90°"),
+        #((0, 0.5, 0), "Move forward to (0, 0.5)"),
+        #((0.5, 0.5, -np.pi/2), "Move to (0.5, 0.5) and turn -90°")
+    ]
+    
+    # Solve for each test case
+    for x_final, description in test_cases:
+        print(f"\nSolving for: {description}")
+        print(f"Target state: p1={x_final[0]}, p2={x_final[1]}, theta={x_final[2]:.2f}")
+        
+        try:
+            x_opt, u_opt = solve_car_control(x_final)
+            
+            if x_opt is not None and u_opt is not None:
+                print("Optimization successful!")
+                print(f"Final position: p1={x_opt[-1, 0]:.3f}, p2={x_opt[-1, 1]:.3f}, theta={x_opt[-1, 2]:.3f}")
+                
+                # Plot the trajectory
+                fig, ax = plot_trajectory(x_opt, u_opt, L=0.1, h=0.1, title=description)
+                plt.show()
+            else:
+                print("Optimization failed!")
+                
+        except Exception as e:
+            print(f"Error: {e}")
+    
+    # Additional analysis: plot control inputs
+    if x_opt is not None and u_opt is not None:
+        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 6))
+        
+        time_steps = np.arange(len(u_opt)) * 0.1  # h = 0.1
+        
+        ax1.plot(time_steps, u_opt[:, 0], 'b-', linewidth=2)
+        ax1.set_ylabel('Speed s(t)')
+        ax1.set_xlabel('Time')
+        ax1.grid(True, alpha=0.3)
+        
+        ax2.plot(time_steps, u_opt[:, 1], 'r-', linewidth=2)
+        ax2.set_ylabel('Steering angle φ(t)')
+        ax2.set_xlabel('Time')
+        ax2.grid(True, alpha=0.3)
+        
+        plt.tight_layout()
+        plt.show()

cvxpy/sandbox/rocket_landing.py

@@ -0,0 +1,120 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+import cvxpy as cp
+
+# Data - all parameters normalized to be dimensionless
+h_0 = 1                      # Initial height
+v_0 = 0                      # Initial velocity
+m_0 = 1.0                    # Initial mass
+m_T = 0.6                    # Final mass
+g_0 = 1                      # Gravity at the surface
+h_c = 500                    # Used for drag
+c = 0.5 * np.sqrt(g_0 * h_0) # Thrust-to-fuel mass
+D_c = 0.5 * 620 * m_0 / g_0  # Drag scaling
+u_t_max = 3.5 * g_0 * m_0    # Maximum thrust
+T_max = 0.2                  # Number of seconds
+T = 10                     # Number of time steps
+dt = 0.2 / T                 # Time per discretized step
+
+# Create variables
+x_v = cp.Variable(T, bounds=[0, np.inf])  # Velocity
+x_h = cp.Variable(T, bounds=[0,np.inf])  # Height  
+x_m = cp.Variable(T)               # Mass
+u_t = cp.Variable(T, bounds=[0, u_t_max])  # Thrust
+
+# Set starting values (equivalent to JuMP's start parameter)
+x_v.value = np.full(T, v_0)        # start = v_0
+x_h.value = np.full(T, h_0)        # start = h_0
+x_m.value = np.full(T, m_0)        # start = m_0
+u_t.value = np.zeros(T)            # start = 0
+
+# Define forces as functions
+def D(x_h_val, x_v_val):
+    """Drag force function"""
+    return D_c * cp.square(x_v_val) * cp.exp(-h_c * (x_h_val - h_0) / h_0)
+
+def g(x_h_val):
+    """Gravity function"""
+    return g_0 * cp.square(h_0 / x_h_val)
+
+# Define time derivative function
+def ddt(x, t):
+    """Time derivative using backward difference"""
+    return (x[t] - x[t-1]) / dt
+
+# Initialize constraints list
+constraints = []
+
+# Boundary conditions
+constraints.append(x_v[0] == v_0)
+constraints.append(x_h[0] == h_0)
+constraints.append(x_m[0] == m_0)
+constraints.append(u_t[T-1] == 0.0)
+
+# Mass constraints
+constraints.append(x_m >= m_T)
+
+# Thrust constraints
+constraints.append(u_t <= u_t_max)
+
+# Dynamical equations as constraints
+for t in range(1, T):
+    # Rate of ascent: dx_h/dt = x_v
+    constraints.append(ddt(x_h, t) == x_v[t-1])
+    
+    # Acceleration: dx_v/dt = (u_t - D(x_h, x_v))/x_m - g(x_h)
+    constraints.append(
+        ddt(x_v, t) == (u_t[t-1] - D(x_h[t-1], x_v[t-1])) / x_m[t-1] - g(x_h[t-1])
+    )
+    
+    # Rate of mass loss: dx_m/dt = -u_t/c
+    constraints.append(ddt(x_m, t) == -u_t[t-1] / c)
+
+# Objective: maximize altitude at end of time of flight
+objective = cp.Maximize(x_h[T-1])
+
+# Create and solve problem
+problem = cp.Problem(objective, constraints)
+
+# Solve the problem
+problem.solve(solver=cp.IPOPT, nlp=True, verbose=True)
+
+# Check if solution was found
+if problem.status == cp.OPTIMAL:
+    print(f"Optimal value: {problem.value}")
+    print(f"Final altitude: {x_h.value[T-1]}")
+    print(f"Final mass: {x_m.value[T-1]}")
+    print(f"Final velocity: {x_v.value[T-1]}")
+else:
+    print(f"Problem status: {problem.status}")
+
+# Plot results if solution found
+if problem.status == cp.OPTIMAL:
+    def plot_trajectory(y, ylabel, subplot_idx):
+        """Plot trajectory over time"""
+        time = np.arange(T) * dt
+        plt.subplot(2, 2, subplot_idx)
+        plt.plot(time, y)
+        plt.xlabel('Time (s)')
+        plt.ylabel(ylabel)
+        plt.grid(True)
+    
+    plt.figure(figsize=(12, 8))
+    
+    plot_trajectory(x_h.value, 'Altitude', 1)
+    plot_trajectory(x_m.value, 'Mass', 2)
+    plot_trajectory(x_v.value, 'Velocity', 3)
+    plot_trajectory(u_t.value, 'Thrust', 4)
+    
+    plt.tight_layout()
+    plt.show()
+    
+    # Print some statistics
+    print("\nSolution Statistics:")
+    print(f"Maximum altitude reached: {np.max(x_h.value):.6f}")
+    print(f"Maximum velocity reached: {np.max(x_v.value):.6f}")
+    print(f"Maximum thrust used: {np.max(u_t.value):.6f}")
+    print(f"Total fuel consumed: {m_0 - x_m.value[T-1]:.6f}")
+else:
+    print("No optimal solution found. Check problem formulation.")