Add option fix_active_set_fe0, add oscilator tests (#27)

* oscilator examples: compare error wrt exact solution

* fix typo

* dont warn on Feasible_Point_Found

* implement option fix_active_set_fe0

* add warnings.simplefilter(always)

* add debug print

* refactor oscilator_example, add oscilator_test

* add test_fix_active_set

* add tests in TestValidation

* fix test_least_squares_problem, quite sensitive to options still

* remove unused init_gamma

* add annotation

examples/oscilator_example.py

@@ -9,6 +9,8 @@
 A2 = np.array([[1, -OMEGA], [OMEGA, 1]])
 R_OSC = 1
 
+TSIM = np.pi / 2
+X_SOL = np.array([np.exp(TSIM-1) * np.cos(2*np.pi * (TSIM-1)), -np.exp(TSIM-1) * np.sin(2*np.pi*(TSIM-1))])
 
 def get_oscilator_model(use_g_Stewart=False):
 
@@ -37,23 +39,26 @@ def get_oscilator_model(use_g_Stewart=False):
 
     return model
 
-
-def main(use_g_Stewart=False):
+def get_default_options():
     opts = nosnoc.NosnocOpts()
-    opts.use_fesd = True
-    comp_tol = 1e-6
+    comp_tol = 1e-8
     opts.comp_tol = comp_tol
     opts.homotopy_update_slope = 0.1  # decrease rate
     opts.N_finite_elements = 2
-    opts.n_s = 2
-    opts.step_equilibration = nosnoc.StepEquilibrationMode.L2_RELAXED_SCALED
+    opts.n_s = 3
+    opts.step_equilibration = nosnoc.StepEquilibrationMode.DIRECT_COMPLEMENTARITY
+    opts.print_level = 1
+    return opts
+
+
+def solve_oscilator(opts=None, use_g_Stewart=False, do_plot=True):
+    if opts is None:
+        opts = get_default_options()
 
     model = get_oscilator_model(use_g_Stewart)
 
-    Tsim = np.pi / 2
     Nsim = 29
-    Tstep = Tsim / Nsim
-
+    Tstep = TSIM / Nsim
     opts.terminal_time = Tstep
 
     solver = nosnoc.NosnocSolver(opts, model)
@@ -63,15 +68,20 @@ def main(use_g_Stewart=False):
     looper.run()
     results = looper.get_results()
 
+    error = np.max(np.abs(X_SOL - results["X_sim"][-1]))
+    print(f"error wrt exact solution {error:.2e}")
+
+    if do_plot:
+        plot_oscilator(results["X_sim"], results["t_grid"])
+    # nosnoc.plot_timings(results["cpu_nlp"])
 
-    plot_oscilator(results["X_sim"], results["t_grid"])
-    nosnoc.plot_timings(results["cpu_nlp"])
     # store solution
     # import json
     # json_file = 'oscilator_results_ref.json'
     # with open(json_file, 'w') as f:
     #     json.dump(results['w_sim'], f, indent=4, sort_keys=True, default=make_object_json_dumpable)
     # print(f"saved results in {json_file}")
+    return results
 
 def main_least_squares():
 
@@ -82,20 +92,17 @@ def main_least_squares():
     #     w_sim_ref = json.load(f)
 
     opts = nosnoc.NosnocOpts()
-    opts.use_fesd = True
     comp_tol = 1e-7
     opts.comp_tol = comp_tol
-    opts.N_finite_elements = 2
-    opts.n_s = 2
-    opts.print_level = 3
+    opts.print_level = 2
 
     # opts.homotopy_update_rule = nosnoc.HomotopyUpdateRule.SUPERLINEAR
     opts.cross_comp_mode = nosnoc.CrossComplementarityMode.COMPLEMENT_ALL_STAGE_VALUES_WITH_EACH_OTHER
-    opts.mpcc_mode = nosnoc.MpccMode.FISCHER_BURMEISTER_IP_AUG
+    opts.mpcc_mode = nosnoc.MpccMode.FISCHER_BURMEISTER
     opts.constraint_handling = nosnoc.ConstraintHandling.LEAST_SQUARES
     opts.step_equilibration = nosnoc.StepEquilibrationMode.DIRECT
     opts.initialization_strategy = nosnoc.InitializationStrategy.ALL_XCURRENT_W0_START
-    # opts.initialization_strategy = nosnoc.InitializationStrategy.RK4_SMOOTHENED
+    opts.initialization_strategy = nosnoc.InitializationStrategy.RK4_SMOOTHENED
     opts.sigma_0 = 1e0
     # opts.gamma_h = np.inf
     # opts.opts_casadi_nlp['ipopt']['max_iter'] = 0
@@ -104,9 +111,8 @@ def main_least_squares():
 
     model = get_oscilator_model()
 
-    Tsim = np.pi / 2
     Nsim = 29
-    Tstep = Tsim / Nsim
+    Tstep = TSIM / Nsim
 
     opts.terminal_time = Tstep
 
@@ -119,41 +125,29 @@ def main_least_squares():
     results = looper.get_results()
     print(f"max cost_val = {max(results['cost_vals']):.2e}")
 
+    error = np.max(np.abs(X_SOL - results["X_sim"][-1]))
+    print(f"error wrt exact solution {error:.2e}")
+
+    breakpoint()
     plot_oscilator(results["X_sim"], results["t_grid"])
-    nosnoc.plot_timings(results["cpu_nlp"])
-    # breakpoint()
+    # nosnoc.plot_timings(results["cpu_nlp"])
 
 
 def main_polishing():
 
-    opts = nosnoc.NosnocOpts()
+    opts = get_default_options()
     opts.comp_tol = 1e-4
-    opts.print_level = 3
+    opts.do_polishing_step = True
 
     opts.cross_comp_mode = nosnoc.CrossComplementarityMode.COMPLEMENT_ALL_STAGE_VALUES_WITH_EACH_OTHER
     opts.step_equilibration = nosnoc.StepEquilibrationMode.DIRECT
-    # opts.constraint_handling = nosnoc.ConstraintHandling.LEAST_SQUARES
-    # opts.mpcc_mode = nosnoc.MpccMode.FISCHER_BURMEISTER
-    opts.do_polishing_step = True
-    opts.homotopy_update_slope = 0.1
-
-    model = get_oscilator_model()
-
-    Tsim = np.pi / 2
-    Nsim = 29
-    Tstep = Tsim / Nsim
-
-    opts.terminal_time = Tstep
+    opts.constraint_handling = nosnoc.ConstraintHandling.LEAST_SQUARES
+    opts.mpcc_mode = nosnoc.MpccMode.FISCHER_BURMEISTER
+    opts.print_level = 3
 
-    solver = nosnoc.NosnocSolver(opts, model)
-    solver.print_problem()
-    # loop
-    looper = nosnoc.NosnocSimLooper(solver, model.x0, Nsim)
-    looper.run()
-    results = looper.get_results()
+    results = solve_oscilator(opts, do_plots=False)
     print(f"max cost_val = {max(results['cost_vals']):.2e}")
 
-    plot_oscilator(results["X_sim"], results["t_grid"])
     nosnoc.plot_timings(results["cpu_nlp"])
 
 
@@ -198,6 +192,6 @@ def make_object_json_dumpable(input):
         raise TypeError(f"Cannot make input of type {type(input)} dumpable.")
 
 if __name__ == "__main__":
-    main(len(argv) > 1)
+    solve_oscilator(use_g_Stewart=False, do_plot=True)
     # main_least_squares()
     # main_polishing()

nosnoc/__init__.py

@@ -5,3 +5,6 @@
 from .utils import casadi_length, print_casadi_vector, flatten_layer
 from .plot_utils import plot_timings, plot_iterates, latexify_plot
 from .rk_utils import rk4
+
+import warnings
+warnings.simplefilter("always")

nosnoc/nosnoc.py

@@ -1065,7 +1065,7 @@ def __init__(self, opts: NosnocOpts, model: NosnocModel, ocp: Optional[NosnocOcp
             all_h = [fe.h() for stage in self.stages for fe in stage]
             self.add_constraint(sum(all_h) - opts.terminal_time)
 
-        # CasADi ca.Functions
+        # CasADi Functions
         self.cost_fun = ca.Function('cost_fun', [self.w], [self.cost])
         self.comp_res = ca.Function('comp_res', [self.w, self.p], [J_comp])
         self.g_fun = ca.Function('g_fun', [self.w, self.p], [self.g])
@@ -1351,7 +1351,7 @@ def solve(self) -> dict:
         self.initialize()
         opts = self.opts
         prob = self.problem
-        w0 = prob.w0
+        w0 = prob.w0.copy()
 
         w_all = [w0.copy()]
         n_iter_polish = opts.max_iter_homotopy + 1
@@ -1370,6 +1370,43 @@ def solve(self) -> dict:
         p0 = prob.model.p_val_ctrl_stages[0]
         lambda00 = self.model.lambda00_fun(x0, p0).full().flatten()
 
+        if opts.fix_active_set_fe0 and opts.pss_mode == PssMode.STEWART:
+            lbw = prob.lbw.copy()
+            ubw = prob.ubw.copy()
+
+            # lambda00 != 0.0 -> corresponding thetas on first fe are zero
+            I_active_lam = np.where(lambda00 > 1e1*opts.comp_tol)[0].tolist()
+            ind_theta_fe1 = flatten_layer(prob.ind_theta[0][0], 2) # flatten sys
+            w_zero_indices = []
+            for i in range(opts.n_s):
+                tmp = flatten(ind_theta_fe1[i])
+                try:
+                    w_zero_indices += [tmp[i] for i in I_active_lam]
+                except:
+                    breakpoint()
+
+            # if all but one lambda are zero: this theta can be fixed to 1.0, all other thetas are 0.0
+            w_one_indices = []
+            # I_lam_zero = set(range(len(lambda00))).difference( I_active_lam )
+            # n_lam = sum(prob.model.dims.n_f_sys)
+            # if len(I_active_lam) == n_lam - 1:
+            #     for i in range(opts.n_s):
+            #         tmp = flatten(ind_theta_fe1[i])
+            #         w_one_indices += [tmp[i] for i in I_lam_zero]
+            if opts.print_level > 1:
+                print(f"fixing {prob.w[w_one_indices]} = 1. and {prob.w[w_zero_indices]} = 0.")
+                print(f"Since lambda00 = {lambda00}")
+            w0[w_zero_indices] = 0.0
+            lbw[w_zero_indices] = 0.0
+            ubw[w_zero_indices] = 0.0
+            w0[w_one_indices] = 1.0
+            lbw[w_one_indices] = 1.0
+            ubw[w_one_indices] = 1.0
+
+        else:
+            lbw = prob.lbw
+            ubw = prob.ubw
+
         # homotopy loop
         for ii in range(opts.max_iter_homotopy):
             tau_val = sigma_k
@@ -1383,8 +1420,8 @@ def solve(self) -> dict:
             sol = self.solver(x0=w0,
                               lbg=prob.lbg,
                               ubg=prob.ubg,
-                              lbx=prob.lbw,
-                              ubx=prob.ubw,
+                              lbx=lbw,
+                              ubx=ubw,
                               p=p_val)
             cpu_time_nlp[ii] = time.time() - t
 
@@ -1404,7 +1441,7 @@ def solve(self) -> dict:
             if opts.print_level:
                 self._print_iter_stats(sigma_k, complementarity_residual, nlp_res, cost_val,
                                        cpu_time_nlp[ii], nlp_iter[ii], status)
-            if status not in ['Solve_Succeeded', 'Solved_To_Acceptable_Level']:
+            if status not in ['Solve_Succeeded', 'Solved_To_Acceptable_Level', 'Feasible_Point_Found']:
                 print(f"Warning: IPOPT exited with status {status}")
 
             if complementarity_residual < opts.comp_tol:
@@ -1439,6 +1476,7 @@ def solve(self) -> dict:
                 for ii in range(len(g_val)):
                     if g_val[ii] > threshold:
                         print(f"g_val[{ii}] = {g_val[ii]} expr: {prob.g_lsq[ii]}")
+                print(f"h values: {w_opt[prob.ind_h]}")
 
         if opts.initialization_strategy == InitializationStrategy.ALL_XCURRENT_WOPT_PREV:
             prob.w0[:] = w_opt[:]

nosnoc/nosnoc_opts.py

@@ -34,6 +34,7 @@ class NosnocOpts:
 
     smoothing_parameter: float = 1e1  #: used for smoothed Step representation
     # used in InitializationStrategy.RK4_smoothed
+    fix_active_set_fe0: bool = False
 
     # initialization - Stewart
     init_theta: float = 1.0
@@ -42,7 +43,6 @@ class NosnocOpts:
     # initialization - Step
     init_alpha: float = 1.0  # for step only
     init_beta: float = 1.0
-    init_gamma: float = 1.0
 
     N_finite_elements: int = 2  #
     Nfe_list: list = field(default_factory=list)  #: list of length N_stages, Nfe per stage
@@ -61,7 +61,7 @@ class NosnocOpts:
     rho_h: float = 1.0
 
     # polishing step
-    do_polishing_step = False
+    do_polishing_step: bool = False
 
     # OCP only
     N_stages: int = 1

test/oscilator_test.py

@@ -0,0 +1,94 @@
+from examples.oscilator_example import (
+    get_default_options,
+    TSIM,
+    X_SOL,
+    solve_oscilator,
+)
+import unittest
+import nosnoc
+import numpy as np
+
+EXACT_SWITCH_TIME = 1
+X_SWITCH_EXACT = np.array([1.0, 0.0])
+
+
+def compute_errors(results) -> dict:
+    X_sim = results["X_sim"]
+    switch_diff = np.abs(results["t_grid"] - EXACT_SWITCH_TIME)
+    err_t_switch = np.min(switch_diff)
+
+    switch_index = np.where(switch_diff == err_t_switch)[0][0]
+    err_x_switch = np.max(np.abs(X_sim[switch_index] - X_SWITCH_EXACT))
+
+    err_t_end = np.abs(results["t_grid"][-1] - TSIM)
+
+    err_x_end = np.max(np.abs(X_sim[-1] - X_SOL))
+    return {
+        "t_switch": err_t_switch,
+        "t_end": err_t_end,
+        "x_switch": err_x_switch,
+        "x_end": err_x_end,
+    }
+
+
+class OscilatorTests(unittest.TestCase):
+
+    def test_default(self):
+        opts = get_default_options()
+        opts.print_level = 0
+        results = solve_oscilator(opts, do_plot=False)
+        errors = compute_errors(results)
+
+        print(errors)
+        tol = 1e-5
+        assert errors["t_switch"] < tol
+        assert errors["t_end"] < tol
+        assert errors["x_switch"] < tol
+        assert errors["x_end"] < tol
+
+    def test_polishing(self):
+        opts = get_default_options()
+        opts.print_level = 0
+        opts.comp_tol = 1e-4
+        opts.do_polishing_step = True
+
+        opts.cross_comp_mode = nosnoc.CrossComplementarityMode.COMPLEMENT_ALL_STAGE_VALUES_WITH_EACH_OTHER
+        opts.step_equilibration = nosnoc.StepEquilibrationMode.DIRECT
+        # opts.constraint_handling = nosnoc.ConstraintHandling.LEAST_SQUARES
+        # opts.mpcc_mode = nosnoc.MpccMode.FISCHER_BURMEISTER
+
+        results = solve_oscilator(opts, do_plot=False)
+        errors = compute_errors(results)
+
+        print(errors)
+        tol = 1e-5
+        assert errors["t_switch"] < tol
+        assert errors["t_end"] < tol
+        assert errors["x_switch"] < tol
+        assert errors["x_end"] < tol
+
+
+    def test_fix_active_set(self):
+        opts = get_default_options()
+        opts.print_level = 0
+        opts.fix_active_set_fe0 = True
+
+        opts.cross_comp_mode = nosnoc.CrossComplementarityMode.COMPLEMENT_ALL_STAGE_VALUES_WITH_EACH_OTHER
+        opts.step_equilibration = nosnoc.StepEquilibrationMode.L2_RELAXED
+
+        results = solve_oscilator(opts, do_plot=False)
+        errors = compute_errors(results)
+
+        print(errors)
+        tol = 1e-5
+        assert errors["t_switch"] < tol
+        assert errors["t_end"] < tol
+        assert errors["x_switch"] < tol
+        assert errors["x_end"] < tol
+
+
+if __name__ == "__main__":
+    unittest.main()
+    # uncomment to run single test locally
+    # oscilator_test = OscilatorTests()
+    # oscilator_test.test_least_squares_problem()