Search code examples
pythonlistloopscolorsjupyter-notebook

Iterate list of colors regardless of length of index of dataframes - Jupyter Notebook


I am using Jupyter notebook. I have DataFrames that when created are never the same length and would like to apply a list of colors I define on a unique value of a column. The DataFrames do share the same column I intend to color based on unique values. Where I am getting stuck is if the length of the DataFrame index is greater than the number of colors listed in my list I receive an error saying "IndexError: list index out of range". I would like to continue to iterate with my list of colors regardless the length of unique values in the below example, column x1.

Below I have a one off MRE that should be able to test. You will note that df3 works correctly due to column x1 having the same length of values as colors in my list whereas df4 throws an error due to column x1 having more values than colors in my list.

df1 = pd.DataFrame({'x1':['a', 'b', 'b', 'c', 'd', 'd'],
                     'x2':['3', '4', '8', '0', '11', '1']})


df2 = pd.DataFrame({'x1':['a', 'a', 'c', 'd', 'e', 'e', 'f', 'g', 'g', 'g', 'h'],
                     'x2':['0', '41', '22', '5', '19', '21', '5', '7', '8', '24', '15']})
def Color_Unique(s):
    df = s.copy()
    color_map = {}

    Trade_Cusip_Combo_Key = df['x1'].unique()
    colors_to_use = ['background-color: #ADD8E6', 'background-color: #90ee90', 'background-color: #FFD580', 'background-color: #CBC3E3', 'background-color: #D3D3D3', 'background-color: #C4A484']

    for Trade_Cusip_Combo in Trade_Cusip_Combo_Key:
        color_map[Trade_Cusip_Combo] = colors_to_use[0]
        colors_to_use.pop(0)

    for index, row in df.iterrows():
        if row['x1'] in Trade_Cusip_Combo_Key:
            Trade_Cusip_Combo = row['x1']
            my_color = color_map[Trade_Cusip_Combo]
            df.loc[index,:] = my_color
        else:
            df.loc[index,:] = 'background-color: '        
    return df
df3 = df1.style.apply(Color_Unique, axis=None)
df3
df4 = df2.style.apply(Color_Unique, axis=None)
df4

Solution

  • You can easily loop through your list of colors, going back to the start once you reach the end, by using itertools.cycle that will cycle through the list indefinitely.

    We first create the cycle object with:

    from itertools import cycle
    
    
    colors_to_use = ['background-color: #ADD8E6', 'background-color: #90ee90', 
                     'background-color: #FFD580', 'background-color: #CBC3E3',
                     'background-color: #D3D3D3', 'background-color: #C4A484']
    colors_cycle = cycle(colors_to_use)
    

    Then you can get the next color with the next function:

    print(next(colors_cycle))
    print(next(colors_cycle))
    

    outputs:

    background-color: #ADD8E6
    background-color: #90ee90
    

    So, the coloring part of your code becomes: (don't forget from itertools import cycle at the beginning of your script)

    colors_to_use = ['background-color: #ADD8E6', 'background-color: #90ee90', 'background-color: #FFD580', 'background-color: #CBC3E3', 'background-color: #D3D3D3', 'background-color: #C4A484']
    colors_cycle = cycle(colors_to_use)
    
    for Trade_Cusip_Combo in Trade_Cusip_Combo_Key:
        color_map[Trade_Cusip_Combo] = next(colors_cycle)
    

    and you get this output, where colors start looping again from blue after brown:

    df4 with cycling colors

    The complete code (with imports) of the Jupyter cell:

    from itertools import cycle
    import pandas as pd
    
    
    df2 = pd.DataFrame({'x1':['a', 'a', 'c', 'd', 'e', 'e', 'f', 'g', 'g', 'g', 'h'],
                         'x2':['0', '41', '22', '5', '19', '21', '5', '7', '8', '24', '15']})
    
    def Color_Unique(s):
        df = s.copy()
        color_map = {}
    
        Trade_Cusip_Combo_Key = df['x1'].unique()
        colors_to_use = ['background-color: #ADD8E6', 'background-color: #90ee90', 'background-color: #FFD580', 'background-color: #CBC3E3', 'background-color: #D3D3D3', 'background-color: #C4A484']
        colors_cycle = cycle(colors_to_use)
    
        for Trade_Cusip_Combo in Trade_Cusip_Combo_Key:
            color_map[Trade_Cusip_Combo] = next(colors_cycle)
    
        for index, row in df.iterrows():
            if row['x1'] in Trade_Cusip_Combo_Key:
                Trade_Cusip_Combo = row['x1']
                my_color = color_map[Trade_Cusip_Combo]
                df.loc[index,:] = my_color
            else:
                df.loc[index,:] = 'background-color: '        
        return df
    
    df4 = df2.style.apply(Color_Unique, axis=None)
    df4