Legend in Matplotlib -- Subplotting by a For loop-CodePudding

I'm new to Python and Matplotlib, I would appreciate any help on how to create legends for every subplot that I have created with a FOR loop. Here is the code and the best that I could get close to labeling my figures.

import matplotlib.pyplot as plt
import numpy as np
n_rows=2
n_cols=2
leg=['A','B','C','D']
fig, axes = plt.subplots(n_rows,n_cols)
for row_num in range(n_rows):
    for col_num in range (n_cols):
    ax = axes[row_num][col_num]
    ax.plot(np.random.rand(20))
    ax.set_title(f'Plot ({row_num 1}, {col_num 1})')
    ax.legend(leg[row_num col_num])   
fig.suptitle('Main Title')
fig.tight_layout()               
plt.show()

here is the output of the code: Image with incorrect legends

CodePudding user response：

You are using ax.legend(leg[row_num col_num]) but row_num col_num is not a correct representation for an index of a list.

This is what is happening

row_num | col_num | idx=row_num col_num |  leg[idx]
    0   |     0   |  0                  |    A
    0   |     1   |  1                  |    B
    1   |     0   |  1                  |    B
    1   |     1   |  2                  |    C

If you use leg[row_num col_num] you get the incorrect legend entry.

There are many ways to fix this. A simple one is to introduce a counter (the variable j in the code below), which increments at every loop.

import matplotlib.pyplot as plt
import numpy as np
n_rows=2
n_cols=2
leg=['A','B','C','D']
fig, axes = plt.subplots(n_rows,n_cols)
j = 0
for row_num in range(n_rows):
    for col_num in range(n_cols):
        ax = axes[row_num][col_num]
        ax.plot(np.random.rand(20))
        ax.set_title(f'Plot ({row_num 1}, {col_num 1})')
        ax.legend(leg[j]) 
        j  = 1
fig.suptitle('Main Title')
fig.tight_layout()               
plt.show()