implemented linear dynamics in single shooting for scenario MPC

(cherry picked from commit 66683d2)

examples/optimal_control/scenario_approach_for_mpc.py

@@ -69,7 +69,7 @@ def sample_disturbances(K: int, N: int) -> npt.NDArray[np.floating]:
 F = get_dynamics()
 nx, nu, nd = F.size1_in(0), F.size1_in(1), F.size1_in(2)
 
-scmpc = ScenarioBasedMpc[cs.SX](Nlp(), K, N)
+scmpc = ScenarioBasedMpc[cs.SX](Nlp(), K, N, shooting="single")
 x, _, _ = scmpc.state("x", nx, bound_initial=False)
 u, _ = scmpc.action("u", nu, lb=-5, ub=+5)
 d, _ = scmpc.disturbance("d", nd)

src/csnlp/wrappers/mpc/mpc.py

@@ -32,6 +32,37 @@ def _callable2csfunc(
     return cs.Function("F", sym_in, (sym_out,), {"allow_free": True, "cse": True})
 
 
+def _create_qp_mats(
+    N: int, A: MatType, B: MatType, D: Optional[MatType]
+) -> tuple[MatType, MatType, Optional[MatType]]:
+    """Internal utility to build the linear MPC matrices for building a QP."""
+    ns, na = B.shape
+
+    F = cs.vcat([cs.mpower(A, m) for m in range(1, N + 1)])
+
+    zero = cs.DM.zeros(ns, na)
+    Nnx = ns * (N - 1)
+    G_col = cs.vertcat(B, F[:Nnx, :] @ B)
+    G_cols = [G_col]
+    for _ in range(1, N):
+        G_col = cs.vertcat(zero, G_col[:Nnx, :])
+        G_cols.append(G_col)
+    G = cs.hcat(G_cols)
+
+    D_not_none = D is not None
+    if D_not_none:
+        zero = cs.DM.zeros(ns, D.shape[1])
+        H_col = cs.vertcat(D, F[:Nnx, :] @ D)
+        H_cols = [H_col]
+        for _ in range(1, N):
+            H_col = cs.vertcat(zero, H_col[:Nnx, :])
+            H_cols.append(H_col)
+        H = cs.hcat(H_cols)
+    else:
+        H = None
+    return F, G, H
+
+
 class Mpc(NonRetroactiveWrapper[SymType]):
     """A wrapper to easily turn an NLP scheme into an MPC controller. Most of the theory
     for MPC is taken from :cite:`rawlings_model_2017`.
@@ -500,34 +531,9 @@ def _set_singleshooting_linear_dynamics(
     ) -> tuple[MatType, MatType, Optional[MatType]]:
         """Internal utility to create linear dynamics constraints and states in
         single shooting mode."""
-        ns, na = B.shape
+        ns = A.shape[0]
         N = self._prediction_horizon
-
-        # create the QP matrices
-        F = cs.vcat([cs.mpower(A, m) for m in range(1, N + 1)])
-        #
-        zero = cs.DM.zeros(ns, na)
-        Nnx = ns * (N - 1)
-        G_col = cs.vertcat(B, F[:Nnx, :] @ B)
-        G_cols = [G_col]
-        for _ in range(1, N):
-            G_col = cs.vertcat(zero, G_col[:Nnx, :])
-            G_cols.append(G_col)
-        G = cs.hcat(G_cols)
-        #
-        D_not_none = D is not None
-        if D_not_none:
-            zero = cs.DM.zeros(ns, D.shape[1])
-            H_col = cs.vertcat(D, F[:Nnx, :] @ D)
-            H_cols = [H_col]
-            for _ in range(1, N):
-                H_col = cs.vertcat(zero, H_col[:Nnx, :])
-                H_cols.append(H_col)
-            H = cs.hcat(H_cols)
-        else:
-            H = None
-
-        # compute the next states
+        F, G, H = _create_qp_mats(self._prediction_horizon, A, B, D)
         x_0 = cs.vcat(self._initial_states.values())
         U = cs.vec(cs.vcat(self._actions_exp.values()))  # NOTE: different from vvcat!
         X_next = F @ x_0 + G @ U

src/csnlp/wrappers/mpc/scenario_based_mpc.py

@@ -8,7 +8,7 @@
 from csnlp.multistart.multistart_nlp import _chained_subevalf, _n
 
 from ..wrapper import Nlp
-from .mpc import Mpc, _callable2csfunc
+from .mpc import Mpc, _callable2csfunc, _create_qp_mats
 from .mpc import _n as _name_init_state
 
 SymType = TypeVar("SymType", cs.SX, cs.MX)
@@ -368,9 +368,31 @@ def set_nonlinear_dynamics(
     def _set_singleshooting_linear_dynamics(
         self, A: MatType, B: MatType, D: MatType
     ) -> tuple[MatType, MatType, Optional[MatType]]:
-        # TODO: recall what I did for the nonlinear single shooting case
-        # and check if the creation of F,G,H has to be extracted from the base class
-        raise NotImplementedError("not implemented yet.")
+        disturbance_names = self.single_disturbances.keys()
+        X0 = cs.vcat(self._initial_states.values())
+        U = cs.vec(cs.vcat(self._actions_exp.values()))  # NOTE: different from vvcat!
+        D_all = cs.hcat(
+            [
+                cs.vec(
+                    cs.vcat([self._disturbances[_n(n, i)] for n in disturbance_names])
+                )
+                for i in range(self._n_scenarios)
+            ]
+        )
+
+        F, G, H = _create_qp_mats(self._prediction_horizon, A, B, D)
+        X_next_pred = F @ X0 + G @ U + H @ D_all
+
+        state_names = self.single_states.keys()
+        N = self._prediction_horizon
+        ns = A.shape[0]
+        cumsizes = np.cumsum([0] + [s.shape[0] for s in self._initial_states.values()])
+        for i in range(self._n_scenarios):
+            X_i = cs.vertcat(X0, X_next_pred[:, i]).reshape((ns, N + 1))
+            X_i_split = cs.vertsplit(X_i, cumsizes)
+            for n, x in zip(state_names, X_i_split):
+                self._states[_n(n, i)] = x
+        return F, G, H
 
     def _set_multishooting_nonlinear_dynamics(
         self,