on interleaved/rolling variables

src/csnlp/wrappers/mpc/mpc.py

@@ -1,7 +1,7 @@
-from collections.abc import Sequence
+from collections.abc import Collection, Generator
 from inspect import signature
 from math import ceil
-from typing import Callable, Literal, Optional, TypeVar, Union
+from typing import Any, Callable, Literal, Optional, TypeVar, Union
 
 import casadi as cs
 import numpy as np
@@ -589,3 +589,138 @@ def _set_singleshooting_nonlinear_dynamics(
         X = cs.horzcat(X0, X_next)
         cumsizes = np.cumsum([0] + [s.shape[0] for s in self._initial_states.values()])
         self._states = dict(zip(self._states.keys(), cs.vertsplit(X, cumsizes)))
+
+    def rolling_quantities(
+        self,
+        dynamics: Union[
+            cs.Function,
+            Callable[[tuple[npt.ArrayLike, ...]], tuple[npt.ArrayLike, ...]],
+        ],
+        states_kwargs: Union[dict[str, Any], Collection[dict[str, Any]]],
+        actions_kwargs: Union[dict[str, Any], Collection[dict[str, Any]]],
+        disturbances_kwargs: Union[
+            None, dict[str, Any], Collection[dict[str, Any]]
+        ] = None,
+    ) -> Generator[
+        tuple[dict[str, Union[SymType, tuple[SymType, ...]]], ...], None, None
+    ]:
+        if self._dynamics is not None:
+            raise RuntimeError("Dynamics were already set.")
+        X = []
+        if isinstance(states_kwargs, dict):
+            states_kwargs = (states_kwargs,)
+        U, U_exp = [], []
+        if isinstance(actions_kwargs, dict):
+            actions_kwargs = (actions_kwargs,)
+        disturbed = disturbances_kwargs is not None
+        if disturbed:
+            if isinstance(disturbances_kwargs, dict):
+                disturbances_kwargs = (disturbances_kwargs,)
+            D = []
+
+        # create initial states' parameters
+        names = [kw["name"] for kw in states_kwargs]
+        sizes = [kw.get("size", 1) for kw in states_kwargs]
+        cumsizes = np.cumsum([0] + sizes)
+        states = {n: self.nlp.parameter(_n(n), (s, 1)) for n, s in zip(names, sizes)}
+        x = cs.vcat(states.values())
+        initial_states = {_n(n): s for n, s in states.items()}
+        if not self._is_multishooting:
+            X.append(x)
+
+        for k in range(self._prediction_horizon):
+            # create states - if single shooting, the state is given by the last state
+            # propagated via the dynamics; otherwise, if multiple shooting, create a new
+            # one and impose the dynamics via an equality constraint
+            if self._is_multishooting:
+                states = {}
+                for kw in states_kwargs:
+                    name = kw["name"]
+                    if k == 0 and not kw.get("bound_initial", True):
+                        lb, ub = -np.inf, +np.inf
+                    else:
+                        lb, ub = kw.get("lb", -np.inf), kw.get("ub", +np.inf)
+                    states[name] = self.nlp.variable(
+                        f"{name}{k}", (kw.get("size", 1), 1), lb, ub
+                    )
+                x_prev = x
+                x = cs.vcat([s[0] for s in states.values()])
+                self.nlp.constraint(f"dyn{k}", x_prev, "==", x)
+                X.append(x)
+
+            # create actions - do so only if we are still within the control horizon,
+            # and on a multiple of the spacing; otherwise, use last actions
+            if k % self._input_spacing == 0 and k < self._control_horizon:
+                actions = {
+                    kw["name"]: self.nlp.variable(
+                        f"{kw['name']}{k}",
+                        (kw.get("size", 1), 1),
+                        kw.get("lb", -np.inf),
+                        kw.get("ub", +np.inf),
+                    )
+                    for kw in actions_kwargs
+                }
+                u = cs.vcat([a[0] for a in actions.values()])
+                U.append(u)
+            U_exp.append(u)
+
+            # create disturbances
+            if disturbed:
+                disturbances = {
+                    kw["name"]: self.nlp.parameter(
+                        f"{kw['name']}{k}", (kw.get("size", 1), 1)
+                    )
+                    for kw in disturbances_kwargs
+                }
+                d = cs.vcat(disturbances.values())
+                D.append(d)
+
+            yield (states, actions, disturbances) if disturbed else (states, actions)
+
+            # compute next state via the dynamics
+            x = dynamics(x, u, d) if disturbed else dynamics(x, u)
+            if isinstance(x, (tuple, list)):
+                x = x[0]
+            if not self._is_multishooting:
+                states = dict(zip(names, cs.vertsplit(x, cumsizes)))
+                X.append(x)
+
+        # yield also the terminal state
+        if self._is_multishooting:
+            states = {}
+            for kw in states_kwargs:
+                name = kw["name"]
+                if not kw.get("bound_terminal", True):
+                    lb, ub = -np.inf, +np.inf
+                else:
+                    lb, ub = kw.get("lb", -np.inf), kw.get("ub", +np.inf)
+                states[name] = self.nlp.variable(
+                    f"{name}{k + 1}", (kw.get("size", 1), 1), lb, ub
+                )
+            x_prev = x
+            x = cs.vcat([s[0] for s in states.values()])
+            self.nlp.constraint(f"dyn{k + 1}", x_prev, "==", x)
+            X.append(x)
+
+        yield (states, {}, {}) if disturbed else (states, {})
+
+        # save all rolled-out lists of variables to internal dictionaries
+        X = cs.hcat(X)
+        self._states.update(zip(names, cs.vertsplit(X, cumsizes)))
+        self._initial_states.update(initial_states)
+
+        U, U_exp = cs.hcat(U), cs.hcat(U_exp)
+        names = [kw["name"] for kw in actions_kwargs]
+        cumsizes = np.cumsum([0] + [kw.get("size", 1) for kw in actions_kwargs])
+        self._actions.update(zip(names, cs.vertsplit(U, cumsizes)))
+        self._actions_exp.update(zip(names, cs.vertsplit(U_exp, cumsizes)))
+
+        if disturbed:
+            D = cs.hcat(D)
+            names = (kw["name"] for kw in disturbances_kwargs)
+            cumsizes = np.cumsum(
+                [0] + [kw.get("size", 1) for kw in disturbances_kwargs]
+            )
+            self._disturbances.update(zip(names, cs.vertsplit(D, cumsizes)))
+
+        self._dynamics = dynamics

tests/test_wrappers_mpc.py

@@ -354,6 +354,135 @@ def test_can_be_pickled(self, sym_type: str):
 
         self.assertIn(repr(mpc), repr(mpc2))
 
+    @parameterized.expand([("single",), ("multi",)])
+    def test_rolling_quantities(self, shooting: str):
+        Np = np.random.randint(10, 20)
+        Nc = np.random.randint((Np - 2) // 2, Np - 2)
+        space = np.random.randint(1, Nc // 2)
+        mpc = Mpc[cs.SX](
+            nlp=Nlp(sym_type="SX"),
+            prediction_horizon=Np,
+            control_horizon=Nc,
+            input_spacing=space,
+            shooting=shooting,
+        )
+        nxs = (2, 3)
+        nus = (1, 2)
+        nd = sum(nxs)
+        A = np.full((nd, nd), 1.0)
+        B = np.full((nd, sum(nus)), 1.0)
+        F = lambda x, u, d: (A @ x + B @ u + d, object())
+        x_kws = [
+            {
+                "name": "x0",
+                "size": nxs[0],
+                "lb": -1,
+                "ub": 2,
+                "bound_initial": False,
+                "bound_terminal": True,
+            },
+            {
+                "name": "x1",
+                "size": nxs[1],
+                "lb": -3,
+                "ub": 4,
+                "bound_initial": True,
+                "bound_terminal": False,
+            },
+        ]
+        u_kws = [
+            {"name": "u0", "size": nus[0], "lb": -5, "ub": 6},
+            {"name": "u1", "size": nus[1], "lb": -7, "ub": 8},
+        ]
+        d_kws = {"name": "d", "size": nd}
+
+        # checkout rolling outputs
+        generator = mpc.rolling_quantities(F, x_kws, u_kws, d_kws)
+        for k, (states, actions, disturbances) in enumerate(generator):
+            for i in range(len(nxs)):
+                name = f"x{i}"
+                self.assertTrue(name in states)
+                if shooting == "single":
+                    state = states[name]
+                    self.assertTupleEqual(state.shape, (nxs[i], 1))
+                else:
+                    state, lam_lbx, lam_ubx = states[name]
+                    expected_shape = (
+                        (0, 1)
+                        if (k == 0 and not x_kws[i]["bound_initial"])
+                        or (k == Np and not x_kws[i]["bound_terminal"])
+                        else (nxs[i], 1)
+                    )
+                    self.assertTupleEqual(state.shape, (nxs[i], 1))
+                    self.assertTupleEqual(lam_lbx.shape, expected_shape)
+                    self.assertTupleEqual(lam_ubx.shape, expected_shape)
+
+            if k < Np:
+                for i in range(len(nus)):
+                    name = f"u{i}"
+                    self.assertTrue(name in actions)
+                    action, lam_lbu, lam_ubu = actions[name]
+                    expected_shape = (nus[i], 1)
+                    self.assertTupleEqual(action.shape, expected_shape)
+                    self.assertTupleEqual(lam_lbu.shape, expected_shape)
+                    self.assertTupleEqual(lam_ubu.shape, expected_shape)
+                self.assertTrue("d" in disturbances)
+                self.assertTupleEqual(disturbances["d"].shape, (nd, 1))
+            else:
+                self.assertEqual(len(actions), 0)
+                self.assertEqual(len(disturbances), 0)
+
+        # check dict properties
+        for i in range(len(nxs)):
+            name = f"x{i}"
+            self.assertTrue(name in mpc.states)
+            self.assertTupleEqual(mpc.states[name].shape, (nxs[i], Np + 1))
+        for i in range(len(nus)):
+            name = f"u{i}"
+            self.assertTrue(name in mpc.actions)
+            self.assertTupleEqual(mpc.actions[name].shape, (nus[i], ceil(Nc / space)))
+            self.assertTrue(name in mpc.actions_expanded)
+            self.assertTupleEqual(mpc.actions_expanded[name].shape, (nus[i], Np))
+        self.assertTrue("d" in mpc.disturbances)
+        self.assertTupleEqual(mpc.disturbances["d"].shape, (nd, Np))
+
+        # check lb and ub on primal variables
+        lbx = [np.full(nx, kw["lb"]) for nx, kw in zip(nxs, x_kws)]
+        ubx = [np.full(nx, kw["ub"]) for nx, kw in zip(nxs, x_kws)]
+        lbu = np.concatenate([np.full(nu, kw["lb"]) for nu, kw in zip(nus, u_kws)])
+        ubu = np.concatenate([np.full(nu, kw["ub"]) for nu, kw in zip(nus, u_kws)])
+        if shooting == "single":
+            np.testing.assert_array_equal(np.tile(lbu, ceil(Nc / space)), mpc.lbx.data)
+            np.testing.assert_array_equal(np.tile(ubu, ceil(Nc / space)), mpc.ubx.data)
+        else:
+            LBX, UBX = [], []
+            for k in range(Np):
+                if k == 0:
+                    LBX.extend(
+                        np.full(nx, kw["lb"] if kw["bound_initial"] else -np.inf)
+                        for nx, kw in zip(nxs, x_kws)
+                    )
+                    UBX.extend(
+                        np.full(nx, kw["ub"] if kw["bound_initial"] else +np.inf)
+                        for nx, kw in zip(nxs, x_kws)
+                    )
+                else:
+                    LBX.extend(lbx)
+                    UBX.extend(ubx)
+                if k % space == 0 and k < Nc:
+                    LBX.append(lbu)
+                    UBX.append(ubu)
+            LBX.extend(
+                np.full(nx, kw["lb"] if kw["bound_terminal"] else -np.inf)
+                for nx, kw in zip(nxs, x_kws)
+            )
+            UBX.extend(
+                np.full(nx, kw["ub"] if kw["bound_terminal"] else +np.inf)
+                for nx, kw in zip(nxs, x_kws)
+            )
+            np.testing.assert_array_equal(np.concatenate(LBX), mpc.lbx.data)
+            np.testing.assert_array_equal(np.concatenate(UBX), mpc.ubx.data)
+
 
 class TestPwaMpc(unittest.TestCase):
     def test_pwa_dynamics__raises__if_dynamics_already_set(self):