Cleanup casadi import and run formatter (#26)

* import casadi as ca * run formatter on nosnoc and test
nosnoc · Jan 21, 2023 · 3d18ca7 · 3d18ca7
1 parent 4bdad94
commit 3d18ca7
Show file tree

Hide file tree

Showing 13 changed files with 357 additions and 296 deletions.
diff --git a/nosnoc/helpers.py b/nosnoc/helpers.py
@@ -6,7 +6,12 @@
 
 class NosnocSimLooper:
 
-    def __init__(self, solver: NosnocSolver, x0: np.ndarray, Nsim: int, p_values: Optional[np.ndarray]=None, w_init: Optional[list]=None):
+    def __init__(self,
+                 solver: NosnocSolver,
+                 x0: np.ndarray,
+                 Nsim: int,
+                 p_values: Optional[np.ndarray] = None,
+                 w_init: Optional[list] = None):
         """
         :param solver: NosnocSolver to be called in a loop
         :param x0: np.ndarray: initial state
@@ -39,7 +44,7 @@ def __init__(self, solver: NosnocSolver, x0: np.ndarray, Nsim: int, p_values: Op
         self.cost_vals = []
         self.w_init = w_init
 
-        self.cpu_nlp = np.zeros((Nsim, solver.opts.max_iter_homotopy+1))
+        self.cpu_nlp = np.zeros((Nsim, solver.opts.max_iter_homotopy + 1))
 
     def run(self) -> None:
         """Run the simulation loop."""
@@ -64,7 +69,6 @@ def run(self) -> None:
             self.w_all += [results["w_all"]]
             self.cost_vals.append(results["cost_val"])
 
-
     def get_results(self) -> dict:
         self.t_grid = np.concatenate((np.array([0.0]), np.cumsum(self.time_steps)))
         results = {

diff --git a/nosnoc/nosnoc.py b/nosnoc/nosnoc.py
diff --git a/nosnoc/nosnoc_opts.py b/nosnoc/nosnoc_opts.py
@@ -140,14 +140,22 @@ def preprocess(self):
             self.right_boundary_point_explicit = False
 
         # checks:
-        if (self.cross_comp_mode == CrossComplementarityMode.SUM_LAMBDAS_COMPLEMENT_WITH_EVERY_THETA and
-            self.mpcc_mode in [MpccMode.FISCHER_BURMEISTER, MpccMode.FISCHER_BURMEISTER_IP_AUG]):
-            Warning("UNSUPPORTED option combination comp_mode: SUM_LAMBDAS_COMPLEMENT_WITH_EVERY_THETA and mpcc_mode: MpccMode.FISCHER_BURMEISTER")
+        if (self.cross_comp_mode == CrossComplementarityMode.SUM_LAMBDAS_COMPLEMENT_WITH_EVERY_THETA
+                and self.mpcc_mode
+                in [MpccMode.FISCHER_BURMEISTER, MpccMode.FISCHER_BURMEISTER_IP_AUG]):
+            Warning(
+                "UNSUPPORTED option combination comp_mode: SUM_LAMBDAS_COMPLEMENT_WITH_EVERY_THETA and mpcc_mode: MpccMode.FISCHER_BURMEISTER"
+            )
         if self.mpcc_mode == MpccMode.FISCHER_BURMEISTER and self.constraint_handling != ConstraintHandling.LEAST_SQUARES:
-            Warning("UNSUPPORTED option combination comp_mode: mpcc_mode == MpccMode.FISCHER_BURMEISTER and constraint_handling != ConstraintHandling.LEAST_SQUARES")
-        if (self.step_equilibration in
-            [StepEquilibrationMode.DIRECT, StepEquilibrationMode.DIRECT_COMPLEMENTARITY] and self.constraint_handling != ConstraintHandling.LEAST_SQUARES):
-            Warning("UNSUPPORTED option combination: StepEquilibrationMode.DIRECT* and constraint_handling != ConstraintHandling.LEAST_SQUARES")
+            Warning(
+                "UNSUPPORTED option combination comp_mode: mpcc_mode == MpccMode.FISCHER_BURMEISTER and constraint_handling != ConstraintHandling.LEAST_SQUARES"
+            )
+        if (self.step_equilibration
+                in [StepEquilibrationMode.DIRECT, StepEquilibrationMode.DIRECT_COMPLEMENTARITY] and
+                self.constraint_handling != ConstraintHandling.LEAST_SQUARES):
+            Warning(
+                "UNSUPPORTED option combination: StepEquilibrationMode.DIRECT* and constraint_handling != ConstraintHandling.LEAST_SQUARES"
+            )
         return
 
     ## Options in matlab..

diff --git a/nosnoc/nosnoc_types.py b/nosnoc/nosnoc_types.py
@@ -20,7 +20,7 @@ class IrkSchemes(Enum):
 class InitializationStrategy(Enum):
     ALL_XCURRENT_W0_START = auto()
     ALL_XCURRENT_WOPT_PREV = auto()
-    EXTERNAL = auto() # let user do from outsid
+    EXTERNAL = auto()  # let user do from outsid
     RK4_SMOOTHENED = auto()  # experimental
     # Other ideas
     # OLD_SOLUTION = auto()
@@ -59,6 +59,7 @@ class ConstraintHandling(Enum):
     EXACT = auto()
     LEAST_SQUARES = auto()
 
+
 class PssMode(Enum):
     """
     Mode to represent the Piecewise Smooth System (PSS).

diff --git a/nosnoc/utils.py b/nosnoc/utils.py
@@ -55,7 +55,7 @@ def casadi_sum_list(input: list):
 
 
 # Note this is not generalized, it expects equivalent depth, greater than `layer`
-def flatten_layer(L: list, layer: int=0):
+def flatten_layer(L: list, layer: int = 0):
     if layer == 0:
         # Check if already flat
         if any(isinstance(e, list) for e in L):
@@ -73,6 +73,7 @@ def flatten(L):
     else:
         return [L]
 
+
 def flatten_outer_layers(L: list, n_layers: int):
     for _ in range(n_layers):
         L = flatten_layer(L, 0)

diff --git a/test/simple_sim_tests.py b/test/simple_sim_tests.py
@@ -56,7 +56,6 @@ def test_opts(opts, model):
     assert errors["x_end"] < tol
 
 
-
 class SimpleTests(unittest.TestCase):
 
     def test_default(self):
@@ -166,7 +165,6 @@ def test_least_squares_problem(self):
             # TODO: Fix test!
             pass
 
-
     def test_initializations(self):
         model = get_simplest_model_switch()
 

diff --git a/test/test_initializing.py b/test/test_initializing.py
@@ -24,30 +24,23 @@ def test_initialization(self):
         u2_est[-2] = -2.0
         u = np.concatenate((u1_est, u2_est), axis=1)
         X0 = X0.reshape((1, -1))
-        x_est = np.concatenate((
-            X0 * 2 / 3, X0 * 1 / 3, np.zeros((opts.N_stages - 2, 4))
-        ))
+        x_est = np.concatenate((X0 * 2 / 3, X0 * 1 / 3, np.zeros((opts.N_stages - 2, 4))))
 
         opts.initialization_strategy = nosnoc.InitializationStrategy.EXTERNAL
         solver = nosnoc.NosnocSolver(opts, model, ocp)
         solver2 = nosnoc.NosnocSolver(opts2, model2, ocp2)
         solver.set("x", x_est)
         for i in range(opts.N_stages):
-            self.assertTrue(np.array_equal(
-                solver.problem.w0[solver.problem.ind_x[i][0][0]], x_est[i, :]
-            ))
+            self.assertTrue(
+                np.array_equal(solver.problem.w0[solver.problem.ind_x[i][0][0]], x_est[i, :]))
 
         solver.set("u", u)
-        self.assertTrue(np.array_equal(
-            solver.problem.w0[solver.problem.ind_u], u
-        ))
+        self.assertTrue(np.array_equal(solver.problem.w0[solver.problem.ind_u], u))
 
         res_initialized = solver.solve()
         res_uninitialized = solver2.solve()
         # If initialized, the iteration should be faster
-        self.assertLessEqual(
-            res_initialized["nlp_iter"][0], res_uninitialized["nlp_iter"][0]
-        )
+        self.assertLessEqual(res_initialized["nlp_iter"][0], res_uninitialized["nlp_iter"][0])
 
 
 if __name__ == "__main__":

diff --git a/test/test_ocp.py b/test/test_ocp.py
@@ -16,21 +16,24 @@
 EQUIDISTANT_CONTROLS = [True, False]
 PSS_MODES = [nosnoc.PssMode.STEWART]
 options = [
-    (equidistant_control_grid, step_equilibration, irk_representation, irk_scheme, pss_mode, nosnoc.HomotopyUpdateRule.LINEAR)
-    for equidistant_control_grid in EQUIDISTANT_CONTROLS
-    for step_equilibration in [nosnoc.StepEquilibrationMode.HEURISTIC_MEAN, nosnoc.StepEquilibrationMode.HEURISTIC_DELTA, nosnoc.StepEquilibrationMode.L2_RELAXED, nosnoc.StepEquilibrationMode.L2_RELAXED_SCALED]
-    for irk_representation in nosnoc.IrkRepresentation
-    for irk_scheme in nosnoc.IrkSchemes
+    (equidistant_control_grid, step_equilibration, irk_representation, irk_scheme, pss_mode,
+     nosnoc.HomotopyUpdateRule.LINEAR) for equidistant_control_grid in EQUIDISTANT_CONTROLS
+    for step_equilibration in [
+        nosnoc.StepEquilibrationMode.HEURISTIC_MEAN, nosnoc.StepEquilibrationMode.HEURISTIC_DELTA,
+        nosnoc.StepEquilibrationMode.L2_RELAXED, nosnoc.StepEquilibrationMode.L2_RELAXED_SCALED
+    ] for irk_representation in nosnoc.IrkRepresentation for irk_scheme in nosnoc.IrkSchemes
     for pss_mode in PSS_MODES
     # Ignore the following cases that currently fail:
-    if (equidistant_control_grid, step_equilibration, irk_representation, irk_scheme, pss_mode) not in [
+    if (equidistant_control_grid, step_equilibration, irk_representation, irk_scheme, pss_mode)
+    not in [
         # (True, nosnoc.StepEquilibrationMode.DIRECT, nosnoc.IrkRepresentation.DIFFERENTIAL_LIFT_X, nosnoc.IrkSchemes.RADAU_IIA, nosnoc.PssMode.STEWART, nosnoc.HomotopyUpdateRule.LINEAR),
     ]
 ]
 
 # test HomotopyUpdateRule.SUPERLINEAR separately without cartesian product
 options += [
-        (True, nosnoc.StepEquilibrationMode.L2_RELAXED, nosnoc.IrkRepresentation.DIFFERENTIAL, nosnoc.IrkSchemes.RADAU_IIA, nosnoc.PssMode.STEWART, nosnoc.HomotopyUpdateRule.SUPERLINEAR),
+    (True, nosnoc.StepEquilibrationMode.L2_RELAXED, nosnoc.IrkRepresentation.DIFFERENTIAL,
+     nosnoc.IrkSchemes.RADAU_IIA, nosnoc.PssMode.STEWART, nosnoc.HomotopyUpdateRule.SUPERLINEAR),
 ]
 
 
@@ -40,10 +43,8 @@ def test_default(self):
         example(plot=False)
 
     @parameterized.expand(options)
-    def test_combination(
-        self, equidistant_control_grid, step_equilibration, irk_representation,
-        irk_scheme, pss_mode, homotopy_update_rule
-    ):
+    def test_combination(self, equidistant_control_grid, step_equilibration, irk_representation,
+                         irk_scheme, pss_mode, homotopy_update_rule):
         opts = get_default_options()
         opts.comp_tol = 1e-5
         opts.N_stages = 5
@@ -55,16 +56,20 @@ def test_combination(
         opts.pss_mode = pss_mode
         opts.homotopy_update_rule = homotopy_update_rule
 
-        print(f"test setting: equidistant_control_grid {equidistant_control_grid}" +
-              f"\n{step_equilibration}\n{irk_representation}\n{irk_scheme}\n{pss_mode}\n{homotopy_update_rule}")
+        print(
+            f"test setting: equidistant_control_grid {equidistant_control_grid}" +
+            f"\n{step_equilibration}\n{irk_representation}\n{irk_scheme}\n{pss_mode}\n{homotopy_update_rule}"
+        )
         results = solve_ocp(opts)
 
         x_traj = results["x_traj"]
         u_traj = results["u_traj"]
         t_grid = results["t_grid"]
 
-        message = (f"For parameters: control grid {equidistant_control_grid} step_equilibration {step_equilibration}, irk_representation {irk_representation}, "
-                   f"irk_scheme {irk_scheme}, pss_mode {pss_mode}, homotopy_update_rule {homotopy_update_rule}")
+        message = (
+            f"For parameters: control grid {equidistant_control_grid} step_equilibration {step_equilibration}, irk_representation {irk_representation}, "
+            f"irk_scheme {irk_scheme}, pss_mode {pss_mode}, homotopy_update_rule {homotopy_update_rule}"
+        )
 
         self.assertTrue(np.allclose(x_traj[0], X0, atol=1e-4), message)
         self.assertTrue(np.allclose(x_traj[-1][:2], X_TARGET, atol=1e-4), message)

diff --git a/test/test_ocp_motor.py b/test/test_ocp_motor.py
@@ -10,13 +10,11 @@
 import nosnoc
 import numpy as np
 
+options = [(step_equilibration, pss_mode)
+           for pss_mode in nosnoc.PssMode
+           for step_equilibration in nosnoc.StepEquilibrationMode
+           if step_equilibration != nosnoc.StepEquilibrationMode.DIRECT]
 
-options = [
-    (step_equilibration, pss_mode)
-    for pss_mode in nosnoc.PssMode
-    for step_equilibration in nosnoc.StepEquilibrationMode
-    if step_equilibration != nosnoc.StepEquilibrationMode.DIRECT
-]
 
 class TestOcpMotor(unittest.TestCase):
 

diff --git a/test/test_parametric_ocp.py b/test/test_parametric_ocp.py
@@ -17,7 +17,10 @@ def test_one_parametric_ocp(self):
 
         results_with_global_var = solve_paramteric_example(with_global_var=True)
         self.assertTrue(np.allclose(np.ones((1,)), results_with_global_var["v_global"], atol=1e-7))
-        self.assertTrue(np.allclose(np.array(ref_results["x_list"]), np.array(results_with_global_var["x_list"]), atol=1e-7))
+        self.assertTrue(
+            np.allclose(np.array(ref_results["x_list"]),
+                        np.array(results_with_global_var["x_list"]),
+                        atol=1e-7))
 
 
 if __name__ == "__main__":

diff --git a/test/test_problem_dimensions.py b/test/test_problem_dimensions.py
@@ -11,6 +11,7 @@
 
 # TODO: add control stages instead of just simulation
 class TestProblemDimension(unittest.TestCase):
+
     def test_w(self):
         model = get_simplest_model_sliding()
 

diff --git a/test/test_timefreezing.py b/test/test_timefreezing.py
@@ -2,18 +2,13 @@
 from examples.time_freezing_hysteresis_temperature_control import control
 
 
-
 class TestTimeFreezing(unittest.TestCase):
 
     def test_control_example(self):
         model, opts, solver, results = control(with_plot=False)
         end_time = model.t_fun(results["x_list"][-1])
-        self.assertLessEqual(
-            opts.terminal_time - opts.time_freezing_tolerance, end_time
-        )
-        self.assertLessEqual(
-            end_time, opts.terminal_time + opts.time_freezing_tolerance
-        )
+        self.assertLessEqual(opts.terminal_time - opts.time_freezing_tolerance, end_time)
+        self.assertLessEqual(end_time, opts.terminal_time + opts.time_freezing_tolerance)
 
 
 if __name__ == "__main__":

diff --git a/test/test_validation.py b/test/test_validation.py
@@ -42,7 +42,7 @@ def test_switching_system_wrong_g_Stewart(self):
 
     def test_switching_system_wrong_S_shape(self):
         x = SX.sym('x')
-        F = [horzcat(-x, 2*x)]
+        F = [horzcat(-x, 2 * x)]
         c = [x + x**2]
         S = [np.array([-1])]
         X0 = np.array([0])