Retrieve data from pandas dataframe to plot in seaborn

Viewed 80

This question has probably been answered a million times; and yet it is still difficult to begin working with pandas and to get used to how the dataframe works.

Thus, if someone could propose a better version of my method to retrieve the data of interest from the dataframe, I would be much obliged.

I have a dataframe which stores scores, which can be simulated with:

import random
import pandas as pd


columns = ['ID', 'Number', 'Score 1', 'Score 2', 'Score 3', 'Score 4']
df = pd.DataFrame(columns=columns)

for k in range(100):
    idx = random.choice([10, 11, 12, 13, 14])
    number = random.choice([1, 2, 3, 4])
    scores = [random.choice(range(1000)) for _ in range (4)]
    entry = pd.DataFrame([[idx, number] + scores], columns=df.columns)
    df = df.append(entry)

The trick is that for a given ID, a given number N can appear X times, with X non-constant. However, for each line in the dataframe, the number of scores is fixed, i.e. in this example 4.

I want to select the score 4 for a given ID and group them by number. e.g.

# Retrieve scores of ID 10
id_scores = df.loc[df['ID'] == 10]
numbers = id_scores.Number.unique()

data = list()
for n in numbers:
    data.append(id_scores.loc[id_scores['Number'] == n]['Score 4'].values)

The code above works.. but I don't need to be proficient with pandas to understand that this is not the correct way to do it. My goal is then to plot the score 4 for a given number (x-Axis) with boxplots. Thus, every box might not contain the same number of scores because the value number for a given ID is not fixed.

3 Answers

You can use pandas groupby for this:

df.loc[df['ID'] ==10].groupby('Number')['Score 4'].apply(list)

gives (e.g):

Number
1     [528, 501, 69, 251, 387]
2              [549, 321, 302]
3    [409, 720, 629, 699, 803]
4    [524, 228, 682, 852, 525]
Name: Score 4, dtype: object

let me know if this answers your question.

What you're doing here is a typical SAC pattern

But first, a little advice for performance optimization on your df creation:

# I would use np.random to avoid any explicit looping with *100 func calls
# (np.random is vectorized)

import numpy as np
import pandas as pd

idx = np.random.randint(10, 15, 100)
number = np.random.randint(0, 5, 100)
scores = np.random.randint(0, 1000, (4, 100))  # 2D array of shape (4, 100)
columns = ['ID', 'Number', 'Score 1', 'Score 2', 'Score 3', 'Score 4']

# create dataframe from series representing columns and transpose it 
# idx will be the df index to allow for fast df.loc[ID] indexing instead of
# using bool mask (e.g. df[df["ID"] == 10])
df = pd.DataFrame([number, *scores], index=columns, columns=idx).T

Now you can get the matching rows for a specific ID and group them by Number while returning the Score 4 values as a sequence of your preference

df.loc[10].groupby("Number")["Score 4"].apply(tuple)

You can even chain .to_dict() if you prefer a Py native data structure.

The answers above did result in a similar output to the one with my looping code, but neither was actually the correct format for seaborn. So just in case, here is a different approach:

df = pd.read_csv(csv)
df = df.loc[df['ID'] == 10]
df = pd.melt(
    df, id_vars='Number', value_vars='Score 4',
    var_name='Score ID', value_name='Score')

Which can then be used with seaborn:

sns.boxplot(x='Number', y='Score', hue='Score ID', data=df)
Related