Speed of time vars and z output (#52)

Added support for speed of time variables as in matlab, as well as added
z to output struct.

nosnoc/__init__.py

@@ -4,7 +4,7 @@
 from .model import NosnocModel
 from .ocp import NosnocOcp
 from .nosnoc_opts import NosnocOpts
-from .nosnoc_types import MpccMode, IrkSchemes, StepEquilibrationMode, CrossComplementarityMode, IrkRepresentation, PssMode, IrkRepresentation, HomotopyUpdateRule, InitializationStrategy, ConstraintHandling, Status
+from .nosnoc_types import MpccMode, IrkSchemes, StepEquilibrationMode, CrossComplementarityMode, IrkRepresentation, PssMode, IrkRepresentation, HomotopyUpdateRule, InitializationStrategy, ConstraintHandling, Status, SpeedOfTimeVariableMode
 from .helpers import NosnocSimLooper
 from .utils import casadi_length, casadi_vertcat_list, print_casadi_vector, flatten_layer, make_object_json_dumpable
 from .plot_utils import plot_timings, plot_iterates, latexify_plot

nosnoc/helpers.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 from .solver import NosnocSolver
+from .nosnoc_types import SpeedOfTimeVariableMode
 
 
 class NosnocSimLooper:
@@ -42,6 +43,8 @@ def __init__(self,
         self.theta_sim = []
         self.lambda_sim = []
         self.alpha_sim = []
+        self.z_sim = []
+        self.sot = []
         self.w_sim = []
         self.w_all = []
         self.cost_vals = []
@@ -80,10 +83,13 @@ def run(self, stop_on_failure=False) -> None:
             self.theta_sim.append(results["theta_list"])
             self.lambda_sim.append(results["lambda_list"])
             self.alpha_sim.append(results["alpha_list"])
+            self.z_sim.append(results["z_list"])
             self.w_sim += [results["w_sol"]]
             self.w_all += [results["w_all"]]
             self.cost_vals.append(results["cost_val"])
             self.status.append(results["status"])
+            if self.solver.opts.speed_of_time_variables != SpeedOfTimeVariableMode.NONE:
+                self.sot.append(results["sot"])
             if self.print_level > 0:
                 print(f"Sim step {i + 1}/{self.Nsim}\t status: {results['status']}")
 
@@ -102,6 +108,8 @@ def get_results(self) -> dict:
             "theta_sim": self.theta_sim,
             "lambda_sim": self.lambda_sim,
             "alpha_sim": self.alpha_sim,
+            "z_sim": self.z_sim,
+            "sot": self.sot,
             "w_sim": self.w_sim,
             "w_all": self.w_all,
             "cost_vals": self.cost_vals,

nosnoc/nosnoc_opts.py

@@ -5,7 +5,7 @@
 
 from .rk_utils import generate_butcher_tableu, generate_butcher_tableu_integral
 from .utils import validate
-from .nosnoc_types import MpccMode, IrkSchemes, StepEquilibrationMode, CrossComplementarityMode, IrkRepresentation, PssMode, IrkRepresentation, HomotopyUpdateRule, InitializationStrategy, ConstraintHandling
+from .nosnoc_types import MpccMode, IrkSchemes, StepEquilibrationMode, CrossComplementarityMode, IrkRepresentation, PssMode, IrkRepresentation, HomotopyUpdateRule, InitializationStrategy, ConstraintHandling, SpeedOfTimeVariableMode
 
 
 def _assign(dictionary, keys, value):
@@ -104,6 +104,11 @@ class NosnocOpts:
     time_freezing: bool = False
     time_freezing_tolerance: float = 1e-3
 
+    # Speed of time handling
+    speed_of_time_variables: SpeedOfTimeVariableMode = SpeedOfTimeVariableMode.NONE
+    speed_of_time_min: float = 1.0
+    speed_of_time_max: float = 25.0
+
     rootfinder_for_initial_z: bool = False
 
     # Usabillity:

nosnoc/nosnoc_types.py

@@ -82,6 +82,12 @@ class ConstraintHandling(Enum):
     LEAST_SQUARES = auto()
 
 
+class SpeedOfTimeVariableMode(Enum):
+    NONE = auto()    # No speed of time variables
+    LOCAL = auto()   # Speed of time variables as control stage discontinous variables
+    GLOBAL = auto()  # Single speed of time variable used across whole problem
+
+
 class PssMode(Enum):
     """
     Mode to represent the Piecewise Smooth System (PSS).

nosnoc/problem.py

@@ -7,7 +7,7 @@
 
 from nosnoc.model import NosnocModel
 from nosnoc.nosnoc_opts import NosnocOpts
-from nosnoc.nosnoc_types import MpccMode, CrossComplementarityMode, StepEquilibrationMode, PssMode, IrkRepresentation, ConstraintHandling
+from nosnoc.nosnoc_types import MpccMode, CrossComplementarityMode, StepEquilibrationMode, PssMode, IrkRepresentation, ConstraintHandling, SpeedOfTimeVariableMode
 from nosnoc.ocp import NosnocOcp
 from nosnoc.utils import casadi_length, casadi_vertcat_list, casadi_sum_list, flatten, increment_indices, create_empty_list_matrix
 
@@ -432,7 +432,7 @@ def X_fe(self) -> List[ca.SX]:
 
         return X_fe
 
-    def forward_simulation(self, ocp: NosnocOcp, Uk: ca.SX) -> None:
+    def forward_simulation(self, ocp: NosnocOcp, Uk: ca.SX, sot: ca.SX) -> None:
         opts = self.opts
         model = self.model
 
@@ -453,8 +453,8 @@ def forward_simulation(self, ocp: NosnocOcp, Uk: ca.SX) -> None:
 
         for j in range(opts.n_s):
             # Dynamics excluding complementarities
-            fj = model.f_x_fun(X_fe[j], self.rk_stage_z(j), Uk, self.p, model.v_global)
-            qj = ocp.f_q_fun(X_fe[j], Uk, self.p, model.v_global)
+            fj = sot*model.f_x_fun(X_fe[j], self.rk_stage_z(j), Uk, self.p, model.v_global)
+            qj = sot*ocp.f_q_fun(X_fe[j], Uk, self.p, model.v_global)
             gqj = ocp.g_path_fun(X_fe[j], Uk, self.p, model.v_global)
             self.add_constraint(gqj, ocp.lbg, ocp.ubg)
 
@@ -569,6 +569,11 @@ def __create_control_stage(self, ctrl_idx, prev_fe):
         Uk = ca.SX.sym(f'U_{ctrl_idx}', self.model.dims.n_u)
         self.add_variable(Uk, self.ind_u, self.ocp.lbu, self.ocp.ubu, self.ocp.u_guess)
 
+        # Create stage local speed of time variables
+        if self.opts.speed_of_time_variables == SpeedOfTimeVariableMode.LOCAL:
+            sot = ca.SX.sym(f'sot_{ctrl_idx}', 1)
+            self.add_variable(sot, self.ind_sot, [self.opts.speed_of_time_min], [self.opts.speed_of_time_max], [1.0])
+
         # Create Finite elements in this control stage
         control_stage = []
         for ii in range(self.opts.Nfe_list[ctrl_idx]):
@@ -597,6 +602,11 @@ def __create_primal_variables(self):
         self.add_variable(self.model.v_global, self.ind_v_global, self.ocp.lbv_global,
                           self.ocp.ubv_global, self.ocp.v_global_guess)
 
+        # add global speed of time if necessary
+        if self.opts.speed_of_time_variables == SpeedOfTimeVariableMode.GLOBAL:
+            sot = ca.SX.sym('sot', 1)
+            self.add_variables(sot, self.ind_sot, self.opts.speed_of_time_min, self.opts.speed_of_time_max, 1.0)
+
         # Generate control_stages
         prev_fe = self.fe0
         for ii in range(self.opts.N_stages):
@@ -682,6 +692,7 @@ def __init__(self, opts: NosnocOpts, model: NosnocModel, ocp: Optional[NosnocOcp
         self.ind_elastic = []
 
         self.ind_u = []
+        self.ind_sot = []
         self.ind_h = []
         self.ind_v_global = []
 
@@ -714,10 +725,17 @@ def __init__(self, opts: NosnocOpts, model: NosnocModel, ocp: Optional[NosnocOcp
 
         for ctrl_idx, stage in enumerate(self.stages):
             Uk = self.w[self.ind_u[ctrl_idx]]
+            if self.opts.speed_of_time_variables == SpeedOfTimeVariableMode.GLOBAL:
+                sot = self.w[self.ind_sot[0]]
+            elif self.opts.speed_of_time_variables == SpeedOfTimeVariableMode.LOCAL:
+                sot = self.w[self.ind_sot[ctrl_idx]]
+            else:
+                sot = ca.SX.eye(1)
+
             for _, fe in enumerate(stage):
 
                 # 1) Stewart Runge-Kutta discretization
-                fe.forward_simulation(ocp, Uk)
+                fe.forward_simulation(ocp, Uk, sot)
 
                 # 2) Complementarity Constraints
                 fe.create_complementarity_constraints(sigma_p, tau, Uk, s_elastic)

nosnoc/solver.py

@@ -1,4 +1,3 @@
-
 from abc import ABC, abstractmethod
 from typing import Optional
 
@@ -8,7 +7,7 @@
 
 from nosnoc.model import NosnocModel
 from nosnoc.nosnoc_opts import NosnocOpts
-from nosnoc.nosnoc_types import InitializationStrategy, PssMode, HomotopyUpdateRule, ConstraintHandling, Status
+from nosnoc.nosnoc_types import InitializationStrategy, PssMode, HomotopyUpdateRule, ConstraintHandling, Status, SpeedOfTimeVariableMode
 from nosnoc.ocp import NosnocOcp
 from nosnoc.problem import NosnocProblem
 from nosnoc.rk_utils import rk4_on_timegrid
@@ -482,6 +481,8 @@ def get_results_from_primal_vector(prob: NosnocProblem, w_opt: np.ndarray) -> di
     ind_x_all = flatten_outer_layers(prob.ind_x, 2)
     results["x_all_list"] = [x0] + [w_opt[np.array(ind)] for ind in ind_x_all]
     results["u_list"] = [w_opt[ind] for ind in prob.ind_u]
+    if opts.speed_of_time_variables != SpeedOfTimeVariableMode.NONE:
+        results["sot"] = [w_opt[ind] for ind in prob.ind_sot]
 
     results["theta_list"] = [w_opt[ind] for ind in get_cont_algebraic_indices(prob.ind_theta)]
     results["lambda_list"] = [w_opt[ind] for ind in get_cont_algebraic_indices(prob.ind_lam)]
@@ -496,6 +497,7 @@ def get_results_from_primal_vector(prob: NosnocProblem, w_opt: np.ndarray) -> di
     results["lambda_p_list"] = [
         w_opt[flatten_layer(ind)] for ind in get_cont_algebraic_indices(prob.ind_lambda_p)
     ]
+    results["z_list"] = [w_opt[ind] for ind in get_cont_algebraic_indices(prob.ind_z)]
 
     if opts.use_fesd:
         time_steps = w_opt[prob.ind_h]