Add checks, docstrings, parameter updates (#14)

* check if problem is pure simulation

* add n_p* to dims, remove default values

* add option to update p_global

* add p_values and docstrings to NosnocSimLooper

* make p_time_var optional for model, allow update

* update parametric OCP example

examples/parametric_cart_pole_with_friction.py

@@ -103,7 +103,7 @@ def solve_paramteric_example(with_global_var=False):
 
     model = nosnoc.NosnocModel(x=x, F=F, S=S, c=c, x0=x0, u=u,
                                p_global=p_global, p_global_val=p_global_val,
-                               p_time_var=p_time_var, p_time_var_val=p_time_var_val,
+                               p_time_var=p_time_var,
                                v_global=v_global
     )
 
@@ -113,13 +113,16 @@ def solve_paramteric_example(with_global_var=False):
     ocp = nosnoc.NosnocOcp(lbu=lbu, ubu=ubu, f_q=f_q, f_terminal=f_terminal, g_terminal=g_terminal,
                            lbx=lbx, ubx=ubx, lbv_global=lbv_global, ubv_global=ubv_global, v_global_guess=v_global_guess)
 
-    ## Solve OCP
+    # create solver
     solver = nosnoc.NosnocSolver(opts, model, ocp)
+    # set / update parameters
+    solver.set('p_time_var', p_time_var_val)
+    solver.set('p_global', p_global_val)
 
+    # solve OCP
     results = solver.solve()
     return results
 
-
 def main():
     results = solve_paramteric_example()
     plot_results(results)

nosnoc/helpers.py

@@ -1,10 +1,29 @@
+from typing import Optional
+
 import numpy as np
 from .nosnoc import NosnocSolver
 
 
 class NosnocSimLooper:
 
-    def __init__(self, solver: NosnocSolver, x0: np.ndarray, Nsim: int):
+    def __init__(self, solver: NosnocSolver, x0: np.ndarray, Nsim: int, p_values: Optional[np.ndarray]=None):
+        """
+        :param solver: NosnocSolver to be called in a loop
+        :param x0: np.ndarray: initial state
+        :param Nsim: int: number of simulation steps
+        :param: p_values: Optional np.ndarray of shape (Nsim, n_p_glob), parameter values p_glob are updated at each simulation step accordingly.
+        """
+        # check that NosnocSolver solves a pure simulation problem.
+        if not solver.problem.is_sim_problem():
+            raise Exception("NosnocSimLooper can only be used with pure simulation problem")
+
+        # p values
+        self.p_values = p_values
+        if self.p_values is not None:
+            if self.p_values.shape != (Nsim, solver.problem.model.dims.n_p_glob):
+                raise ValueError("p_values should have shape (Nsim, n_p_glob).")
+
+        # create
         self.solver = solver
         self.Nsim = Nsim
 
@@ -21,7 +40,11 @@ def __init__(self, solver: NosnocSolver, x0: np.ndarray, Nsim: int):
 
     def run(self) -> None:
         for i in range(self.Nsim):
+            # set values
             self.solver.set("x", self.xcurrent)
+            if self.p_values is not None:
+                self.solver.set("p_global", self.p_values[i, :])
+            # solve
             results = self.solver.solve()
             # collect
             self.X_sim += results["x_list"]

nosnoc/nosnoc.py

@@ -33,7 +33,7 @@ def __init__(self,
                  u: SX = SX.sym('u_dummy', 0, 1),
                  p_time_var: SX = SX.sym('p_tim_var_dummy', 0, 1),
                  p_global: SX = SX.sym('p_global_dummy', 0, 1),
-                 p_time_var_val: np.ndarray = None,
+                 p_time_var_val: Optional[np.ndarray] = None,
                  p_global_val: np.ndarray = np.array([]),
                  v_global: SX = SX.sym('v_global_dummy', 0, 1),
                  name: str = 'nosnoc'):
@@ -66,8 +66,6 @@ def preprocess_model(self, opts: NosnocOpts):
         n_c_sys = [casadi_length(self.c[i]) for i in range(n_sys)]
         n_f_sys = [self.F[i].shape[1] for i in range(n_sys)]
 
-        self.dims = NosnocDims(n_x=n_x, n_u=n_u, n_sys=n_sys, n_c_sys=n_c_sys, n_f_sys=n_f_sys)
-
         # sanity checks
         if not isinstance(self.F, list):
             raise ValueError("model.F should be a list.")
@@ -98,6 +96,9 @@ def preprocess_model(self, opts: NosnocOpts):
             self.p_val_ctrl_stages[i, :n_p_time_var] = self.p_time_var_val[i, :]
             self.p_val_ctrl_stages[i, n_p_time_var:] = self.p_global_val
 
+        self.dims = NosnocDims(n_x=n_x, n_u=n_u, n_sys=n_sys, n_c_sys=n_c_sys, n_f_sys=n_f_sys,
+                               n_p_time_var=n_p_time_var, n_p_glob=n_p_glob)
+
         # g_Stewart
         g_Stewart_list = [-self.S[i] @ self.c[i] for i in range(n_sys)]
         g_Stewart = casadi_vertcat_list(g_Stewart_list)
@@ -301,12 +302,13 @@ class NosnocDims:
     """
     detected automatically
     """
-    n_x: int = 0
-    n_u: int = 0
-    n_sys: int = 0
-    n_c_sys: list = field(default_factory=list)
-    n_f_sys: list = field(default_factory=list)
-
+    n_x: int
+    n_u: int
+    n_sys: int
+    n_p_time_var: int
+    n_p_glob: int
+    n_c_sys: list
+    n_f_sys: list
 
 class NosnocFormulationObject(ABC):
 
@@ -801,8 +803,14 @@ def __init__(self, opts: NosnocOpts, model: NosnocModel, ocp: Optional[NosnocOcp
         super().__init__()
 
         self.model = model
-        self.ocp = ocp
         self.opts = opts
+        if ocp is None:
+            self.ocp_trivial = True
+            ocp = NosnocOcp()
+        else:
+            self.ocp_trivial = False
+        ocp.preprocess_ocp(model)
+        self.ocp = ocp
 
         h_ctrl_stage = opts.terminal_time / opts.N_stages
         self.stages: list[list[FiniteElementBase]] = []
@@ -891,6 +899,15 @@ def print(self):
         print(f"cost:\n{self.cost}")
 
 
+    def is_sim_problem(self):
+        if self.model.dims.n_u != 0:
+            return False
+        if self.opts.N_stages != 1:
+            return False
+        if not self.ocp_trivial:
+            return False
+        return True
+
 def get_results_from_primal_vector(prob: NosnocProblem, w_opt: np.ndarray) -> dict:
     opts = prob.opts
 
@@ -940,11 +957,8 @@ class NosnocSolver():
     def __init__(self, opts: NosnocOpts, model: NosnocModel, ocp: Optional[NosnocOcp] = None):
 
         # preprocess inputs
-        if ocp is None:
-            ocp = NosnocOcp()
         opts.preprocess()
         model.preprocess_model(opts)
-        ocp.preprocess_ocp(model)
 
         if opts.initialization_strategy == InitializationStrategy.RK4_SMOOTHENED:
             model.add_smooth_step_representation(smoothing_parameter=opts.smoothing_parameter)
@@ -1118,13 +1132,24 @@ def solve(self) -> dict:
         return results
 
     # TODO: move this to problem?
-    def set(self, field: str, value):
+    def set(self, field: str, value: np.ndarray) -> None:
+        """
+        :param field: in ["x", "p_global", "p_time_var"]
+        :param value: np.ndarray: numerical value of appropriate size
+        """
         prob = self.problem
+        dims = prob.model.dims
         if field == 'x':
             ind_x0 = prob.ind_x[0]
             prob.w0[ind_x0] = value
             prob.lbw[ind_x0] = value
             prob.ubw[ind_x0] = value
+        elif field == 'p_global':
+            for i in range(self.opts.N_stages):
+                self.model.p_val_ctrl_stages[i, dims.n_p_time_var:] = value
+        elif field == 'p_time_var':
+            for i in range(self.opts.N_stages):
+                self.model.p_val_ctrl_stages[i, :dims.n_p_time_var] = value[i, :]
         else:
             raise NotImplementedError()