openscvx 0.1.3__py3-none-any.whl → 0.2.1__py3-none-any.whl

openscvx/io.py

@@ -5,17 +5,134 @@ import queue
 import time
 
 from termcolor import colored
+from openscvx._version import version
+import jax
+from openscvx.results import OptimizationResults
 
 # Define colors for printing
 col_main = "blue"
 col_pos = "green"
 col_neg = "red"
 
+def print_summary_box(lines, title="Summary"):
+    """
+    Print a centered summary box with the given lines.
+    
+    Args:
+        lines (list): List of strings to display in the box
+        title (str): Title for the box (default: "Summary")
+    """
+    # Find the longest line (excluding the title which will be handled separately)
+    content_lines = lines[1:] if len(lines) > 1 else []
+    max_content_width = max(len(line) for line in content_lines) if content_lines else 0
+    title_width = len(title)
+    
+    # Box width should accommodate both title and content
+    box_width = max(max_content_width, title_width) + 4  # Add padding for the box borders
+    
+    # Center with respect to the 89-character horizontal lines in io.py
+    total_width = 89
+    if box_width <= total_width:
+        indent = (total_width - box_width) // 2
+    else:
+        # If box is wider than 89 chars, use a smaller fixed indentation
+        indent = 2
+    
+    # Print the box with dynamic width and centering
+    print(f"\n{' ' * indent}╭{'─' * box_width}╮")
+    print(f"{' ' * indent}│ {title:^{box_width-2}} │")
+    print(f"{' ' * indent}├{'─' * box_width}┤")
+    for line in content_lines:
+        print(f"{' ' * indent}│ {line:<{box_width-2}} │")
+    print(f"{' ' * indent}╰{'─' * box_width}╯\n")
+
+def print_problem_summary(settings):
+    """
+    Print the problem summary box.
+    
+    Args:
+        settings: Configuration settings containing problem information
+    """
+    n_nodal_convex = sum(1 for c in settings.sim.constraints_nodal if c.convex)
+    n_nodal_nonconvex = sum(1 for c in settings.sim.constraints_nodal if not c.convex)
+    n_ctcs = len(settings.sim.constraints_ctcs)
+    n_augmented = settings.sim.n_states - settings.sim.idx_x_true.stop
+    
+    # Get JAX backend information
+    jax_backend = jax.devices()[0].platform.upper()
+    jax_version = jax.__version__
+    
+    # Build weights string conditionally
+    weights_parts = [
+        f"λ_cost={settings.scp.lam_cost:4.1f}",
+        f"λ_tr={settings.scp.w_tr:4.1f}",
+        f"λ_vc={settings.scp.lam_vc:4.1f}"
+    ]
+    
+    # Add λ_vb only if there are nodal nonconvex constraints
+    if n_nodal_nonconvex > 0:
+        weights_parts.append(f"λ_vb={settings.scp.lam_vb:4.1f}")
+    
+    weights_str = ", ".join(weights_parts)
+    
+    lines = [
+        "Problem Summary",
+        f"Dimensions: {settings.sim.n_states} states ({n_augmented} aug), {settings.sim.n_controls} controls, {settings.scp.n} nodes",
+        f"Constraints: {n_nodal_convex} conv, {n_nodal_nonconvex} nonconv, {n_ctcs} ctcs",
+        f"Weights: {weights_str}",
+        f"CVX Solver: {settings.cvx.solver}, Discretization Solver: {settings.dis.solver}",
+        f"JAX Backend: {jax_backend} (v{jax_version})"
+    ]
+    
+    print_summary_box(lines, "Problem Summary")
+
+def print_results_summary(result: OptimizationResults, timing_post, timing_init, timing_solve):
+    """
+    Print the results summary box.
+    
+    Args:
+        result (OptimizationResults): Optimization results object
+        timing_post (float): Post-processing time
+        timing_init (float): Initialization time
+        timing_solve (float): Solve time
+    """
+    cost = result.get("cost", 0.0)
+    ctcs_violation = result.get("ctcs_violation", 0.0)
+    
+    # Convert numpy arrays to scalars for formatting
+    if hasattr(cost, 'item'):
+        cost = cost.item()
+    
+    # Handle CTCS violation - display as 1D array
+    if hasattr(ctcs_violation, 'size'):
+        if ctcs_violation.size == 1:
+            ctcs_violation_str = f"[{ctcs_violation.item():.2e}]"
+        else:
+            # Display as 1D array
+            ctcs_violation_str = f"[{', '.join([f'{v:.2e}' for v in ctcs_violation])}]"
+    else:
+        ctcs_violation_str = f"[{ctcs_violation:.2e}]"
+    
+    # Calculate total computation time
+    total_time = (timing_init or 0.0) + (timing_solve or 0.0) + timing_post
+    
+    lines = [
+        "Results Summary",
+        f"Cost: {cost:.6f}",
+        f"CTCS Constraint Violation: {ctcs_violation_str}",
+        f"Preprocessing Time: {timing_init or 0.0:.3f}s",
+        f"Main Solve Time: {timing_solve or 0.0:.3f}s",
+        f"Post-processing Time: {timing_post:.3f}s",
+        f"Total Computation Time: {total_time:.3f}s"
+    ]
+    
+    print_summary_box(lines, "Results Summary")
+
 def intro():
     # Silence syntax warnings
     warnings.filterwarnings("ignore")
     # fmt: off
-    ascii_art = '''
+    ascii_art = f'''
                              
                             ____                    _____  _____           
                            / __ \                  / ____|/ ____|          
@@ -25,19 +142,21 @@ def intro():
                            \____/| .__/ \___|_| |_|_____/ \_____\_/  /_/\_\ 
                                  | |                                       
                                  |_|                                       
----------------------------------------------------------------------------------------------------------
+─────────────────────────────────────────────────────────────────────────────────────────────────────────
                                 Author: Chris Hayner and Griffin Norris
                                     Autonomous Controls Laboratory
                                        University of Washington
----------------------------------------------------------------------------------------------------------
+                                         Version: {version}
+─────────────────────────────────────────────────────────────────────────────────────────────────────────
 '''
     # fmt: on
     print(ascii_art)
 
 def header():
-    print("{:^4} | {:^7} | {:^7} | {:^7} | {:^7} | {:^7} | {:^7} |  {:^7} | {:^14}".format(
+    print(colored("─────────────────────────────────────────────────────────────────────────────────────────────────────────"))
+    print("{:^4} │ {:^7} │ {:^7} │ {:^7} │ {:^7} │ {:^7} │ {:^7} │  {:^7} │ {:^14}".format(
         "Iter", "Dis Time (ms)", "Solve Time (ms)", "J_total", "J_tr", "J_vb", "J_vc", "Cost", "Solver Status"))
-    print(colored("---------------------------------------------------------------------------------------------------------"))
+    print(colored("─────────────────────────────────────────────────────────────────────────────────────────────────────────"))
 
 def intermediate(print_queue, params):
     hz = 30.0
@@ -60,22 +179,15 @@ def intermediate(print_queue, params):
             cost_colored = colored("{:.1e}".format(data["cost"]))
             prob_stat_colored = colored(data["prob_stat"], col_pos if data["prob_stat"] == 'optimal' else col_neg)
 
-            print("{:^4} |     {:^6.2f}    |      {:^6.2F}     | {:^7} | {:^7} | {:^7} | {:^7} |  {:^7} | {:^14}".format(
+            print("{:^4} │     {:^6.2f}    │      {:^6.2F}     │ {:^7} │ {:^7} │ {:^7} │ {:^7} │  {:^7} │ {:^14}".format(
                 iter_colored, data["dis_time"], data["subprop_time"], J_tot_colored, J_tr_colored, J_vb_colored, J_vc_colored, cost_colored, prob_stat_colored))
 
-            print(colored("---------------------------------------------------------------------------------------------------------"))
-            print("{:^4} | {:^7} | {:^7} | {:^7} | {:^7} | {:^7} | {:^7} |  {:^7} | {:^14}".format(
+            print(colored("─────────────────────────────────────────────────────────────────────────────────────────────────────────"))
+            print("{:^4} │ {:^7} │ {:^7} │ {:^7} │ {:^7} │ {:^7} │ {:^7} │  {:^7} │ {:^14}".format(
                 "Iter", "Dis Time (ms)", "Solve Time (ms)", "J_total", "J_tr", "J_vb", "J_vc", "Cost", "Solver Status"))
         except queue.Empty:
             pass
         time.sleep(max(0.0, 1.0/hz - (time.time() - t_start)))
 
-def footer(computation_time):
-    print(colored("---------------------------------------------------------------------------------------------------------"))
-    # Define ANSI color codes
-    BOLD = "\033[1m"
-    RESET = "\033[0m"
-
-    # Print with bold text
-    print("------------------------------------------------ " + BOLD + "RESULTS" + RESET + " ------------------------------------------------")
-    print("Total Computation Time: ", computation_time)
+def footer():
+    print(colored("─────────────────────────────────────────────────────────────────────────────────────────────────────────"))

openscvx/ocp.py

@@ -3,12 +3,18 @@ import numpy.linalg as la
 from numpy import block
 import numpy as np
 import cvxpy as cp
-from cvxpygen import cpg
 from openscvx.config import Config
-from cvxpygen import cpg
 
+# Optional cvxpygen import
+try:
+    from cvxpygen import cpg
+    CVXPYGEN_AVAILABLE = True
+except ImportError:
+    CVXPYGEN_AVAILABLE = False
+    cpg = None
 
-def OptimalControlProblem(params: Config):
+
+def OptimalControlProblem(settings: Config):
     ########################
     # VARIABLES & PARAMETERS
     ########################
@@ -18,49 +24,51 @@ def OptimalControlProblem(params: Config):
     lam_cost = cp.Parameter(nonneg=True, name='lam_cost')
 
     # State
-    x = cp.Variable((params.scp.n, params.sim.n_states), name='x') # Current State
-    dx = cp.Variable((params.scp.n, params.sim.n_states), name='dx') # State Error
-    x_bar = cp.Parameter((params.scp.n, params.sim.n_states), name='x_bar') # Previous SCP State
+    x = cp.Variable((settings.scp.n, settings.sim.n_states), name='x') # Current State
+    dx = cp.Variable((settings.scp.n, settings.sim.n_states), name='dx') # State Error
+    x_bar = cp.Parameter((settings.scp.n, settings.sim.n_states), name='x_bar') # Previous SCP State
+    x_init = cp.Parameter(settings.sim.n_states, name='x_init') # Initial State
+    x_term = cp.Parameter(settings.sim.n_states, name='x_term') # Final State
 
     # Affine Scaling for State
-    S_x = params.sim.S_x
-    inv_S_x = params.sim.inv_S_x
-    c_x = params.sim.c_x
+    S_x = settings.sim.S_x
+    inv_S_x = settings.sim.inv_S_x
+    c_x = settings.sim.c_x
 
     # Control
-    u = cp.Variable((params.scp.n, params.sim.n_controls), name='u') # Current Control
-    du = cp.Variable((params.scp.n, params.sim.n_controls), name='du') # Control Error
-    u_bar = cp.Parameter((params.scp.n, params.sim.n_controls), name='u_bar') # Previous SCP Control
+    u = cp.Variable((settings.scp.n, settings.sim.n_controls), name='u') # Current Control
+    du = cp.Variable((settings.scp.n, settings.sim.n_controls), name='du') # Control Error
+    u_bar = cp.Parameter((settings.scp.n, settings.sim.n_controls), name='u_bar') # Previous SCP Control
 
     # Affine Scaling for Control
-    S_u = params.sim.S_u
-    inv_S_u = params.sim.inv_S_u
-    c_u = params.sim.c_u
+    S_u = settings.sim.S_u
+    inv_S_u = settings.sim.inv_S_u
+    c_u = settings.sim.c_u
 
     # Discretized Augmented Dynamics Constraints
-    A_d = cp.Parameter((params.scp.n - 1, (params.sim.n_states)*(params.sim.n_states)), name='A_d')
-    B_d = cp.Parameter((params.scp.n - 1, params.sim.n_states*params.sim.n_controls), name='B_d')
-    C_d = cp.Parameter((params.scp.n - 1, params.sim.n_states*params.sim.n_controls), name='C_d')
-    z_d = cp.Parameter((params.scp.n - 1, params.sim.n_states), name='z_d')
-    nu  = cp.Variable((params.scp.n - 1, params.sim.n_states), name='nu') # Virtual Control
+    A_d = cp.Parameter((settings.scp.n - 1, (settings.sim.n_states)*(settings.sim.n_states)), name='A_d')
+    B_d = cp.Parameter((settings.scp.n - 1, settings.sim.n_states*settings.sim.n_controls), name='B_d')
+    C_d = cp.Parameter((settings.scp.n - 1, settings.sim.n_states*settings.sim.n_controls), name='C_d')
+    z_d = cp.Parameter((settings.scp.n - 1, settings.sim.n_states), name='z_d')
+    nu  = cp.Variable((settings.scp.n - 1, settings.sim.n_states), name='nu') # Virtual Control
 
     # Linearized Nonconvex Nodal Constraints
-    if params.sim.constraints_nodal:
+    if settings.sim.constraints_nodal:
         g = []
         grad_g_x = []
         grad_g_u = []
         nu_vb = []
-        for idx_ncvx, constraint in enumerate(params.sim.constraints_nodal):
+        for idx_ncvx, constraint in enumerate(settings.sim.constraints_nodal):
             if not constraint.convex:
-                g.append(cp.Parameter(params.scp.n, name = 'g_' + str(idx_ncvx)))
-                grad_g_x.append(cp.Parameter((params.scp.n, params.sim.n_states), name='grad_g_x_' + str(idx_ncvx)))
-                grad_g_u.append(cp.Parameter((params.scp.n, params.sim.n_controls), name='grad_g_u_' + str(idx_ncvx)))
-                nu_vb.append(cp.Variable(params.scp.n, name='nu_vb_' + str(idx_ncvx))) # Virtual Control for VB
+                g.append(cp.Parameter(settings.scp.n, name = 'g_' + str(idx_ncvx)))
+                grad_g_x.append(cp.Parameter((settings.scp.n, settings.sim.n_states), name='grad_g_x_' + str(idx_ncvx)))
+                grad_g_u.append(cp.Parameter((settings.scp.n, settings.sim.n_controls), name='grad_g_u_' + str(idx_ncvx)))
+                nu_vb.append(cp.Variable(settings.scp.n, name='nu_vb_' + str(idx_ncvx))) # Virtual Control for VB
 
     # Applying the affine scaling to state and control
     x_nonscaled = []
     u_nonscaled = []
-    for k in range(params.scp.n):
+    for k in range(settings.scp.n):
         x_nonscaled.append(S_x @ x[k] + c_x)
         u_nonscaled.append(S_u @ u[k] + c_u)
 
@@ -71,74 +79,77 @@ def OptimalControlProblem(params: Config):
     # CONSTRAINTS
     #############
     idx_ncvx = 0
-    if params.sim.constraints_nodal:
-        for constraint in params.sim.constraints_nodal:
+    if settings.sim.constraints_nodal:
+        for constraint in settings.sim.constraints_nodal:
             if constraint.nodes is None:
-                nodes = range(params.scp.n)
+                nodes = range(settings.scp.n)
             else:
                 nodes = constraint.nodes
 
-            if constraint.convex:
+            if constraint.convex and constraint.vectorized:
+                constr += [constraint(x_nonscaled, u_nonscaled)]
+                
+            elif constraint.convex:
                 constr += [constraint(x_nonscaled[node], u_nonscaled[node]) for node in nodes]
 
             elif not constraint.convex:
                 constr += [((g[idx_ncvx][node] + grad_g_x[idx_ncvx][node] @ dx[node] + grad_g_u[idx_ncvx][node] @ du[node])) == nu_vb[idx_ncvx][node] for node in nodes]
                 idx_ncvx += 1
 
-    for i in range(params.sim.idx_x_true.start, params.sim.idx_x_true.stop):
-        if params.sim.initial_state.type[i] == 'Fix':
-            constr += [x_nonscaled[0][i] == params.sim.initial_state.value[i]]  # Initial Boundary Conditions
-        if params.sim.final_state.type[i] == 'Fix':
-            constr += [x_nonscaled[-1][i] == params.sim.final_state.value[i]]   # Final Boundary Conditions
-        if params.sim.initial_state.type[i] == 'Minimize':
-            cost += lam_cost * x_nonscaled[0][i]
-        if params.sim.final_state.type[i] == 'Minimize':
-            cost += lam_cost * x_nonscaled[-1][i]
-        if params.sim.initial_state.type[i] == 'Maximize':
+    for i in range(settings.sim.idx_x_true.start, settings.sim.idx_x_true.stop):
+        if settings.sim.x.initial_type[i] == 'Fix':
+            constr += [x_nonscaled[0][i] == x_init[i]]  # Initial Boundary Conditions
+        if settings.sim.x.final_type[i] == 'Fix':
+            constr += [x_nonscaled[-1][i] == x_term[i]]   # Final Boundary Conditions
+        if settings.sim.x.initial_type[i] == 'Minimize':
             cost += lam_cost * x_nonscaled[0][i]
-        if params.sim.final_state.type[i] == 'Maximize':
+        if settings.sim.x.final_type[i] == 'Minimize':
             cost += lam_cost * x_nonscaled[-1][i]
+        if settings.sim.x.initial_type[i] == 'Maximize':
+            cost -= lam_cost * x_nonscaled[0][i]
+        if settings.sim.x.final_type[i] == 'Maximize':
+            cost -= lam_cost * x_nonscaled[-1][i]
 
-    if params.scp.uniform_time_grid:
-        constr += [x_nonscaled[i][params.sim.idx_t] - x_nonscaled[i-1][params.sim.idx_t] == x_nonscaled[i-1][params.sim.idx_t] - x_nonscaled[i-2][params.sim.idx_t] for i in range(2, params.scp.n)] # Uniform Time Step
+    if settings.scp.uniform_time_grid:
+        constr += [u_nonscaled[i][settings.sim.idx_s] == u_nonscaled[i-1][settings.sim.idx_s] for i in range(1, settings.scp.n)]
 
-    constr += [0 == la.inv(S_x) @ (x_nonscaled[i] - x_bar[i] - dx[i]) for i in range(params.scp.n)] # State Error
-    constr += [0 == la.inv(S_u) @ (u_nonscaled[i] - u_bar[i] - du[i]) for i in range(params.scp.n)] # Control Error
+    constr += [0 == la.inv(S_x) @ (x_nonscaled[i] - x_bar[i] - dx[i]) for i in range(settings.scp.n)] # State Error
+    constr += [0 == la.inv(S_u) @ (u_nonscaled[i] - u_bar[i] - du[i]) for i in range(settings.scp.n)] # Control Error
 
     constr += [x_nonscaled[i] == \
-                      cp.reshape(A_d[i-1], (params.sim.n_states, params.sim.n_states)) @ x_nonscaled[i-1] \
-                    + cp.reshape(B_d[i-1], (params.sim.n_states, params.sim.n_controls)) @ u_nonscaled[i-1] \
-                    + cp.reshape(C_d[i-1], (params.sim.n_states, params.sim.n_controls)) @ u_nonscaled[i] \
+                      cp.reshape(A_d[i-1], (settings.sim.n_states, settings.sim.n_states)) @ x_nonscaled[i-1] \
+                    + cp.reshape(B_d[i-1], (settings.sim.n_states, settings.sim.n_controls)) @ u_nonscaled[i-1] \
+                    + cp.reshape(C_d[i-1], (settings.sim.n_states, settings.sim.n_controls)) @ u_nonscaled[i] \
                     + z_d[i-1] \
-                    + nu[i-1] for i in range(1, params.scp.n)] # Dynamics Constraint
+                    + nu[i-1] for i in range(1, settings.scp.n)] # Dynamics Constraint
     
-    constr += [u_nonscaled[i] <= params.sim.max_control for i in range(params.scp.n)]
-    constr += [u_nonscaled[i] >= params.sim.min_control for i in range(params.scp.n)] # Control Constraints
+    constr += [u_nonscaled[i] <= settings.sim.u.max for i in range(settings.scp.n)]
+    constr += [u_nonscaled[i] >= settings.sim.u.min for i in range(settings.scp.n)] # Control Constraints
 
-    constr += [x_nonscaled[i][params.sim.idx_x_true] <= params.sim.max_state[params.sim.idx_x_true] for i in range(params.scp.n)]
-    constr += [x_nonscaled[i][params.sim.idx_x_true] >= params.sim.min_state[params.sim.idx_x_true] for i in range(params.scp.n)] # State Constraints (Also implemented in CTCS but included for numerical stability)
+    constr += [x_nonscaled[i][settings.sim.idx_x_true] <= settings.sim.x.true.max for i in range(settings.scp.n)]
+    constr += [x_nonscaled[i][settings.sim.idx_x_true] >= settings.sim.x.true.min for i in range(settings.scp.n)] # State Constraints (Also implemented in CTCS but included for numerical stability)
 
     ########
     # COSTS
     ########
     
     inv = block([[inv_S_x, np.zeros((S_x.shape[0], S_u.shape[1]))], [np.zeros((S_u.shape[0], S_x.shape[1])), inv_S_u]])
-    cost += sum(w_tr * cp.sum_squares(inv @ cp.hstack((dx[i], du[i]))) for i in range(params.scp.n))  # Trust Region Cost
-    cost += sum(params.scp.lam_vc * cp.sum(cp.abs(nu[i-1])) for i in range(1, params.scp.n)) # Virtual Control Slack
+    cost += sum(w_tr * cp.sum_squares(inv @ cp.hstack((dx[i], du[i]))) for i in range(settings.scp.n))  # Trust Region Cost
+    cost += sum(settings.scp.lam_vc * cp.sum(cp.abs(nu[i-1])) for i in range(1, settings.scp.n)) # Virtual Control Slack
     
     idx_ncvx = 0
-    if params.sim.constraints_nodal:
-        for constraint in params.sim.constraints_nodal:
+    if settings.sim.constraints_nodal:
+        for constraint in settings.sim.constraints_nodal:
             if not constraint.convex:
-                cost += params.scp.lam_vb * cp.sum(cp.pos(nu_vb[idx_ncvx]))
+                cost += settings.scp.lam_vb * cp.sum(cp.pos(nu_vb[idx_ncvx]))
                 idx_ncvx += 1
 
-    for idx, nodes in zip(np.arange(params.sim.idx_y.start, params.sim.idx_y.stop), params.sim.ctcs_node_intervals):  
+    for idx, nodes in zip(np.arange(settings.sim.idx_y.start, settings.sim.idx_y.stop), settings.sim.ctcs_node_intervals):  
         if nodes[0] == 0:
             start_idx = 1
         else:
             start_idx = nodes[0]
-        constr += [cp.abs(x_nonscaled[i][idx] - x_nonscaled[i-1][idx]) <= params.sim.max_state[idx] for i in range(start_idx, nodes[1])]
+        constr += [cp.abs(x_nonscaled[i][idx] - x_nonscaled[i-1][idx]) <= settings.sim.x.max[idx] for i in range(start_idx, nodes[1])]
         constr += [x_nonscaled[0][idx] == 0]
 
     
@@ -146,15 +157,23 @@ def OptimalControlProblem(params: Config):
     # PROBLEM
     #########
     prob = cp.Problem(cp.Minimize(cost), constr)
-    if params.cvx.cvxpygen:
+    if settings.cvx.cvxpygen:
+        if not CVXPYGEN_AVAILABLE:
+            raise ImportError(
+                "cvxpygen is required for code generation but not installed. "
+                "Install it with: pip install openscvx[cvxpygen] or pip install cvxpygen"
+            )
         # Check to see if solver directory exists
         if not os.path.exists('solver'):
-            cpg.generate_code(prob, solver = params.cvx.solver, code_dir='solver', wrapper = True)
+            cpg.generate_code(prob, solver = settings.cvx.solver, code_dir='solver', wrapper = True)
         else:
             # Prompt the use to indicate if they wish to overwrite the solver directory or use the existing compiled solver
-            overwrite = input("Solver directory already exists. Overwrite? (y/n): ")
-            if overwrite.lower() == 'y':
-                cpg.generate_code(prob, solver = params.cvx.solver, code_dir='solver', wrapper = True)
+            if settings.cvx.cvxpygen_override:
+                cpg.generate_code(prob, solver = settings.cvx.solver, code_dir='solver', wrapper = True)
             else:
-                pass
+                overwrite = input("Solver directory already exists. Overwrite? (y/n): ")
+                if overwrite.lower() == 'y':
+                    cpg.generate_code(prob, solver = settings.cvx.solver, code_dir='solver', wrapper = True)
+                else:
+                    pass
     return prob