PythonLinearNonlinearControl/ilqr/animation.py

297 lines
9.2 KiB
Python

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as ani
import matplotlib.font_manager as fon
import sys
import math
# default setting of figures
plt.rcParams["mathtext.fontset"] = 'stix' # math fonts
plt.rcParams['xtick.direction'] = 'in' # x axis in
plt.rcParams['ytick.direction'] = 'in' # y axis in
plt.rcParams["font.size"] = 10
plt.rcParams['axes.linewidth'] = 1.0 # axis line width
plt.rcParams['axes.grid'] = True # make grid
def coordinate_transformation_in_angle(positions, base_angle):
'''
Transformation the coordinate in the angle
Parameters
-------
positions : numpy.ndarray
this parameter is composed of xs, ys
should have (2, N) shape
base_angle : float [rad]
Returns
-------
traslated_positions : numpy.ndarray
the shape is (2, N)
'''
if positions.shape[0] != 2:
raise ValueError('the input data should have (2, N)')
positions = np.array(positions)
positions = positions.reshape(2, -1)
rot_matrix = [[np.cos(base_angle), np.sin(base_angle)],
[-1*np.sin(base_angle), np.cos(base_angle)]]
rot_matrix = np.array(rot_matrix)
translated_positions = np.dot(rot_matrix, positions)
return translated_positions
def square_make_with_angles(center_x, center_y, size, angle):
'''
Create square matrix with angle line matrix(2D)
Parameters
-------
center_x : float in meters
the center x position of the square
center_y : float in meters
the center y position of the square
size : float in meters
the square's half-size
angle : float in radians
Returns
-------
square xs : numpy.ndarray
lenght is 5 (counterclockwise from right-up)
square ys : numpy.ndarray
length is 5 (counterclockwise from right-up)
angle line xs : numpy.ndarray
angle line ys : numpy.ndarray
'''
# start with the up right points
# create point in counterclockwise
square_xys = np.array([[size, 0.5 * size], [-size, 0.5 * size], [-size, -0.5 * size], [size, -0.5 * size], [size, 0.5 * size]])
trans_points = coordinate_transformation_in_angle(square_xys.T, -angle) # this is inverse type
trans_points += np.array([[center_x], [center_y]])
square_xs = trans_points[0, :]
square_ys = trans_points[1, :]
angle_line_xs = [center_x, center_x + math.cos(angle) * size]
angle_line_ys = [center_y, center_y + math.sin(angle) * size]
return square_xs, square_ys, np.array(angle_line_xs), np.array(angle_line_ys)
def circle_make_with_angles(center_x, center_y, radius, angle):
'''
Create circle matrix with angle line matrix
Parameters
-------
center_x : float
the center x position of the circle
center_y : float
the center y position of the circle
radius : float
angle : float [rad]
Returns
-------
circle xs : numpy.ndarray
circle ys : numpy.ndarray
angle line xs : numpy.ndarray
angle line ys : numpy.ndarray
'''
point_num = 100 # 分解能
circle_xs = []
circle_ys = []
for i in range(point_num + 1):
circle_xs.append(center_x + radius * math.cos(i*2*math.pi/point_num))
circle_ys.append(center_y + radius * math.sin(i*2*math.pi/point_num))
angle_line_xs = [center_x, center_x + math.cos(angle) * radius]
angle_line_ys = [center_y, center_y + math.sin(angle) * radius]
return np.array(circle_xs), np.array(circle_ys), np.array(angle_line_xs), np.array(angle_line_ys)
class AnimDrawer():
"""create animation of path and robot
Attributes
------------
cars :
anim_fig : figure of matplotlib
axis : axis of matplotlib
"""
def __init__(self, objects):
"""
Parameters
------------
objects : list of objects
Notes
---------
lead_history_states, lead_history_predict_states, traj_ref, history_traj_ref, history_angle_ref
"""
self.car_history_state = objects[0]
self.target = objects[1]
self.history_target_x = [self.target[:, 0]]
self.history_target_y = [self.target[:, 1]]
self.history_xs = [self.car_history_state[:, 0]]
self.history_ys = [self.car_history_state[:, 1]]
self.history_ths = [self.car_history_state[:, 2]]
# setting up figure
self.anim_fig = plt.figure()
self.axis = self.anim_fig.add_subplot(111)
# imgs
self.car_imgs = []
self.traj_imgs = []
def draw_anim(self, interval=10):
"""draw the animation and save
Parameteres
-------------
interval : int, optional
animation's interval time, you should link the sampling time of systems
default is 50 [ms]
"""
self._set_axis()
self._set_img()
self.skip_num = 2
frame_num = int((len(self.history_xs[0])-1) / self.skip_num)
animation = ani.FuncAnimation(self.anim_fig, self._update_anim, interval=interval, frames=frame_num)
# self.axis.legend()
print('save_animation?')
shuold_save_animation = int(input())
if shuold_save_animation:
print('animation_number?')
num = int(input())
animation.save('animation_{0}.mp4'.format(num), writer='ffmpeg')
# animation.save("Sample.gif", writer = 'imagemagick') # gif保存
plt.show()
def _set_axis(self):
""" initialize the animation axies
"""
# (1) set the axis name
self.axis.set_xlabel(r'$\it{x}$ [m]')
self.axis.set_ylabel(r'$\it{y}$ [m]')
self.axis.set_aspect('equal', adjustable='box')
LOW_MARGIN = 2
HIGH_MARGIN = 2
self.axis.set_xlim(np.min(self.history_xs) - LOW_MARGIN, np.max(self.history_xs) + HIGH_MARGIN)
self.axis.set_ylim(np.min(self.history_ys) - LOW_MARGIN, np.max(self.history_ys) + HIGH_MARGIN)
def _set_img(self):
""" initialize the imgs of animation
this private function execute the make initial imgs for animation
"""
# object imgs
obj_color_list = ["k", "k", "m", "m"]
obj_styles = ["solid", "solid", "solid", "solid"]
for i in range(len(obj_color_list)):
temp_img, = self.axis.plot([], [], color=obj_color_list[i], linestyle=obj_styles[i])
self.car_imgs.append(temp_img)
traj_color_list = ["k", "b"]
for i in range(len(traj_color_list)):
temp_img, = self.axis.plot([],[], color=traj_color_list[i], linestyle="dashed")
self.traj_imgs.append(temp_img)
temp_img, = self.axis.plot([],[], "*", color="b")
self.traj_imgs.append(temp_img)
def _update_anim(self, i):
"""the update animation
this function should be used in the animation functions
Parameters
------------
i : int
time step of the animation
the sampling time should be related to the sampling time of system
Returns
-----------
object_imgs : list of img
traj_imgs : list of img
"""
i = int(i * self.skip_num)
# self._draw_set_axis(i)
self._draw_car(i)
self._draw_traj(i)
# self._draw_prediction(i)
return self.car_imgs, self.traj_imgs,
def _draw_set_axis(self, i):
"""
"""
# (2) set the xlim and ylim
LOW_MARGIN = 20
HIGH_MARGIN = 20
OVER_LOOK = 50
self.axis.set_xlim(np.min(self.history_xs[0][i : i + OVER_LOOK]) - LOW_MARGIN, np.max(self.history_xs[0][i : i + OVER_LOOK]) + HIGH_MARGIN)
self.axis.set_ylim(np.min(self.history_ys[0][i : i + OVER_LOOK]) - LOW_MARGIN, np.max(self.history_ys[0][i : i + OVER_LOOK]) + HIGH_MARGIN)
def _draw_car(self, i):
"""
This private function is just divided thing of
the _update_anim to see the code more clear
Parameters
------------
i : int
time step of the animation
the sampling time should be related to the sampling time of system
"""
# cars
object_x, object_y, angle_x, angle_y = square_make_with_angles(self.history_xs[0][i],
self.history_ys[0][i],
1.0,
self.history_ths[0][i])
self.car_imgs[0].set_data([object_x, object_y])
self.car_imgs[1].set_data([angle_x, angle_y])
def _draw_traj(self, i):
"""
This private function is just divided thing of
the _update_anim to see the code more clear
Parameters
------------
i : int
time step of the animation
the sampling time should be related to the sampling time of system
"""
# car
self.traj_imgs[0].set_data(self.history_xs[0][:i], self.history_ys[0][:i])
# goal
self.traj_imgs[-1].set_data(self.history_target_x[0][i-1], self.history_target_y[0][i-1])
# traj_ref
self.traj_imgs[1].set_data(self.history_target_x[0][:i], self.history_target_y[0][:i])