CODEWITHSUNDEEP
pythonpandasdataframeplot
Fourier
Fourier

Reputation: 2993

Create grouped/stacked bar plots from multiple categories containing several labels inside a pandas dataframe

I have the following pandas dataframe (df) [only an excerpt of the full dataframe]:

   Name    Cat_1    Cat_2
0   foo        P    Apples, Pears, Cats
1   bar     R, M    Apples
2   bla        E    Pears
3   blu        F    Cats, Pears
4   boo        G    Apples, Pears
5   faa     P, E    Apples, Cats

I would like to create bar plots that are build from Cat_1 and Cat_2. These columns contain multiple tags, which have to be use for plotting.

Currently, I am running this simple code to plot Cat_1:

import pandas as pd
from matplotlib import pyplot as plt

fig, ax = plt.subplots(figsize = (4,4))
s = df["Cat_1"].str.split(", ", expand = True).stack()
s.value_counts().plot(kind = 'bar', ax = ax)

This returns a nice bar plot for each of the different labels in Cat_1 allowing multiple assignments (as intended).

One could apply the same to Cat_2 and obtain a separate plot with the respective labels.

However, I want to have a single plot that is first "stacked" by Cat_1 and subsequently the values are counted for Cat_2.

I guess a way to think of this is to build a nested dictionary that would look like the following:

{"P": {"Apples": 2, "Pears": 1, "Cats": 2}, "R": {"Apples": 1}, ....}

but at the same time keep track of the total count of Cat_1.

It does not matter whether its a grouped or stacked bar chart in the end.

Please take a look a the enclosed figure for a more visual idea:

enter image description here

Upvotes: 0

Views: 817

Answers (1)

ak_slick
ak_slick

Reputation: 1016

This should get you pretty close if I understand correctly.

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

df = pd.DataFrame(columns=['Name', 'Cat_1', 'Cat_2'])

df['Name'] = ['foo', 'bar', 'bla', 'blu', 'boo', 'faa']
df['Cat_1'] = ['P', 'R, M', 'E', 'F', 'G', 'P, E']
df['Cat_2'] = ['Apples, Pears, Cats', 'Apples', 'Pears', 'Cats, Pears', 'Apples, Pears', 'Apples, Cats']

# arrange data simply prepopulate with zero
df_pl = pd.DataFrame(columns=df["Cat_1"].str.split(", ", expand=True).stack().unique().tolist(),
                     index=df["Cat_2"].str.split(", ", expand=True).stack().unique().tolist(),
                     data=0)

# get chunk size for each combination
for x in df_pl.columns:
    ind = df.Cat_1.str.contains(x)
    for name in df_pl.index:
        df_pl.set_value(name, x, df.loc[ind, 'Cat_2'].str.contains(name).sum())

N = len(df_pl.columns)
ind = np.arange(N)    # the x locations for the groups
width = 0.35       # the width of the bars: can also be len(x) sequence

plotted = []
p = {}
for name in df_pl.index:
    bottoms = df_pl.index.isin(plotted).sum()
    p[name] = plt.bar(ind, df_pl.loc[name].values.tolist(), bottom=bottoms)
    plotted.append(name)

plt.ylabel('y_label')
plt.title('some plot')
plt.xticks(ind, df_pl.columns.tolist())
plt.legend(p.values(), p.keys())

plt.show()

Example Output

Upvotes: 1

Related Questions