openscvx 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

openscvx/discretization.py

@@ -0,0 +1,170 @@
+import jax.numpy as jnp
+import numpy as np
+
+from openscvx.dynamics import Dynamics
+from openscvx.integrators import solve_ivp_rk45, solve_ivp_diffrax
+
+
+def dVdt(
+    tau: float,
+    V: jnp.ndarray,
+    u_cur: np.ndarray,
+    u_next: np.ndarray,
+    state_dot: callable,
+    A: callable,
+    B: callable,
+    n_x: int,
+    n_u: int,
+    N: int,
+    dis_type: str,
+) -> jnp.ndarray:
+    # Define the nodes
+    nodes = jnp.arange(0, N-1)
+
+    # Define indices for slicing the augmented state vector
+    i0 = 0
+    i1 = n_x
+    i2 = i1 + n_x * n_x
+    i3 = i2 + n_x * n_u
+    i4 = i3 + n_x * n_u
+    i5 = i4 + n_x
+
+    # Unflatten V
+    V = V.reshape(-1, i5)
+
+    # Compute the interpolation factor based on the discretization type
+    if dis_type == "ZOH":
+        beta = 0.0
+    elif dis_type == "FOH":
+        beta = (tau) * 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[:, :n_x]
+    u = u[: x.shape[0]]
+
+    # Compute the nonlinear propagation term
+    f = state_dot(x, u[:, :-1], nodes)
+    F = s[:, None] * f
+
+    # Evaluate the State Jacobian
+    dfdx = A(x, u[:, :-1], nodes)
+    sdfdx = s[:, None, None] * dfdx
+
+    # Evaluate the Control Jacobian
+    dfdu_veh = B(x, u[:, :-1], nodes)
+    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
+    # fmt: off
+    dVdt = jnp.zeros_like(V)
+    dVdt = dVdt.at[:, i0:i1].set(F)
+    dVdt = dVdt.at[:, i1:i2].set(jnp.matmul(sdfdx, V[:, i1:i2].reshape(-1, n_x, n_x)).reshape(-1, n_x * n_x))
+    dVdt = dVdt.at[:, i2:i3].set((jnp.matmul(sdfdx, V[:, i2:i3].reshape(-1, n_x, n_u)) + dfdu * alpha).reshape(-1, n_x * n_u))
+    dVdt = dVdt.at[:, i3:i4].set((jnp.matmul(sdfdx, V[:, i3:i4].reshape(-1, n_x, n_u)) + dfdu * beta).reshape(-1, n_x * n_u))
+    dVdt = dVdt.at[:, i4:i5].set((jnp.matmul(sdfdx, V[:, i4:i5].reshape(-1, n_x)[..., None]).squeeze(-1) + z).reshape(-1, n_x))
+    # fmt: on
+    return dVdt.flatten()
+
+
+def calculate_discretization(
+    x,
+    u,
+    state_dot: callable,
+    A: callable,
+    B: callable,
+    n_x: int,
+    n_u: int,
+    N: int,
+    custom_integrator: bool,
+    debug: bool,
+    solver: str,
+    rtol,
+    atol,
+    dis_type: str,
+):
+
+    # Define indices for slicing the augmented state vector
+    i0 = 0
+    i1 = n_x
+    i2 = i1 + n_x * n_x
+    i3 = i2 + n_x * n_u
+    i4 = i3 + n_x * n_u
+    i5 = i4 + n_x
+
+    # initial augmented state
+    V0 = jnp.zeros((N - 1, i5))
+    V0 = V0.at[:, :n_x].set(x[:-1].astype(float))
+    V0 = V0.at[:, n_x : n_x + n_x * n_x].set(
+        jnp.eye(n_x).reshape(1, -1).repeat(N - 1, axis=0)
+    )
+
+    # choose integrator
+    if custom_integrator:
+        # fmt: off
+        sol = solve_ivp_rk45(
+            lambda t,y,*a: dVdt(t, y, *a),
+            1.0/(N-1),
+            V0.reshape(-1),
+            args=(u[:-1].astype(float), u[1:].astype(float),
+                  state_dot, A, B, n_x, n_u, N, dis_type),
+            is_not_compiled=debug,
+        )
+        # fmt: on
+    else:
+        # fmt: off
+        sol = solve_ivp_diffrax(
+            lambda t,y,*a: dVdt(t, y, *a),
+            1.0/(N-1),
+            V0.reshape(-1),
+            args=(u[:-1].astype(float), u[1:].astype(float),
+                  state_dot, A, B, n_x, n_u, N, dis_type),
+            solver_name=solver,
+            rtol=rtol,
+            atol=atol,
+            extra_kwargs=None,
+        )
+        # fmt: on
+
+    Vend = sol[-1].T.reshape(-1, i5)
+    Vmulti = sol.T
+
+    # fmt: off
+    A_bar = Vend[:, i1:i2].reshape(N-1, n_x, n_x).transpose(1,2,0).reshape(n_x*n_x, -1, order='F').T
+    B_bar = Vend[:, i2:i3].reshape(N-1, n_x, n_u).transpose(1,2,0).reshape(n_x*n_u, -1, order='F').T
+    C_bar = Vend[:, i3:i4].reshape(N-1, n_x, n_u).transpose(1,2,0).reshape(n_x*n_u, -1, order='F').T
+    z_bar = Vend[:, i4:i5]
+    # fmt: on
+
+    return A_bar, B_bar, C_bar, z_bar, Vmulti
+
+
+def get_discretization_solver(dyn: Dynamics, params):
+    return lambda x, u: calculate_discretization(
+        x=x,
+        u=u,
+        state_dot=dyn.f,
+        A=dyn.A,
+        B=dyn.B,
+        n_x=params.sim.n_states,
+        n_u=params.sim.n_controls,
+        N=params.scp.n,
+        custom_integrator=params.dis.custom_integrator,
+        debug=params.dev.debug,
+        solver=params.dis.solver,
+        rtol=params.dis.rtol,
+        atol=params.dis.atol,
+        dis_type=params.dis.dis_type,
+    )

openscvx/dynamics.py

@@ -0,0 +1,41 @@
+from dataclasses import dataclass
+from typing import Callable, Optional
+import functools
+
+import jax.numpy as jnp
+
+
+@dataclass
+class Dynamics:
+    f: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]
+    A: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None
+    B: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None
+
+def dynamics(
+    _func=None,
+    *,
+    A: Optional[Callable] = None,
+    B: Optional[Callable] = None,):
+    """Decorator to mark a function as defining the system dynamics.
+
+    Use as:
+    @dynamics(A=my_grad_f_x, B=my_grad_f_u)')
+    def my_dynamics(x,u): ...
+    """
+
+    def decorator(f: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]):
+        # wrap so name, doc, signature stay on f
+        wrapped = functools.wraps(f)(f)
+        return Dynamics(
+            f=wrapped,
+            A=A,
+            B=B,
+        )
+
+    # if called as @dynamics or @dynamics(...), _func will be None and we return decorator
+    if _func is None:
+        return decorator
+    # if called as dynamics(func), we immediately decorate
+    else:
+        return decorator(_func)
+

openscvx/integrators.py

@@ -0,0 +1,139 @@
+import jax
+import jax.numpy as jnp
+import diffrax as dfx
+
+SOLVER_MAP = {
+    "Tsit5": dfx.Tsit5,
+    "Euler": dfx.Euler,
+    "Heun": dfx.Heun,
+    "Midpoint": dfx.Midpoint,
+    "Ralston": dfx.Ralston,
+    "Dopri5": dfx.Dopri5,
+    "Dopri8": dfx.Dopri8,
+    "Bosh3": dfx.Bosh3,
+    "ReversibleHeun": dfx.ReversibleHeun,
+    "ImplicitEuler": dfx.ImplicitEuler,
+    "KenCarp3": dfx.KenCarp3,
+    "KenCarp4": dfx.KenCarp4,
+    "KenCarp5": dfx.KenCarp5,
+}
+
+# fmt: off
+def rk45_step(f, t, y, h, *args):
+    k1 = f(t, y, *args)
+    k2 = f(t + h/4, y + h*k1/4, *args)
+    k3 = f(t + 3*h/8, y + 3*h*k1/32 + 9*h*k2/32, *args)
+    k4 = f(t + 12*h/13, y + 1932*h*k1/2197 - 7200*h*k2/2197 + 7296*h*k3/2197, *args)
+    k5 = f(t + h, y + 439*h*k1/216 - 8*h*k2 + 3680*h*k3/513 - 845*h*k4/4104, *args)
+    y_next = y + h * (25*k1/216 + 1408*k3/2565 + 2197*k4/4104 - k5/5)
+    return y_next
+# fmt: on
+
+
+def solve_ivp_rk45(
+    f,
+    tau_final: float,
+    y_0,
+    args,
+    tau_0: float = 0.0,
+    num_substeps: int = 50,
+    is_not_compiled: bool = False,
+):
+    substeps = jnp.linspace(tau_0, tau_final, num_substeps)
+
+    h = (tau_final - tau_0) / (len(substeps) - 1)
+    solution = jnp.zeros((len(substeps), len(y_0)))
+    solution = solution.at[0].set(y_0)
+
+    if is_not_compiled:
+        for i in range(1, len(substeps)):
+            t = tau_0 + i * h
+            solution = solution.at[i].set(rk45_step(f, t, solution[i - 1], h, *args))
+    else:
+
+        def body_fun(i, val):
+            t, y, V_result = val
+            y_next = rk45_step(f, t, y, h, *args)
+            V_result = V_result.at[i].set(y_next)
+            return (t + h, y_next, V_result)
+
+        _, _, solution = jax.lax.fori_loop(
+            1, len(substeps), body_fun, (tau_0, y_0, solution)
+        )
+
+    return solution
+
+
+def solve_ivp_diffrax(
+    f,
+    tau_final,
+    y_0,
+    args,
+    tau_0: float = 0.0,
+    num_substeps: int = 50,
+    solver_name="Dopri8",
+    rtol: float = 1e-3,
+    atol: float = 1e-6,
+    extra_kwargs=None,
+):
+    substeps = jnp.linspace(tau_0, tau_final, num_substeps)
+
+    solver_class = SOLVER_MAP.get(solver_name)
+    if solver_class is None:
+        raise ValueError(f"Unknown solver: {solver_name}")
+    solver = solver_class()
+
+    term = dfx.ODETerm(lambda t, y, args: f(t, y, *args))
+    stepsize_controller = dfx.PIDController(rtol=rtol, atol=atol)
+    solution = dfx.diffeqsolve(
+        term,
+        solver=solver,
+        t0=tau_0,
+        t1=tau_final,
+        dt0=(tau_final - tau_0) / (len(substeps) - 1),
+        y0=y_0,
+        args=args,
+        stepsize_controller=stepsize_controller,
+        saveat=dfx.SaveAt(ts=substeps),
+        **(extra_kwargs or {}),
+    )
+
+    return solution.ys
+
+
+# TODO: (norrisg) this function is basically identical to `solve_ivp_diffrax`, could combine, but requires returning solution and getting `.ys` wherever the `solve_ivp_diffrax` is called
+def solve_ivp_diffrax_prop(
+    f,
+    tau_final,
+    y_0,
+    args,
+    tau_0: float = 0.0,
+    num_substeps: int = 50,
+    solver_name="Dopri8",
+    rtol: float = 1e-3,
+    atol: float = 1e-6,
+    extra_kwargs=None,
+):
+    substeps = jnp.linspace(tau_0, tau_final, num_substeps)
+
+    solver_class = SOLVER_MAP.get(solver_name)
+    if solver_class is None:
+        raise ValueError(f"Unknown solver: {solver_name}")
+    solver = solver_class()
+
+    term = dfx.ODETerm(lambda t, y, args: f(t, y, *args))
+    stepsize_controller = dfx.PIDController(rtol=rtol, atol=atol)
+    solution = dfx.diffeqsolve(
+        term,
+        solver=solver,
+        t0=tau_0,
+        t1=tau_final,
+        dt0=(tau_final - tau_0) / (len(substeps) - 1),
+        y0=y_0,
+        args=args,
+        stepsize_controller=stepsize_controller,
+        saveat=dfx.SaveAt(dense=True, ts=substeps),
+        **(extra_kwargs or {}),
+    )
+
+    return solution

openscvx/io.py

@@ -0,0 +1,81 @@
+import sys
+import warnings
+warnings.filterwarnings("ignore")
+import queue
+import time
+
+from termcolor import colored
+
+# Define colors for printing
+col_main = "blue"
+col_pos = "green"
+col_neg = "red"
+
+def intro():
+    # Silence syntax warnings
+    warnings.filterwarnings("ignore")
+    # fmt: off
+    ascii_art = '''
+                             
+                            ____                    _____  _____           
+                           / __ \                  / ____|/ ____|          
+                          | |  | |_ __   ___ _ __ | (___ | |  __   ____  __
+                          | |  | | '_ \ / _ \ '_ \ \___ \| |  \ \ / /\ \/ /
+                          | |__| | |_) |  __/ | | |____) | |___\ V /  >  < 
+                           \____/| .__/ \___|_| |_|_____/ \_____\_/  /_/\_\ 
+                                 | |                                       
+                                 |_|                                       
+---------------------------------------------------------------------------------------------------------
+                                Author: Chris Hayner and Griffin Norris
+                                    Autonomous Controls Laboratory
+                                       University of Washington
+---------------------------------------------------------------------------------------------------------
+'''
+    # fmt: on
+    print(ascii_art)
+
+def header():
+    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("---------------------------------------------------------------------------------------------------------"))
+
+def intermediate(print_queue, params):
+    hz = 30.0
+    while True:
+        t_start = time.time()
+        try:
+            data = print_queue.get(timeout=1.0/hz)
+            # remove bottom labels and line
+            if not data["iter"] == 1:
+                sys.stdout.write('\x1b[1A\x1b[2K\x1b[1A\x1b[2K')
+            if data["prob_stat"][3] == 'f':
+                # Only show the first element of the string
+                data["prob_stat"] = data["prob_stat"][0]
+
+            iter_colored = colored("{:4d}".format(data["iter"]))
+            J_tot_colored = colored("{:.1e}".format(data["J_total"]))
+            J_tr_colored = colored("{:.1e}".format(data["J_tr"]), col_pos if data["J_tr"] <= params.scp.ep_tr else col_neg)
+            J_vb_colored = colored("{:.1e}".format(data["J_vb"]), col_pos if data["J_vb"] <= params.scp.ep_vb else col_neg)
+            J_vc_colored = colored("{:.1e}".format(data["J_vc"]), col_pos if data["J_vc"] <= params.scp.ep_vc else col_neg)
+            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(
+                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(
+                "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)

openscvx/ocp.py

@@ -0,0 +1,160 @@
+import os
+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
+
+
+def OptimalControlProblem(params: Config):
+    ########################
+    # VARIABLES & PARAMETERS
+    ########################
+
+    # Parameters
+    w_tr = cp.Parameter(nonneg = True, name='w_tr')
+    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
+
+    # Affine Scaling for State
+    S_x = params.sim.S_x
+    inv_S_x = params.sim.inv_S_x
+    c_x = params.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
+
+    # Affine Scaling for Control
+    S_u = params.sim.S_u
+    inv_S_u = params.sim.inv_S_u
+    c_u = params.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
+
+    # Linearized Nonconvex Nodal Constraints
+    if params.sim.constraints_nodal:
+        g = []
+        grad_g_x = []
+        grad_g_u = []
+        nu_vb = []
+        for idx_ncvx, constraint in enumerate(params.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
+
+    # Applying the affine scaling to state and control
+    x_nonscaled = []
+    u_nonscaled = []
+    for k in range(params.scp.n):
+        x_nonscaled.append(S_x @ x[k] + c_x)
+        u_nonscaled.append(S_u @ u[k] + c_u)
+
+    constr = []
+    cost = lam_cost * 0
+
+    #############
+    # CONSTRAINTS
+    #############
+    idx_ncvx = 0
+    if params.sim.constraints_nodal:
+        for constraint in params.sim.constraints_nodal:
+            if constraint.nodes is None:
+                nodes = range(params.scp.n)
+            else:
+                nodes = constraint.nodes
+
+            if 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':
+            cost += lam_cost * x_nonscaled[0][i]
+        if params.sim.final_state.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
+
+    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 += [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] \
+                    + z_d[i-1] \
+                    + nu[i-1] for i in range(1, params.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 += [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)
+
+    ########
+    # 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
+    
+    idx_ncvx = 0
+    if params.sim.constraints_nodal:
+        for constraint in params.sim.constraints_nodal:
+            if not constraint.convex:
+                cost += params.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):  
+        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 += [x_nonscaled[0][idx] == 0]
+
+    
+    #########
+    # PROBLEM
+    #########
+    prob = cp.Problem(cp.Minimize(cost), constr)
+    if params.cvx.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)
+        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)
+            else:
+                pass
+    return prob