Subplotting after groupby with pandas plot

Question

I have a dataframe like so:

df = pd.DataFrame({"idx":[1,2,3]*2,"a":[1]*3+[2]*3,'b':[3]*3+[4]*3,'grp':[4]*3+[5]*3})
df = df.set_index("idx")
df
     a  b  grp
idx           
1    1  3    4
2    1  3    4
3    1  3    4
1    2  4    5
2    2  4    5
3    2  4    5

and I would like to plot the values of a and b as function of idx . Making one subplot per column and one line per group.

I manage to do this creating axis separately and iterating over groups as proposed here . But I would like to use the subplots parameter of the plot function to avoid looping. I tried solutions like

df.groupby("grp").plot(subplots=True)

But it plot the groups in different subplots and removing the groupby does not make appear the two separated lines as in the example.

Is it possible? Also is it better to iterate and use matplotlib plot or use pandas plot function?

Answer 1

IIUC, you can do something like this:

axs = df.set_index('grp', append=True)\
  .stack()\
  .unstack('grp')\
  .rename_axis(['idx','title'])\
  .reset_index('title').groupby('title').plot()

[v.set_title(f'{i}') for i, v in axs.items()]

Output:

Maybe eaiser to simple loop and plot:

fig, ax = plt.subplots(1,2, figsize=(10,5))
ax = iter(ax)
for n, g in df.set_index('grp', append=True)\
              .stack()\
              .unstack('grp')\
              .rename_axis(['idx','title'])\
              .reset_index('title').groupby('title'):
    g.plot(ax=next(ax), title=f'{n}')

Output:

Answer 2

If i understod your question correct, you can access columns and rows in a pandas dataframe. An example can be like this:

import numpy as np
import matplotlib.pyplot as plt
x = np.array(df['idx']) 

a = np.array(df['a']) 

b = np.array(df['b']) 

plt.subplot(1,2,1)#(121) will also work
#fill inn title etc for the first plot under here
plt.plot(x,a)
plt.subplot(1,2,2)
#fill inn title etc for the second plot under here
plt.plot(x,b)
plt.show()

edit: Sorry now changed for subplot.