How to pivot a dataframe to plot bars

Question

I am trying to create a stacked bar chart, and color the bars using color names I have stored in a variable party_color.

This is my attempt:

import pandas as pd
import matplotlib.pyplot as plt

marginal_electorates_2016 = {'margin': {0: 'Fairly safe', 1: 'Fairly safe', 2: 'Fairly safe', 3: 'Marginal', 4: 'Marginal', 5: 'Marginal', 6: 'Marginal', 7: 'Safe', 8: 'Safe', 9: 'Safe', 10: 'Safe'},
                             'PartyNm': {0: 'Australian Labor Party', 1: "Katter's Australian Party", 2: 'Liberal/National Coalition', 3: 'Australian Labor Party', 4: 'Independent', 5: 'Liberal/National Coalition', 6: 'Nick Xenophon Team', 7: 'Australian Labor Party', 8: 'Independent', 9: 'Liberal/National Coalition', 10: 'The Greens'},
                             'count': {0: 32, 1: 1, 2: 29, 3: 24, 4: 1, 5: 28, 6: 1, 7: 13, 8: 1, 9: 19, 10: 1},
                             'party_color': {0: 'red', 1: 'yellow', 2: 'blue', 3: 'red', 4: 'pink', 5: 'blue', 6: 'orange', 7: 'red', 8: 'pink', 9: 'blue', 10: 'green'}}

marginal_electorates_2016 = pd.DataFrame(marginal_electorates_2016)

         margin                     PartyNm  count party_color
0   Fairly safe      Australian Labor Party     32         red
1   Fairly safe   Katters Australian Party       1      yellow
2   Fairly safe  Liberal/National Coalition     29        blue
3      Marginal      Australian Labor Party     24         red
4      Marginal                 Independent      1        pink
5      Marginal  Liberal/National Coalition     28        blue
6      Marginal          Nick Xenophon Team      1      orange
7          Safe      Australian Labor Party     13         red
8          Safe                 Independent      1        pink
9          Safe  Liberal/National Coalition     19        blue
10         Safe                  The Greens      1       green

plt.figure(figsize=(16, 6))

marginal_electorates_2016.plot(
    kind = 'bar',
    x = ['margin', 'PartyNm'],
    y = 'count',
    stacked = False,
    subplots = True,
    figsize = [10,15],
    sharey = True,
    c = 'party_colour'
    )
plt.tight_layout

Trenton McKinney · Accepted Answer · 2021-12-04T20:00:04.713

The easiest way to plot groups of bars is to reshape the dataframe from a long to wide format with pandas.DataFrame.pivot
Create a dict of colors with the party as the key and the color as the value, which can be passed to the color parameter when plotting.
- There are a number of ways to convert two column to a dict with one column as the keys and another as the values. It doesn't matter which way is used. How to create a dictionary of two pandas DataFrame columns
Plot the dataframe directly with pandas.DataFrame.plot and kind='bar'. Use stacked=True for stacked bars if desired, but I do not recommend it, because it makes the data more difficult to compare.

import pandas as pd
import matplotlib.pyplot as plt

# create the dataframe
df = pd.DataFrame(marginal_electorates_2016)

# transform the shape
dfp = df.pivot(index='margin', columns='PartyNm', values='count')

# display(dfp)
PartyNm      Australian Labor Party  Independent   Katters Australian Party  Liberal/National Coalition  Nick Xenophon Team  The Greens
margin                                                                                                                                 
Fairly safe                    32.0          NaN                        1.0                        29.0                 NaN         NaN
Marginal                       24.0          1.0                        NaN                        28.0                 1.0         NaN
Safe                           13.0          1.0                        NaN                        19.0                 NaN         1.0

# create a colors dict
colors = dict(df[['PartyNm', 'party_color']].drop_duplicates().to_numpy())

# plot
ax = dfp.plot(kind='bar', rot=0, color=colors, figsize=(16, 6))
ax.legend(bbox_to_anchor=(1, 1.02), loc='upper left')

with stacked=True

How to pivot a dataframe to plot bars

1 Answers1

Related