Search code examples
pythonmatplotliblegendmathematical-expressions

Writing variables as subscripts in math mode

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?

Solution

  • 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()

    Plot