Gráficos de dispersión con codificación de colores por valores de columna en Python

En la visualización de datos, asignar colores a diferentes categorías puede mejorar la claridad y revelar patrones. Esta funcionalidad está disponible en ggplot2 para R, pero ¿cómo podemos lograr lo mismo en Python usando pandas y matplotlib?

Actualización: Mejoras de Seaborn

Desde la respuesta original , Seaborn se ha convertido en una poderosa biblioteca para crear tramas informativas y visualmente atractivas. Sus actualizaciones recientes ofrecen funciones convenientes para colorear diagramas de dispersión basados ​​en valores de columnas:

• Uso de seaborn.relplot: esta función de alto nivel combina aspectos de matplotlib.pyplot.scatter y Seaborn. FacetaGrid. Maneja automáticamente la codificación de colores según los parámetros de tono y orden especificados.
• Asignación de matplotlib.pyplot.scatter a seaborn.FacetGrid: similar al enfoque original, puede asignar la función de dispersión a un FacetGrid y personaliza los colores según el tono.

Original Pandas y Matplotlib Enfoque

Para aquellos que buscan un enfoque directo con Matplotlib, aquí hay una función personalizada que asigna colores a puntos según una columna categórica:

import matplotlib.pyplot as plt
import pandas as pd

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

Esta función crea un diccionario de colores a partir de valores de categorías únicos y asigna los colores correspondientes a los puntos de datos. Luego, el diagrama de dispersión se genera con puntos codificados por colores.

Ejemplo

Usando el marco de datos de muestra proporcionado:

df = pd.DataFrame({'Height': np.append(np.random.normal(6, 0.25, size=5), np.random.normal(5.4, 0.25, size=5)),
                   'Weight': np.append(np.random.normal(180, 20, size=5), np.random.normal(140, 20, size=5)),
                   'Gender': ["Male", "Male", "Male", "Male", "Male",
                              "Female", "Female", "Female", "Female", "Female"]})

Llamando a la función dfScatter con el marco de datos:

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

Produce un diagrama de dispersión donde los puntos se colorean por género:

[Imagen del diagrama de dispersión coloreado por género]

Avanzado de Seaborn Las características y la función dfScatter personalizada brindan opciones flexibles para agregar codificación de colores a los diagramas de dispersión en Python, lo que hace que la visualización de datos sea más informativa y visual. atractivo.