Merge branch 'main' into feature/support-ltv

Author: shaoanlu

README.md

@@ -15,6 +15,56 @@ Refer to [examples](examples/):
 - `examples/afti16_example.py` (linear AFTI-16 aircraft model)
 - `examples/unicycle_nmpc_example.py` (nonlinear unicycle with time-varying linearization)
 
+Create a model, configure with `Modelsetup()`, and solve with `Model.solve()`.
+
+```python
+model = Model()
+model.setup(A, B, Np, kwargs...)
+results = model.solve(x0, u0, yref, uref, w, verbose=False)
+```
+
+**Required Arguments:**
+- `A`: State matrix
+- `B`: Input matrix
+- `Np`: Prediction horizon
+
+**Optional Problem Arguments:**
+- `C`: Output matrix (default: I, meaning y = x)
+- `e`: Affine term in dynamics (default: zeros)
+- `Wy`: Output weight matrix (default: I)
+- `Wu`: Input weight matrix (default: I)
+- `Wdu`: Input increment weight matrix (default: I)
+- `Wf`: Terminal cost weight matrix (default: Wy)
+- `xmin, xmax`: State constraints (default: ±Inf)
+- `umin, umax`: Input constraints (default: ±Inf)
+- `dumin, dumax`: Input increment constraints (default: ±Inf)
+
+**Solver Settings:**
+- `rho`: ADMM penalty parameter (default: 1.0)
+- `tol`: Convergence tolerance (default: 1e-4)
+- `eta`: Acceleration restart factor (default: 0.999)
+- `maxiter`: Maximum iterations (default: 100)
+- `precond`: Use preconditioning (default: false)
+- `accel`: Use Nesterov acceleration (default: false)
+- `device`: Compute device, `cpu` or `gpu` (default: cpu)
+
+**System Model:**
+```
+x_{k+1} = A * x_k + B * u_k + e + w
+y_k = C * x_k
+```
+
+where:
+- `e`: Constant affine term (defined in `setup!`)
+- `w`: Known/measured disturbance (provided at solve time)
+
+**Arguments:**
+- `x0`: Current state (nx)
+- `u0`: Previous input (nu)
+- `yref`: Output reference (ny)
+- `uref`: Input reference (nu)
+- `w`: Known disturbance (nx), use `zeros(nx)` if none
+
 ![](assets/example_afti16.png)
 
 ![](assets/example_unicycle.png)

pimpc/solver.py

@@ -381,16 +381,19 @@ def _solve_gpu(model, x0, u0, yref, uref, w, *, warm_vars=None, verbose=False):
         E_inv = cp.ones(nx_bar, dtype=dtype)
 
     C_part = C_bar[:, :nx]
-    Q_bar = cp.block(
+    Q11 = C_part.T @ cp.asarray(model.Wy, dtype=dtype) @ C_part
+    Q22 = cp.asarray(model.Wu, dtype=dtype)
+    Q_bar = cp.vstack(
         [
-            [C_part.T @ cp.asarray(model.Wy, dtype=dtype) @ C_part, cp.zeros((nx, nu), dtype=dtype)],
-            [cp.zeros((nu, nx), dtype=dtype), cp.asarray(model.Wu, dtype=dtype)],
+            cp.hstack([Q11, cp.zeros((nx, nu), dtype=dtype)]),
+            cp.hstack([cp.zeros((nu, nx), dtype=dtype), Q22]),
         ]
     )
-    Q_bar_N = cp.block(
+    Q11_N = C_part.T @ cp.asarray(model.Wf, dtype=dtype) @ C_part
+    Q_bar_N = cp.vstack(
         [
-            [C_part.T @ cp.asarray(model.Wf, dtype=dtype) @ C_part, cp.zeros((nx, nu), dtype=dtype)],
-            [cp.zeros((nu, nx), dtype=dtype), cp.asarray(model.Wu, dtype=dtype)],
+            cp.hstack([Q11_N, cp.zeros((nx, nu), dtype=dtype)]),
+            cp.hstack([cp.zeros((nu, nx), dtype=dtype), Q22]),
         ]
     )
     q_bar = cp.concatenate(