Color by Column Values in Matplotlib

后端 未结 4 1528
情话喂你
情话喂你 2020-11-28 07:50

One of my favorite aspects of using the ggplot2 library in R is the ability to easily specify aesthetics. I can quickly make a scatterplot and apply color assoc

4条回答
  •  隐瞒了意图╮
    2020-11-28 08:36

    Update October 2015

    Seaborn handles this use-case splendidly:

    import numpy 
    import pandas
    from  matplotlib import pyplot
    import seaborn
    seaborn.set(style='ticks')
    
    numpy.random.seed(0)
    N = 37
    _genders= ['Female', 'Male', 'Non-binary', 'No Response']
    df = pandas.DataFrame({
        'Height (cm)': numpy.random.uniform(low=130, high=200, size=N),
        'Weight (kg)': numpy.random.uniform(low=30, high=100, size=N),
        'Gender': numpy.random.choice(_genders, size=N)
    })
    
    fg = seaborn.FacetGrid(data=df, hue='Gender', hue_order=_genders, aspect=1.61)
    fg.map(pyplot.scatter, 'Weight (kg)', 'Height (cm)').add_legend()
    

    Which immediately outputs:

    Old Answer

    In this case, I would use matplotlib directly.

    import numpy as np
    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
    
    if 1:
        df = pd.DataFrame({'Height':np.random.normal(size=10),
                           'Weight':np.random.normal(size=10),
                           'Gender': ["Male","Male","Unknown","Male","Male",
                                      "Female","Did not respond","Unknown","Female","Female"]})    
        fig = dfScatter(df)
        fig.savefig('fig1.png')
    

    And that gives me:

    scalle plot with categorized colors As far as I know, that color column can be any matplotlib compatible color (RBGA tuples, HTML names, hex values, etc).

    I'm having trouble getting anything but numerical values to work with the colormaps.

提交回复
热议问题