Add plotting tools. (#34)

examples/hysteresis_car_control/hysteresis_car_control_time_optimal_2d.py

@@ -9,6 +9,7 @@
 
 
 import nosnoc
+from nosnoc.plot_utils import plot_voronoi_2d
 import casadi as ca
 import numpy as np
 from math import ceil, log
@@ -196,10 +197,10 @@ def create_gearbox_voronoi(u=None, q_goal=None, mode=Stages.STAGE_2,
             x=X, F=F, g_Stewart=g_ind, x0=X0, t_var=t,
             name="gearbox"
         )
-    return model, lbx, ubx, lbu, ubu, f_q, f_terminal, g_path, g_terminal
+    return model, lbx, ubx, lbu, ubu, f_q, f_terminal, g_path, g_terminal, Z
 
 
-def plot(x_list, t_grid, u_list, t_grid_u):
+def plot(x_list, t_grid, u_list, t_grid_u, Z):
     """Plot."""
     q = [x[0] for x in x_list]
     v = [x[1] for x in x_list]
@@ -224,20 +225,22 @@ def plot(x_list, t_grid, u_list, t_grid_u):
     plt.plot(t, q, label="Position vs actual time")
     plt.xlabel("Actual Time [$s$]")
 
-    plt.figure()
+    fig = plt.figure()
+    ax = fig.add_subplot()
+    plot_voronoi_2d(Z, ax=ax, show=False, annotate=True)
     psi = [(vi - v1) / (v2 - v1) for vi in v]
-    im = plt.scatter(psi, aux, c=t_grid, cmap=plt.hot())
+    im = ax.scatter(psi, aux, c=t_grid, cmap=plt.hot())
     im.set_label('Time')
-    plt.colorbar(im)
-    plt.xlabel("$\\psi(x)$")
-    plt.ylabel("$w$")
+    plt.colorbar(im, ax=ax)
+    ax.set_xlabel("$\\psi(x)$")
+    ax.set_ylabel("$w$")
     plt.show()
 
 
 def simulation(u=25, Tsim=6, Nsim=30, with_plot=True):
     """Simulate the temperature control system with a fixed input."""
     opts = create_options()
-    model, lbx, ubx, lbu, ubu, f_q, f_terminal, g_path, g_terminal = create_gearbox_voronoi(u=u, q_goal=q_goal)
+    model, lbx, ubx, lbu, ubu, f_q, f_terminal, g_path, g_terminal, Z = create_gearbox_voronoi(u=u, q_goal=q_goal)
     Tstep = Tsim / Nsim
     opts.N_finite_elements = 2
     opts.N_stages = 1
@@ -252,13 +255,13 @@ def simulation(u=25, Tsim=6, Nsim=30, with_plot=True):
     results = looper.get_results()
     print(f"Ends in zone: {np.argmax(results['theta_sim'][-1][-1])}")
     print(results['theta_sim'][-1][-1])
-    plot(results["X_sim"], results["t_grid"], None, None)
+    plot(results["X_sim"], results["t_grid"], None, None, Z)
 
 
 def control():
     """Execute one Control step."""
     N = 15
-    model, lbx, ubx, lbu, ubu, f_q, f_terminal, g_path, g_terminal = create_gearbox_voronoi(
+    model, lbx, ubx, lbu, ubu, f_q, f_terminal, g_path, g_terminal, Z = create_gearbox_voronoi(
         q_goal=q_goal,
     )
     opts = create_options()
@@ -279,9 +282,9 @@ def control():
     results = solver.solve()
     plot(
         results["x_traj"], results["t_grid"],
-        results["u_list"], results["t_grid_u"]
+        results["u_list"], results["t_grid_u"], Z
     )
 
 
 if __name__ == "__main__":
-    control()
+    simulation()

examples/hysteresis_car_control/hysteresis_car_control_time_optimal_3d.py

@@ -8,6 +8,7 @@
 """
 
 import nosnoc
+from nosnoc.plot_utils import plot_colored_line_3d
 import casadi as ca
 import numpy as np
 from math import ceil, log
@@ -264,6 +265,10 @@ def plot(x_list, t_grid, u_list, t_grid_u):
     plt.colorbar(im)
     plt.xlabel("$\\psi(x)$")
     plt.ylabel("$w$")
+
+    plot_colored_line_3d(
+        psi, [x[-3] for x in x_list], [x[-2] for x in x_list], t
+    )
     plt.show()
 
 

nosnoc/plot_utils.py

@@ -4,7 +4,7 @@
 import matplotlib
 
 from nosnoc import NosnocProblem, get_results_from_primal_vector
-from .utils import flatten_layer, flatten
+from .utils import flatten_layer
 from .nosnoc_types import PssMode
 
 
@@ -98,9 +98,9 @@ def plot_iterates(problem: NosnocProblem,
         # plot lambda, mu
         plt.subplot(n_row, n_iterates, n_iterates * 0 + it + 1)
         for i in range(n_lam):
-            plt.plot([x[i] for x in lambdas], label=f'$\lambda_{i+1}$')
+            plt.plot([x[i] for x in lambdas], label=f'$\\lambda_{i+1}$')
         for i in range(n_mu):
-            plt.plot([x[i] for x in mus], label=f'$\mu_{i+1}$')
+            plt.plot([x[i] for x in mus], label=f'$\\mu_{i+1}$')
         plt.grid(alpha=0.3)
         plt.title(title_list[it])
         plt.legend()
@@ -109,7 +109,7 @@ def plot_iterates(problem: NosnocProblem,
         # TODO: make this step plot?
         plt.subplot(n_row, n_iterates, n_iterates * 1 + it + 1)
         for i in range(n_lam):
-            plt.plot([x[i] for x in thetas], label=r'$\theta_' + f'{i+1}$')
+            plt.plot([x[i] for x in thetas], label=r'$\\theta_' + f'{i+1}$')
         plt.grid(alpha=0.3)
         plt.legend()
 
@@ -127,3 +127,91 @@ def plot_iterates(problem: NosnocProblem,
 
     # plt.show()
     return
+
+
+def plot_voronoi_2d(Z, show=True, annotate=False, ax=None):
+    """Plot voronoi regions."""
+    from scipy.spatial import Voronoi, voronoi_plot_2d
+    if not isinstance(Z, np.ndarray):
+        Z = np.array(Z)
+    vor = Voronoi(Z)
+    fig = voronoi_plot_2d(vor, ax=ax)
+    if ax is None:
+        ax = fig.axes[0]
+
+    if annotate:
+        for i in range(Z.shape[0]):
+            ax.text(Z[i, 0], Z[i, 1], f"p{i+1}")
+
+    if show:
+        plt.show()
+
+    return ax
+
+
+def scatter_3d(Z, show=True, ax=None, annotate=False):
+    """
+    3D scatter points.
+
+    :param Z: List of points
+    :param show: Show the scatterplot after drawing
+    :param ax: Optional axis to draw the points onto
+    :param annotate: Annotate the points (with p_x)
+    :return: ax
+    """
+    if ax is None:
+        fig = plt.figure()
+        ax = fig.add_subplot(projection='3d')
+
+    ax.scatter(
+        [z[0] for z in Z],
+        [z[1] for z in Z],
+        [z[2] for z in Z],
+    )
+
+    if annotate:
+        for i, pi in enumerate(Z):
+            ax.text(pi[0], pi[1], pi[2], f"p{i+1}", zdir=(1, 0, 0))
+
+    if show:
+        plt.show()
+
+    return ax
+
+
+def _plot_color_hack_3d(ax, x, y, z, t):
+    """Color hack for 3d plot."""
+    t_min = min(t)
+    dt = max(t) - t_min
+    for i in range(np.size(x)-1):
+        ax.plot(x[i:i+2], y[i:i+2], z[i:i+2], color=plt.cm.hot(
+            (t[i] - t_min) / dt
+        ))
+
+
+def plot_colored_line_3d(x, y, z, t, ax=None, label="Trajectory", label_4d="Time"):
+    """
+    Plot colored line in 3D.
+
+    :param x: x values
+    :param y: y values
+    :param z: z values
+    :param t: time values (fourth dimension)
+    :param ax: axis
+    :param label: Label of the line
+    :param label_4d: Label of the 4th dimension (None = don't add)
+    :return ax
+    """
+    if ax is None:
+        fig = plt.figure()
+        ax = fig.add_subplot(projection='3d')
+
+    _plot_color_hack_3d(ax, x, y, z, t)
+    im = ax.scatter(x, y, z, c=t,
+                    label=label, cmap=plt.hot())
+
+    if label_4d is not None:
+        im.set_label(label_4d)
+        plt.colorbar(im, ax=ax)
+
+    return ax