wonderboy3489
wonderboy3489

Reputation: 37

Python legend labels as expression

I'm trying to plot 9 lines, each of which belongs to 1 of 3 categories, and I want to show a legend with just 3 labels instead of 9. Right now my code looks like

tau = np.array([.1, .2, .3])
num_repeats = 3
plot_colors = ['r-','b-','k-']
labels = ['tau = 0.1','tau = 0.2','tau = 0.3']
plot_handles=[None]*3

for k in np.arange(tau.size):
  for ell in np.arange(num_repeats):
    ( _, _, n_iter, history,_ ) = opti.minimise (theta0, f_tol=f_tol, theta_tol = theta_tol, tau=tau[k], N=3,report=50000, m=1)
    true_energy = -59.062
    error = np.absolute(true_energy - history[:,num_p])
    plot_handles[k] = plt.plot(np.arange(0,n_iter),error,plot_colors[k],'label'=labels[k])

plt.xlabel('Iteration')
plt.ylabel('Absolute Error')
plt.yscale('log')
plt.legend(handles=plot_handles)
plt.show()

I get an error saying the 'label' keyword can't be an expression. Does anyone know of a way to do this? Thanks in advance!

-Jim

Upvotes: 0

Views: 179

Answers (1)

David Zwicker
David Zwicker

Reputation: 24268

One trick to achieve this is to only set a label for one plot in the inner loop. In the following code, label is only non-empty if ell==0:

tau = np.array([.1, .2, .3])
num_repeats = 3
plot_colors = ['r-','b-','k-']
labels = ['tau = 0.1','tau = 0.2','tau = 0.3']
plot_handles=[None]*3

for k in np.arange(tau.size):
  for ell in np.arange(num_repeats):
    ( _, _, n_iter, history,_ ) = opti.minimise (theta0, f_tol=f_tol, theta_tol = theta_tol, tau=tau[k], N=3,report=50000, m=1)
    true_energy = -59.062
    error = np.absolute(true_energy - history[:,num_p])
    if ell == 0:
        label = labels[k] 
    else:
        label = '' 
    plot_handles[k] = plt.plot(np.arange(0, n_iter), error, plot_colors[k], label=label)

plt.xlabel('Iteration')
plt.ylabel('Absolute Error')
plt.yscale('log')
plt.legend(handles=plot_handles)
plt.show()

Upvotes: 1

Related Questions