def cols(choice, its, noise, miniters = 0, highlight = -1): '''Returns an array of colors mapped to iteration by index''' colors = zeros((its, 3), uint8) # switches between highlighting individual iterations or just giving a color map if highlight == -1: colors[0] = (0, 0, 0) if choice == 0: for i in range(miniters, its): colors[i] = (int(remap(i, miniters, its, 0.0, 255.0)), int(remap(i, miniters, its, 0.0, 255.0)), int(remap(i, miniters, its, 0.0, 255.0))) return colors if choice == 1: for i in range(1, its): colors[i] = (colmap(cm.gnuplot2(abs(remap(((i + 256) % 512), 0, 512, -1.0, 1.0)))[0]), colmap(cm.gnuplot2(abs(remap(((i + 256) % 512), 0, 512, -1.0, 1.0)))[1]), colmap(cm.gnuplot2(abs(remap(((i + 256) % 512), 0, 512, -1.0, 1.0)))[2])) return colors if choice == 2: for i in range(1, its): colors[i] = noise.gencolor(i) return colors else: if choice == 0: colors[highlight] = (255, 255, 255) return colors if choice == 1: colors[highlight] = (colmap(cm.gnuplot2(abs(remap(((highlight + 256) % 512), 0, 512, -1.0, 1.0)))[0]), colmap(cm.gnuplot2(abs(remap(((highlight + 256) % 512), 0, 512, -1.0, 1.0)))[1]), colmap(cm.gnuplot2(abs(remap(((highlight + 256) % 512), 0, 512, -1.0, 1.0)))[2])) return colors if choice == 2: colors[highlight] = noise.gencolor(highlight) return colors
def get_rgb(temp): # NOTE: this takes considerable time. Use binning later on. # Or even give everything to the shader. # Good ones are afmhot, CMRmap, gist_heat, gnuplot, gnuplot2. # See http://matplotlib.org/users/colormaps.html for more. color = cm.gnuplot2(int(temp), bytes=True) return '{r},{g},{b}'.format(r=color[0], g=color[1], b=color[2])
def get_colors(self, qty): qty = np.power(qty / qty.max(), 1.0 / CONTRAST) if COLORMAP == 0: rgba = cm.gray(qty, alpha=ALPHA) elif COLORMAP == 1: rgba = cm.afmhot(qty, alpha=ALPHA) elif COLORMAP == 2: rgba = cm.hot(qty, alpha=ALPHA) elif COLORMAP == 3: rgba = cm.gist_heat(qty, alpha=ALPHA) elif COLORMAP == 4: rgba = cm.copper(qty, alpha=ALPHA) elif COLORMAP == 5: rgba = cm.gnuplot2(qty, alpha=ALPHA) elif COLORMAP == 6: rgba = cm.gnuplot(qty, alpha=ALPHA) elif COLORMAP == 7: rgba = cm.gist_stern(qty, alpha=ALPHA) elif COLORMAP == 8: rgba = cm.gist_earth(qty, alpha=ALPHA) elif COLORMAP == 9: rgba = cm.spectral(qty, alpha=ALPHA) return rgba
def draw_agent(agent: ConspiracyAgent) -> dict: portrayal = { 'Shape': 'circle', 'r': 0.5, 'Filled': 'true', 'Layer': 0, 'Color': c.to_hex(cm.gnuplot2(agent.prior_value)), } return portrayal
def main(): f = open('label.txt','w') #target_names = np.array(args.names) X, target_names, y = getXY(args.image_dir) X = np.asfarray(X,dtype='float') colors = cm.gnuplot2(np.linspace(0, 1, len(target_names))) #X_pca = PCA(n_components=128).fit_transform(X) X_pca = X tsne = TSNE(n_components=2, init='random', random_state=0) X_r = tsne.fit_transform(X_pca) for c, i, target_name in zip(colors, list(range(0, len(target_names))), target_names): plt.scatter(X_r[y[i], 0], X_r[y[i], 1], c=c, label=str(i+1)) f.write(target_name+'\n') plt.legend() plt.savefig("{}/10crop1.png".format('./')) f.close()
def make_plots_3to7(dict_in, param_varied, var_name): from matplotlib.ticker import MultipleLocator import matplotlib.cm as cm fig, ((ax1, ax2), (ax3, ax4)) = pl.subplots(2, 2, figsize=[12, 12], sharex=True) dashes = [(None, None), (2, 2), (5, 5), (5, 3, 1, 3), (5, 3, 1, 3, 1, 3), (10, 4)] colors = cm.gnuplot2( np.union1d([0], np.linspace(0.3, 0.8, len(dict_in.keys()) - 1))) i = 0 ks = [float(val) for val in dict_in.keys()] ks = [str(val) for val in np.sort(ks)] for val in ks: log_mL = dict_in[val]['log_mL'] VRcolor = dict_in[val]['VRcolor'] UVcolor = dict_in[val]['UVcolor'] VKcolor = dict_in[val]['VKcolor'] times = dict_in[val]['times'] ax1.plot(times, log_mL, marker='', dashes=dashes[i], color=colors[i], lw=2, label=param_varied + ' = ' + str(val)) ax2.plot(times, UVcolor, marker='', dashes=dashes[i], color=colors[i], lw=2, label=param_varied + ' = ' + str(val)) ax3.plot(times, VRcolor, marker='', dashes=dashes[i], color=colors[i], lw=2, label=param_varied + ' = ' + str(val)) ax4.plot(times, VKcolor, marker='', dashes=dashes[i], color=colors[i], lw=2, label=param_varied + ' = ' + str(val)) i += 1 # Make figures look nice # Upper Left: M/L vs time ax1.set_xlim([0, 15]) ax1.set_ylim([-0.3, 1.0]) ax1.set_xticks([0, 5, 10, 15]) minor_locator = MultipleLocator(1) ax1.xaxis.set(minor_locator=minor_locator) ax1.set_yticks(np.arange(-0.2, 1.1, 0.2)) minor_locator = MultipleLocator(0.2 / 4) ax1.yaxis.set(minor_locator=minor_locator) ax1.set_ylabel(r'log(M/$\rm L_V$)') ax1.legend(loc='lower right') # Upper Right: U - V vs time ax2.set_ylim([0.6, 1.8]) ax2.set_yticks(np.arange(0.6, 1.81, 0.2)) minor_locator = MultipleLocator(0.2 / 4) ax2.yaxis.set(minor_locator=minor_locator) ax2.set_ylabel(r'U - V') # Lower Left: V - R vs time ax3.set_ylim(0.4, 0.7) ax3.set_yticks(np.arange(0.4, 0.71, 0.05)) minor_locator = MultipleLocator(0.05 / 4) ax3.yaxis.set(minor_locator=minor_locator) ax3.set_ylabel(r'V - R') ax3.set_xlabel(r'Time (Gyr)') # Lower Right: V - K vs time ax4.set_ylim([2.0, 4.0]) ax4.set_yticks(np.arange(2, 4.1, 0.5)) minor_locator = MultipleLocator(0.5 / 4) ax4.yaxis.set(minor_locator=minor_locator) ax4.set_xlabel('Time (Gyr)') ax4.set_ylabel('V - K') fig.savefig('test_%s.png' % var_name, bbox='tight') #pl.show() pl.close() return [fig]
import matplotlib.pyplot as plt import numpy as np import array as arr import matplotlib.cm as cm from itertools import cycle from matplotlib.figure import figaspect my_colors = iter(cm.gnuplot2(np.linspace(0, 1, 8))) thread_default = np.array([88]) runtime_default = np.array([0.0264344]) threads = np.array([1, 2, 4, 8, 16, 22, 44, 88]) af_none = np.array([ 0.324832, 0.163288, 0.0842142, 0.0435799, 0.0230997, 0.0194625, 0.0210078, 0.0264214 ]) af_sockets_false = np.array([ 0.325076, 0.163253, 0.0838667, 0.0436435, 0.023046, 0.0190545, 0.0206504, 0.0238985 ]) af_cores_false = np.array([ 0.324648, 0.16354, 0.0837889, 0.0434537, 0.0231496, 0.0194563, 0.0212046, 0.0237055 ]) af_threads_false = np.array([ 0.325007, 0.163234, 0.0840898, 0.0434851, 0.0231797, 0.0194611, 0.0210849, 0.0239731 ]) af_sockets_true = np.array([ 0.326206, 0.163408, 0.0837538, 0.0521883, 0.0292147, 0.0231815, 0.0199503, 0.0232362
Note: This example assumes that `name i` corresponds to `label i` in `labels.csv`. """) parser = argparse.ArgumentParser() parser.add_argument('workDir', type=str) parser.add_argument('--names', type=str, nargs='+', required=True) args = parser.parse_args() y = pd.read_csv("{}/labels.csv".format(args.workDir)).as_matrix()[:, 0] X = pd.read_csv("{}/reps.csv".format(args.workDir)).as_matrix() target_names = np.array(args.names) colors = cm.gnuplot2(np.linspace(0, 0.7, len(target_names))) X_pca = PCA(n_components=50).fit_transform(X, X) tsne = TSNE(n_components=2, init='random', random_state=0) X_r = tsne.fit_transform(X_pca) for c, i, target_name in zip(colors, list(range(1, len(target_names) + 1)), target_names): plt.scatter(X_r[y == i, 0], X_r[y == i, 1], c=c, label=target_name) plt.legend() out = "{}/tsne.pdf".format(args.workDir) plt.savefig(out) print("Saved to: {}".format(out))
#plt.plot(dG_FEP[1,:],-dG_cumulant[0,1,:]-dG_cumulant[1,1,:],'g.') #plt.plot(dG_FEP[2,:],-dG_cumulant[0,2,:]-dG_cumulant[1,2,:],'b.') #plt.plot([0],[0],'r.',label="U") #plt.plot([0],[0],'g.',label="TS") #plt.plot([0],[0],'b.',label="N") plt.show() raise SystemExit Y = dcontact_E/(0.5*(stdEij_U**2 + stdEij_N**2)) plt.figure() for i in range(Eij.shape[1]): plt.plot(Qavg,Y[:,i],color=cm.gnuplot2((relative_TS_fluct[i] - relative_TS_fluct.min())/relative_TS_fluct.max())) plt.xlabel("$Q$",fontsize=18) plt.ylabel("$\\frac{2\\langle\\delta\\epsilon_{ij}^2\\rangle(Q)}{\\langle\\delta\\epsilon_{ij}^2\\rangle_U + \\langle\\delta\\epsilon_{ij}^2\\rangle_N}$",fontsize=22) #plt.title("$\\frac{2\\langle\\delta\\epsilon_{ij}^2\\rangle}{\\left(\\langle\\delta\\epsilon_{ij}^2\\rangle_U + \\langle\\delta\\epsilon_{ij}^2\\rangle_N\\right)}$") plt.title("Energy Fluctuations Relative to average of U and N") plt.savefig("eps_fluct_relative_U_and_N.png") plt.savefig("eps_fluct_relative_U_and_N.pdf") plt.savefig("eps_fluct_relative_U_and_N.eps") Z = dcontact_E/(stdEij_N**2) plt.figure() for i in range(Eij.shape[1]): plt.plot(Qavg,Z[:,i],color=cm.gnuplot2((relative_TS_fluct[i] - relative_TS_fluct.min())/relative_TS_fluct.max())) plt.xlabel("Q") #plt.ylabel("$\\langle\\delta\\epsilon_{ij}^2\\rangle / \\langle\\delta\\epsilon_{ij}^2\\rangle_N $")
np.savetxt("contact_Q.dat", contact_Q) np.savetxt("dQ2.dat", dQ2) np.savetxt("contact_avgE.dat", contact_avgE) np.savetxt("dcontact_E.dat", dcontact_E) np.savetxt("dcontact_S.dat", dcontact_E) np.savetxt("ddG_vs_Q.dat", ddG_vs_Q) np.savetxt("ddG_deriv_vs_Q.dat", ddG_deriv_vs_Q) Y = dcontact_E / (0.5 * (stdEij_U**2 + stdEij_N**2)) plt.figure() for i in range(Eij.shape[1]): plt.plot(Qavg, Y[:, i], color=cm.gnuplot2( (relative_TS_fluct[i] - relative_TS_fluct.min()) / relative_TS_fluct.max())) plt.xlabel("$Q$", fontsize=18) plt.ylabel( "$\\frac{2\\langle\\delta\\epsilon_{ij}^2\\rangle(Q)}{\\langle\\delta\\epsilon_{ij}^2\\rangle_U + \\langle\\delta\\epsilon_{ij}^2\\rangle_N}$", fontsize=22) #plt.title("$\\frac{2\\langle\\delta\\epsilon_{ij}^2\\rangle}{\\left(\\langle\\delta\\epsilon_{ij}^2\\rangle_U + \\langle\\delta\\epsilon_{ij}^2\\rangle_N\\right)}$") plt.title("Energy Fluctuations Relative to average of U and N") plt.savefig("eps_fluct_relative_U_and_N.png") plt.savefig("eps_fluct_relative_U_and_N.pdf") plt.savefig("eps_fluct_relative_U_and_N.eps") Z = dcontact_E / (stdEij_N**2) plt.figure() for i in range(Eij.shape[1]): plt.plot(Qavg,
S = S - 1 if j == 1: S == S b = np.random.choice([2, 0], I, p=[γ, 1 - γ]) #infected individuals have chances to recover or stay ill for e in b: if e == 2: R = +1 #to append the newly recovered people and delete it in the infectious number I = I - 1 if e == 0: I == I B.append(I) #store the result of I for a vaccination rate BB.append(B) #to store infectious population of different vaccination rate B = [1] #plot #print('the number of infectious',B) plt.figure(figsize=(6, 4), dpi=150) for j in range(0, 11): X = range(0, 1001) f = plt.plot( X, BB[j], label=str(j * 10) + '%', color=cm.gnuplot2(j * 20) ) #the plot function is inside the loop, so it can be ploted in one figure #plt.savefig ("markflow" ,type="png") #plt.plot(x, y, format_string, **kwargs): #x is the X-axis data, can be a list or array;Y in the same way;Format string for control curve, **kwargs second group or more (x, y, format_string) plt.xlabel("times", fontsize=14) plt.ylabel("number of people", fontsize=14) plt.title("SIR model with different vaccination rates") plt.legend() plt.show()
import os import sys # sys.path.insert(0, r'./') import synthwave_parts as synth matplotlib.matplotlib_fname() rcParams['animation.convert_path'] = '/usr/local/bin/convert' # import synthwave_parts as synth background = '#000000' # black to purple gradient horizon = {'grad': cm.gnuplot(range(28)), 'dist': 16} ground = {'grad': cm.gnuplot2(range(20)), 'dist': 50} matplotlib.rcParams['axes.facecolor'] = background # key paramters perspective = (0, 5) motion_lines = 10 frames = 200 xlim = (-100, 100) ylim = (-50, 70) # figure and axes setup fig = plt.figure(figsize=(20, 12)) ax = plt.axes(xlim=xlim, ylim=ylim) exp_debug = False
rho = 28.0 beta = 8.0 / 3.0 """ Create PyGnosis System """ references = {} variables = { 'x': [[1.0]], # rate of convection, formally 'x' 'y': [[1.0]], # horizontal temperature variation, formally 'y' 'z': [[1.0]] # vertical temperature variation, formally 'z' } sys = pygnosis.System(references, variables) """ Define system dynamics """ dx_dt = lambda: sigma * (sys.y - sys.x) dy_dt = lambda: sys.x * (rho - sys.z) - sys.y dz_dt = lambda: sys.x * sys.y - beta * sys.z dynamics = {'x': dx_dt, 'y': dy_dt, 'z': dz_dt} sys.setup_dynamics(dynamics) """ Build simulation """ sim = pygnosis.Simulation(sys, timestep=0.01, timespan=40.0) """ Run simulation """ states = sim() """ 3-D Scatter """ from matplotlib import cm t = sim.time ax, fig = sim.scatter('x', 'y', 'z', s=1, c=cm.gnuplot2((t - t.min()) / (t.ptp()))) """ Plot lines """ ax, fig = sim.plot(t, 'x', t, 'y', t, 'z')
for t in range(cfg['T']): nmis[i * cfg['J'] + j, t] = MI[t] / np.sqrt(Hz[t] * Hgt[t]) # nmis[i*cfg['J']+j,t] = MI*2. / (Hz+Hgt) mis[i * cfg['J'] + j, t] = MI[t] vMeasures[i * cfg['J'] + j, t] = 2. * MI[t] / (Hz[t] + Hgt[t]) print nmis[i * cfg['J'] + j, t], 2. * MI[t], Hz[t], Hgt[t] baseMap={'spkm':'spkm','kmeans':'k-means','NiwSphere':'DirSNIW', \ 'DpNiw':'DP-GMM','DpNiwSphere':'DpSNIW opt','DpNiwSphereFull':'DP-TGMM', \ 'DPvMFmeans':'DPvMFmeans',"DirNiwSphereFull":"Dir-TGMM"} #cl = cm.gnuplot2(np.arange(len(bases))) cl = cm.hsv(np.arange(255)) cl = cm.brg(np.arange(255)) cl = cm.gist_rainbow(np.arange(255)) cl = cm.gnuplot2(np.arange(255)) cl = cm.gnuplot(np.arange(255)) cl = cm.spectral(np.arange(255)) #print cltlib I = len(bases) + 1 print nmis.shape fig = plt.figure(figsize=figSize, dpi=80, facecolor='w', edgecolor='k') for i, base in enumerate(bases): if not re.search('_', base) is None: info = base.split('_') base = info[0] Ki = int(info[1]) plt.plot(np.arange(cfg['T']), nmis[i * cfg['J'], :], label=baseMap[base] + ' ($K_0={}$)'.format(Ki),
def plot(args): limit = 1000 env_name = args.env_name if args.env_id == 9: limit = 500 """ get config about trajectory file""" if args.env_id != 21: discriminator_updater = IRL(1, 1, args, False) traj_name = discriminator_updater.traj_name m_return_list = discriminator_updater.m_return_list else: traj_name = "traj10" m_return_list = None seed_list = [1, 2, 3, 4, 5] seed_num = len(seed_list) print(seed_list) load_robo = 0 plot_methods = [ "vild_per0", "vild_per2", "meirl", "airl", "gail", "vail", "infogail_mean", "infogail_best" ] if args.env_id == 9: args.info_loss_type = "linear" args.rl_method = "sac" args.hidden_size = (100, 100) args.activation = "relu" args.q_step = 1 if args.env_id == 0: plot_methods = [ "vild_per0", "vild_per2", "gail", "infogail_mean", "infogail_best" ] args.vild_loss_type = "bce" if args.env_id == -3: plot_methods = [ "vild_per0", "vild_per2", "gail", "infogail_mean", "infogail_best" ] args.vild_loss_type = "bce" args.rl_method = "ppo" args.bce_negative = 1 args.hidden_size = (100, 100) args.activation = "relu" args.gp_lambda = 0 if args.env_id == 21: plot_methods = [ "vild_per0", "vild_per2", "gail", "infogail_mean", "infogail_best", "meirl", "airl", "vail" ] # for paper # plot_methods = ["vild_per2", "gail", "meirl", "airl", "vail"] # for slide args.vild_loss_type = "bce" env_name += "_reach" args.hidden_size = (100, 100) args.activation = "relu" load_robo = 1 if args.env_id == 15: plot_methods = ["vild_per2", "gail", "meirl"] R_test_all = [] gail_legend = [] c_tmp = [] l_tmp = [] max_len = 10000 if any("vild_per2" in s for s in plot_methods): """ VILD """ R_test_avg = [] cat = 1 noise_array_is = [] args.il_method = "vild" args.per_alpha = 2 args.clip_discriminator = 5 method_name = args.il_method.upper() + "_" + args.rl_method.upper() if args.vild_loss_type.lower() != "linear": method_name += "_" + args.vild_loss_type.upper() args.clip_discriminator = 0 hypers = rl_hypers_parser(args) + "_" + irl_hypers_parser(args) for seed in seed_list: exp_name = "%s-%s-%s_s%d" % (traj_name, method_name, hypers, seed) if load_robo: filename = "./results_IL/test/%s/%s-%s" % (env_name, env_name, exp_name) else: filename = "./results_IL/%s/%s/%s-%s" % (method_name, env_name, env_name, exp_name) R_test_avg_i = load(filename, limit, load_robo) if not load_robo and args.c_data == 7: noise_array_is += [load_noise(filename)] if R_test_avg_i[0, 0] == -999: cat = 0 print("cannot load %s" % exp_name) else: load_iter = R_test_avg_i.shape[0] if load_iter < max_len: max_len = load_iter for i in range(0, seed - 1): R_test_avg[i] = R_test_avg[i][0:max_len, :] R_test_avg_i = R_test_avg_i[0:max_len, :] R_test_avg += [R_test_avg_i] if cat: R_test_avg = np.concatenate( R_test_avg, axis=1) # array [iter, test_episode * seed] R_test_all += [R_test_avg] gail_legend += ["VILD (IS)"] c_tmp += ["r"] l_tmp += ["-"] if any("vild_per0" in s for s in plot_methods): """ VILD no IS """ R_test_avg = [] cat = 1 noise_array = [] args.il_method = "vild" args.per_alpha = 0 args.clip_discriminator = 5 method_name = args.il_method.upper() + "_" + args.rl_method.upper() if args.vild_loss_type.lower() != "linear": method_name += "_" + args.vild_loss_type.upper() args.clip_discriminator = 0 hypers = rl_hypers_parser(args) + "_" + irl_hypers_parser(args) for seed in seed_list: exp_name = "%s-%s-%s_s%d" % (traj_name, method_name, hypers, seed) if load_robo: filename = "./results_IL/test/%s/%s-%s" % (env_name, env_name, exp_name) else: filename = "./results_IL/%s/%s/%s-%s" % (method_name, env_name, env_name, exp_name) R_test_avg_i = load(filename, limit, load_robo) if not load_robo: noise_array += [load_noise(filename)] if R_test_avg_i[0, 0] == -999: cat = 0 print("cannot load %s" % exp_name) else: load_iter = R_test_avg_i.shape[0] if load_iter < max_len: max_len = load_iter for i in range(0, seed - 1): R_test_avg[i] = R_test_avg[i][0:max_len, :] R_test_avg_i = R_test_avg_i[0:max_len, :] R_test_avg += [R_test_avg_i] if cat: R_test_avg = np.concatenate( R_test_avg, axis=1) # array [iter, test_episode * seed] R_test_all += [R_test_avg] gail_legend += ["VILD (w/o IS)"] c_tmp += ["blueviolet"] l_tmp += ["-"] if any("gail" in s for s in plot_methods): """ GAIL """ R_test_avg = [] cat = 1 args.il_method = "gail" args.clip_discriminator = 0 hypers = rl_hypers_parser(args) + "_" + irl_hypers_parser(args) method_name = args.il_method.upper() + "_" + args.rl_method.upper() for seed in seed_list: exp_name = "%s-%s-%s_s%d" % (traj_name, method_name, hypers, seed) if load_robo: filename = "./results_IL/test/%s/%s-%s" % (env_name, env_name, exp_name) else: filename = "./results_IL/%s/%s/%s-%s" % (method_name, env_name, env_name, exp_name) R_test_avg_i = load(filename, limit, load_robo) if R_test_avg_i[0, 0] == -999: cat = 0 print("cannot load %s" % exp_name) else: load_iter = R_test_avg_i.shape[0] if load_iter < max_len: max_len = load_iter for i in range(0, seed - 1): R_test_avg[i] = R_test_avg[i][0:max_len, :] R_test_avg_i = R_test_avg_i[0:max_len, :] R_test_avg += [R_test_avg_i] if cat: R_test_avg = np.concatenate( R_test_avg, axis=1) # array [iter, test_episode * seed] R_test_all += [R_test_avg] gail_legend += ["GAIL"] c_tmp += ["darkgoldenrod"] l_tmp += ["--"] if any("airl" in s for s in plot_methods): """ AIRL """ R_test_avg = [] cat = 1 args.il_method = "airl" args.clip_discriminator = 5 hypers = rl_hypers_parser(args) + "_" + irl_hypers_parser(args) method_name = args.il_method.upper() + "_" + args.rl_method.upper() for seed in seed_list: exp_name = "%s-%s-%s_s%d" % (traj_name, method_name, hypers, seed) if load_robo: filename = "./results_IL/test/%s/%s-%s" % (env_name, env_name, exp_name) else: filename = "./results_IL/%s/%s/%s-%s" % (method_name, env_name, env_name, exp_name) R_test_avg_i = load(filename, limit, load_robo) if R_test_avg_i[0, 0] == -999: cat = 0 print("cannot load %s" % exp_name) else: load_iter = R_test_avg_i.shape[0] if load_iter < max_len: max_len = load_iter for i in range(0, seed - 1): R_test_avg[i] = R_test_avg[i][0:max_len, :] R_test_avg_i = R_test_avg_i[0:max_len, :] R_test_avg += [R_test_avg_i] if cat: R_test_avg = np.concatenate( R_test_avg, axis=1) # array [iter, test_episode * seed] R_test_all += [R_test_avg] gail_legend += ["AIRL"] c_tmp += ["b"] l_tmp += ["-."] if any("meirl" in s for s in plot_methods): """ ME-IRL """ R_test_avg = [] cat = 1 args.il_method = "irl" args.clip_discriminator = 5 hypers = rl_hypers_parser(args) + "_" + irl_hypers_parser(args) method_name = args.il_method.upper() + "_" + args.rl_method.upper() for seed in seed_list: exp_name = "%s-%s-%s_s%d" % (traj_name, method_name, hypers, seed) if load_robo: filename = "./results_IL/test/%s/%s-%s" % (env_name, env_name, exp_name) else: filename = "./results_IL/%s/%s/%s-%s" % (method_name, env_name, env_name, exp_name) R_test_avg_i = load(filename, limit, load_robo) if R_test_avg_i[0, 0] == -999: cat = 0 print("cannot load %s" % exp_name) else: load_iter = R_test_avg_i.shape[0] if load_iter < max_len: max_len = load_iter for i in range(0, seed - 1): R_test_avg[i] = R_test_avg[i][0:max_len, :] R_test_avg_i = R_test_avg_i[0:max_len, :] R_test_avg += [R_test_avg_i] if cat: R_test_avg = np.concatenate( R_test_avg, axis=1) # array [iter, test_episode * seed] R_test_all += [R_test_avg] gail_legend += ["ME-IRL"] c_tmp += ["olive"] l_tmp += [":"] if any("vail" in s for s in plot_methods): """ VAIL """ R_test_avg = [] cat = 1 args.il_method = "vail" args.clip_discriminator = 0 hypers = rl_hypers_parser(args) + "_" + irl_hypers_parser(args) method_name = args.il_method.upper() + "_" + args.rl_method.upper() for seed in seed_list: exp_name = "%s-%s-%s_s%d" % (traj_name, method_name, hypers, seed) if load_robo: filename = "./results_IL/test/%s/%s-%s" % (env_name, env_name, exp_name) else: filename = "./results_IL/%s/%s/%s-%s" % (method_name, env_name, env_name, exp_name) R_test_avg_i = load(filename, limit, load_robo) if R_test_avg_i[0, 0] == -999: cat = 0 print("cannot load %s" % exp_name) else: load_iter = R_test_avg_i.shape[0] if load_iter < max_len: max_len = load_iter for i in range(0, seed - 1): R_test_avg[i] = R_test_avg[i][0:max_len, :] R_test_avg_i = R_test_avg_i[0:max_len, :] R_test_avg += [R_test_avg_i] if cat: R_test_avg = np.concatenate( R_test_avg, axis=1) # array [iter, test_episode * seed] R_test_all += [R_test_avg] gail_legend += ["VAIL"] c_tmp += ["m"] l_tmp += ["--"] if any("infogail_mean" in s for s in plot_methods): """ InfoGAIL """ R_test_avg = [] cat = 1 args.il_method = "infogail" args.clip_discriminator = 5 if args.info_loss_type == "linear" else 0 hypers = rl_hypers_parser(args) + "_" + irl_hypers_parser(args) method_name = args.il_method.upper() + "_" + args.rl_method.upper() if args.il_method == "infogail" and args.info_loss_type.lower( ) != "bce": method_name += "_" + args.info_loss_type.upper() for seed in seed_list: exp_name = "%s-%s-%s_s%d" % (traj_name, method_name, hypers, seed) if load_robo: filename = "./results_IL/test/%s/%s-%s" % (env_name, env_name, exp_name) else: filename = "./results_IL/%s/%s/%s-%s" % (method_name, env_name, env_name, exp_name) R_test_avg_i = load_info( filename, limit, load_robo) # [iter, 10, test_episode (1?) ] array if R_test_avg_i[0, 0, 0] == -999: cat = 0 print("cannot load %s" % exp_name) else: load_iter = R_test_avg_i.shape[0] if load_iter < max_len: max_len = load_iter for i in range(0, seed - 1): R_test_avg[i] = R_test_avg[i][0:max_len, :] R_test_avg_i = R_test_avg_i[0:max_len, :] R_test_avg += [R_test_avg_i] if cat: # R_test_avg_x = np.stack(R_test_avg, axis=2) # array [iter, 10, test_episode * seed] R_test_avg_x = np.concatenate(R_test_avg, axis=2) # average over uniform contexts R_test_avg_mean = np.mean(R_test_avg_x, axis=1) R_test_all += [R_test_avg_mean] gail_legend += ["InfoGAIL"] c_tmp += ["black"] l_tmp += ["-."] if any("infogail_best" in s for s in plot_methods): R_test_avg_sort = [] ## For each seed, sort array along context descending based on performance for i in range(0, len(seed_list)): R_test_x = R_test_avg[ i] # R_test_avg is a list. R_test_x is a np array shape [test_iter, 10, test_epi ] test_last = np.mean( np.mean(R_test_x, axis=2)[-100:, :], axis=0 ) # final performance average over last 100 iterations sort_idx = np.argsort(test_last)[::-1] # descending indices R_test_x_sort = R_test_x[:, sort_idx, :] R_test_avg_sort += [R_test_x_sort] R_test_avg_sort = np.concatenate( R_test_avg_sort, axis=2) # array [iter, 10, test_episode * seed] R_test_avg_sort_best = R_test_avg_sort[:, 0, :] R_test_all += [R_test_avg_sort_best] gail_legend += ["InfoGAIL (best)"] c_tmp += ["dodgerblue"] l_tmp += [":"] if any("infogail_worst" in s for s in plot_methods): R_test_avg_sort_worst = R_test_avg_sort[:, -1, :] R_test_all += [R_test_avg_sort_worst] gail_legend += ["InfoGAIL (worst)"] c_tmp += ["teal"] l_tmp += ["--"] ## infogail_all R_test_info = [] R_test_info_mean = [] R_test_info_std = [] for i in range(0, 10): R_test_info += [R_test_avg_sort[:, i, :]] R_test_info_mean += [np.mean(R_test_info[i], 1)] R_test_info_std += [ np.std(R_test_info[i], 1) / np.sqrt(R_test_info[i].shape[1]) ] """ Compute statistics for plotting """ R_test_mean = [] R_test_std = [] R_test_last = [] R_test_last_mean = [] R_test_last_std = [] best_idx = -1 best_R = -1e6 for i in range(0, len(R_test_all)): R_test = R_test_all[i] """ This is for plotting """ R_test_mean += [np.mean(R_test, 1)] R_test_std += [np.std(R_test, 1) / np.sqrt(R_test.shape[1])] ## Average last final 100 iteration. R_test_last += [np.mean(R_test[-100:, :], 0)] last_mean = np.mean(R_test_last[i], 0) if last_mean > best_R: best_idx = i best_R = last_mean R_test_last_mean += [last_mean] R_test_last_std += [ np.std(R_test_last[i], 0) / np.sqrt(R_test.shape[1]) ] ### Print statistics # print("%s: %0.1f(%0.1f)" % (gail_legend[i], R_test_mean_last, R_test_std_last)) """ For t_test """ ## paired t-test from scipy import stats best_m = R_test_last[best_idx] p_list = [] for i in range(0, len(R_test_all)): # if gail_legend[i] == "InfoGAIL (best)": continue if i != best_idx: _, p = stats.ttest_ind(best_m, R_test_last[i], nan_policy="omit") else: p = 1 p_list += [p] ## copied able latex format latex = "" for i in range(0, len(R_test_all)): # if gail_legend[i] == "InfoGAIL (best)": continue print("%-70s: Sum %0.0f(%0.0f) with p-value %f" % (gail_legend[i], R_test_last_mean[i], R_test_last_std[i], p_list[i])) if p_list[i] > 0.01: # 1 percent confidence latex += " & \\textbf{%0.0f (%0.0f)}" % (R_test_last_mean[i], R_test_last_std[i]) else: latex += " & %0.0f (%0.0f)" % (R_test_last_mean[i], R_test_last_std[i]) print(latex + " \\\\") plot_large = args.plot_large skipper = 10 # Use sliding window to compute running mean and std for clearer figures. skipper 1 = no running. running = 1 plot_name = "%s_%s_%s" % (args.rl_method, env_name, args.noise_type) title = "%s" % (env_name[:-3]) x_max_iter = len(R_test_mean[0]) """ Plot """ if plot_large: linewidth = 2 fontsize = 21 f = plt.figure(figsize=(8, 6)) # for Figure 2, ... else: linewidth = 1 # for Figure 2, ... fontsize = 14 # for Figure 2, ... f = plt.figure(figsize=(4, 3)) # for Figure 2, ... ax = f.gca() """ Plot """ cc_tmp = ["black"] + ["gray"] + ["darkgray"] + ["dimgray"] + ["silver"] for i in range(0, len(R_test_all)): if running: y_plot = running_mean(R_test_mean[i][:x_max_iter], skipper) y_err = running_mean(R_test_std[i][:x_max_iter], skipper) else: y_plot = R_test_mean[i][:x_max_iter:skipper] y_err = R_test_std[i][:x_max_iter:skipper] x_ax = np.arange(0, len(y_plot)) * 10000 // 2 if args.rl_method == "sac": x_ax = np.arange(0, len(y_plot)) * 10000 // 2 // 5 if args.env_id == 21: x_ax = np.arange(0, len(y_plot)) * 100000 * 5 / 4 if i == 0 and m_return_list is not None: ic_k = 0 for i_k in range(0, len(m_return_list)): opt_plot = (x_ax * 0) + m_return_list[i_k] if i_k == 0 or i_k == 2 or i_k == 4 or i_k == 6 or i_k == 9: # if i_k == 0 or i_k == 4 or i_k == 8: ax.plot(x_ax, opt_plot, color=cc_tmp[ic_k], linestyle=":", linewidth=linewidth + 0.5) # plt.text(x=0, y=m_return_list[i_k], s="k=%d" % (i_k+1), fontsize=fontsize-3, color=cc_tmp[ic_k]) ic_k += 1 errorfill(x_ax, y_plot, yerr=y_err, color=c_tmp[i], linestyle=l_tmp[i], linewidth=linewidth, label=gail_legend[i], shade=1) if plot_large == 1: if args.env_id == 0: ax.legend(prop={"size": fontsize - 3}, frameon=True, loc='lower right', framealpha=1, ncol=2) elif args.env_id == -3: # ax.legend(prop={"size":fontsize-4}, frameon=True, loc = 'lower right', framealpha=1, ncol=2) ax.legend(prop={"size": fontsize - 3}, frameon=True, loc='lower right', framealpha=1, ncol=2) elif args.env_id == 21: ax.legend(prop={"size": fontsize - 6.3}, frameon=True, loc='lower right', framealpha=1, ncol=3) # handles, labels = ax.get_legend_handles_labels() # leg_order = [0, 1, 2, 3, 4, 5, 6] # leg_order = [0, 1, 6, 3, 4, 5, 2] # handles, labels = [handles[i] for i in leg_order], [labels[i] for i in leg_order] # ax.legend(handles, labels, prop={"size":fontsize-4}, frameon=True, loc = 'lower right', framealpha=1, ncol=3) else: ax.legend(prop={"size": fontsize - 3}, frameon=True, loc='upper left', framealpha=1) else: if args.env_id == -3: ax.legend(prop={"size": fontsize - 3}, frameon=True, loc='lower right', framealpha=1, ncol=2) if args.env_id == 21: ax.legend(prop={"size": fontsize - 6}, frameon=True, loc='lower right', framealpha=1, ncol=3) if args.env_id == -3: title = "%s" % ("LunarLander") # ax.set_ylim([-200,250]) # ax.set_ylim([-200,300]) # f.legend(prop={"size":fontsize-3}, frameon=True, loc="lower right", framealpha=1, ncol=3) if args.env_id == 21: title = "%s" % ("RobosuiteReacher") plot_name = "%s_%s" % (args.rl_method, env_name) ax.set_ylim([-10, 45]) # paper # ax.set_ylim([-2.4, 29]) # slide if args.env_id == 15: title = "%s" % ("AntBullet") fig_name = "./figures/%s" % (plot_name) if plot_large: fig_name += "_large" # fig_name += "_slide" print(fig_name) if plot_large: plt.title(title, fontsize=fontsize + 1) else: plt.title(title, fontsize=fontsize) plt.xlabel("Transition samples", fontsize=fontsize - 2) plt.ylabel('Cumulative rewards', fontsize=fontsize - 2) plt.xticks(fontsize=fontsize - 2) plt.yticks(fontsize=fontsize - 2) plt.tight_layout() plt.ticklabel_format(style='sci', axis='y', scilimits=(0, 0)) plt.ticklabel_format(style='sci', axis='x', scilimits=(0, 0)) ## plot legend. if plot_large == 2: figlegend = pylab.figure(figsize=(18, 1)) #pylab.figlegend(*ax.get_legend_handles_labels(), prop={"size":fontsize-3}, ncol=7, handlelength=2.8, ) pylab.figlegend( *ax.get_legend_handles_labels(), prop={"size": fontsize - 3}, ncol=10, handlelength=2.0, ) figlegend.savefig("./figures/legend.pdf", bbox_inches='tight', pad_inches=-0.25) if args.plot_save: f.savefig(fig_name + ".pdf", bbox_inches='tight', pad_inches=0) ## plot covariance. if not load_robo and args.c_data == 7: f_cov = plt.figure(figsize=(4, 3)) # for Figure 2, ... ax_cov = plt.gca() noise_mean_is = np.mean(np.array(noise_array_is), axis=0) noise_mean = np.mean(np.array(noise_array), axis=0) true_variance = np.array( args.noise_level_list)**2 # convert to variance if args.env_id == -3: 1 # true_variance = [0.021, 0.022, 0.023, 0.024, 0.025, 0.026, 0.027, 0.028, 0.029, 0.030] # true_variance = true_variance * 0.034 x = ['1', '2', '3', '4', '5', '6', '7', '8', '9', '10'] x_pos = [i for i, _ in enumerate(x)] plt.ylabel('Covariance', fontsize=fontsize - 2) plt.xlabel('Demonstrator number k', fontsize=fontsize - 2) plt.plot(x_pos, noise_mean_is, 'rx', markersize=8, label="VILD (IS)") plt.plot(x_pos, noise_mean, 'm+', markersize=8, label="VILD (w/o IS)") if args.env_id == -3: plt.legend(prop={"size": fontsize - 3}, frameon=True, loc='lower right', framealpha=1) elif args.noise_type == "SDNTN": plt.legend(prop={"size": fontsize - 3}, frameon=True, loc='upper left', framealpha=1) else: plt.plot(x_pos, true_variance, 'ko', markersize=5, label="Ground-truth") plt.legend(prop={"size": fontsize - 3}, frameon=True, loc='upper left', framealpha=1) plt.yticks(fontsize=fontsize - 2) plt.tight_layout() plt.xticks(x_pos, x, fontsize=fontsize - 2) fig_name_cov = "./figures/%s_cov" % (plot_name) if plot_large: plt.title(title + " (Covariance)", fontsize=fontsize + 1) else: plt.title(title + " (Covariance)", fontsize=fontsize) print(fig_name_cov) if args.plot_save: f_cov.savefig(fig_name_cov + ".pdf", bbox_inches='tight', pad_inches=0) ## plot bc. if args.env_id == 2 or args.env_id == 5 or args.env_id == 7 or args.env_id == 9: # mujoco tasks. c_tmp_2 = [] l_tmp_2 = [] R_test_all_2 = [] gail_legend_2 = [] """ BC """ R_test_avg = [] cat = 1 method_name = "BC" for seed in seed_list: hypers_tmp = bc_hypers_parser(args) exp_name = "%s-%s-%s_s%d" % (traj_name, method_name, hypers_tmp, seed) filename = "./results_IL/%s/%s/%s-%s" % (method_name, env_name, env_name, exp_name) R_test_avg_i = load(filename, limit, load_robo) if R_test_avg_i[0, 0] == -999: cat = 0 print("cannot load %s" % exp_name) else: load_iter = R_test_avg_i.shape[0] if load_iter < max_len: max_len = load_iter for i in range(0, seed - 1): R_test_avg[i] = R_test_avg[i][0:max_len, :] R_test_avg_i = R_test_avg_i[0:max_len, :] R_test_avg += [R_test_avg_i] if cat: R_test_avg = np.concatenate( R_test_avg, axis=1) # array [iter, test_episode * seed] R_test_all_2 += [R_test_avg] gail_legend_2 += ["BC"] c_tmp_2 += ["blue"] l_tmp_2 += ["--"] """ D-BC """ R_test_avg = [] cat = 1 method_name = "DBC" for seed in seed_list: hypers_tmp = bc_hypers_parser(args) exp_name = "%s-%s-%s_s%d" % (traj_name, method_name, hypers_tmp, seed) filename = "./results_IL/%s/%s/%s-%s" % (method_name, env_name, env_name, exp_name) R_test_avg_i = load(filename, limit, load_robo) if R_test_avg_i[0, 0] == -999: cat = 0 print("cannot load %s" % exp_name) else: load_iter = R_test_avg_i.shape[0] if load_iter < max_len: max_len = load_iter for i in range(0, seed - 1): R_test_avg[i] = R_test_avg[i][0:max_len, :] R_test_avg_i = R_test_avg_i[0:max_len, :] R_test_avg += [R_test_avg_i] if cat: R_test_avg = np.concatenate( R_test_avg, axis=1) # array [iter, test_episode * seed] R_test_all_2 += [R_test_avg] gail_legend_2 += ["BC-D"] c_tmp_2 += ["orange"] l_tmp_2 += ["-."] """ Co-BC """ R_test_avg = [] cat = 1 method_name = "COBC" for seed in seed_list: hypers_tmp = bc_hypers_parser(args) exp_name = "%s-%s-%s_s%d" % (traj_name, method_name, hypers_tmp, seed) filename = "./results_IL/%s/%s/%s-%s" % (method_name, env_name, env_name, exp_name) R_test_avg_i = load(filename, limit, load_robo) if R_test_avg_i[0, 0] == -999: cat = 0 print("cannot load %s" % exp_name) else: load_iter = R_test_avg_i.shape[0] if load_iter < max_len: max_len = load_iter for i in range(0, seed - 1): R_test_avg[i] = R_test_avg[i][0:max_len, :] R_test_avg_i = R_test_avg_i[0:max_len, :] R_test_avg += [R_test_avg_i] if cat: R_test_avg = np.concatenate( R_test_avg, axis=1) # array [iter, test_episode * seed] R_test_all_2 += [R_test_avg] gail_legend_2 += ["Co-teaching"] c_tmp_2 += ["darkgreen"] l_tmp_2 += [":"] R_test_mean_2 = [] R_test_std_2 = [] for i in range(0, len(R_test_all_2)): R_test = R_test_all_2[i] """ This is for plotting """ R_test_mean_2 += [np.mean(R_test, 1)] R_test_std_2 += [np.std(R_test, 1) / np.sqrt(R_test.shape[1])] if plot_large: linewidth = 2 fontsize = 21 f_2 = plt.figure(figsize=(8, 6)) # for Figure 2, ... else: linewidth = 1 # for Figure 2, ... fontsize = 14 # for Figure 2, ... f_2 = plt.figure(figsize=(4, 3)) # for Figure 2, ... ax_2 = f_2.gca() from matplotlib import cm for i in range(0, len(R_test_all_2)): if running: y_plot = running_mean(R_test_mean_2[i][:x_max_iter], skipper) y_err = running_mean(R_test_std_2[i][:x_max_iter], skipper) else: y_plot = R_test_mean_2[i][:x_max_iter:skipper] y_err = R_test_std_2[i][:x_max_iter:skipper] x_ax = np.arange(0, len(y_plot)) * 10000 // 2 // 5 if i == 0: errorfill(x_ax, (x_ax * 0) + R_test_last_mean[0], yerr=(x_ax * 0) + R_test_last_std[0], color=c_tmp[0], linestyle="-", linewidth=linewidth, shade=1) errorfill(x_ax, y_plot, yerr=y_err, color=c_tmp_2[i], linestyle=l_tmp_2[i], linewidth=linewidth, label=gail_legend_2[i], shade=1) # if plot_large: # ax_2.legend(prop={"size":fontsize-3}, frameon=True, loc = 'lower right', framealpha=1, ncol=2) if args.env_id == 2: ax_2.legend(prop={"size": fontsize - 3}, frameon=True, loc='lower right', framealpha=1, ncol=2) if args.env_id == 2: ax_2.set_ylim([-999, 5000]) fig_name_bc = "./figures/%s_bc" % (plot_name) if plot_large: fig_name_bc += "_large" print(fig_name_bc) # title += " (Supervised)" if plot_large: plt.title(title, fontsize=fontsize + 1) else: plt.title(title, fontsize=fontsize) plt.xlabel("Gradient steps", fontsize=fontsize - 2) plt.ylabel('Cumulative rewards', fontsize=fontsize - 2) plt.xticks(fontsize=fontsize - 2) plt.yticks(fontsize=fontsize - 2) plt.tight_layout() plt.ticklabel_format(style='sci', axis='y', scilimits=(0, 0)) plt.ticklabel_format(style='sci', axis='x', scilimits=(0, 0)) ## plot legend. if plot_large == 2: figlegend = pylab.figure(figsize=(18, 1)) pylab.figlegend( *ax_2.get_legend_handles_labels(), prop={"size": fontsize - 3}, ncol=10, handlelength=2.0, ) figlegend.savefig("./figures/legend_bc.pdf", bbox_inches='tight', pad_inches=-0.25) if args.plot_save: f_2.savefig(fig_name_bc + ".pdf", bbox_inches='tight', pad_inches=0) ## plot infogail sorted performance. if any("infogail_mean" in s for s in plot_methods): if plot_large: linewidth = 2 fontsize = 21 f_3 = plt.figure(figsize=(8, 6)) # for Figure 2, ... else: linewidth = 1 # for Figure 2, ... fontsize = 14 # for Figure 2, ... f_3 = plt.figure(figsize=(4, 3)) # for Figure 2, ... ax_3 = f_3.gca() # info_color = ["black"] + ["gray"] + ["darkgray"] + ["dimgray"] + ["silver"] \ # + ["r"] + ["g"] + ["b"] + ["m"] + ["c"] from matplotlib import cm info_color = cm.gnuplot2(np.linspace(0, 1, 12)) for i in range(0, 10): if running: y_plot = running_mean(R_test_info_mean[i][:x_max_iter], skipper) y_err = running_mean(R_test_info_std[i][:x_max_iter], skipper) else: y_plot = R_test_info_mean[i][:x_max_iter:skipper] y_err = R_test_info_std[i][:x_max_iter:skipper] x_ax = np.arange(0, len(y_plot)) * 10000 // 2 if args.env_id == 21: x_ax = np.arange(0, len(y_plot)) * 100000 * 5 / 4 if i == 0 and m_return_list is not None: ic_k = 0 for i_k in range(0, len(m_return_list)): opt_plot = (x_ax * 0) + m_return_list[i_k] if i_k == 0 or i_k == 2 or i_k == 4 or i_k == 6 or i_k == 9: # if i_k == 0 or i_k == 4 or i_k == 8: ax_3.plot(x_ax, opt_plot, color=cc_tmp[ic_k], linestyle=":", linewidth=linewidth + 0.5) # plt.text(x=0, y=m_return_list[i_k], s="k=%d" % (i_k+1), fontsize=fontsize-3, color=cc_tmp[ic_k]) ic_k += 1 if i != 8: errorfill(x_ax, y_plot, yerr=y_err, color=info_color[i], linestyle="-", linewidth=linewidth, shade=1) else: errorfill(x_ax, y_plot, yerr=y_err, color="goldenrod", linestyle="-", linewidth=linewidth, shade=1) fig_name_info = "./figures/%s_info" % (plot_name) if plot_large: fig_name_info += "_large" print(fig_name_info) if args.env_id == 21: title = "%s" % ("RobosuiteReacher") if "infogail_best" in plot_methods: ax_3.set_ylim([-10, 45]) # for paper else: ax_3.set_ylim([0, 35]) # for slide if plot_large: plt.title(title + " (InfoGAIL)", fontsize=fontsize + 1) else: plt.title(title + " (InfoGAIL)", fontsize=fontsize) plt.xlabel("Transition samples", fontsize=fontsize - 2) plt.ylabel('Cumulative rewards', fontsize=fontsize - 2) plt.xticks(fontsize=fontsize - 2) plt.yticks(fontsize=fontsize - 2) plt.tight_layout() plt.ticklabel_format(style='sci', axis='y', scilimits=(0, 0)) plt.ticklabel_format(style='sci', axis='x', scilimits=(0, 0)) if args.plot_save: f_3.savefig(fig_name_info + ".pdf", bbox_inches='tight', pad_inches=0) if args.plot_show: plt.show()
else: sp = [N - vc - 1] ift = [1] rcv = [0] for t in range(1, 1000): count1 = np.random.choice( range(2), sp[t - 1], p=[1 - beta * ift[t - 1] / N, beta * ift[t - 1] / N]).sum() #yield infected people count2 = np.random.choice(range(2), ift[t - 1], p=[1 - gamma, gamma]).sum() #yield recovered people sp.append(sp[t - 1] - count1) rcv.append(rcv[t - 1] + count2) ift.append(ift[t - 1] + count1 - count2) count1, count2 = 0, 0 #reset list_inf.append(ift) #========================plot================================================== plt.figure(figsize=(6, 4), dpi=150) x = [y for y in range(1, 1001)] for u in range(0, 11): plt.plot(x, list_inf[u], label=str(u * 10) + '%', color=cm.gnuplot2(u * 20)) #plot graphs in one figure with loops plt.xlabel('time') plt.ylabel('number of people') plt.title('SIR model with different vaccination rates') plt.legend() plt.savefig("SIR model with different vaccination rates", type="png")