def plot(obj, *args, **kwargs): if isinstance(obj, TimeSeries): is_time_series = True else: is_time_series = False if not is_time_series and isinstance(obj, OnlineFunction.Type()): is_truth_solution = False elif is_time_series and isinstance(obj[0], OnlineFunction.Type()): assert all(isinstance(obj_, OnlineFunction.Type()) for obj_ in obj) is_truth_solution = False else: is_truth_solution = True if is_time_series: if "every" in kwargs: obj = obj[::kwargs["every"]] del kwargs["every"] if "interval" in kwargs: anim_interval = kwargs["interval"] del kwargs["interval"] else: anim_interval = None if not is_truth_solution: assert "reduced_problem" in kwargs, ( "Please use this method as plot(reduced_solution, reduced_problem=my_reduced_problem)" + " when plotting a reduced solution") if not is_time_series: N = obj.N else: N = obj[0].N assert all(obj_.N == N for obj_ in obj) reduced_problem = kwargs["reduced_problem"] del kwargs["reduced_problem"] basis_functions = reduced_problem.basis_functions[:N] truth_problem = reduced_problem.truth_problem if not is_time_series: obj = basis_functions * obj else: obj = [basis_functions * obj_ for obj_ in obj] elif "truth_problem" in kwargs: truth_problem = kwargs["truth_problem"] del kwargs["truth_problem"] else: truth_problem = None if "component" in kwargs: component = kwargs["component"] del kwargs["component"] if not is_time_series: obj = obj.sub(component) else: obj = [obj_.sub(component) for obj_ in obj] if truth_problem is not None and hasattr(truth_problem, "mesh_motion"): truth_problem.mesh_motion.move_mesh() if not is_time_series: output = original_plot(obj, *args, **kwargs) else: def animate(i): return original_plot(obj[i], *args, **kwargs).collections fig = plt.figure() output = matplotlib.animation.FuncAnimation( fig, animate, frames=len(obj), interval=anim_interval, repeat=False) try: from IPython.display import HTML except ImportError: pass else: output = HTML(output.to_html5_video()) plt.close() if truth_problem is not None and hasattr(truth_problem, "mesh_motion"): truth_problem.mesh_motion.reset_reference() return output
def animate_planar_quad(t, x, y, θ, title_string=None, display_in_notebook=False): """Animate the planar quadrotor system from given position data. All arguments are assumed to be 1-D NumPy arrays, where `x`, `y`, and `θ` are the degrees of freedom of the planar quadrotor over time `t`. Example usage: import matplotlib.pyplot as plt from animations import animate_planar_quad fig, ani = animate_planar_quad(t, x, θ) ani.save('planar_quad.mp4', writer='ffmpeg') plt.show() """ # Geometry rod_width = 2. rod_height = 0.15 axle_height = 0.2 axle_width = 0.05 prop_width = 0.5 * rod_width prop_height = 1.5 * rod_height hub_width = 0.3 * rod_width hub_height = 2.5 * rod_height # Figure and axis fig, ax = plt.subplots(dpi=100) x_min, x_max = np.min(x), np.max(x) x_pad = (rod_width + prop_width) / 2 + 0.1 * (x_max - x_min) y_min, y_max = np.min(y), np.max(y) y_pad = (rod_width + prop_width) / 2 + 0.1 * (y_max - y_min) ax.set_xlim([x_min - x_pad, x_max + x_pad]) ax.set_ylim([y_min - y_pad, y_max + y_pad]) ax.set_aspect(1.) if title_string is not None: plt.title(title_string) # Artists rod = mpatches.Rectangle((-rod_width / 2, -rod_height / 2), rod_width, rod_height, facecolor='tab:blue', edgecolor='k') hub = mpatches.FancyBboxPatch((-hub_width / 2, -hub_height / 2), hub_width, hub_height, facecolor='tab:blue', edgecolor='k', boxstyle='Round,pad=0.,rounding_size=0.05') axle_left = mpatches.Rectangle((-rod_width / 2, rod_height / 2), axle_width, axle_height, facecolor='tab:blue', edgecolor='k') axle_right = mpatches.Rectangle( (rod_width / 2 - axle_width, rod_height / 2), axle_width, axle_height, facecolor='tab:blue', edgecolor='k') prop_left = mpatches.Ellipse( ((axle_width - rod_width) / 2, rod_height / 2 + axle_height), prop_width, prop_height, facecolor='tab:gray', edgecolor='k', alpha=0.7) prop_right = mpatches.Ellipse( ((rod_width - axle_width) / 2, rod_height / 2 + axle_height), prop_width, prop_height, facecolor='tab:gray', edgecolor='k', alpha=0.7) patches = (rod, hub, axle_left, axle_right, prop_left, prop_right) for patch in patches: ax.add_patch(patch) trace = ax.plot([], [], '--', linewidth=2, color='tab:orange')[0] timestamp = ax.text(0.1, 0.9, '', transform=ax.transAxes) def animate(k, t, x, y, θ): transform = mtransforms.Affine2D().rotate_around(0., 0., θ[k]) transform += mtransforms.Affine2D().translate(x[k], y[k]) transform += ax.transData for patch in patches: patch.set_transform(transform) trace.set_data(x[:k + 1], y[:k + 1]) timestamp.set_text('t = {:.1f} s'.format(t[k])) artists = patches + (trace, timestamp) return artists dt = t[1] - t[0] step = max(int(np.floor((1 / 30) / dt)), 1) # max out at 30Hz for faster rendering ani = animation.FuncAnimation(fig, animate, t[::step].size, fargs=(t[::step], x[::step], y[::step], θ[::step]), interval=step * dt * 1000, blit=True) if display_in_notebook: try: get_ipython() from IPython.display import HTML ani = HTML(ani.to_html5_video()) except (NameError, ImportError): raise RuntimeError( "`display_in_notebook = True` requires this code to be run in jupyter/colab." ) return fig, ani