Cours 7

cours/inverted-pendulum/main.py

@@ -0,0 +1,143 @@
+#!/usr/bin/env python3
+
+import pendulum
+import math, time, argparse, sys
+
+from tkinter import *
+from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
+from matplotlib.figure import Figure
+import matplotlib.animation as animation
+
+parser = argparse.ArgumentParser(
+    usage = "%(prog)s [OPTION]",
+    description = "Simulate an inverted pendulum"
+)
+parser.add_argument('-p', '--plot', dest='plot', action='store_true')
+parser.set_defaults(plot=False)
+args = parser.parse_args()
+
+c = pendulum.Cart()
+input_x0 = pendulum.max_w / 2
+mouse_diff = False
+mode_diff = False
+pause = False
+last_time = {0: time.time()}
+
+theta_samples = [0]
+error_samples = [0]
+
+root = Tk()
+if sys.platform == 'linux':
+    root.attributes('-type', 'dialog')
+root.title("Inverted pendulum")
+
+w = Canvas(root, width=pendulum.max_w, height=pendulum.max_h)
+
+f = Frame(root)
+f.pack(side=TOP, fill=X)
+
+vmode = StringVar()
+lmode = Label(f, textvariable=vmode)
+lmode.pack(side=LEFT)
+vinfo = StringVar()
+linfo = Label(f, textvariable=vinfo)
+linfo.pack(side=RIGHT)
+
+def create_line(x1, y1, x2, y2, **kwargs):
+    w.create_line(x1, pendulum.max_h - y1, x2, pendulum.max_h - y2, **kwargs)
+
+def create_oval(x1, y1, x2, y2, **kwargs):
+    w.create_oval(x1, pendulum.max_h - y1, x2, pendulum.max_h - y2, **kwargs)
+
+def push_sample(data, sample):
+    data.append(sample)
+
+def repaint():
+    global input_x0, out, pause, error_samples
+    joint_color = "#4760B2"
+    mass_color = "#EF1010"
+    radius = 4
+    if not pause:
+        out = c.step(input_x0, mouse_diff, mode_diff)
+        push_sample(theta_samples, math.fmod(out.p.theta, 2 * math.pi))
+        push_sample(error_samples, out.p.error)
+    w.delete("all")
+    create_line(0, out.p.y0, pendulum.max_w, out.p.y0, width=5);
+    create_line(out.p.x0, out.p.y0, out.p.x, out.p.y, width=3)
+    create_oval(out.p.x - radius, out.p.y - radius,
+                out.p.x + radius, out.p.y + radius, fill=mass_color)
+    create_oval(out.p.x0 - 2 * radius, out.p.y0 - 2 * radius,
+                out.p.x0 + 2 * radius, out.p.y0 + 2 * radius, fill=joint_color)
+    input_x0 = out.p.x0
+    vmode.set(f"Current mode: {out.p.mode_name}"
+              + (" (paused)" if pause else ""))
+    new_time = time.time()
+    vinfo.set("x: {:07.2f}, theta: {:06.2f}, error: {:06.2f}, fps: {:06.2f}" \
+              .format(out.p.x,
+                      out.p.theta,
+                      out.p.error,
+                      1. / (new_time - last_time[0])))
+    last_time.update({0:new_time})
+
+    w.after(int(1000 * pendulum.dt), repaint)
+
+def mouse_change_callback(event):
+    global mouse_diff, input_x0
+    input_x0 = event.x
+    mouse_diff = not mouse_diff
+
+def key_pressed(event):
+    global input_x0, mouse_diff, mode_diff, pause
+    x_step = 3
+    if event.keysym == 'q':
+        quit()
+    elif event.keysym == 'r':
+        c.reset()
+        input_x0 = pendulum.max_w / 2
+        theta_samples.clear()
+        error_samples.clear()
+    elif event.keysym == 'm':
+        mode_diff = not mode_diff
+    elif event.keysym == 'p':
+        pause = not pause
+    elif event.keysym == 'Right' and not pause:
+        input_x0 = input_x0 + x_step
+        mouse_diff = not mouse_diff
+    elif event.keysym == 'Left' and not pause:
+        input_x0 = input_x0 - x_step
+        mouse_diff = not mouse_diff
+
+w.pack(expand = YES, fill = BOTH)
+root.bind("<B1-Motion>", mouse_change_callback)
+root.bind("<Button-1>", mouse_change_callback)
+root.bind("<Key>", key_pressed)
+
+fig = Figure(figsize=(6, 4.5), dpi=100)
+ax = fig.add_subplot(111)
+text = ax.text(0.97,0.97, "", transform=ax.transAxes, ha="left", va="top")
+line1, = ax.plot([])
+line1.set_label("theta")
+line2, = ax.plot([])
+line2.set_label("error")
+ax.legend(loc = 'upper left')
+
+def animate(i):
+    line1.set_data(range(len(theta_samples)), theta_samples)
+    line2.set_data(range(len(error_samples)), error_samples)
+    ax.set_xlim([0, len(theta_samples)])
+    ax.set_ylim([min(min(theta_samples), min(error_samples)),
+                 max(max(theta_samples), max(error_samples)) + 0.0001])
+
+    return line1, text
+
+if args.plot:
+    canfig = FigureCanvasTkAgg(fig, master=root)
+    canfig.draw()
+    canfig.get_tk_widget().pack(side=BOTTOM, fill=X)
+    ani = animation.FuncAnimation(fig,
+                                  animate,
+                                  interval=int(1000 * pendulum.dt))
+
+root.focus_set()
+repaint()
+mainloop()