openscvx 0.5.3.dev45__tar.gz → 0.5.3.dev47__tar.gz

{openscvx-0.5.3.dev45/openscvx.egg-info → openscvx-0.5.3.dev47}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.5.3.dev45
+Version: 0.5.3.dev47
 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.5.3.dev45 → openscvx-0.5.3.dev47}/docs/Foundations/control_parameterization.md

@@ -44,8 +44,9 @@ discretizer:
 import openscvx as ox
 
 problem = ox.Problem(
-    dynamics, ...,
-    discretizer={"dis_type": "ZOH"},   # or "FOH" (default)
+    dynamics,
+    ...,
+    discretizer={"dis_type": "ZOH"},  # or "FOH" (default)
 )
 ```
 

openscvx-0.5.3.dev47/docs/Foundations/discretization.md

@@ -0,0 +1,94 @@
+---
+title: Exact Discretization
+description: >-
+  How OpenSCvx discretizes and linearizes continuous-time dynamics by
+  integrating an augmented state vector to build the ZOH/FOH transition matrices.
+---
+
+# Exact Discretization
+
+``` py title="dVdt.py"
+def dVdt(self, tau: float, V: jnp.ndarray, u_cur: np.ndarray, u_next: np.ndarray) -> jnp.ndarray:
+    """
+    Computes the time derivative of the augmented state vector for the system for a sequence of states.
+
+    Parameters:
+    tau (float): Current time.
+    V (np.ndarray): Sequence of augmented state vectors.
+    u_cur (np.ndarray): Sequence of current control inputs.
+    u_next (np.ndarray): Sequence of next control inputs.
+    A: Function that computes the Jacobian of the system dynamics with respect to the state.
+    B: Function that computes the Jacobian of the system dynamics with respect to the control input.
+    obstacles: List of obstacles in the environment.
+    params (dict): Parameters of the system.
+
+    Returns:
+    np.ndarray: Time derivatives of the augmented state vectors.
+    """
+
+    # Extract the number of states and controls from the parameters
+    n_x = self.params.sim.n_states
+    n_u = self.params.sim.n_controls
+
+    # Unflatten V
+    V = V.reshape(-1, self.i5)
+
+    # Compute the interpolation factor based on the discretization type
+    if self.dis_type == "ZOH":
+        beta = 0.0
+    elif self.dis_type == "FOH":
+        beta = (tau) * self.params.sim.n
+    alpha = 1 - beta
+
+    # Interpolate the control input
+    u = u_cur + beta * (u_next - u_cur)
+    s = u[:, -1]
+
+    # Initialize the augmented Jacobians
+    dfdx = jnp.zeros((V.shape[0], n_x, n_x))
+    dfdu = jnp.zeros((V.shape[0], n_x, n_u))
+
+    # Ensure x_seq and u have the same batch size
+    x = V[:, : self.params.sim.n_states]
+    u = u[: x.shape[0]]
+
+    # Compute the nonlinear propagation term
+    f = self.params.dyn.state_dot(x, u[:, :-1])
+    F = s[:, None] * f
+
+    # Evaluate the State Jacobian
+    dfdx = self.params.dyn.A(x, u[:, :-1])
+    sdfdx = s[:, None, None] * dfdx
+
+    # Evaluate the Control Jacobian
+    dfdu_veh = self.params.dyn.B(x, u[:, :-1])
+    dfdu = dfdu.at[:, :, :-1].set(s[:, None, None] * dfdu_veh)
+    dfdu = dfdu.at[:, :, -1].set(f)
+
+    # Compute the defect
+    z = F - jnp.einsum("ijk,ik->ij", sdfdx, x) - jnp.einsum("ijk,ik->ij", dfdu, u)
+
+    # Stack up the results into the augmented state vector
+    dVdt = jnp.zeros_like(V)
+    dVdt = dVdt.at[:, self.i0 : self.i1].set(F)
+    dVdt = dVdt.at[:, self.i1 : self.i2].set(
+        jnp.matmul(sdfdx, V[:, self.i1 : self.i2].reshape(-1, n_x, n_x)).reshape(-1, n_x * n_x)
+    )
+    dVdt = dVdt.at[:, self.i2 : self.i3].set(
+        (jnp.matmul(sdfdx, V[:, self.i2 : self.i3].reshape(-1, n_x, n_u)) + dfdu * alpha).reshape(
+            -1, n_x * n_u
+        )
+    )
+    dVdt = dVdt.at[:, self.i3 : self.i4].set(
+        (jnp.matmul(sdfdx, V[:, self.i3 : self.i4].reshape(-1, n_x, n_u)) + dfdu * beta).reshape(
+            -1, n_x * n_u
+        )
+    )
+    dVdt = dVdt.at[:, self.i4 : self.i5].set(
+        (
+            jnp.matmul(sdfdx, V[:, self.i4 : self.i5].reshape(-1, n_x)[..., None]).squeeze(-1) + z
+        ).reshape(-1, n_x)
+    )
+    return dVdt.flatten()
+```
+    

{openscvx-0.5.3.dev45 → openscvx-0.5.3.dev47}/docs/Foundations/ocp.md

@@ -18,21 +18,21 @@ The problem is defined as a function that takes in a `Config` object, which cont
 The state, constrol and additional parameters are defined as follows:
 
 ```python
+lam_prox = cp.Parameter(nonneg=True, name="lam_prox")  # Weight on the Trust Region
+lam_cost = cp.Parameter(nonneg=True, name="lam_cost")  # Weight on the Nonlinear Cost
+lam_vc = cp.Parameter(
+    (params.sim.n - 1, params.sim.n_states), nonneg=True, name="lam_vc"
+)  # Weight on Virtual Control
+lam_vb = cp.Parameter(nonneg=True, name="lam_vb")  # Weight on Virtual Buffer
 
-lam_prox = cp.Parameter(nonneg = True, name='lam_prox')       # Weight on the Trust Region
-lam_cost = cp.Parameter(nonneg=True, name='lam_cost') # Weight on the Nonlinear Cost
-lam_vc = cp.Parameter((params.sim.n - 1, params.sim.n_states), nonneg=True, name='lam_vc')  # Weight on Virtual Control
-lam_vb = cp.Parameter(nonneg=True, name='lam_vb')     # Weight on Virtual Buffer
+x = cp.Variable((params.sim.n, params.sim.n_states), name="x")  # State
+dx = cp.Variable((params.sim.n, params.sim.n_states), name="dx")  # State Trust Region
+x_bar = cp.Parameter((params.sim.n, params.sim.n_states), name="x_bar")  # Previous SCP State
 
-x = cp.Variable((params.sim.n, params.sim.n_states), name='x')   # State
-dx = cp.Variable((params.sim.n, params.sim.n_states), name='dx') # State Trust Region
-x_bar = cp.Parameter((params.sim.n, params.sim.n_states), name='x_bar') # Previous SCP State
-
-
-u = cp.Variable((params.sim.n, params.sim.n_controls), name='u')   # Control
-du = cp.Variable((params.sim.n, params.sim.n_controls), name='du') # Control Trust Region
-u_bar = cp.Parameter((params.sim.n, params.sim.n_controls), name='u_bar') # Previous SCP Control
 
+u = cp.Variable((params.sim.n, params.sim.n_controls), name="u")  # Control
+du = cp.Variable((params.sim.n, params.sim.n_controls), name="du")  # Control Trust Region
+u_bar = cp.Parameter((params.sim.n, params.sim.n_controls), name="u_bar")  # Previous SCP Control
 ```
 
 ## Scaling Definitions
@@ -104,15 +104,15 @@ def OptimalControlProblem(params: Config):
     ########################
 
     # Parameters
-    lam_prox = cp.Parameter(nonneg = True, name='lam_prox')
-    lam_cost = cp.Parameter(nonneg=True, name='lam_cost')
-    lam_vc = cp.Parameter((params.sim.n - 1, params.sim.n_states), nonneg=True, name='lam_vc')
-    lam_vb = cp.Parameter(nonneg=True, name='lam_vb')
+    lam_prox = cp.Parameter(nonneg=True, name="lam_prox")
+    lam_cost = cp.Parameter(nonneg=True, name="lam_cost")
+    lam_vc = cp.Parameter((params.sim.n - 1, params.sim.n_states), nonneg=True, name="lam_vc")
+    lam_vb = cp.Parameter(nonneg=True, name="lam_vb")
 
     # State
-    x = cp.Variable((params.sim.n, params.sim.n_states), name='x') 
-    dx = cp.Variable((params.sim.n, params.sim.n_states), name='dx') # State Trust Region
-    x_bar = cp.Parameter((params.sim.n, params.sim.n_states), name='x_bar') # Previous SCP State
+    x = cp.Variable((params.sim.n, params.sim.n_states), name="x")
+    dx = cp.Variable((params.sim.n, params.sim.n_states), name="dx")  # State Trust Region
+    x_bar = cp.Parameter((params.sim.n, params.sim.n_states), name="x_bar")  # Previous SCP State
 
     # Affine Scaling for State
     S_x = params.sim.S_x
@@ -120,9 +120,11 @@ def OptimalControlProblem(params: Config):
     c_x = params.sim.c_x
 
     # Control
-    u = cp.Variable((params.sim.n, params.sim.n_controls), name='u') 
-    du = cp.Variable((params.sim.n, params.sim.n_controls), name='du') # Control Trust Region
-    u_bar = cp.Parameter((params.sim.n, params.sim.n_controls), name='u_bar') # Previous SCP Control
+    u = cp.Variable((params.sim.n, params.sim.n_controls), name="u")
+    du = cp.Variable((params.sim.n, params.sim.n_controls), name="du")  # Control Trust Region
+    u_bar = cp.Parameter(
+        (params.sim.n, params.sim.n_controls), name="u_bar"
+    )  # Previous SCP Control
 
     # Affine Scaling for Control
     S_u = params.sim.S_u
@@ -130,11 +132,11 @@ def OptimalControlProblem(params: Config):
     c_u = params.sim.c_u
 
     # Discretized Augmented Dynamics Constraints
-    A_d = cp.Parameter((params.sim.n - 1, params.sim.n_states, params.sim.n_states), name='A_d')
-    B_d = cp.Parameter((params.sim.n - 1, params.sim.n_states, params.sim.n_controls), name='B_d')
-    C_d = cp.Parameter((params.sim.n - 1, params.sim.n_states, params.sim.n_controls), name='C_d')
-    z_d = cp.Parameter((params.sim.n - 1, params.sim.n_states), name='z_d')
-    nu  = cp.Variable((params.sim.n - 1, params.sim.n_states), name='nu') # Virtual Control
+    A_d = cp.Parameter((params.sim.n - 1, params.sim.n_states, params.sim.n_states), name="A_d")
+    B_d = cp.Parameter((params.sim.n - 1, params.sim.n_states, params.sim.n_controls), name="B_d")
+    C_d = cp.Parameter((params.sim.n - 1, params.sim.n_states, params.sim.n_controls), name="C_d")
+    z_d = cp.Parameter((params.sim.n - 1, params.sim.n_states), name="z_d")
+    nu = cp.Variable((params.sim.n - 1, params.sim.n_states), name="nu")  # Virtual Control
 
     # Linearized Nonconvex Nodal Constraints
     if params.sim.constraints_nodal:
@@ -144,10 +146,20 @@ def OptimalControlProblem(params: Config):
         nu_vb = []
         for idx_ncvx, constraint in enumerate(params.sim.constraints_nodal):
             if not constraint.convex:
-                g.append(cp.Parameter(params.sim.n, name = 'g_' + str(idx_ncvx)))
-                grad_g_x.append(cp.Parameter((params.sim.n, params.sim.n_states), name='grad_g_x_' + str(idx_ncvx)))
-                grad_g_u.append(cp.Parameter((params.sim.n, params.sim.n_controls), name='grad_g_u_' + str(idx_ncvx)))
-                nu_vb.append(cp.Variable(params.sim.n, name='nu_vb_' + str(idx_ncvx))) # Virtual Control for VB
+                g.append(cp.Parameter(params.sim.n, name="g_" + str(idx_ncvx)))
+                grad_g_x.append(
+                    cp.Parameter(
+                        (params.sim.n, params.sim.n_states), name="grad_g_x_" + str(idx_ncvx)
+                    )
+                )
+                grad_g_u.append(
+                    cp.Parameter(
+                        (params.sim.n, params.sim.n_controls), name="grad_g_u_" + str(idx_ncvx)
+                    )
+                )
+                nu_vb.append(
+                    cp.Variable(params.sim.n, name="nu_vb_" + str(idx_ncvx))
+                )  # Virtual Control for VB
 
     # Applying the affine scaling to state and control
     x_nonscaled = []
@@ -174,50 +186,90 @@ def OptimalControlProblem(params: Config):
                 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]
+                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':
+        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':
+        if params.sim.final_state.type[i] == "Minimize":
             cost += lam_cost * x_nonscaled[-1][i]
-        if params.sim.initial_state.type[i] == 'Maximize':
+        if params.sim.initial_state.type[i] == "Maximize":
             cost += lam_cost * x_nonscaled[0][i]
-        if params.sim.final_state.type[i] == 'Maximize':
+        if params.sim.final_state.type[i] == "Maximize":
             cost += lam_cost * x_nonscaled[-1][i]
 
     if params.sim._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.sim.n)] # Uniform Time Step
-
-    constr += [0 == la.inv(S_x) @ (x_nonscaled[i] - x_bar[i] - dx[i]) for i in range(params.sim.n)] # State Error
-    constr += [0 == la.inv(S_u) @ (u_nonscaled[i] - u_bar[i] - du[i]) for i in range(params.sim.n)] # Control Error
-
-    constr += [x_nonscaled[i] == \
-                      A_d[i-1] @ x_nonscaled[i-1] \
-                    + B_d[i-1] @ u_nonscaled[i-1] \
-                    + C_d[i-1] @ u_nonscaled[i] \
-                    + z_d[i-1] \
-                    + nu[i-1] for i in range(1, params.sim.n)] # Dynamics Constraint
-    
-    constr += [u_nonscaled[i] <= params.sim.max_control for i in range(params.sim.n)]
-    constr += [u_nonscaled[i] >= params.sim.min_control for i in range(params.sim.n)] # Control Constraints
+        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.sim.n)
+        ]  # Uniform Time Step
+
+    constr += [
+        0 == la.inv(S_x) @ (x_nonscaled[i] - x_bar[i] - dx[i]) for i in range(params.sim.n)
+    ]  # State Error
+    constr += [
+        0 == la.inv(S_u) @ (u_nonscaled[i] - u_bar[i] - du[i]) for i in range(params.sim.n)
+    ]  # Control Error
+
+    constr += [
+        x_nonscaled[i]
+        == A_d[i - 1] @ x_nonscaled[i - 1]
+        + B_d[i - 1] @ u_nonscaled[i - 1]
+        + C_d[i - 1] @ u_nonscaled[i]
+        + z_d[i - 1]
+        + nu[i - 1]
+        for i in range(1, params.sim.n)
+    ]  # Dynamics Constraint
 
-    constr += [x_nonscaled[i][params.sim.idx_x_true] <= params.sim.max_state[params.sim.idx_x_true] for i in range(params.sim.n)]
-    constr += [x_nonscaled[i][params.sim.idx_x_true] >= params.sim.min_state[params.sim.idx_x_true] for i in range(params.sim.n)] # State Constraints (Also implemented in CTCS but included for numerical stability)
+    constr += [u_nonscaled[i] <= params.sim.max_control for i in range(params.sim.n)]
+    constr += [
+        u_nonscaled[i] >= params.sim.min_control for i in range(params.sim.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.sim.n)
+    ]
+    constr += [
+        x_nonscaled[i][params.sim.idx_x_true] >= params.sim.min_state[params.sim.idx_x_true]
+        for i in range(params.sim.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(lam_prox * cp.sum_squares(inv @ cp.hstack((dx[i], du[i]))) for i in range(params.sim.n))  # Trust Region Cost
-    cost += sum(lam_vc * cp.sum(cp.abs(nu[i-1])) for i in range(1, params.sim.n)) # Virtual Control Slack
-    
+
+    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(
+        lam_prox * cp.sum_squares(inv @ cp.hstack((dx[i], du[i]))) for i in range(params.sim.n)
+    )  # Trust Region Cost
+    cost += sum(
+        lam_vc * cp.sum(cp.abs(nu[i - 1])) for i in range(1, params.sim.n)
+    )  # Virtual Control Slack
+
     idx_ncvx = 0
     if params.sim.constraints_nodal:
         for constraint in params.sim.constraints_nodal:
@@ -225,28 +277,32 @@ def OptimalControlProblem(params: Config):
                 cost += 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(params.sim.idx_y.start, params.sim.idx_y.stop), params.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]) <= params.sim.max_state[idx]
+            for i in range(start_idx, nodes[1])
+        ]
         constr += [x_nonscaled[0][idx] == 0]
 
-    
     #########
     # PROBLEM
     #########
     prob = cp.Problem(cp.Minimize(cost), constr)
     if solver.cvxpygen:
         # Check to see if solver directory exists
-        if not os.path.exists('solver'):
-            cpg.generate_code(prob, solver = solver.cvx_solver, code_dir='solver', wrapper = True)
+        if not os.path.exists("solver"):
+            cpg.generate_code(prob, solver=solver.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 = solver.cvx_solver, code_dir='solver', wrapper = True)
+            if overwrite.lower() == "y":
+                cpg.generate_code(prob, solver=solver.cvx_solver, code_dir="solver", wrapper=True)
             else:
                 pass
     return prob

{openscvx-0.5.3.dev45 → openscvx-0.5.3.dev47}/docs/UnderTheHood/batching_jit_grad.md

@@ -25,9 +25,7 @@ result = problem.solve()
 result = problem.solve_jax()
 
 # Batched — standard jax.vmap, no library-specific API.
-batched = jax.vmap(problem.solve_jax, in_axes=(0, 0, None))(
-    x_initial_stack, x_final_stack, params
-)
+batched = jax.vmap(problem.solve_jax, in_axes=(0, 0, None))(x_initial_stack, x_final_stack, params)
 
 # JIT-compile once, solve forever (e.g. MPC inner loop).
 fast_solve = jax.jit(problem.solve_jax)
@@ -101,10 +99,10 @@ batched inputs agree.)
 
 ```python
 results = problem.solve_batched(
-    x_initial=x0_batch,                  # (B, n_x): rank = declared+1 -> batched
+    x_initial=x0_batch,  # (B, n_x): rank = declared+1 -> batched
     parameters={
-        "gate_center": centers,          # (B, 3) vs declared (3,)  -> batched
-        "gate_radius": 0.5,              # declared shape           -> shared
+        "gate_center": centers,  # (B, 3) vs declared (3,)  -> batched
+        "gate_radius": 0.5,  # declared shape           -> shared
     },
 )
 ```
@@ -117,7 +115,7 @@ fall back to the rank rule):
 
 ```python
 results = problem.solve_batched(
-    parameters={"weights": w},           # (B,) is also w's declared shape...
+    parameters={"weights": w},  # (B,) is also w's declared shape...
     in_axes={"parameters": {"weights": 0}},  # ...so say "batched" explicitly
 )
 
@@ -157,7 +155,7 @@ kwarg's field. Construction-only keys of the constructor dict — `autotuner`,
 
 ```python
 B = 8  # batch size. Just an int for building the arrays below — there is no
-       # B argument; solve_batched infers it from the leading-axis lengths.
+# B argument; solve_batched infers it from the leading-axis lengths.
 sweep = problem.solve_batched(algorithm={"ep_tr": jnp.logspace(-6, -3, B)})
 weights = problem.solve_batched(algorithm={"lam_prox": jnp.array([0.5, 1.0, 4.0])})
 budgets = problem.solve_batched(max_iters=jnp.array([5, 10, 20]))  # per-element caps

{openscvx-0.5.3.dev45 → openscvx-0.5.3.dev47}/docs/UnderTheHood/custom_algorithms_autotuners.md

@@ -183,7 +183,7 @@ the rule:
 
 ```python
 def converged(self, state):
-    return state.J_vc < state.ep_vc   # custom: virtual control alone
+    return state.J_vc < state.ep_vc  # custom: virtual control alone
 ```
 
 The default — every metric below its threshold — is algorithm-agnostic, since

{openscvx-0.5.3.dev45 → openscvx-0.5.3.dev47}/docs/UnderTheHood/lowering_architecture.md

@@ -71,12 +71,12 @@ After preprocessing, all symbolic definitions are collected into a `SymbolicProb
 ```python
 @dataclass
 class SymbolicProblem:
-    dynamics: Expr              # Symbolic dx/dt = f(x, u)
-    states: List[State]         # All state variables (including augmented)
-    controls: List[Control]     # All control variables (including virtual)
+    dynamics: Expr  # Symbolic dx/dt = f(x, u)
+    states: List[State]  # All state variables (including augmented)
+    controls: List[Control]  # All control variables (including virtual)
     constraints: ConstraintSet  # Categorized constraints
-    parameters: dict            # User-defined parameters
-    N: int                      # Discretization nodes
+    parameters: dict  # User-defined parameters
+    N: int  # Discretization nodes
     # ... plus propagation variants
 ```
 
@@ -88,20 +88,20 @@ Lowering produces a `LoweredProblem` containing everything needed for optimizati
 @dataclass
 class LoweredProblem:
     # JAX dynamics (callable functions with Jacobians)
-    dynamics: Dynamics           # f, A=df/dx, B=df/du
-    dynamics_prop: Dynamics      # For forward propagation
+    dynamics: Dynamics  # f, A=df/dx, B=df/du
+    dynamics_prop: Dynamics  # For forward propagation
 
     # Lowered constraints (by backend)
     jax_constraints: LoweredJaxConstraints
     cvxpy_constraints: LoweredCvxpyConstraints
 
     # Unified state/control interfaces
-    x_unified: UnifiedState      # Aggregates all states into one vector
-    u_unified: UnifiedControl    # Aggregates all controls into one vector
+    x_unified: UnifiedState  # Aggregates all states into one vector
+    u_unified: UnifiedControl  # Aggregates all controls into one vector
 
     # CVXPy optimization variables
-    ocp_vars: OCPVariables       # x, u, dx, du, nu, etc.
-    cvxpy_params: dict           # User parameters as cp.Parameter
+    ocp_vars: OCPVariables  # x, u, dx, du, nu, etc.
+    cvxpy_params: dict  # User parameters as cp.Parameter
 ```
 
 ## Constraint Routing
@@ -123,9 +123,9 @@ Non-convex constraints become JAX functions with gradients:
 ```python
 @dataclass
 class LoweredJaxConstraints:
-    nodal: List[LoweredNodalConstraint]      # func, grad_g_x, grad_g_u
+    nodal: List[LoweredNodalConstraint]  # func, grad_g_x, grad_g_u
     cross_node: List[LoweredCrossNodeConstraint]  # func, grad_g_X, grad_g_U
-    ctcs: List[CTCS]                         # Handled via dynamics augmentation
+    ctcs: List[CTCS]  # Handled via dynamics augmentation
 ```
 
 ### CVXPy-Lowered Constraints

{openscvx-0.5.3.dev45 → openscvx-0.5.3.dev47}/docs/UnderTheHood/vectorization_and_vmapping.md

@@ -80,16 +80,16 @@ The unification process (in `openscvx/symbolic/unified.py`) sorts variables (use
 For a problem with `N` decision nodes:
 
 ```python
-x_unified.shape = (n_x,)          # Sum of all state dimensions
-u_unified.shape = (n_u,)          # Sum of all control dimensions
+x_unified.shape = (n_x,)  # Sum of all state dimensions
+u_unified.shape = (n_u,)  # Sum of all control dimensions
 x_unified.guess.shape = (N, n_x)  # State trajectory
 u_unified.guess.shape = (N, n_u)  # Control trajectory
 ```
 
 **Concrete example** (brachistochrone with N=10, no CTCS constraints):
 ```python
-x_unified.shape = (4,)        # position(2) + velocity(1) + time(1)
-u_unified.shape = (2,)        # theta(1) + _time_dilation(1)
+x_unified.shape = (4,)  # position(2) + velocity(1) + time(1)
+u_unified.shape = (2,)  # theta(1) + _time_dilation(1)
 x_unified.guess.shape = (10, 4)
 u_unified.guess.shape = (10, 2)
 ```
@@ -116,9 +116,9 @@ dyn_fn = lower_to_jax(dynamics_aug)
 
 # Create Dynamics object with Jacobians
 dynamics_augmented = Dynamics(
-    f=dyn_fn,                      # State derivative function
-    A=jacfwd(dyn_fn, argnums=0),   # Jacobian df/dx
-    B=jacfwd(dyn_fn, argnums=1),   # Jacobian df/du
+    f=dyn_fn,  # State derivative function
+    A=jacfwd(dyn_fn, argnums=0),  # Jacobian df/dx
+    B=jacfwd(dyn_fn, argnums=1),  # Jacobian df/du
 )
 ```
 
@@ -158,10 +158,10 @@ constraints_nodal_fns = lower_to_jax(constraints_nodal)
 # Create LoweredNodalConstraint objects with Jacobians
 for i, fn in enumerate(constraints_nodal_fns):
     constraint = LoweredNodalConstraint(
-        func=fn,                          # Constraint function
+        func=fn,  # Constraint function
         grad_g_x=jacfwd(fn, argnums=0),  # Jacobian dg/dx
         grad_g_u=jacfwd(fn, argnums=1),  # Jacobian dg/du
-        nodes=constraints_nodal[i].nodes, # Node indices where constraint applies
+        nodes=constraints_nodal[i].nodes,  # Node indices where constraint applies
     )
 ```
 
@@ -239,18 +239,9 @@ Dynamics functions are vmapped to process all intervals in parallel (in `Problem
 
 ```python
 # Vectorize dynamics functions across decision nodes
-self.dynamics_augmented.f = jax.vmap(
-    self.dynamics_augmented.f,
-    in_axes=(0, 0, 0, None)
-)
-self.dynamics_augmented.A = jax.vmap(
-    self.dynamics_augmented.A,
-    in_axes=(0, 0, 0, None)
-)
-self.dynamics_augmented.B = jax.vmap(
-    self.dynamics_augmented.B,
-    in_axes=(0, 0, 0, None)
-)
+self.dynamics_augmented.f = jax.vmap(self.dynamics_augmented.f, in_axes=(0, 0, 0, None))
+self.dynamics_augmented.A = jax.vmap(self.dynamics_augmented.A, in_axes=(0, 0, 0, None))
+self.dynamics_augmented.B = jax.vmap(self.dynamics_augmented.B, in_axes=(0, 0, 0, None))
 ```
 
 **Dynamics Vmap Configuration: `in_axes=(0, 0, 0, None)`**
@@ -372,9 +363,9 @@ Cross-node constraints are **not vmapped** because they already operate on full
 
 ```python
 # Each LoweredCrossNodeConstraint operates on full trajectories
-residual = constraint.func(X, U, params)      # scalar
-grad_X = constraint.grad_g_X(X, U, params)    # (N, n_x) - sparse, mostly zeros
-grad_U = constraint.grad_g_U(X, U, params)    # (N, n_u) - sparse, mostly zeros
+residual = constraint.func(X, U, params)  # scalar
+grad_X = constraint.grad_g_X(X, U, params)  # (N, n_x) - sparse, mostly zeros
+grad_U = constraint.grad_g_U(X, U, params)  # (N, n_u) - sparse, mostly zeros
 ```
 
 The Jacobians are dense arrays but exhibit sparsity patterns determined by which nodes the constraint couples.
@@ -385,23 +376,23 @@ The vmapped dynamics functions are called during discretization (in `calculate_d
 
 ```python
 # Setup batch inputs
-x = V[:, :n_x]                          # Shape: (N-1, n_x) - States at interval starts
-u = u[: x.shape[0]]                     # Shape: (N-1, n_u) - Controls (includes time dilation)
-nodes = jnp.arange(0, N-1)              # Shape: (N-1,) - Node indices
+x = V[:, :n_x]  # Shape: (N-1, n_x) - States at interval starts
+u = u[: x.shape[0]]  # Shape: (N-1, n_u) - Controls (includes time dilation)
+nodes = jnp.arange(0, N - 1)  # Shape: (N-1,) - Node indices
 
 # Extract time dilation (last control dimension)
-s = u[:, -1]                            # Shape: (N-1,) - Time dilation values
+s = u[:, -1]  # Shape: (N-1,) - Time dilation values
 
 # Call vmapped dynamics - evaluates all intervals in parallel
 # Note: dynamics receive u[:, :-1] (vehicle controls only, excluding time dilation)
 f = state_dot(x, u[:, :-1], nodes, params)  # Shape: (N-1, n_x)
-dfdx = A(x, u[:, :-1], nodes, params)       # Shape: (N-1, n_x, n_x)
-dfdu_veh = B(x, u[:, :-1], nodes, params)   # Shape: (N-1, n_x, n_u-1)
+dfdx = A(x, u[:, :-1], nodes, params)  # Shape: (N-1, n_x, n_x)
+dfdu_veh = B(x, u[:, :-1], nodes, params)  # Shape: (N-1, n_x, n_u-1)
 
 # Build full control Jacobian including time dilation
 dfdu = jnp.zeros((x.shape[0], n_x, n_u))
 dfdu = dfdu.at[:, :, :-1].set(s[:, None, None] * dfdu_veh)  # Vehicle control derivatives
-dfdu = dfdu.at[:, :, -1].set(f)                              # Time dilation derivative = f
+dfdu = dfdu.at[:, :, -1].set(f)  # Time dilation derivative = f
 ```
 
 **Why exclude time dilation from dynamics?** Time dilation is a meta-control that scales the entire dynamics (used for time-optimal problems). The actual vehicle dynamics are defined without it, and time dilation is applied as a scaling factor during discretization. This is why `n_u-1` appears in the vehicle dynamics Jacobians.

{openscvx-0.5.3.dev45 → openscvx-0.5.3.dev47}/docs/UsersGuide/01_hello_world_brachistochrone.md

@@ -214,12 +214,7 @@ To mark these as CTCS constraints we simply wrap them in `ox.ctcs(...)`
 # Generate box constraints for all states
 constraints = []
 for state in states:
-    constraints.extend(
-        [
-            ox.ctcs(state <= state.max),
-            ox.ctcs(state.min <= state)
-        ]
-    )
+    constraints.extend([ox.ctcs(state <= state.max), ox.ctcs(state.min <= state)])
 ```
 
 This style of constraint definition as a list should feel familiar to CVXPY users.