simplified API for development

Author: ishaanamahajan

examples/cartpole_example_code_generation.py

@@ -15,11 +15,9 @@
 N = 20
 
 prob = tinympc.TinyMPC()
-prob.setup(A, B, Q, R, N, rho=1, max_iter=10)
 
-# Optional bounds via new API
 u_min = np.array([-0.5])
 u_max = np.array([0.5])
-prob.set_bound_constraints([], [], u_min, u_max)
+prob.setup(A, B, Q, R, N, rho=1, max_iter=10, u_min=u_min, u_max=u_max)
 
 prob.codegen("out", verbose=1)

examples/cartpole_example_mpc_reference_constrained.py

@@ -18,12 +18,11 @@
 N = 20
 
 prob = tinympc.TinyMPC()
-prob.setup(A, B, Q, R, N, rho=1.0)
 
 # Define input constraints
 u_min = np.array([-0.45])
 u_max = np.array([0.45])
-prob.set_bound_constraints([], [], u_min, u_max)
+prob.setup(A, B, Q, R, N, rho=1.0, u_min=u_min, u_max=u_max)
 
 # Goal must be another equilibrium position
 prob.set_x_ref(np.array([1.0, 0, 0, 0]))

examples/quadrotor_hover_code_generation.py

@@ -4,7 +4,7 @@
 import autograd.numpy as anp
 
 # Toggle switch for adaptive rho
-ENABLE_ADAPTIVE_RHO = False   # Set to True to enable adaptive rho
+ENABLE_ADAPTIVE_RHO = False  # Set to True to enable adaptive rho
 
 # Quadrotor system matrices (12 states, 4 inputs)
 rho_value = 5.0

examples/rocket_landing_constraints.py

@@ -51,15 +51,16 @@
 qcu = np.array([3])     # dimensions for input cones
 
 # Setup solver
+cone_constraints = {
+    'Acu': Acu, 'qcu': qcu, 'cu': cu,
+    'Acx': Acx, 'qcx': qcx, 'cx': cx
+}
+
 solver = tinympc.TinyMPC()
 solver.setup(A, B, Q, R, NHORIZON, rho=1.0, fdyn=fdyn,
-             max_iter=100, abs_pri_tol=2e-3, verbose=True)
-
-# Set bounds explicitly 
-solver.set_bound_constraints(x_min, x_max, u_min, u_max)
-
-# Set cone constraints (inputs first)
-solver.set_cone_constraints(Acu, qcu, cu, Acx, qcx, cx)
+             max_iter=100, abs_pri_tol=2e-3, verbose=True,
+             x_min=x_min, x_max=x_max, u_min=u_min, u_max=u_max,
+             cone_constraints=cone_constraints)
 
 # Initial and goal states
 xinit = np.array([4.0, 2.0, 20.0, -3.0, 2.0, -4.5])

src/tinympc/interface.py

@@ -81,19 +81,23 @@ def expand_ndarray(self, array_, expected_rows, expected_cols, fallback):
         return array_
 
     # Setup the problem data and solver options
-    def setup(self, A, B, Q, R, N, rho=1.0, fdyn=None, verbose=False, **settings):
+    def setup(self, A, B, Q, R, N, rho=1.0, fdyn=None, 
+              x_min=None, x_max=None, u_min=None, u_max=None, 
+              cone_constraints=None, verbose=False, **settings):
         """Instantiate necessary algorithm variables and parameters
         
         :param A (np.ndarray): State transition matrix of the linear system, size nx x nx
         :param B (np.ndarray): Input matrix of the linear system, size nx x nu
         :param Q (np.ndarray): Stage cost for state, must be diagonal and positive semi-definite, size nx x nx
         :param R (np.ndarray): Stage cost for input, must be diagonal and positive definite, size nu x nu
+        :param N (int): Prediction horizon length
         :param rho (float, optional): Penalty term used in ADMM, default 1.0
         :param fdyn (np.ndarray or None, optional): Affine offset vector for dynamics, size nx x 1. If None, defaults to zeros (linear system), default None
-        :param x_min (list[float] or None, optional): Lower bound state constraints of the same length as nx, default None
-        :param x_max (list[float] or None, optional): Upper bound state constraints of the same length as nx, default None
-        :param u_min (list[float] or None, optional): Lower bound input constraints of the same length as nu, default None
-        :param u_max (list[float] or None, optional): Upper bound input constraints of the same length as nu, default None
+        :param x_min (np.ndarray or None, optional): Lower bound state constraints, size nx or nx x N, default None
+        :param x_max (np.ndarray or None, optional): Upper bound state constraints, size nx or nx x N, default None  
+        :param u_min (np.ndarray or None, optional): Lower bound input constraints, size nu or nu x (N-1), default None
+        :param u_max (np.ndarray or None, optional): Upper bound input constraints, size nu or nu x (N-1), default None
+        :param cone_constraints (dict or None, optional): Cone constraints dict with keys {Acu, qcu, cu, Acx, qcx, cx}, default None
         :param verbose (bool): Whether or not to print data to console during setup, default False
         :params settings: Dictionary of optional settings
             :param abs_pri_tol (float): Solution tolerance for primal variables
@@ -155,6 +159,14 @@ def setup(self, A, B, Q, R, N, rho=1.0, fdyn=None, verbose=False, **settings):
                                            self.nx, self.nu, self.N,
                                            self.settings, self.verbose)
 
+        # Handle constraints if provided 
+        if any(x is not None for x in [x_min, x_max, u_min, u_max]):
+            self.set_bound_constraints(x_min, x_max, u_min, u_max)
+        
+        if cone_constraints is not None:
+            # cone_constraints should be a dict with keys: Acu, qcu, cu, Acx, qcx, cx
+            self.set_cone_constraints(**cone_constraints)
+
     def set_x0(self, x0):
         assert len(x0.shape) == 1
         assert len(x0) == self.nx