hbaromega
hbaromega

Reputation: 2374

Plotting using for loops in Python

I have defined a 1D array x and a matrix y. Now I want to plot all rows or a specific row of y as function(s) of x. My code below shows 3 cases where case 1 successfully executes.

Can I rectify the syntaxes of the other two cases without using Pandas DataFrame? Also, what if I want to do the same using list instead of array?

import numpy as np
import matplotlib.pyplot as plt

x = np.array([0, 1, 2, 3, 4])
y = np.array([[0, 1, 2, 3, 4],
     [0, 1, 4, 9, 16],
     [0, 1, 8, 27, 64],
     [0, 1, 16, 81, 256]])
print('x =',x)
print('y =',y)

''' Case 1: Want to plot all rows of y-matrix as different plots '''
plt.plot(x, y.T) # works

print('y[2,:]=',y[2,:])
y2 = y[2,:]  # Cube of x

''' Case 2: Want to plot row 2 of y-matrix '''
for i in range(5):  # Doesn't work
  print(x[i], y2[i])
  plt.plot(x[i], y2[i])

''' Case 3: Want to plot all rows of y-matrix as different plots'''
for i in range(5):  # Doesn't work
  for j in range(4):
    plt.plot(x[i], y[i][j])
plt.show()

Upvotes: 0

Views: 2515

Answers (1)

JohanC
JohanC

Reputation: 80279

If you want your points connected with lines, the complete curve needs to be given to plt.plot in one go. Therefore, your second case should be plt.plot(x, y2), and similar for the third.

Note that in Python, the use of explicit indices is discouraged. Instead of writing the third case as

for j in range(4):
    plt.plot(x, y[j])

The "pythonic" way to write the loop would be:

for y_j in y:
    plt.plot(x, y_j)

Here is the updated code:

import numpy as np
import matplotlib.pyplot as plt

x = np.array([0, 1, 2, 3, 4])
y = np.array([[0, 1, 2, 3, 4],
              [0, 1, 4, 9, 16],
              [0, 1, 8, 27, 64],
              [0, 1, 16, 81, 256]])

''' Case 1: Want to plot all rows of y-matrix as different plots '''
plt.plot(x, y.T)

y2 = y[2, :]  # Cube of x
''' Case 2: Want to plot row 2 of y-matrix '''
plt.plot(x, y2)

''' Case 3: Want to plot all rows of y-matrix as different plots'''
for y_j in y:
    plt.plot(x, y_j)
plt.show()

The pure python version would be:

import numpy as np
import matplotlib.pyplot as plt

x = [0, 1, 2, 3, 4]
y = [[0, 1, 2, 3, 4],
     [0, 1, 4, 9, 16],
     [0, 1, 8, 27, 64],
     [0, 1, 16, 81, 256]]

''' Case 1: Want to plot all rows of y-matrix as different plots '''
y_trans = list(map(list, zip(*y)))
plt.plot(x, y_trans)

y2 = y[2]  # Cube of x
''' Case 2: Want to plot row 2 of y-matrix '''
plt.plot(x, y2)

''' Case 3: Want to plot all rows of y-matrix as different plots'''
for y_j in y:
    plt.plot(x, y_j)
plt.show()

In general, it is recommended to use numpy arrays, especially if you want to make more complicated calculations, or when speed is important for larger lists. Numpy's popularity for plotting (and almost all other standard precision calculations) is due to broadcasting and vectorization. Broadcasting and vectorization often make the code easier to read and write, as well as much faster.

For example, to get an array of squares, you simply write y2 = x**2 with x = np.arange(5).

Upvotes: 1

Related Questions