add z and sot output to solver and looper

nosnoc · Aug 31, 2023 · 6dc0c23 · 6dc0c23
1 parent 7be6c16
commit 6dc0c23
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 3 deletions.
diff --git a/nosnoc/helpers.py b/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,

diff --git a/nosnoc/problem.py b/nosnoc/problem.py
@@ -313,7 +313,7 @@ def __init__(self,
                         ca.SX.sym(f'alpha_{ctrl_idx}_{fe_idx}_{ii+1}_{ij+1}',
                                   dims.n_c_sys[ij]), self.ind_alpha,
                         lb_dual * np.ones(dims.n_c_sys[ij]), np.ones(dims.n_c_sys[ij]),
-                        opts.init_theta * np.ones(dims.n_c_sys[ij]), ii, ij)
+                        opts.init_alpha * np.ones(dims.n_c_sys[ij]), ii, ij)
                 # add lambda_n
                 for ij in range(dims.n_sys):
                     self.add_variable(

diff --git a/nosnoc/solver.py b/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]