Irma integration order experiment (#40)

.gitignore

@@ -4,6 +4,7 @@ env
 *.egg-info
 
 *.pdf
+*.gif
 
 private*
 external

examples/Acary2014/irma.py

@@ -35,6 +35,9 @@ def get_default_options():
     opts.n_s = 2
     opts.step_equilibration = nosnoc.StepEquilibrationMode.HEURISTIC_MEAN
     opts.pss_mode = nosnoc.PssMode.STEP
+    opts.print_level = 0
+    opts.homotopy_update_rule = nosnoc.HomotopyUpdateRule.LINEAR
+
     return opts
 
 
@@ -90,30 +93,67 @@ def solve_irma(opts=None, model=None):
     return results
 
 
-def plot_results(results):
+def plot_trajectory(results, figure_filename=None):
     nosnoc.latexify_plot()
 
-    plt.figure()
-    for i in range(len(X0)):
-        plt.subplot(5, 1, i+1)
-        plt.plot(results["t_grid"], results["X_sim"][:, i])
-        plt.hlines(thresholds[i], xmin=0, xmax=1000, linestyles='dotted')
-        plt.xlim(0, 1000)
-        plt.ylim(0, 1.1*max(results["X_sim"][:, i]))
-        plt.ylabel(f'$x_{i+1}$')
-        plt.xlabel('$t$')
+    n_subplot = len(X0)
+    fig, axs = plt.subplots(n_subplot, 1)
+    for i in range(n_subplot):
+        axs[i].plot(results["t_grid"], results["X_sim"][:, i])
+        axs[i].hlines(thresholds[i], xmin=0, xmax=TSIM, linestyles='dotted')
+        axs[i].set_xlim(0, TSIM)
+        axs[i].set_ylim(0, 1.1*max(results["X_sim"][:, i]))
+        axs[i].set_ylabel(f'$x_{i+1}$')
+        axs[i].grid()
+        if i == n_subplot - 1:
+            plt.xlabel('$t$ [min]')
+        else:
+            axs[i].xaxis.set_ticklabels([])
+
+    if figure_filename is not None:
+        plt.savefig(figure_filename)
+        print(f'stored figure as {figure_filename}')
+
     plt.show()
 
+def plot_algebraic_traj(results, figure_filename=None):
+    nosnoc.latexify_plot()
+    alpha_sim = np.array([results['alpha_sim'][0][0]] + nosnoc.flatten_layer(results['alpha_sim']))
+    n_subplot = len(alpha_sim[0])
+
+    fig, axs = plt.subplots(n_subplot, 1)
+    for i in range(n_subplot):
+        axs[i].plot(results["t_grid"], alpha_sim[:,i])
+        # axs[i].hlines(thresholds[i], xmin=0, xmax=TSIM, linestyles='dotted')
+        axs[i].set_xlim(0, TSIM)
+        axs[i].set_ylim(0, 1.1*max(alpha_sim[:, i]))
+        axs[i].set_ylabel(r'$\alpha_' + f'{i+1}$')
+        # axs[i].set_xlabel('$t$ [min]')
+        axs[i].grid()
+        if i == n_subplot - 1:
+            axs[i].set_xlabel('$t$ [min]')
+        else:
+            axs[i].xaxis.set_ticklabels([])
+
+    if figure_filename is not None:
+        plt.savefig(figure_filename)
+        print(f'stored figure as {figure_filename}')
+
+    plt.show()
 
 # EXAMPLE
 def example():
     opts = get_default_options()
     opts.print_level = 1
-    results = []
+
     model = get_irma_model(SWITCH_ON, LIFTING)
     results = solve_irma(opts=opts, model=model)
 
-    plot_results(results)
+    plot_algebraic_traj(results)
+    plot_trajectory(results)
+
+    # plot_algebraic_traj(results, figure_filename='irma_algebraic_traj.pdf')
+    # plot_trajectory(results, figure_filename='irma_traj.pdf')
 
 
 if __name__ == "__main__":

examples/Acary2014/irma_integration_order_experiment.py

@@ -0,0 +1,176 @@
+from matplotlib import pyplot as plt
+import numpy as np
+from irma import get_irma_model, get_default_options, SWITCH_ON, LIFTING, X0
+import nosnoc
+import pickle
+import os
+import itertools
+
+BENCHMARK_DATA_PATH = 'private_irma_benchmark_data'
+REF_RESULTS_FILENAME = 'irma_benchmark_results.pickle'
+SCHEMES = [nosnoc.IrkSchemes.GAUSS_LEGENDRE, nosnoc.IrkSchemes.RADAU_IIA]
+
+NS_VALUES = [1, 2, 3, 4]
+NFE_VALUES = [3]
+# NSIM_VALUES = [1, 3, 10, 20, 50, 100, 300] # convergence issues for Legendre
+NSIM_VALUES = [1, 3, 9, 18, 50, 100, 300]
+
+
+# # NOTE: this has convergence issues
+# NS_VALUES = [2]
+# NSIM_VALUES = [10]
+# SCHEME = nosnoc.IrkSchemes.GAUSS_LEGENDRE
+
+
+TOL = 1e-12
+TSIM = 100
+
+
+def pickle_results(results, filename):
+    file = os.path.join(BENCHMARK_DATA_PATH, filename)
+    with open(file, 'wb') as f:
+        pickle.dump(results, f)
+
+def unpickle_results(filename):
+    file = os.path.join(BENCHMARK_DATA_PATH, filename)
+    with open(file, 'rb') as f:
+        results = pickle.load(f)
+    return results
+
+def generate_reference_solution():
+    opts = get_default_options()
+    opts.n_s = 5
+    opts.N_finite_elements = 3
+
+    Nsim = 500
+    Tstep = TSIM / Nsim
+    opts.terminal_time = Tstep
+    opts.comp_tol = TOL * 1e-2
+    opts.sigma_N = TOL * 1e-2
+    opts.do_polishing_step = False
+    opts.step_equilibration = nosnoc.StepEquilibrationMode.HEURISTIC_MEAN
+
+    model = get_irma_model(SWITCH_ON, LIFTING)
+
+    solver = nosnoc.NosnocSolver(opts, model)
+
+    # loop
+    looper = nosnoc.NosnocSimLooper(solver, model.x0, Nsim, print_level=1)
+    looper.run()
+    results = looper.get_results()
+    results['opts'] = opts
+
+    pickle_results(results, REF_RESULTS_FILENAME)
+
+
+def get_results_filename(opts):
+    filename = 'irma_bm_results_'
+    filename += 'Nfe_' + str(opts.N_finite_elements) + '_'
+    filename += 'ns' + str(opts.n_s) + '_'
+    filename += 'tol' + str(opts.comp_tol) + '_'
+    filename += 'dt' + str(opts.terminal_time) + '_'
+    filename += 'Tsim' + str(TSIM) + '_'
+    filename += opts.irk_scheme.name
+    filename += '.pickle'
+    return filename
+
+
+def get_opts(Nsim, n_s, N_fe, scheme):
+    opts = get_default_options()
+    Tstep = TSIM / Nsim
+    opts.terminal_time = Tstep
+    opts.comp_tol = TOL
+    opts.sigma_N = TOL
+    opts.irk_scheme = scheme
+    opts.print_level = 1
+
+    opts.n_s = n_s
+    opts.N_finite_elements = N_fe
+    # opts.step_equilibration = nosnoc.StepEquilibrationMode.DIRECT
+    # opts.irk_representation = nosnoc.IrkRepresentation.DIFFERENTIAL
+    return opts
+
+
+def run_benchmark():
+    for n_s, N_fe, Nsim, scheme in itertools.product(NS_VALUES, NFE_VALUES, NSIM_VALUES, SCHEMES):
+
+        model = get_irma_model(SWITCH_ON, LIFTING)
+        opts = get_opts(Nsim, n_s, N_fe, scheme)
+        solver = nosnoc.NosnocSolver(opts, model)
+
+        # loop
+        looper = nosnoc.NosnocSimLooper(solver, model.x0, Nsim, print_level=1)
+        looper.run()
+        results = looper.get_results()
+        results['opts'] = opts
+        filename = get_results_filename(opts)
+        pickle_results(results, filename)
+        del solver, looper, results, model, opts
+
+
+def get_reference_solution():
+    return unpickle_results(REF_RESULTS_FILENAME)
+
+def count_failures(results):
+    status_list: list = results['status']
+    return len([x for x in status_list if x != nosnoc.Status.SUCCESS])
+
+def evaluate_reference_solution():
+    results = get_reference_solution()
+    n_fail = count_failures(results)
+
+    print(f"Reference solution got {n_fail} failing subproblems")
+
+
+def order_plot():
+    nosnoc.latexify_plot()
+    N_fe = 3
+    ref_results = get_reference_solution()
+    x_end_ref = ref_results['X_sim'][-1]
+
+    linestyles = ['-', '--', '-.', ':']
+    SCHEME = nosnoc.IrkSchemes.RADAU_IIA
+
+    plt.figure()
+    for i, n_s in enumerate(NS_VALUES):
+        errors = []
+        step_sizes = []
+        for Nsim in NSIM_VALUES:
+            opts = get_opts(Nsim, n_s, N_fe, SCHEME)
+            filename = get_results_filename(opts)
+            results = unpickle_results(filename)
+            x_end = results['X_sim'][-1]
+            n_fail = count_failures(results)
+            error = np.max(np.abs(x_end - x_end_ref))
+            print("opts.n_s: ", opts.n_s, "opts.terminal_time: ", opts.terminal_time, "error: ", error, "n_fail: ", n_fail)
+            errors.append(error)
+            step_sizes.append(opts.terminal_time)
+
+        label = r'$n_s=' + str(n_s) +'$'
+        if results['opts'].irk_scheme == nosnoc.IrkSchemes.RADAU_IIA:
+            if n_s == 1:
+                label = 'implicit Euler: 1'
+            else:
+                label = 'Radau IIA: ' + str(2*n_s-1)
+        elif results['opts'].irk_scheme == nosnoc.IrkSchemes.GAUSS_LEGENDRE:
+            label = 'Gauss-Legendre: ' + str(2*n_s)
+        plt.plot(step_sizes, errors, label=label, marker='o', linestyle=linestyles[i])
+    plt.grid()
+    plt.xlabel('Step size')
+    plt.ylabel('Error')
+    plt.yscale('log')
+    plt.xscale('log')
+    plt.legend()
+
+    plt.savefig(f'irma_benchmark_{SCHEME.name}.pdf', bbox_inches='tight')
+    plt.show()
+
+
+if __name__ == "__main__":
+    # generate data
+    run_benchmark()
+    generate_reference_solution()
+
+    # evalute
+    evaluate_reference_solution()
+    order_plot()

nosnoc/__init__.py

@@ -4,9 +4,9 @@
 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
+from .nosnoc_types import MpccMode, IrkSchemes, StepEquilibrationMode, CrossComplementarityMode, IrkRepresentation, PssMode, IrkRepresentation, HomotopyUpdateRule, InitializationStrategy, ConstraintHandling, Status
 from .helpers import NosnocSimLooper
-from .utils import casadi_length, casadi_vertcat_list, print_casadi_vector, flatten_layer
+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
 from .rk_utils import rk4
 

nosnoc/helpers.py

@@ -11,7 +11,9 @@ def __init__(self,
                  x0: np.ndarray,
                  Nsim: int,
                  p_values: Optional[np.ndarray] = None,
-                 w_init: Optional[list] = None):
+                 w_init: Optional[list] = None,
+                 print_level: Optional[int] = None
+                ):
         """
         :param solver: NosnocSolver to be called in a loop
         :param x0: np.ndarray: initial state
@@ -44,6 +46,11 @@ def __init__(self,
         self.w_all = []
         self.cost_vals = []
         self.w_init = w_init
+        if print_level is not None:
+            self.print_level = print_level
+        else:
+            self.print_level = solver.opts.print_level
+        self.status = []
 
         self.cpu_nlp = np.zeros((Nsim, solver.opts.max_iter_homotopy + 1))
 
@@ -69,6 +76,9 @@ def run(self) -> None:
             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.print_level > 0:
+                print(f"Sim step {i + 1}/{self.Nsim}\t status: {results['status']}")
 
     def get_results(self) -> dict:
         self.t_grid = np.concatenate((np.array([0.0]), np.cumsum(self.time_steps)))
@@ -79,8 +89,10 @@ def get_results(self) -> dict:
             "t_grid": self.t_grid,
             "theta_sim": self.theta_sim,
             "lambda_sim": self.lambda_sim,
+            "alpha_sim": self.alpha_sim,
             "w_sim": self.w_sim,
             "w_all": self.w_all,
-            "cost_vals": self.cost_vals
+            "cost_vals": self.cost_vals,
+            "status": self.status,
         }
         return results

nosnoc/nosnoc_types.py

@@ -86,3 +86,8 @@ class PssMode(Enum):
     lambda_p_i >= 0;    for all i = 1,..., n_sys
     alpha_i >= 0;     for all i = 1,..., n_sys
     """
+
+
+class Status(Enum):
+    SUCCESS = auto()
+    INFEASIBLE = auto()

nosnoc/plot_utils.py

@@ -24,8 +24,7 @@ def latexify_plot():
     matplotlib.rcParams.update(params)
 
 
-def plot_timings(timings: np.ndarray, latexify=True, title='', figure_filename=''):
-    # latexify plot
+def plot_timings(timings: np.ndarray, latexify=True, title=None, figure_filename=None):
     if latexify:
         latexify_plot()
 
@@ -51,8 +50,9 @@ def plot_timings(timings: np.ndarray, latexify=True, title='', figure_filename='
     plt.legend()
     plt.xlim(x_range)
     plt.grid(alpha=0.3)
-    plt.title(title)
-    if figure_filename != '':
+    if title is not None:
+        plt.title(title)
+    if figure_filename is not None:
         plt.savefig(figure_filename)
         print(f'stored figure as {figure_filename}')
 
@@ -63,9 +63,8 @@ def plot_iterates(problem: NosnocProblem,
                   iterates: list,
                   latexify=False,
                   title_list=[],
-                  figure_filename=''):
+                  figure_filename=None):
 
-    # latexify plot
     if latexify:
         latexify_plot()
 
@@ -121,7 +120,7 @@ def plot_iterates(problem: NosnocProblem,
         plt.grid(alpha=0.3)
         plt.legend()
 
-    if figure_filename != '':
+    if figure_filename is not None:
         plt.savefig(figure_filename)
         print(f'stored figure as {figure_filename}')
 

nosnoc/problem.py

@@ -598,6 +598,7 @@ def _add_finite_element(self, fe: FiniteElement, ctrl_idx: int):
         self.ind_alpha[ctrl_idx].append(increment_indices(fe.ind_alpha, w_len))
         self.ind_lambda_n[ctrl_idx].append(increment_indices(fe.ind_lambda_n, w_len))
         self.ind_lambda_p[ctrl_idx].append(increment_indices(fe.ind_lambda_p, w_len))
+        self.ind_z[ctrl_idx].append(increment_indices(fe.ind_z, w_len))
 
     # TODO: can we just use add_variable? It is a bit involved, since index vectors here have different format.
     def _add_primal_vector(self, symbolic: ca.SX, lb: np.array, ub, initial):
@@ -656,6 +657,7 @@ def __init__(self, opts: NosnocOpts, model: NosnocModel, ocp: Optional[NosnocOcp
         self.ind_alpha = create_empty_list_matrix((opts.N_stages,))
         self.ind_lambda_n = create_empty_list_matrix((opts.N_stages,))
         self.ind_lambda_p = create_empty_list_matrix((opts.N_stages,))
+        self.ind_z = create_empty_list_matrix((opts.N_stages,))
 
         self.ind_u = []
         self.ind_h = []