How to plot multiple columns with a legend using python matplotlib?

user_17548328 0
Viewed 367

After transposing a dataframe with the following code:

groupedt = grouped.set_index('StateName').T
groupedt = groupedt.reset_index()

I have this resulting dataframe:

enter image description here

Is there a way for me to plot a line graph with every state on it that includes a legend of the states? I currently have the following code:

plt.figure(figsize=(16,6))
plt.plot(groupedt['index'], groupedt[groupedt.columns[2:41]])
plt.show()

which plots the time series for every state but I am not sure how I can create a legend that shows which state is which color?

matplotlib python
1 Answers

One way would be to:

  • Plot each column separately (with a label)
  • Add a legend

Not a MatPlotLib expert, there might be easier ways (maybe plot directly from the DataFrame?).

code00.py:

#!/usr/bin/env python

import sys

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


def main(*argv):
    shp = (4, 12)
    arr = np.arange(np.prod(shp)).reshape(shp).transpose()  # Generate some data for the dataframe
    df = pd.DataFrame(arr, columns=("Column {:d}".format(i + 1) for i in range(shp[0])))
    print(df)
    col_x = 0  # Column index for x axis
    plt.xlabel(df.columns[col_x])
    col_y_min = 1  # Minimum column index for lines
    col_y_max = len(df.columns)  # Maximum column index for lines
    for i, c in enumerate(df.columns[col_y_min:col_y_max], start=col_y_min):
        plt.plot(df[df.columns[col_x]], df[df.columns[i]], label=c)
    plt.legend()
    plt.show()


if __name__ == "__main__":
    print("Python {:s} {:03d}bit on {:s}\n".format(" ".join(elem.strip() for elem in sys.version.split("\n")),
                                                   64 if sys.maxsize > 0x100000000 else 32, sys.platform))
    rc = main(*sys.argv[1:])
    print("\nDone.")
    sys.exit(rc)

Output:

[cfati@CFATI-5510-0:e:\Work\Dev\StackOverflow\q071930291]> "e:\Work\Dev\VEnvs\py_pc064_03.09_test0\Scripts\python.exe" code00.py
Python 3.9.9 (tags/v3.9.9:ccb0e6a, Nov 15 2021, 18:08:50) [MSC v.1929 64 bit (AMD64)] 064bit on win32

    Column 1  Column 2  Column 3  Column 4
0          0        12        24        36
1          1        13        25        37
2          2        14        26        38
3          3        15        27        39
4          4        16        28        40
5          5        17        29        41
6          6        18        30        42
7          7        19        31        43
8          8        20        32        44
9          9        21        33        45
10        10        22        34        46
11        11        23        35        47

Done.

And the plot:

Img00

user_47885460
Related