generate_missing_figures.py

#!/usr/bin/env python3
"""
Generate missing figures for optimization-control-systems template
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend
from scipy.optimize import minimize, linprog
from scipy.linalg import solve_continuous_are
import os

# Set random seed for reproducibility
np.random.seed(42)

# Define colors for plots (matching LaTeX color definitions)
optblue = '#0077BB'    # RGB(0,119,187)
optred = '#CC3311'     # RGB(204,51,17)
optgreen = '#009966'   # RGB(0,153,102)
optorange = '#FF9933'  # RGB(255,153,51)
optpurple = '#990099'  # RGB(153,0,153)

# Ensure assets directory exists
os.makedirs('assets', exist_ok=True)

def generate_linear_constrained_optimization():
    """Generate linear programming and constrained optimization figure"""

    # Linear programming setup
    c = np.array([-3, -2])  # Maximize 3x1 + 2x2 (negate for minimize)
    A_ub = np.array([[1, 1], [2, 1]])  # Constraints matrix
    b_ub = np.array([4, 6])  # Constraints bounds
    bounds = [(0, None), (0, None)]  # x1, x2 >= 0

    # Solve linear programming problem
    result = linprog(c, A_ub=A_ub, b_ub=b_ub, bounds=bounds, method='highs')

    # Constrained optimization setup
    def objective_function(x):
        return x[0]**2 + x[1]**2 - 2*x[0] - 4*x[1] + 5

    def constraint_function(x):
        return x[0] + x[1] - 3

    # Solve constrained optimization
    constraints = {'type': 'ineq', 'fun': lambda x: -constraint_function(x)}
    result_constrained = minimize(objective_function, [0.5, 0.5], constraints=constraints)

    # Create visualization
    fig, axes = plt.subplots(2, 2, figsize=(15, 12))
    fig.suptitle('Linear Programming and Constrained Optimization', fontsize=16, fontweight='bold')

    # Linear programming feasible region
    ax1 = axes[0, 0]
    x1_range = np.linspace(0, 4, 100)
    x2_constraint1 = 4 - x1_range  # x1 + x2 <= 4
    x2_constraint2 = 6 - 2*x1_range  # 2x1 + x2 <= 6

    # Plot constraints
    ax1.fill_between(x1_range, 0, np.minimum(x2_constraint1, x2_constraint2),
                    where=(x2_constraint1 >= 0) & (x2_constraint2 >= 0),
                    alpha=0.3, color=optblue, label='Feasible region')
    ax1.plot(x1_range, x2_constraint1, 'b-', linewidth=2, label='x1 + x2 <= 4')
    ax1.plot(x1_range, x2_constraint2, 'r-', linewidth=2, label='2x1 + x2 <= 6')

    # Plot corner points
    corner_points = np.array([[0, 0], [0, 4], [2, 2], [3, 0]])
    ax1.plot(corner_points[:, 0], corner_points[:, 1], 'ko', markersize=8)
    ax1.plot(result.x[0], result.x[1], 'r*', markersize=15, label='Optimal solution')

    # Objective function contours
    x1_obj, x2_obj = np.meshgrid(np.linspace(0, 4, 50), np.linspace(0, 5, 50))
    obj_values = 3*x1_obj + 2*x2_obj
    ax1.contour(x1_obj, x2_obj, obj_values, levels=10, colors='gray', alpha=0.5)

    ax1.set_xlim(0, 4)
    ax1.set_ylim(0, 5)
    ax1.set_xlabel('x1')
    ax1.set_ylabel('x2')
    ax1.set_title('Linear Programming: Feasible Region')
    ax1.legend()
    ax1.grid(True, alpha=0.3)

    # Constrained optimization visualization
    ax2 = axes[0, 1]
    x_range_const = np.linspace(-1, 4, 100)
    y_range_const = np.linspace(-1, 4, 100)
    X_const, Y_const = np.meshgrid(x_range_const, y_range_const)
    Z_obj = X_const**2 + Y_const**2 - 2*X_const - 4*Y_const + 5

    # Plot objective function contours
    contours = ax2.contour(X_const, Y_const, Z_obj, levels=15, colors='blue', alpha=0.6)
    ax2.clabel(contours, inline=True, fontsize=8)

    # Plot constraint
    x_constraint = np.linspace(0, 3, 100)
    y_constraint = 3 - x_constraint
    ax2.fill_between(x_constraint, 0, y_constraint, alpha=0.2, color=optred, label='Feasible region')
    ax2.plot(x_constraint, y_constraint, 'r-', linewidth=3, label='x1 + x2 <= 3')

    # Plot solutions
    ax2.plot(1, 2, 'go', markersize=10, label='Unconstrained optimum')
    ax2.plot(result_constrained.x[0], result_constrained.x[1], 'r*', markersize=15, label='Constrained optimum')

    ax2.set_xlim(-0.5, 3.5)
    ax2.set_ylim(-0.5, 3.5)
    ax2.set_xlabel('x1')
    ax2.set_ylabel('x2')
    ax2.set_title('Constrained Optimization')
    ax2.legend()
    ax2.grid(True, alpha=0.3)

    # Algorithm convergence comparison
    ax3 = axes[1, 0]
    algorithms = ['Gradient\nDescent', 'Newton\nMethod', 'Interior\nPoint', 'Simplex']
    iterations = [1000, 5, 20, 8]
    accuracy = [1e-6, 1e-12, 1e-8, 1e-15]

    ax3_twin = ax3.twinx()
    bars1 = ax3.bar([x - 0.2 for x in range(len(algorithms))], iterations,
                   width=0.4, color=optblue, alpha=0.7, label='Iterations')
    bars2 = ax3_twin.semilogy([x + 0.2 for x in range(len(algorithms))], accuracy,
                             'ro', markersize=10, label='Final accuracy')

    ax3.set_xlabel('Algorithm')
    ax3.set_ylabel('Iterations to convergence', color=optblue)
    ax3_twin.set_ylabel('Final accuracy', color=optred)
    ax3.set_title('Algorithm Performance Comparison')
    ax3.set_xticks(range(len(algorithms)))
    ax3.set_xticklabels(algorithms, rotation=45)
    ax3.grid(True, alpha=0.3)

    # Optimization landscape
    ax4 = axes[1, 1]
    # Create a 3D-like visualization using contours and shading
    levels = np.linspace(np.min(Z_obj), np.max(Z_obj), 20)
    cs = ax4.contourf(X_const, Y_const, Z_obj, levels=levels, cmap='viridis', alpha=0.8)
    ax4.contour(X_const, Y_const, Z_obj, levels=levels, colors='black', alpha=0.3, linewidths=0.5)

    # Add colorbar
    cbar = plt.colorbar(cs, ax=ax4, shrink=0.8)
    cbar.set_label('Objective function value')

    ax4.plot(result_constrained.x[0], result_constrained.x[1], 'r*', markersize=15, label='Optimum')
    ax4.set_xlabel('x1')
    ax4.set_ylabel('x2')
    ax4.set_title('Optimization Landscape')
    ax4.legend()

    plt.tight_layout()
    plt.savefig('assets/linear_constrained_optimization.pdf', dpi=300, bbox_inches='tight')
    plt.close()

    print("Generated linear_constrained_optimization.pdf")

def generate_control_systems_analysis():
    """Generate control systems analysis figure"""

    # Define LQR continuous function
    def lqr_continuous(A, B, Q, R):
        """Solve continuous-time LQR problem"""
        P = solve_continuous_are(A, B, Q, R)
        K = np.linalg.inv(R) @ B.T @ P
        return K, P

    # System matrices for double integrator
    A = np.array([[0, 1], [0, 0]])
    B = np.array([[0], [1]])
    C = np.array([[1, 0]])

    # LQR design parameters
    Q = np.array([[1, 0], [0, 1]])
    R = np.array([[1]])

    # Design LQR controller
    K_lqr, P_lqr = lqr_continuous(A, B, Q, R)
    A_cl = A - B @ K_lqr

    # Simulate step response
    t = np.linspace(0, 10, 1000)
    dt = t[1] - t[0]

    # Initial condition response
    x0 = np.array([[1], [0]])
    x_open = np.zeros((2, len(t)))
    x_closed = np.zeros((2, len(t)))
    x_open[:, [0]] = x0
    x_closed[:, [0]] = x0

    for i in range(1, len(t)):
        # Open-loop (unstable system would blow up, so we'll simulate a different scenario)
        x_open[:, [i]] = x_open[:, [i-1]] + dt * A @ x_open[:, [i-1]]

        # Closed-loop
        x_closed[:, [i]] = x_closed[:, [i-1]] + dt * A_cl @ x_closed[:, [i-1]]

    # Generate Kalman filter data
    # Process and measurement noise
    Q_kf = 0.01 * np.eye(2)
    R_kf = 0.1

    # Generate true trajectory with noise
    x_true = np.zeros((2, len(t)))
    y_meas = np.zeros(len(t))
    x_true[:, [0]] = x0

    for i in range(1, len(t)):
        process_noise = np.random.multivariate_normal([0, 0], Q_kf).reshape(2, 1)
        x_true[:, [i]] = x_true[:, [i-1]] + dt * A_cl @ x_true[:, [i-1]] + process_noise
        y_meas[i] = (C @ x_true[:, [i]])[0, 0] + np.random.normal(0, np.sqrt(R_kf))

    # Simple Kalman filter estimate (simplified)
    x_est = np.zeros((2, len(t)))
    x_est[:, [0]] = x0
    P_est = np.eye(2)

    for i in range(1, len(t)):
        # Prediction
        x_pred = x_est[:, [i-1]] + dt * A_cl @ x_est[:, [i-1]]
        P_pred = P_est + dt * Q_kf

        # Update
        K_kf = P_pred @ C.T / (C @ P_pred @ C.T + R_kf)
        x_est[:, [i]] = x_pred + K_kf * (y_meas[i] - (C @ x_pred)[0, 0])
        P_est = (np.eye(2) - K_kf @ C) @ P_pred

    # Create visualization
    fig, axes = plt.subplots(2, 2, figsize=(15, 12))
    fig.suptitle('Optimal Control Theory and State Estimation', fontsize=16, fontweight='bold')

    # LQR pole placement
    ax1 = axes[0, 0]

    # Open-loop poles
    open_poles = np.linalg.eigvals(A)
    closed_poles = np.linalg.eigvals(A_cl)

    ax1.plot(np.real(open_poles), np.imag(open_poles), 'ro', markersize=10, label='Open-loop poles')
    ax1.plot(np.real(closed_poles), np.imag(closed_poles), 'bo', markersize=10, label='Closed-loop poles')
    ax1.axvline(0, color='black', linestyle='--', alpha=0.5)
    ax1.fill_between([-2, 0], [-3, -3], [3, 3], alpha=0.2, color='green', label='Stable region')
    ax1.set_xlabel('Real Part')
    ax1.set_ylabel('Imaginary Part')
    ax1.set_title('LQR Pole Placement')
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    ax1.set_xlim(-2, 1)
    ax1.set_ylim(-2, 2)

    # Step response comparison
    ax2 = axes[0, 1]
    ax2.plot(t, x_open[0, :], 'r-', linewidth=2, label='Open-loop')
    ax2.plot(t, x_closed[0, :], 'b-', linewidth=2, label='Closed-loop (LQR)')
    ax2.set_xlabel('Time (s)')
    ax2.set_ylabel('Position')
    ax2.set_title('Step Response Comparison')
    ax2.legend()
    ax2.grid(True, alpha=0.3)
    ax2.set_xlim(0, 10)

    # Kalman filter performance
    ax3 = axes[1, 0]
    ax3.plot(t, x_true[0, :], 'g-', linewidth=2, label='True state')
    ax3.plot(t, y_meas, 'r.', markersize=2, alpha=0.5, label='Measurements')
    ax3.plot(t, x_est[0, :], 'b-', linewidth=2, label='Kalman estimate')

    # Add uncertainty bounds (simplified)
    uncertainty = 0.1 * np.ones_like(t)
    ax3.fill_between(t, x_est[0, :] - uncertainty, x_est[0, :] + uncertainty,
                     alpha=0.3, color='blue', label='Uncertainty bounds')

    ax3.set_xlabel('Time (s)')
    ax3.set_ylabel('Position')
    ax3.set_title('Kalman Filter State Estimation')
    ax3.legend()
    ax3.grid(True, alpha=0.3)
    ax3.set_xlim(0, 10)

    # LQR design trade-offs
    ax4 = axes[1, 1]
    Q_weights = [0.1, 1, 10, 100]
    settling_times = []
    control_efforts = []

    for q_weight in Q_weights:
        Q_test = q_weight * np.eye(2)
        K_test, _ = lqr_continuous(A, B, Q_test, R)
        A_cl_test = A - B @ K_test

        # Approximate settling time (time to reach 2% of steady state)
        settling_times.append(4 / abs(np.real(np.linalg.eigvals(A_cl_test)).max()))
        control_efforts.append(np.linalg.norm(K_test))

    ax4_twin = ax4.twinx()

    lines1 = ax4.plot(Q_weights, settling_times, 'bo-', linewidth=2, markersize=8, label='Settling time')
    lines2 = ax4_twin.plot(Q_weights, control_efforts, 'ro-', linewidth=2, markersize=8, label='Control effort')

    ax4.set_xlabel('Q weight')
    ax4.set_ylabel('Settling time (s)', color='blue')
    ax4_twin.set_ylabel('Control effort', color='red')
    ax4.set_title('LQR Design Trade-offs')
    ax4.set_xscale('log')
    ax4.grid(True, alpha=0.3)

    # Combine legends
    lines1, labels1 = ax4.get_legend_handles_labels()
    lines2, labels2 = ax4_twin.get_legend_handles_labels()
    ax4.legend(lines1 + lines2, labels1 + labels2, loc='center right')

    plt.tight_layout()
    plt.savefig('assets/control_systems_analysis.pdf', dpi=300, bbox_inches='tight')
    plt.close()

    print("Generated control_systems_analysis.pdf")

if __name__ == "__main__":
    print("Generating missing optimization and control systems figures...")

    try:
        generate_linear_constrained_optimization()
        generate_control_systems_analysis()
        print("All missing figures generated successfully!")
    except Exception as e:
        print(f"Error generating figures: {e}")
        import traceback
        traceback.print_exc()