# # viz-attractor.py: A basic set of functions + main() for visualization and animiation of (strange) attractors of a # 3D non linear dynamical system. It shows: the attractor in (x,y,z), the two dimensional projections, the # time evolution of the individual variables, the motion of two solutions initially close. # A few non-linear systems (Lorenz, Rossler, Lotka-Volterra) are included as examples. # # Used in 15-382, Collective Intelligence # Author: Gianni A. Di Caro, CMU-Q, 2018 # Free to use acknowledging the source # # import numpy as np import matplotlib.pyplot as plt from mpl_toolkits.mplot3d import Axes3D from matplotlib import animation from scipy.integrate import odeint # Lorenz system, fluid convection dynamics def lorenz(state, t): # parameters of the model. To observer different attractors / behaviors, change the value of rho # (interesting ranges: [0,1], ]1,28[, [28, ]) rho = 28.0 sigma = 10.0 beta = 8.0 / 3.0 x, y, z = state dxdt = sigma * (y - x) dydt = x * (rho - z) - y dzdt = x * y - beta * z return dxdt, dydt, dzdt # Roessler system, from chemical reactions def rossler(state, t): a = 0.2 b = 0.2 c = 5.7 x, y, z = state dxdt = -(y + z) dydt = x + a*y dzdt = b + x * z - c * z return dxdt, dydt, dzdt def lotka_volterra_3d(state, t): a = 1 b = 1 c = 1 d = 1 e = 1 f = 1 g = 1 x, y, z = state dxdt = x * (a - b*y) dydt = y*(-c + d*x - e*z) dzdt = z*( -f + g*y) return dxdt, dydt, dzdt # plot trajectories based on a number of selected initial conditions def plot_trajectories(dynamical_sys, initial_conditions, colors, seconds, time_step): fig = plt.figure("Attractor") ax = fig.gca(projection='3d') ax.set_xlabel('$X$') ax.set_ylabel('$Y$') ax.set_zlabel('$Z$') t = np.arange(0.0, seconds, time_step) for i in range(len(initial_conditions)): state0 = initial_conditions[i] #states = odeint(lorenz, state0, t) states = odeint(dynamical_sys, state0, t) ax.plot(states[:,0], states[:,1], states[:,2], color=colors[i], alpha=0.95, linewidth=0.4) plt.ioff() plt.gcf().show() # keep the window open raw_input("\n\t\t--------------- Press Enter to continue ... ------------------") show_two_dimensional_projections(states[:, 0], states[:, 1], states[:, 2]) show_individual_variables(states[:, 0], states[:, 1], states[:, 2], t) # plot the x-y, x-z, and y-z projections of a trajectory def show_two_dimensional_projections(x, y, z): fig, ax = plt.subplots(1, 3, sharex=False, sharey=False, figsize=(15, 4)) fig.canvas.set_window_title('2d projections') # # plot the x values vs the y values ax[0].plot(x, y, color='r', linewidth=0.8) ax[0].set_title('x-y phase plane') # # plot the x values vs the z values ax[1].plot(x, z, color='m', linewidth=0.8) ax[1].set_title('x-z phase plane') # # plot the y values vs the z values ax[2].plot(y, z, color='b', linewidth=0.8) ax[2].set_title('y-z phase plane') # fig.savefig('{}/lorenz-attractor-phase-plane.png'.format(save_folder), dpi=180, bbox_inches='tight') plt.ioff() plt.gcf().show() # keep the window open raw_input("\n\t\t--------------- Press Enter to continue ... ------------------") # plot the time evolution of a single state variable def show_individual_variables(x, y, z, t): fig, ax = plt.subplots(1, 3, sharex=False, sharey=False, figsize=(15, 4)) fig.canvas.set_window_title('Single state variable, time evolution') # # plot the x values vs time ax[0].plot(t, x, color='r', alpha=0.7, linewidth=0.4) ax[0].set_title('x(t)') # # plot the x values vs the z values ax[1].plot(t, y, color='m', alpha=0.7, linewidth=0.4) ax[1].set_title('y(t)') # # plot the y values vs the z values ax[2].plot(t, z, color='b', alpha=0.7, linewidth=0.4) ax[2].set_title('z(t)') plt.ioff() plt.gcf().show() # keep the window open raw_input("\n\t\t--------------- Press Enter to continue ... ------------------") # set the data for the point to be move along a trajectory # this function is called sequentially by FuncAnimation() # def set_point_on_trajectory(i, states, states_near, attractor): # set the plot data for the moving point on the figure "attractor" #attractor.set_data(states[i,0], states[i,1]) #attractor.set_3d_properties(states[i,2]) # set the data for two moving points on the attractor attractor.set_data([states[i,0], states_near[i,0]], [states[i,1], states_near[i,1]]) attractor.set_3d_properties([states[i,2], states_near[i,2]]) return attractor, # show the moving point over a selected trajectory, make use of FuncAnimation() that calls set_point_on_trajectory() def move_trajectory_point(dynamical_sys, initial_condition, seconds, time_step): fig = plt.figure("Motion on the attractor") ax1 = fig.gca(projection='3d') ax1.set_xlabel('$X$') ax1.set_ylabel('$Y$') ax1.set_zlabel('$Z$') attractor, = ax1.plot([], [], [], 'ro', alpha=0.7, linewidth=0.4) t = np.arange(0.0, seconds, time_step) frame_num = int((seconds / time_step)) state0 = initial_condition states = odeint(dynamical_sys, state0, t) state0_near = [initial_condition[0] + 0.11, initial_condition[1]+ 0.11, initial_condition[2] + 0.01 ] states_near = odeint(dynamical_sys, state0_near, t) ax1.set_xlim(min(states[:,0]), max(states[:,0])) ax1.set_ylim(min(states[:,1]), max(states[:,1])) ax1.set_zlim(min(states[:,2]), max(states[:,2])) # the "background" image, with the full trajectory ax1.plot(states[:,0], states[:,1], states[:,2], color="g", alpha=0.7, linewidth=0.4) # call the animator anim_speed = 1 anim = animation.FuncAnimation(fig, set_point_on_trajectory, fargs=(states, states_near, attractor,), frames=frame_num, interval=anim_speed) plt.ioff() plt.gcf().show() # keep the window open raw_input("\n\t\t--------------- Press Enter to continue ... ------------------") # Uncomment the following lines to save an .mp4 of the animation #writer = animation.FFMpegWriter(bitrate=500) #anim.save('attractor.mp4', writer=writer,fps=5) # -------------------- main() -------------------------------------------- # # if __name__ == '__main__': # Set of (arbitrary) four initial conditions, that will be plotted as four trajectories # in different colors. One of the trajectories is also plotted in the two dimensional planes (xy, xz, yz) # and will be used for animate the motion on the attractor. # Initial conditions and evolution times are provided for the three included examples non-linear systems # initial_conditions_lorenz = [[1.0, 1.0, 1.0], [3.0, 3.0, 3.0], [-10, -10, 10], [-5,-2,0]] seconds_lorenz = 40.0 time_step_lorenz = 0.01 initial_conditions_rossler = [[10, -10, 5], [5, -20, 10], [5, 5, 5], [10,-2,10]] seconds_rossler = 200.0 time_step_rossler = 0.01 initial_conditions_lv = [[0.25, 0.5, 2.5], [0.5, 0.5, 2], [1, 0.5, 1.5], [1.7, 0.5, 1.2]] seconds_lv = 200.0 time_step_lv = 0.01 # select the system to show and get its parameters dynamical_system = lorenz initial_conditions = initial_conditions_lorenz seconds = seconds_lorenz time_step = time_step_lorenz colors = ["r", "g", "b", "y"] plot_trajectories(dynamical_system, initial_conditions, colors, seconds, time_step) move_trajectory_point(dynamical_system, initial_conditions[0], seconds, time_step) plt.close('all')