def plot_both_things(sol, masses, plot_times, cell_start, cell_end, eps, y0, y1, gridpoints, title, save_location, T): fig, (ax1, ax2) = plt.subplots(1, 2) ax1.set_xlabel('time') ax1.set_ylabel('position') ax2.set_xlabel('position') ax2.set_ylabel('density') ax1.set_title('particle trajectories colored by mass') ax2.set_title('profiles colored by timestep') fig.suptitle(title) t = sol['t'] traj = sol['y'] cmap = plt.get_cmap('jet') cNorm = colors.Normalize(vmin=0, vmax=max(masses)) ScalarMap1 = cmx.ScalarMappable(norm=cNorm, cmap=cmap) for i in range(len(masses)): linecolor = ScalarMap1.to_rgba(masses[i]) ax1.plot(t, traj[i], color=linecolor) ax1.plot([0, 0], [y0, y1], 'ro') ax3 = fig.add_axes([0.01, 0.12, 0.02, 0.7]) cb = matplotlib.colorbar.ColorbarBase(ax3, cmap=cmap, norm=cNorm) cNorm = colors.Normalize(vmin=0, vmax=T) ScalarMap2 = cmx.ScalarMappable(norm=cNorm, cmap=cmap) grid = np.linspace( np.min(np.array(sol['y'])) - .1, np.max(np.array(sol['y'])) + .1, gridpoints) for i in plot_times: particles = sol['y'][:, i] profile = blob(particles, grid, masses, eps) true_time = (T / (max(plot_times) - 1)) * i linecolor = ScalarMap2.to_rgba(true_time) ax2.plot(grid, profile, color=linecolor) ax2.plot([y0, y1], [0, 0], 'ro') ax4 = fig.add_axes([0.94, 0.15, 0.02, 0.7]) cb = matplotlib.colorbar.ColorbarBase(ax4, cmap=cmap, norm=cNorm) mng = plt.get_current_fig_manager() mng.full_screen_toggle() fig.set_size_inches((10, 5), forward=False) fig.savefig(save_location + '.png', dpi=500) return fig
def compute_and_plot_profiles(sol, masses, plot_times, cell_start, cell_end,eps,y0,y1,gridpoints, title): cmap = plt.get_cmap('viridis') fig, ax = plt.subplots(figsize = (7,5)) ax.set_title(title) ax.set_xlabel('position') ax.set_ylabel('density') cNorm = colors.Normalize(vmin = 0, vmax = max(plot_times)) ScalarMap = cmx.ScalarMappable(norm = cNorm, cmap = cmap) grid = np.linspace(cell_start, cell_end, gridpoints) for i in plot_times: particles = sol['y'][:,i] profile = blob(particles, grid, masses, eps) linecolor = ScalarMap.to_rgba(i) ax.plot(grid, profile, color = linecolor) ax.plot([y0,y1],[0,0],'ro') ax2 = fig.add_axes([0.94,0.1,0.02,0.7]) cb = matplotlib.colorbar.ColorbarBase(ax2, cmap=cmap, norm=cNorm) plt.savefig(title)
def plot_four_profiles(x_label, y_label, sol_dict_path, param_dict_path, plot_number, title, param_set = "drift", custom_params = None, save = False, save_path = None): fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2,2, sharex = True, sharey = True) #ax1.set_xlabel('position') ax1.set_ylabel(y_label) #ax2.set_xlabel('position') #ax2.set_ylabel('density') ax3.set_xlabel(x_label) ax3.set_ylabel(y_label) ax4.set_xlabel(x_label) #ax4.set_ylabel('density') fig.suptitle(title) ax4_box = fig.axes[3].get_position() ax2_box = fig.axes[1].get_position() ax4_box.x1 = ax4_box.x1 - .04 ax4_box.x0 = ax4_box.x0 - .04 ax2_box.x1 = ax2_box.x1 - .04 ax2_box.x0 = ax2_box.x0 - .04 fig.axes[3].set_position(ax4_box) fig.axes[1].set_position(ax2_box) if param_set == 'drift': param_1 = (10,5,1) param_2 = (10,1,5) param_3 = (10,1,1) param_4 = (10,5,5) elif param_set == 'interaction': param_1 = (5,10,1) param_2 = (1,10,5) param_3 = (1,10,1) param_4 = (5,10,5) else: param_1 = custom_params[0] param_2 = custom_params[1] param_3 = custom_params[2] param_4 = custom_params[3] ax1.set_title('$A= {}$, $B ={} $, $C = {}$'.format(param_1[0], param_1[1], param_1[2]), pad = 3) ax2.set_title('$A= {}$, $B ={} $, $C = {}$'.format(param_2[0], param_2[1], param_2[2]), pad = 3) ax3.set_title('$A= {}$, $B ={} $, $C = {}$'.format(param_3[0], param_3[1], param_3[2]), pad = 3) ax4.set_title('$A= {}$, $B ={} $, $C = {}$'.format(param_4[0], param_4[1], param_4[2]), pad = 3) sol_dict = pickle.load(open(sol_dict_path, "rb")) param_dict = pickle.load(open(param_dict_path, "rb")) sol_dict.keys() N = param_dict['N'] masses = param_dict['M'] eps = param_dict['eps'] y0 = param_dict['y0'] y1 = param_dict['y1'] T = param_dict['T'] sols = [sol_dict[param_1], sol_dict[param_2], sol_dict[param_3], sol_dict[param_4]] plot_times = [compute_even_plot_times(len(sol['t']), plot_number) for sol in sols] axess = [ax1, ax2, ax3, ax4] cmap = plt.get_cmap('viridis') cNorm = colors.Normalize(vmin = 0, vmax = T) ScalarMap = cmx.ScalarMappable(norm = cNorm, cmap = cmap) colorbar = fig.add_axes([0.90,0.15,0.02,0.7]) colorbar.set_title('time') cb = matplotlib.colorbar.ColorbarBase(colorbar, cmap = cmap, norm = cNorm) for i in range(4): ax = axess[i] plot_time = plot_times[i] sol = sols[i] grid = np.linspace(-2.1, 2.1, 100) for i in plot_time: particles = sol['y'][:,i] profile = blob(particles, grid, masses, eps) true_time = (T / (max(plot_time) -1)) * i linecolor = ScalarMap.to_rgba(true_time) ax.plot(grid, profile, color = linecolor) ax.plot([y0,y1],[0,0],'ro') return fig
def plot_six_profiles(x_label, y_label, sol_dict_path, param_dict_path, plot_number, title, moll_eps, params=None): fig, ((ax1, ax2), (ax3, ax4), (ax5, ax6)) = plt.subplots(3, 2, sharex=True, sharey=True) #ax1.set_xlabel('position') ax1.set_ylabel(y_label) #ax2.set_xlabel('position') #ax2.set_ylabel('density') ax3.set_ylabel(y_label) ax5.set_ylabel(y_label) ax5.set_xlabel(x_label) ax6.set_xlabel(x_label) #ax4.set_ylabel('density') fig.suptitle(title) ax4_box = fig.axes[3].get_position() ax2_box = fig.axes[1].get_position() ax6_box = fig.axes[5].get_position() ax4_box.x1 = ax4_box.x1 - .04 ax4_box.x0 = ax4_box.x0 - .04 ax2_box.x1 = ax2_box.x1 - .04 ax2_box.x0 = ax2_box.x0 - .04 ax6_box.x1 = ax6_box.x1 - .04 ax6_box.x0 = ax6_box.x0 - .04 fig.axes[3].set_position(ax4_box) fig.axes[1].set_position(ax2_box) fig.axes[5].set_position(ax6_box) fig.set_figheight(fig.get_figheight() + fig.get_figheight() * .3) param_1 = params[0] param_2 = params[1] param_3 = params[2] param_4 = params[3] param_5 = params[4] param_6 = params[5] ax1.set_title('$A= {}$, $B ={} $, $C = {}$'.format(param_1[0], param_1[1], param_1[2]), pad=3) ax2.set_title('$A= {}$, $B ={} $, $C = {}$'.format(param_2[0], param_2[1], param_2[2]), pad=3) ax3.set_title('$A= {}$, $B ={} $, $C = {}$'.format(param_3[0], param_3[1], param_3[2]), pad=3) ax4.set_title('$A= {}$, $B ={} $, $C = {}$'.format(param_4[0], param_4[1], param_4[2]), pad=3) ax5.set_title('$A= {}$, $B ={} $, $C = {}$'.format(param_5[0], param_5[1], param_5[2]), pad=3) ax6.set_title('$A= {}$, $B ={} $, $C = {}$'.format(param_6[0], param_6[1], param_6[2]), pad=3) sol_dict = pickle.load(open(sol_dict_path, "rb")) param_dict = pickle.load(open(param_dict_path, "rb")) sol_dict.keys() N = param_dict['N'] masses = param_dict['M'] eps = param_dict['eps'] y0 = param_dict['y0'] y1 = param_dict['y1'] T = param_dict['T'] sols = [ sol_dict[param_1], sol_dict[param_2], sol_dict[param_3], sol_dict[param_4], sol_dict[param_5], sol_dict[param_6] ] plot_times = [ compute_even_plot_times(len(sol['t']), plot_number) for sol in sols ] axess = [ax1, ax2, ax3, ax4, ax5, ax6] cmap = plt.get_cmap('jet') cNorm = colors.Normalize(vmin=0, vmax=T) ScalarMap = cmx.ScalarMappable(norm=cNorm, cmap=cmap) colorbar = fig.add_axes([0.90, 0.15, 0.02, 0.7]) colorbar.set_title('time') cb = matplotlib.colorbar.ColorbarBase(colorbar, cmap=cmap, norm=cNorm) for i in range(6): ax = axess[i] plot_time = plot_times[i] sol = sols[i] grid = np.linspace(-2.1, 2.1, 100) for i in plot_time: particles = sol['y'][:, i] profile = blob(particles, grid, masses, moll_eps) true_time = (T / (max(plot_time) - 1)) * i linecolor = ScalarMap.to_rgba(true_time) ax.plot(grid, profile, color=linecolor) ax.plot([y0, y1], [0, 0], 'ro') return fig