1

I have been making a chart that plots the energy levels of the hydrogen atom and while technically everything shown in the first graph attached is correct I was trying to switch the order of the legend on the right so that the red line is on the bottom and still labeled as n=1. When I switch the range on the legend it makes the graph look like the second picture where n=1 is on the bottom but none of the colors correspond correctly. Any suggestions for how to change the order while keeping the order correct? My code is as follows

import numpy
import math
import os
import matplotlib.pyplot as plt
import pandas as pd
colors = ['#e6194b', '#3cb44b', '#ffe119', '#4363d8', '#f58231', '#911eb4', '#46f0f0', '#f032e6', '#bcf60c', '#fabebe', '#008080', '#e6beff', '#9a6324', '#fffac8', '#800000', '#aaffc3', '#808000', '#ffd8b1', '#000075', '#808080', '#ffffff', '#000000']
#Values the user can edit
n_int = 6 #The inital energy level (must be greater than n_final)
n_final = 5 #The final energy level 
n_max = 11 #This is the highest energy level you want to look at plus 1 (ie: for energy levels 1-10 "n" would be 11)

# Energy level diagram for the Hydrogen Atom (all energy levels)
m_e = 9.1093837015* 10 ** -31 # kg
e = 1.602176634* 10 ** -19 # C
vp = 1.113* 10 ** -10 # C^2 / Jm
h = 6.62607015 * 10 ** -34 # J*s
h_bar = h/(2 * math.pi)  
n_min = 1
n = numpy.linspace(n_min, n_max, n_max-n_min)

#Equation broken down
p1 = m_e * e **4
p2 = 2 * vp ** 2 
p3 = h_bar ** 2
p4 = n ** 2
p5 = p2*p3*p4
e_lv = - p1/p5 #Outputs the values for the energy levels you inputed

max_bound = -numpy.log10(numpy.abs(y_big))
e_lv_math = -numpy.log10(numpy.abs(e_lv_i))
e_lv_math2 = -numpy.log10(numpy.abs(e_lv_f))
d_e_math = e_lv_math2-e_lv_math

inc = 0
for big_int in e_lv:
    x_big = range(0,n_max)
    y_big = [big_int]*n_max
    plt.plot(x_big, -numpy.log10(numpy.abs(y_big)), color = colors[inc])
    inc+=1
    plt.xticks([])
    plt.yticks([])
    plt.ylabel('Relative Energy Levels')
    plt.arrow(4.5, e_lv_math, 0, d_e_math, width = 0)
    plt.tight_layout()
plt.legend(range(n_max,0,-1), loc="center right", bbox_to_anchor = (1.2,.5), title = "n=")

enter image description here

enter image description here

Zephyr
  • 11,891
  • 53
  • 45
  • 80
EJA1065
  • 25
  • 3
  • 1
    If I understand correctly you want the red label '1' to be at the bottom of the legend. Why not draw in reverse `for inc, big_int in enumerate(reversed(e_lv))` and apply your colors like `color=colors[len(e_lv)-1-inc]`. – Marc Aug 29 '21 at 18:24

1 Answers1

1

If you want the invert the order of legend lines and keeping labels as they are, you have to invert the order of legend handles.
You should get handles and labels from the axis where you are plotting:

handles, labels = ax.get_legend_handles_labels()

then pass them to the legend, inverting the order of handles:

ax.legend(handles = handles[::-1], labels = labels)

Example of working code

import matplotlib.pyplot as plt
    

fig, ax = plt.subplots(1, 2)

ax[0].plot([0, 1], [1, 1], color = 'red', label = 'line 1')
ax[0].plot([0, 1], [2, 2], color = 'blue', label = 'line 2')
ax[0].plot([0, 1], [3, 3], color = 'gold', label = 'line 3')

ax[0].legend(frameon = True, loc = 'upper left', bbox_to_anchor = (1.05, 1))

ax[1].plot([0, 1], [1, 1], color = 'red', label = 'line 1')
ax[1].plot([0, 1], [2, 2], color = 'blue', label = 'line 2')
ax[1].plot([0, 1], [3, 3], color = 'gold', label = 'line 3')

handles, labels = ax[1].get_legend_handles_labels()
ax[1].legend(handles = handles[::-1], labels = labels, frameon = True, loc = 'upper left', bbox_to_anchor = (1.05, 1))

plt.tight_layout()

plt.show()

enter image description here

Zephyr
  • 11,891
  • 53
  • 45
  • 80