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) 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 """ self.lead_car_history_state = objects[0] self.follow_car_history_state = objects[1] self.history_xs = [self.lead_car_history_state[:, 0], self.follow_car_history_state[:, 0]] self.history_ys = [self.lead_car_history_state[:, 1], self.follow_car_history_state[:, 1]] self.history_ths = [self.lead_car_history_state[:, 2], self.follow_car_history_state[:, 2]] # setting up figure self.anim_fig = plt.figure(dpi=150) self.axis = self.anim_fig.add_subplot(111) # imgs self.object_imgs = [] self.traj_imgs = [] def draw_anim(self, interval=50): """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 = 1 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') # (2) set the xlim and ylim self.axis.set_xlim(-5, 50) self.axis.set_ylim(-2, 5) 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.object_imgs.append(temp_img) traj_color_list = ["k", "m"] 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) 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_objects(i) self._draw_traj(i) return self.object_imgs, self.traj_imgs, def _draw_objects(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 for j in range(2): fix_j = j * 2 object_x, object_y, angle_x, angle_y = square_make_with_angles(self.history_xs[j][i], self.history_ys[j][i], 1.0, self.history_ths[j][i]) self.object_imgs[fix_j].set_data([object_x, object_y]) self.object_imgs[fix_j + 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 """ for j in range(2): self.traj_imgs[j].set_data(self.history_xs[j][:i], self.history_ys[j][:i])