Colorindo gráficos de dispersão por valores de coluna usando Pandas e Matplotlib

Matplotlib é uma biblioteca Python popular para criar visualizações estáticas, animadas e interativas em Pitão. Este artigo explora o uso do Matplotlib para colorir gráficos de dispersão com base em valores em uma coluna específica de um DataFrame do Pandas.

Importações e dados

Para começar, importamos as bibliotecas necessárias , incluindo Matplotlib (como plt) e Pandas (como pd). Também geramos um DataFrame de amostra ("df") com três colunas: "Altura", "Peso" e "Gênero".

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

np.random.seed(0)
N = 37
_genders = ["Female", "Male", "Non-binary", "No Response"]
df = pd.DataFrame({
    "Height (cm)": np.random.uniform(low=130, high=200, size=N),
    "Weight (kg)": np.random.uniform(low=30, high=100, size=N),
    "Gender": np.random.choice(_genders, size=N),
})

Atualização em agosto de 2021

Seaborn introduziu novas funções em nível de figura, como seaborn.relplot na versão 0.11.0. Essas funções são recomendadas em vez do uso direto do FacetGrid.

sns.relplot(data=df, x="Weight (kg)", y="Height (cm)", hue="Gender", hue_order=_genders, aspect=1.61)
plt.show()

Resposta antiga (2015)

Se desejar usar o Matplotlib diretamente, você precisará mapear a função de dispersão do matplotlib nas categorias de um Pandas DataFrame . Para fazer isso:

• Crie um dicionário com categorias exclusivas da coluna e cores.
• Adicione uma nova coluna "Cor" ao DataFrame, atribuindo a cada categoria uma cor correspondente.
• Use a função de dispersão para plotar os dados, especificando a coluna de cores como o argumento "c".

def dfScatter(df, xcol='Height', ycol='Weight', catcol='Gender'):
    fig, ax = plt.subplots()
    categories = np.unique(df[catcol])
    colors = np.linspace(0, 1, len(categories))
    colordict = dict(zip(categories, colors))

    df["Color"] = df[catcol].apply(lambda x: colordict[x])
    ax.scatter(df[xcol], df[ycol], c=df.Color)
    return fig

fig = dfScatter(df)
fig.savefig('fig1.png')

Seguindo essas etapas, você pode facilmente colorir gráficos de dispersão com base nos valores das colunas usando Pandas e Matplotlib.