Writing variables as subscripts in math mode

Question

I am trying to plot some data, using a for loop to plot distributions. Now I want to label those distributions according to the loop counter as the subscript in math notation. This is where I am with this at the moment.

import matplotlib.pyplot as plt
import numpy as np
import matplotlib.mlab as mlab

mean = [10,12,16,22,25]
variance = [3,6,8,10,12]
x = np.linspace(0,40,1000)
for i in range(4):
    sigma = np.sqrt(variance[i])
    y = mlab.normpdf(x,mean[i],sigma)
    plt.plot(x,y,label=$v_i$) # where i is the variable i want to use to label. I should also be able to use elements from an array, say array[i] for the same.
    plt.xlabel("X")
    plt.ylabel("P(X)")
plt.legend()
plt.axvline(x=15, ymin=0, ymax=1,ls='--',c='black')    
plt.show()

This doesn't work, and I can't keep the variable between the $ $ signs of the math notation, as it is interpreted as text. Is there a way to put the variable in the $ $ notation?

Answer 1

The original question has been edited, this answer has been updated to reflect this.

When trying to work with LaTeX formatting in matplotlib you must use raw strings, denoted by r"".

The code given below will iterate over range(4) and plot using i'th mean and variance (as you originally have done). It will also set the label for each plot using label=r'$v_{}$'.format(i+1). This string formatting simply replaces the {} with whatever is called inside format, in this case i+1. In this way you can automate the labels for your plots.

I have removed the plt.axvline(...), plt.xlabel(...) and plt.ylabel(...) out of the for loop as you only need to call it once. I've also removed the plt.legend() from the for loop for the same reason and have removed its arguments. If you supply the keyword argument label to plt.plot() then you can label your plots individually as you plot them.

import matplotlib.pyplot as plt
import numpy as np
import matplotlib.mlab as mlab

mean = [10,12,16,22,25]
variance = [3,6,8,10,12]

x = np.linspace(0,40,1000)

for i in range(4):
    sigma = np.sqrt(variance[i])
    y = mlab.normpdf(x,mean[i],sigma)
    plt.plot(x,y, label=r'$v_{}$'.format(i+1))

plt.xlabel("X")
plt.ylabel("P(X)")        
plt.axvline(x=15, ymin=0, ymax=1,ls='--',c='black')

plt.legend()
plt.show()

Answer 2

So it turns out that you edited your question based on my answer. However, you;re still not quite there. If you want to do it the way I think you want to code it, it should be like this:

import matplotlib.pyplot as plt
import numpy as np
import matplotlib.mlab as mlab

mean = [10, 12, 16, 22, 25]
variance = [3, 6, 8, 10, 12]
x = np.linspace(0, 40, 1000)
for i in range(4):
    sigma = np.sqrt(variance[i])
    y = mlab.normpdf(x, mean[i], sigma)
    plt.plot(x, y, label = "$v_{" + str(i) + "}$") 
    plt.xlabel("X")
    plt.ylabel("P(X)")

plt.legend()
plt.axvline(x = 15, ymin = 0, ymax = 1, ls = '--', c = 'black')  

plt.show()

This code generates the following figure:

In case you want the first plot start with v_1 instead of v_0 all you need to change is str(i+1). This way the subscripts are 1, 2, 3, and 4 instead of 0, 1, 2 and 3.

Hope this helps!